Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.en.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# PR #2290: Bound Nemotron-H SSD prefill memory

**Date**: 2026-10-13
**Status**: Complete for the bounded SSD-prefill scope; implementation reviewed and focused validation passed
**Risk**: High at issue scope because request-level transient-memory admission remains unimplemented

## Summary

PR #2290 changes Nemotron-H Mamba prefill from one prompt-length quadratic surrogate-attention graph to checkpoint-sized SSD chunks. Each chunk carries the recurrent state forward and is evaluated before the next chunk is constructed, reducing the measured 6,019-token process peak from 72.2 GiB to 18.7 GiB on the tested M1 Ultra checkpoint while preserving substantive greedy outputs.

## Problem

The prior `nemotron_h` prefill path constructed an `L × L` surrogate-attention matrix in every Mamba layer. Long prompts therefore produced a transient memory peak that grew roughly quadratically and could terminate the process in the Metal driver. Issue #2281 also requires predictable oversized requests to be rejected, but the server does not yet estimate per-request activation or transient memory.

## Design and review

- `NemotronHConfig` now reads `chunk_size`, accepts the legacy `mamba_chunk_size` alias, defaults to 128, and rejects zero before model construction.
- `nemotron_h_ssm_attn` preserves float32 timestep, state, and matrix-product boundaries; computes each quadratic block independently; evaluates its output and next state; and uses that state as the prior state for the following block.
- Focused numerical coverage compares a 2+2+1 partial-tail partition with a full step carrying a nonzero prior state, and compares multi-batch/multi-group f16 and bf16 results with a scalar recurrence oracle.
- The implementation is wired through the production Nemotron-H model builder and is exercised by the real server benchmark. No code-level correctness defect remained within the declared partial scope.

## Benchmark evidence

Three fresh-process repetitions were recorded for each 1,519-, 3,019-, 4,519-, and 6,019-token cell. Arm-specific load and forward markers prevent the harness from reporting a comparison unless the candidate executes the new path and the baseline does not. Candidate peak memory remained between 17.6 and 18.7 GiB across those cells and reached 19.3 GiB for three 8,019-token runs. Three substantive greedy prompts, including a 2,810-token prompt, produced identical choices in both arms. The 8,019-token baseline, the reporter's M4 Pro and 2-bit checkpoint, M5, CUDA, and latency under a quiet host were explicitly left unmeasured.

## Verification

- Release Metal/Accelerate `mlxcel-server` build: passed.
- Nemotron-H module suite: 22 passed, 3 ignored.
- Dedicated serial Metal numerical tests: 2 passed.
- Formatting, diff hygiene, Python harness compilation, benchmark artifact consistency, and cross-repository reference review: passed.
- Current CI formatting, Clippy, CPU-only link, license, Python tooling, and repository contract checks observed during review: passed; remaining jobs were still running or queued.

## Remaining work

Request-level activation/transient-memory admission remains required to address arbitrary within-context Metal command-buffer OOM that can still terminate the process. The user approved tracking and implementing that work as a separate follow-up after this bounded SSD-prefill fix merges; no follow-up issue number exists yet. The PR body's non-closing wording was rewritten after GitHub interpreted the earlier negated phrase as a closing keyword, and `closingIssuesReferences` is now empty.
36 changes: 36 additions & 0 deletions TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.ko.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# PR #2290: Nemotron-H SSD prefill 메모리 제한

**날짜**: 2026-10-13
**상태**: 제한된 SSD-prefill 범위 완료; 구현 리뷰와 집중 검증 통과
**위험도**: 이슈 범위에서는 높음. 요청별 transient-memory admission이 아직 구현되지 않음

## 요약

PR #2290은 Nemotron-H Mamba prefill을 prompt 전체 길이의 quadratic surrogate-attention graph 하나로 처리하던 방식에서 checkpoint 크기의 SSD chunk로 나누는 방식으로 변경합니다. 각 chunk는 recurrent state를 다음 chunk로 전달하고 다음 graph를 만들기 전에 평가됩니다. 테스트한 M1 Ultra checkpoint에서 6,019-token process peak가 72.2 GiB에서 18.7 GiB로 감소했고, 실질적인 greedy output은 유지되었습니다.

## 문제

기존 `nemotron_h` prefill 경로는 모든 Mamba layer에서 `L × L` surrogate-attention matrix를 생성했습니다. 따라서 긴 prompt의 transient memory peak가 대략 quadratic하게 증가했고 Metal driver에서 process가 종료될 수 있었습니다. Issue #2281은 예측 가능한 oversized request의 거부도 요구하지만, 현재 server에는 request별 activation 또는 transient memory 추정 기능이 없습니다.

## 설계와 리뷰

- `NemotronHConfig`는 이제 `chunk_size`를 읽고 기존 `mamba_chunk_size` alias를 허용하며, 기본값 128을 사용하고 model 생성 전에 0을 거부합니다.
- `nemotron_h_ssm_attn`은 float32 timestep, state, matrix-product 경계를 유지하고 각 quadratic block을 독립적으로 계산합니다. 각 block의 output과 next state를 평가한 뒤 그 state를 다음 block의 prior state로 사용합니다.
- 집중 수치 검증은 nonzero prior state를 전달하는 2+2+1 partial-tail partition을 full step과 비교하고, multi-batch/multi-group f16 및 bf16 결과를 scalar recurrence oracle과 비교합니다.
- 구현은 production Nemotron-H model builder에 연결되어 있고 실제 server benchmark에서 실행됩니다. 선언된 부분 범위 안에서는 남은 code-level correctness 결함을 발견하지 못했습니다.

## Benchmark 근거

1,519, 3,019, 4,519, 6,019 token의 각 cell에서 fresh process로 3회 반복했습니다. Arm별 load 및 forward marker를 사용해 candidate가 새 경로를 실행하고 baseline은 실행하지 않은 경우에만 비교를 보고합니다. Candidate peak memory는 이 범위에서 17.6–18.7 GiB였고, 8,019-token 실행 3회는 19.3 GiB였습니다. 2,810-token prompt를 포함한 실질적인 greedy prompt 3개에서 두 arm의 choice가 동일했습니다. 8,019-token baseline, 제보자의 M4 Pro 및 2-bit checkpoint, M5, CUDA, quiet host에서의 latency는 미측정으로 명시했습니다.

## 검증

- Metal/Accelerate release `mlxcel-server` build: 통과.
- Nemotron-H module suite: 22개 통과, 3개 ignored.
- 전용 serial Metal 수치 테스트: 2개 통과.
- Formatting, diff hygiene, Python harness compilation, benchmark artifact consistency, cross-repository reference 검토: 통과.
- 리뷰 중 확인한 현재 CI formatting, Clippy, CPU-only link, license, Python tooling, repository contract 검사: 통과. 나머지 job은 실행 또는 대기 중이었습니다.

## 남은 작업

요청별 activation/transient-memory admission은 임의의 within-context Metal command-buffer OOM이 여전히 process를 종료할 수 있는 문제를 해결하기 위해 필요합니다. 사용자는 이 제한된 SSD-prefill 수정이 merge된 뒤 해당 작업을 별도 후속으로 추적하고 구현하는 것에 동의했으며, 아직 후속 issue 번호는 없습니다. 이전 PR 본문의 부정형 표현을 GitHub가 closing keyword로 해석하여 해당 문구를 다시 작성했고, 이제 `closingIssuesReferences`는 비어 있습니다.
Loading
Loading