diff --git a/TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.en.md b/TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.en.md new file mode 100644 index 000000000..3b9e86a72 --- /dev/null +++ b/TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.en.md @@ -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. diff --git a/TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.ko.md b/TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.ko.md new file mode 100644 index 000000000..ff31f6716 --- /dev/null +++ b/TECHNICAL_REPORTS/2290-nemotron-h-prefill-memory-20261013.ko.md @@ -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`는 비어 있습니다. diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/baseline-memory.json b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/baseline-memory.json new file mode 100644 index 000000000..a981cb3fa --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/baseline-memory.json @@ -0,0 +1,442 @@ +{ + "arm": "baseline", + "binary": "/tmp/mlxcel-chain-2281-artifacts/mlxcel-server.baseline", + "model": "models/mlx/nvidia-nemotron-3-nano-30b-a3b-4bit", + "marker_expectation": "absent", + "started_at": "2026-10-13T22:01:54+0900", + "runs": [ + { + "rep": 1, + "target_tokens": 1500, + "status": 200, + "prompt_tokens": 1519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 22441204121, + "peak_gib": 20.899999999441206, + "ready_s": 3.8478498458862305, + "request_s": 6.278626203536987, + "load_before": [ + 5.537109375, + 6.67529296875, + 5.84716796875 + ], + "load_after": [ + 5.447265625, + 6.61962890625, + 5.83642578125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 1, + "target_tokens": 3000, + "status": 200, + "prompt_tokens": 3019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 33715493273, + "peak_gib": 31.399999999441206, + "ready_s": 2.0463509559631348, + "request_s": 10.466207027435303, + "load_before": [ + 5.447265625, + 6.61962890625, + 5.83642578125 + ], + "load_after": [ + 5.72412109375, + 6.63427734375, + 5.85546875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 1, + "target_tokens": 4500, + "status": 200, + "prompt_tokens": 4519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 52076478464, + "peak_gib": 48.5, + "ready_s": 2.032823085784912, + "request_s": 17.74728512763977, + "load_before": [ + 5.72412109375, + 6.63427734375, + 5.85546875 + ], + "load_after": [ + 5.2587890625, + 6.46337890625, + 5.8115234375 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 1, + "target_tokens": 6000, + "status": 200, + "prompt_tokens": 6019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 77524159692, + "peak_gib": 72.19999999925494, + "ready_s": 2.043994903564453, + "request_s": 25.449132680892944, + "load_before": [ + 5.2587890625, + 6.46337890625, + 5.8115234375 + ], + "load_after": [ + 5.51806640625, + 6.45458984375, + 5.83251953125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 2, + "target_tokens": 1500, + "status": 200, + "prompt_tokens": 1519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 22441204121, + "peak_gib": 20.899999999441206, + "ready_s": 3.81768798828125, + "request_s": 4.3409810066223145, + "load_before": [ + 5.51806640625, + 6.45458984375, + 5.83251953125 + ], + "load_after": [ + 5.130859375, + 6.34033203125, + 5.79931640625 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 2, + "target_tokens": 3000, + "status": 200, + "prompt_tokens": 3019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 33715493273, + "peak_gib": 31.399999999441206, + "ready_s": 2.044321060180664, + "request_s": 9.221702814102173, + "load_before": [ + 5.130859375, + 6.34033203125, + 5.79931640625 + ], + "load_after": [ + 4.76953125, + 6.20703125, + 5.76123046875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 2, + "target_tokens": 4500, + "status": 200, + "prompt_tokens": 4519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 52076478464, + "peak_gib": 48.5, + "ready_s": 2.0435571670532227, + "request_s": 15.412789821624756, + "load_before": [ + 4.76953125, + 6.20703125, + 5.76123046875 + ], + "load_after": [ + 4.5732421875, + 6.06396484375, + 5.7197265625 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 2, + "target_tokens": 6000, + "status": 200, + "prompt_tokens": 6019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 77524159692, + "peak_gib": 72.19999999925494, + "ready_s": 2.0413269996643066, + "request_s": 24.723119974136353, + "load_before": [ + 4.5732421875, + 6.06396484375, + 5.7197265625 + ], + "load_after": [ + 4.0166015625, + 5.81591796875, + 5.63916015625 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 3, + "target_tokens": 1500, + "status": 200, + "prompt_tokens": 1519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 22441204121, + "peak_gib": 20.899999999441206, + "ready_s": 2.037729024887085, + "request_s": 4.327772855758667, + "load_before": [ + 4.0166015625, + 5.81591796875, + 5.63916015625 + ], + "load_after": [ + 3.93994140625, + 5.7392578125, + 5.61376953125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 3, + "target_tokens": 3000, + "status": 200, + "prompt_tokens": 3019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 33715493273, + "peak_gib": 31.399999999441206, + "ready_s": 2.035287857055664, + "request_s": 9.245609045028687, + "load_before": [ + 3.93994140625, + 5.7392578125, + 5.61376953125 + ], + "load_after": [ + 3.71875, + 5.603515625, + 5.56689453125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 3, + "target_tokens": 4500, + "status": 200, + "prompt_tokens": 4519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 52076478464, + "peak_gib": 48.5, + "ready_s": 2.0491247177124023, + "request_s": 15.142429113388062, + "load_before": [ + 3.71875, + 5.603515625, + 5.56689453125 + ], + "load_after": [ + 3.4296875, + 5.44384765625, + 5.509765625 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + }, + { + "rep": 3, + "target_tokens": 6000, + "status": 200, + "prompt_tokens": 6019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 77524159692, + "peak_gib": 72.19999999925494, + "ready_s": 2.043894052505493, + "request_s": 23.746572256088257, + "load_before": [ + 3.4296875, + 5.44384765625, + 5.509765625 + ], + "load_after": [ + 3.5546875, + 5.2734375, + 5.44482421875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": false, + "load_marker": false + } + ], + "finished_at": "2026-10-13T22:05:32+0900" +} diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/baseline-outputs.json b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/baseline-outputs.json new file mode 100644 index 000000000..f5c91ada7 --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/baseline-outputs.json @@ -0,0 +1,67 @@ +{ + "arm": "baseline", + "binary": "/tmp/mlxcel-chain-2281-artifacts/mlxcel-server.baseline", + "load_before": [ + 8.7099609375, + 5.65087890625, + 5.400390625 + ], + "responses": [ + { + "index": 0, + "status": 200, + "prompt_tokens": 37, + "completion_tokens": 81, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "A solar eclipse happens when the Moon passes directly between the Earth and the Sun, casting its shadow onto the Earth’s surface. Because the Moon’s shadow is relatively small, only a narrow path experiences the total eclipse, while most locations see only a partial or no eclipse. The geometry of the Sun, Moon, and Earth must align just right, so the eclipse is visible only from specific regions on Earth." + }, + "finish_reason": "stop", + "logprobs": null + } + }, + { + "index": 1, + "status": 200, + "prompt_tokens": 65, + "completion_tokens": 128, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "**Step‑by‑step arithmetic**\n\n1. **Start:** 240 crates \n\n2. **Monday – ship one quarter** \n \\[\n \\text{Shipped} = \\frac{1}{4}\\times 240 = 60\n \\]\n \\[\n \\text{Remaining after Monday} = 240 - 60 = 180\n \\]\n\n3. **Tuesday – receive 36 crates** \n \\[\n \\text{After receiving} = 180 + 36 = 216" + }, + "finish_reason": "length", + "logprobs": null + } + }, + { + "index": 2, + "status": 200, + "prompt_tokens": 2810, + "completion_tokens": 17, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "The blue archive contains cedar records, and the red archive contains maple records." + }, + "finish_reason": "stop", + "logprobs": null + } + } + ], + "health": [ + 200, + { + "status": "ok" + } + ], + "load_after": [ + 7.4384765625, + 5.51953125, + 5.35791015625 + ] +} diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-context-rejection.json b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-context-rejection.json new file mode 100644 index 000000000..49b74ef23 --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-context-rejection.json @@ -0,0 +1,22 @@ +{ + "status": 400, + "response": { + "error": { + "message": "request (8319 tokens) exceeds the available context size (8192 tokens), try increasing it", + "type": "exceed_context_size_error", + "code": null + } + }, + "health": [ + 200, + { + "status": "ok" + } + ], + "process_alive": true, + "load": [ + 5.47119140625, + 4.74072265625, + 4.93408203125 + ] +} diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-memory-8k.json b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-memory-8k.json new file mode 100644 index 000000000..b9b5ec87a --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-memory-8k.json @@ -0,0 +1,118 @@ +{ + "arm": "fixed-final-8k", + "binary": "/tmp/mlxcel-chain-2281-artifacts/mlxcel-server.final", + "model": "models/mlx/nvidia-nemotron-3-nano-30b-a3b-4bit", + "marker_expectation": "present", + "started_at": "2026-10-13T22:24:57+0900", + "runs": [ + { + "rep": 1, + "target_tokens": 8000, + "status": 200, + "prompt_tokens": 8019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 20723217203, + "peak_gib": 19.299999999813735, + "ready_s": 2.0552399158477783, + "request_s": 22.137491941452026, + "load_before": [ + 6.67138671875, + 5.81103515625, + 5.35400390625 + ], + "load_after": [ + 6.63916015625, + 5.873046875, + 5.388671875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 2, + "target_tokens": 8000, + "status": 200, + "prompt_tokens": 8019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 20723217203, + "peak_gib": 19.299999999813735, + "ready_s": 2.044273853302002, + "request_s": 22.105428218841553, + "load_before": [ + 6.63916015625, + 5.873046875, + 5.388671875 + ], + "load_after": [ + 6.98681640625, + 6.02490234375, + 5.45703125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 3, + "target_tokens": 8000, + "status": 200, + "prompt_tokens": 8019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 20723217203, + "peak_gib": 19.299999999813735, + "ready_s": 2.0479462146759033, + "request_s": 22.11466908454895, + "load_before": [ + 6.98681640625, + 6.02490234375, + 5.45703125 + ], + "load_after": [ + 7.513671875, + 6.2548828125, + 5.5576171875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + } + ], + "finished_at": "2026-10-13T22:26:15+0900" +} diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-memory.json b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-memory.json new file mode 100644 index 000000000..b5ec6dacf --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-memory.json @@ -0,0 +1,442 @@ +{ + "arm": "fixed-final", + "binary": "/tmp/mlxcel-chain-2281-artifacts/mlxcel-server.final", + "model": "models/mlx/nvidia-nemotron-3-nano-30b-a3b-4bit", + "marker_expectation": "present", + "started_at": "2026-10-13T22:19:20+0900", + "runs": [ + { + "rep": 1, + "target_tokens": 1500, + "status": 200, + "prompt_tokens": 1519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 18897856102, + "peak_gib": 17.59999999962747, + "ready_s": 2.561246156692505, + "request_s": 4.554424047470093, + "load_before": [ + 3.0283203125, + 4.4365234375, + 4.92724609375 + ], + "load_after": [ + 3.0908203125, + 4.40478515625, + 4.90966796875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 1, + "target_tokens": 3000, + "status": 200, + "prompt_tokens": 3019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 19327352832, + "peak_gib": 18.0, + "ready_s": 2.053757905960083, + "request_s": 8.864027261734009, + "load_before": [ + 3.0908203125, + 4.40478515625, + 4.90966796875 + ], + "load_after": [ + 3.95703125, + 4.53125, + 4.9453125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 1, + "target_tokens": 4500, + "status": 200, + "prompt_tokens": 4519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 19756849561, + "peak_gib": 18.399999999441206, + "ready_s": 2.0424530506134033, + "request_s": 12.59739089012146, + "load_before": [ + 3.95703125, + 4.53125, + 4.9453125 + ], + "load_after": [ + 3.775390625, + 4.45654296875, + 4.91064453125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 1, + "target_tokens": 6000, + "status": 200, + "prompt_tokens": 6019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 20078972108, + "peak_gib": 18.699999999254942, + "ready_s": 2.0454490184783936, + "request_s": 16.56507921218872, + "load_before": [ + 3.775390625, + 4.45654296875, + 4.91064453125 + ], + "load_after": [ + 4.068359375, + 4.48828125, + 4.9111328125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 2, + "target_tokens": 1500, + "status": 200, + "prompt_tokens": 1519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 18897856102, + "peak_gib": 17.59999999962747, + "ready_s": 2.04069185256958, + "request_s": 4.253753185272217, + "load_before": [ + 4.068359375, + 4.48828125, + 4.9111328125 + ], + "load_after": [ + 4.05078125, + 4.4716796875, + 4.900390625 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 2, + "target_tokens": 3000, + "status": 200, + "prompt_tokens": 3019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 19327352832, + "peak_gib": 18.0, + "ready_s": 2.042436122894287, + "request_s": 8.334737777709961, + "load_before": [ + 4.05078125, + 4.4716796875, + 4.900390625 + ], + "load_after": [ + 3.7412109375, + 4.3896484375, + 4.8662109375 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 2, + "target_tokens": 4500, + "status": 200, + "prompt_tokens": 4519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 19756849561, + "peak_gib": 18.399999999441206, + "ready_s": 2.033221960067749, + "request_s": 12.392516851425171, + "load_before": [ + 3.7412109375, + 4.3896484375, + 4.8662109375 + ], + "load_after": [ + 4.337890625, + 4.49072265625, + 4.8916015625 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 2, + "target_tokens": 6000, + "status": 200, + "prompt_tokens": 6019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 20078972108, + "peak_gib": 18.699999999254942, + "ready_s": 2.0374650955200195, + "request_s": 16.56824803352356, + "load_before": [ + 4.337890625, + 4.49072265625, + 4.8916015625 + ], + "load_after": [ + 4.48095703125, + 4.5078125, + 4.8876953125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 3, + "target_tokens": 1500, + "status": 200, + "prompt_tokens": 1519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 18897856102, + "peak_gib": 17.59999999962747, + "ready_s": 2.041900157928467, + "request_s": 4.239515066146851, + "load_before": [ + 4.48095703125, + 4.5078125, + 4.8876953125 + ], + "load_after": [ + 4.2822265625, + 4.4658203125, + 4.87060546875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 3, + "target_tokens": 3000, + "status": 200, + "prompt_tokens": 3019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 19327352832, + "peak_gib": 18.0, + "ready_s": 2.039336919784546, + "request_s": 8.308146953582764, + "load_before": [ + 4.2822265625, + 4.4658203125, + 4.87060546875 + ], + "load_after": [ + 5.0263671875, + 4.63525390625, + 4.923828125 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 3, + "target_tokens": 4500, + "status": 200, + "prompt_tokens": 4519, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 19756849561, + "peak_gib": 18.399999999441206, + "ready_s": 2.047098159790039, + "request_s": 12.469884872436523, + "load_before": [ + 5.0263671875, + 4.63525390625, + 4.923828125 + ], + "load_after": [ + 4.93310546875, + 4.63525390625, + 4.91796875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + }, + { + "rep": 3, + "target_tokens": 6000, + "status": 200, + "prompt_tokens": 6019, + "completion_tokens": 1, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "OK" + }, + "finish_reason": "stop", + "logprobs": null + }, + "peak_bytes": 20078972108, + "peak_gib": 18.699999999254942, + "ready_s": 2.0449910163879395, + "request_s": 16.522796154022217, + "load_before": [ + 4.93310546875, + 4.63525390625, + 4.91796875 + ], + "load_after": [ + 4.3154296875, + 4.51220703125, + 4.86669921875 + ], + "health_status": 200, + "health": { + "status": "ok" + }, + "forward_marker": true, + "load_marker": true + } + ], + "finished_at": "2026-10-13T22:22:13+0900" +} diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-outputs.json b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-outputs.json new file mode 100644 index 000000000..ec01adfc6 --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/candidate-outputs.json @@ -0,0 +1,67 @@ +{ + "arm": "fixed-final", + "binary": "/tmp/mlxcel-chain-2281-artifacts/mlxcel-server.final", + "load_before": [ + 4.26025390625, + 4.49462890625, + 4.85595703125 + ], + "responses": [ + { + "index": 0, + "status": 200, + "prompt_tokens": 37, + "completion_tokens": 81, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "A solar eclipse happens when the Moon passes directly between the Earth and the Sun, casting its shadow onto the Earth’s surface. Because the Moon’s shadow is relatively small, only a narrow path experiences the total eclipse, while most locations see only a partial or no eclipse. The geometry of the Sun, Moon, and Earth must align just right, so the eclipse is visible only from specific regions on Earth." + }, + "finish_reason": "stop", + "logprobs": null + } + }, + { + "index": 1, + "status": 200, + "prompt_tokens": 65, + "completion_tokens": 128, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "**Step‑by‑step arithmetic**\n\n1. **Start:** 240 crates \n\n2. **Monday – ship one quarter** \n \\[\n \\text{Shipped} = \\frac{1}{4}\\times 240 = 60\n \\]\n \\[\n \\text{Remaining after Monday} = 240 - 60 = 180\n \\]\n\n3. **Tuesday – receive 36 crates** \n \\[\n \\text{After receiving} = 180 + 36 = 216" + }, + "finish_reason": "length", + "logprobs": null + } + }, + { + "index": 2, + "status": 200, + "prompt_tokens": 2810, + "completion_tokens": 17, + "choice": { + "index": 0, + "message": { + "role": "assistant", + "content": "The blue archive contains cedar records, and the red archive contains maple records." + }, + "finish_reason": "stop", + "logprobs": null + } + } + ], + "health": [ + 200, + { + "status": "ok" + } + ], + "load_after": [ + 4.50146484375, + 4.54296875, + 4.86865234375 + ] +} diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/memory_harness.py b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/memory_harness.py new file mode 100755 index 000000000..ea53b45a5 --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/memory_harness.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Lablup Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import json +import os +import pathlib +import re +import signal +import socket +import subprocess +import time +import urllib.error +import urllib.request + +UNITS = {"K": 1024, "M": 1024**2, "G": 1024**3} +PEAK_RE = re.compile(r"Physical footprint \(peak\):\s+([0-9.]+)([KMG])") +MARKER = "[NemotronH] chunked SSD prefill forward" +LOAD_MARKER = "[NemotronH] chunked SSD prefill dispatch enabled" + + +def free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def request_json(url, body=None, timeout=900): + data = None if body is None else json.dumps(body).encode() + req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=timeout) as response: + return response.status, json.loads(response.read()) + except urllib.error.HTTPError as error: + return error.code, json.loads(error.read()) + + +def wait_ready(proc, base, timeout=120): + deadline = time.time() + timeout + while time.time() < deadline: + if proc.poll() is not None: + raise RuntimeError(f"server exited during startup with {proc.returncode}") + try: + status, body = request_json(base + "/v1/models", timeout=3) + if status == 200: + return body["data"][0]["id"] + except Exception: + pass + time.sleep(0.25) + raise TimeoutError("server readiness timeout") + + +def vmmap_peak(pid): + result = subprocess.run(["vmmap", "-summary", str(pid)], text=True, capture_output=True, check=True, timeout=60) + match = PEAK_RE.search(result.stdout) + if not match: + raise RuntimeError("vmmap omitted physical footprint peak") + return int(float(match.group(1)) * UNITS[match.group(2)]) + + +def run_one(args, rep, target): + port = free_port() + base = f"http://127.0.0.1:{port}" + log_path = pathlib.Path(args.output).with_suffix(f".rep{rep}.tokens{target}.log") + command = [ + args.binary, "--model", args.model, "--ctx-size", "8192", "--jinja", + "--host", "127.0.0.1", "--port", str(port), "--alias", "repro", + "--parallel", "1", "--offline", + ] + load_before = os.getloadavg() + started = time.time() + with log_path.open("w") as log: + proc = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, text=True, start_new_session=True, env={**os.environ, "RUST_LOG": "info"}) + try: + model_id = wait_ready(proc, base) + ready_s = time.time() - started + prompt = f"Run {rep}, target {target}. " + "benchmark " * max(1, target - 20) + body = { + "model": model_id, + "messages": [ + {"role": "system", "content": "Read the filler and reply with the single word OK."}, + {"role": "user", "content": prompt}, + ], + "max_tokens": 8, + "temperature": 0.0, + "stream": False, + "chat_template_kwargs": {"enable_thinking": False}, + } + request_started = time.time() + status, response = request_json(base + "/v1/chat/completions", body) + request_s = time.time() - request_started + peak_bytes = vmmap_peak(proc.pid) + health_status, health = request_json(base + "/health", timeout=10) + finally: + if proc.poll() is None: + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=15) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait(timeout=5) + logs = log_path.read_text(errors="replace") + forward_marker = MARKER in logs + load_marker = LOAD_MARKER in logs + if args.marker == "absent" and (forward_marker or load_marker): + raise RuntimeError("baseline unexpectedly emitted fixed-arm dispatch marker") + if args.marker == "present" and not (forward_marker and load_marker): + raise RuntimeError("fixed arm did not execute and log chunked SSD forward dispatch") + if args.marker == "loadonly" and not load_marker: + raise RuntimeError("preliminary fixed arm did not emit load-time chunk marker") + usage = response.get("usage") or {} + choice = (response.get("choices") or [{}])[0] + return { + "rep": rep, + "target_tokens": target, + "status": status, + "prompt_tokens": usage.get("prompt_tokens"), + "completion_tokens": usage.get("completion_tokens"), + "choice": choice, + "peak_bytes": peak_bytes, + "peak_gib": peak_bytes / 1024**3, + "ready_s": ready_s, + "request_s": request_s, + "load_before": load_before, + "load_after": os.getloadavg(), + "health_status": health_status, + "health": health, + "forward_marker": forward_marker, + "load_marker": load_marker, + "log": str(log_path), + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--binary", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--arm", required=True) + parser.add_argument("--marker", choices=["present", "absent", "loadonly"], required=True) + parser.add_argument("--reps", type=int, default=3) + parser.add_argument("--points", default="1500,3000,4500,6000") + args = parser.parse_args() + result = { + "arm": args.arm, + "binary": args.binary, + "model": args.model, + "marker_expectation": args.marker, + "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "runs": [], + } + pathlib.Path(args.output).write_text(json.dumps(result, indent=2)) + for rep in range(1, args.reps + 1): + for target in [int(value) for value in args.points.split(",")]: + run = run_one(args, rep, target) + result["runs"].append(run) + pathlib.Path(args.output).write_text(json.dumps(result, indent=2)) + print(json.dumps({key: run[key] for key in ["rep", "target_tokens", "prompt_tokens", "peak_gib", "request_s", "load_before", "forward_marker"]}), flush=True) + result["finished_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z") + pathlib.Path(args.output).write_text(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/output_harness.py b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/output_harness.py new file mode 100755 index 000000000..a9d404076 --- /dev/null +++ b/docs/benchmark_results/data/nemotron-h-prefill-memory-m1ultra-2026-10-13/output_harness.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Lablup Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse, json, os, pathlib, signal, socket, subprocess, time, urllib.error, urllib.request + +def free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)); return sock.getsockname()[1] + +def call(url, body=None, timeout=900): + data = None if body is None else json.dumps(body).encode() + req = urllib.request.Request(url, data=data, headers={"Content-Type":"application/json"}) + try: + with urllib.request.urlopen(req, timeout=timeout) as response: + return response.status, json.loads(response.read()) + except urllib.error.HTTPError as error: + return error.code, json.loads(error.read()) + +def ready(proc, base): + for _ in range(600): + if proc.poll() is not None: raise RuntimeError(f"server exited {proc.returncode}") + try: + status, data = call(base+"/v1/models", timeout=2) + if status == 200: return data["data"][0]["id"] + except Exception: pass + time.sleep(.25) + raise TimeoutError + +def main(): + parser=argparse.ArgumentParser() + parser.add_argument("--binary", required=True); parser.add_argument("--model", required=True) + parser.add_argument("--arm", required=True); parser.add_argument("--output", required=True) + args=parser.parse_args(); port=free_port(); base=f"http://127.0.0.1:{port}" + log_path=pathlib.Path(args.output).with_suffix(".log") + cmd=[args.binary,"--model",args.model,"--ctx-size","8192","--jinja","--host","127.0.0.1","--port",str(port),"--alias","repro","--parallel","1","--offline"] + prompts=[ + "Explain in three concise sentences how a solar eclipse occurs and why it is not visible everywhere on Earth.", + "A warehouse starts with 240 crates. It ships one quarter of them on Monday, receives 36 on Tuesday, then ships one third of the remaining crates on Wednesday. Show the arithmetic and state the final crate count.", + "Context follows. " + "The blue archive contains cedar records; the red archive contains maple records. "*173 + "Using only the context, state in two complete sentences which archive contains cedar records and which contains maple records.", + ] + result={"arm":args.arm,"binary":args.binary,"load_before":os.getloadavg(),"responses":[]} + with log_path.open("w") as log: + proc=subprocess.Popen(cmd,stdout=log,stderr=subprocess.STDOUT,text=True,start_new_session=True,env={**os.environ,"RUST_LOG":"info"}) + try: + model=ready(proc,base) + for index,prompt in enumerate(prompts): + status, body=call(base+"/v1/chat/completions", {"model":model,"messages":[{"role":"user","content":prompt}],"max_tokens":128,"temperature":0.0,"stream":False,"chat_template_kwargs":{"enable_thinking":False}}) + result["responses"].append({"index":index,"status":status,"prompt_tokens":(body.get("usage") or {}).get("prompt_tokens"),"completion_tokens":(body.get("usage") or {}).get("completion_tokens"),"choice":(body.get("choices") or [{}])[0]}) + pathlib.Path(args.output).write_text(json.dumps(result,indent=2)) + print(json.dumps(result["responses"][-1]),flush=True) + result["health"]=call(base+"/health",timeout=10); result["load_after"]=os.getloadavg() + finally: + if proc.poll() is None: + os.killpg(proc.pid,signal.SIGTERM) + try: proc.wait(timeout=15) + except subprocess.TimeoutExpired: + os.killpg(proc.pid,signal.SIGKILL); proc.wait(timeout=5) + result["log"]=str(log_path); pathlib.Path(args.output).write_text(json.dumps(result,indent=2)) +if __name__=="__main__": main() diff --git a/docs/benchmark_results/nemotron-h-prefill-memory-m1ultra-2026-10-13.md b/docs/benchmark_results/nemotron-h-prefill-memory-m1ultra-2026-10-13.md new file mode 100644 index 000000000..cbad5164c --- /dev/null +++ b/docs/benchmark_results/nemotron-h-prefill-memory-m1ultra-2026-10-13.md @@ -0,0 +1,65 @@ +# Nemotron-H chunked SSD prefill memory — M1 Ultra — 2026-10-13 + +## Decision + +Ship the model-internal SSD chunking for Nemotron-H. The production server path no longer builds prompt-length `L × L` surrogate-attention matrices inside every Mamba layer: it carries recurrent state across the checkpoint's `chunk_size=128` blocks and evaluates each block before constructing the next one. + +This record resolves the quadratic-prefill-memory part of issue #2281. It does **not** establish request-level rejection from a projected transient-memory budget. The current server has context-size and KV-budget admission plus an optional startup `--estimate-memory` check, but no per-request activation/transient-memory admission; a separate host/model combination can therefore still reach a fatal Metal driver OOM. The existing 8,319-token context rejection reported below is evidence only for context admission and is not treated as memory admission. + +## Environment and provenance + +- Host: Mac Studio `Mac13,2`, Apple M1 Ultra, 128 GB unified memory, 20 logical CPUs. +- OS: macOS 27.0.1 (`26A434`). +- Repository base: `61329ee2863afe1120a70f054fd39317fb49877a`; the candidate binary was built from the identical production source in this record's commit before committing the documentation. +- MLX pin: `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a`. +- Baseline binary SHA-256: `f89beb4d08739e82c5c99df01a7ee3c0d8c295fb8c5dc2ffda64c2044a054e82`. +- Candidate binary SHA-256: `5d7f9b65901785e6a71cf8c2b711e5d647beff2871d9655af0bdaad067452f31`. +- Build command for both arms: `RUSTC_WRAPPER= cargo build --release --features metal,accelerate --bin mlxcel-server`. +- Checkpoint: local `models/mlx/nvidia-nemotron-3-nano-30b-a3b-4bit`; config SHA-256 `bee8181e005ea1ca595d926f10ee3c327d65e1e5fd77596e42d1421c806b2b90`; affine 4-bit/group-64, 52 layers, 64 Mamba heads, head dimension 64, state dimension 128, `chunk_size=128`. +- Server command: `mlxcel-server --model models/mlx/nvidia-nemotron-3-nano-30b-a3b-4bit --ctx-size 8192 --jinja --host 127.0.0.1 --port --alias repro --parallel 1 --offline`. + +## Method + +Each measurement launched a fresh server process, waited for `/v1/models`, submitted one non-streaming `/v1/chat/completions` request with `temperature=0`, `max_tokens=8`, and `chat_template_kwargs.enable_thinking=false`, then read the process lifetime high-water mark from `vmmap -summary ` (`Physical footprint (peak)`). The deterministic filler was `benchmark ` repeated to the target length; API usage supplied the actual prompt-token count. There were three fresh-process repetitions per reported cell. + +The exact harnesses and JSON outputs are archived under [`data/nemotron-h-prefill-memory-m1ultra-2026-10-13/`](data/nemotron-h-prefill-memory-m1ultra-2026-10-13/). + +The harness refused to report a candidate result unless the log contained both the load-time `chunked SSD prefill dispatch enabled` marker and the production-forward `chunked SSD prefill forward` marker with `sequence_tokens`, `ssm_chunk_size`, and `ssd_chunks`. It also refused a baseline result if either marker appeared. For example, the 6,019-token candidate request logged `sequence_tokens=6013 ssm_chunk_size=128 ssd_chunks=47`; all 12 paired baseline/candidate cells passed the attribution guard. + +`vmmap` prints GiB high-water values to one decimal place, so the zero ranges below mean that all three readings rounded to the same 0.1 GiB value, not that the underlying byte peaks were identical. The host was not quiet: baseline load-average samples ranged from 3.43 to 6.68 and candidate 1,519–6,019-token samples ranged from 3.03 to 5.03. Because the claim is about per-process high-water memory rather than latency, request wall times are retained in the raw run artifacts but are not used for a throughput claim. + +## Results + +| Prompt tokens | Baseline peak median (range), GiB | Candidate peak median (range), GiB | Peak reduction | +|---:|---:|---:|---:| +| 1,519 | 20.9 (20.9–20.9) | 17.6 (17.6–17.6) | 3.3 GiB / 15.8% | +| 3,019 | 31.4 (31.4–31.4) | 18.0 (18.0–18.0) | 13.4 GiB / 42.7% | +| 4,519 | 48.5 (48.5–48.5) | 18.4 (18.4–18.4) | 30.1 GiB / 62.1% | +| 6,019 | 72.2 (72.2–72.2) | 18.7 (18.7–18.7) | 53.5 GiB / 74.1% | +| 8,019 | not run (baseline projection approached the host wired limit) | 19.3 (19.3–19.3) | not computed | + +Across 1,519–6,019 tokens, the baseline peak grew by 51.3 GiB while the candidate grew by 1.1 GiB. The candidate also completed three 8,019-token runs, within the configured 8,192-token context, at a 19.3 GiB peak. The baseline 8,019-token cell was deliberately left unmeasured because the observed quadratic curve made a driver-level OOM plausible; no result is inferred for that arm. + +## Output correctness + +The baseline and final candidate produced byte-identical greedy responses for three substantive prompts: + +| Prompt | Prompt tokens | Completion tokens | Result | +|---|---:|---:|---| +| Three-sentence solar-eclipse explanation | 37 | 81 | exact response and finish-reason match | +| Warehouse arithmetic with shown work | 65 | 128 | exact response and length finish match | +| Long repeated archive context with a two-fact query | 2,810 | 17 | exact response and finish-reason match | + +The numerical unit coverage separately compares 2+2+1 chunking with a full-step evaluation while carrying a nonzero prior recurrent state, and compares multi-batch/multi-group f16 and bf16 chunked outputs and final states with a direct scalar recurrence oracle. + +## Existing rejection behavior and remaining safety gap + +A final-candidate request containing 8,319 prompt tokens under `--ctx-size 8192` returned HTTP 400 with `type=exceed_context_size_error`; the server process remained alive and `/health` returned 200. This verifies the existing context guard only. No test in this record claims that a within-context request projected to exceed the wired-memory budget is rejected, because the scheduler currently has no such admission input. Arbitrary command-buffer OOM remains fatal and cannot be converted into an HTTP error after Metal terminates the process. + +## Unmeasured cells + +- The reporter's M4 Pro 48 GB host and original 2-bit checkpoint were unavailable; the checked-in decision uses the available M1 Ultra and the 4-bit Nemotron checkpoint above. +- Candidate 8,019 tokens has three repetitions; baseline 8,019 tokens is intentionally unmeasured to avoid a likely high-blast-radius OOM. +- M5 and CUDA behavior is unmeasured. +- No throughput or latency improvement is claimed because the host carried background load and this change targets memory. +- Memory-derived per-request admission is unimplemented and unmeasured; the issue should remain open or receive an explicitly split follow-up until that acceptance item is decided. diff --git a/src/models/nemotron_h.rs b/src/models/nemotron_h.rs index 252bbbe3a..a52a4b898 100644 --- a/src/models/nemotron_h.rs +++ b/src/models/nemotron_h.rs @@ -27,13 +27,16 @@ use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{KVCache, Linear, RMSNorm, UnifiedEmbedding, UnifiedLinear}; use mlxcel_core::utils::{create_causal_mask, relu_squared, silu, slice_axis, stack_arrays}; use mlxcel_core::weights::WeightMap; -use mlxcel_core::{MlxArray, UniquePtr, concatenate}; +use mlxcel_core::{MlxArray, UniquePtr, concatenate, concatenate_many}; use serde::Deserialize; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; use super::model_owned::ModelOwnedSequenceState; use super::recurrent_snapshot::{push_optional, restore_optional}; +static SSD_PREFILL_DISPATCH_LOGGED: AtomicBool = AtomicBool::new(false); + // Configuration. // Custom deserializer for hybrid_override_pattern which can be either: // - A string like "MEMEM*EMEMEM*..." (each char is a block type) @@ -171,6 +174,12 @@ pub struct NemotronHConfig { pub conv_kernel: usize, pub n_groups: usize, + /// Number of sequence tokens materialized by one SSD recurrence step. + /// Hugging Face Nemotron-H checkpoints use `chunk_size`; the alias keeps + /// compatibility with early configs that used the Mamba-specific name. + #[serde(default = "default_mamba_chunk_size", alias = "mamba_chunk_size")] + pub chunk_size: usize, + #[serde(default, deserialize_with = "deserialize_time_step_limit_opt")] pub time_step_limit: Option<(f32, f32)>, @@ -262,6 +271,10 @@ impl NemotronHConfig { /// - Normalizes layers_block_type word names to single-char hybrid_override_pattern. /// - Sets num_hidden_layers from pattern length. pub fn post_init(&mut self) -> Result<(), String> { + if self.chunk_size == 0 { + return Err("chunk_size must be greater than zero".to_string()); + } + // When time_step_limit is absent from the config, default to (0.0, +inf). // This matches upstream mlx-lm PR #1026 (commit 6ddfdda): // if time_step_limit is None: time_step_limit = (0.0, float("inf")) @@ -304,6 +317,10 @@ fn default_layer_norm_epsilon() -> f32 { 1e-5 } +fn default_mamba_chunk_size() -> usize { + 128 +} + fn default_true() -> bool { true } @@ -512,6 +529,7 @@ struct NemotronHMamba2Mixer { n_groups: usize, head_dim: usize, time_step_limit: (f32, f32), + chunk_size: usize, conv_dim: usize, conv_weight: UniquePtr, @@ -533,6 +551,17 @@ impl NemotronHMamba2Mixer { // Use fused C++ forward for single-token decode with existing cache let shape = mlxcel_core::array_shape(hidden_states); let seq_len = shape[1]; + if seq_len > 1 && !SSD_PREFILL_DISPATCH_LOGGED.swap(true, Ordering::Relaxed) { + let sequence_tokens = seq_len as usize; + let ssd_chunks = sequence_tokens / self.chunk_size + + usize::from(!sequence_tokens.is_multiple_of(self.chunk_size)); + tracing::info!( + sequence_tokens, + ssm_chunk_size = self.chunk_size, + ssd_chunks, + "[NemotronH] chunked SSD prefill forward" + ); + } if seq_len == 1 && mlxcel_core::ssm_kernel_available() && cache @@ -705,149 +734,18 @@ impl NemotronHMamba2Mixer { ], ); - let num_heads = self.num_heads as i32; - let n_groups = self.n_groups as i32; - let state_dim = self.ssm_state_size as i32; - let head_dim = self.head_dim as i32; - let repeats = num_heads / n_groups; - - // Compute dt with bias and limits in float32 for numerical precision - // (same as Mamba2's compute_dt — upstream mlx-lm casts to float32 before softplus) - let dt_f32 = mlxcel_core::astype(dt, mlxcel_core::dtype::FLOAT32); - let dt_bias_f32 = mlxcel_core::astype(&self.dt_bias, mlxcel_core::dtype::FLOAT32); - let dt_biased = mlxcel_core::add(&dt_f32, &dt_bias_f32); - let dt_soft = mlxcel_core::softplus(&dt_biased); - let min_val = - mlxcel_core::full_f32(&[1], self.time_step_limit.0, mlxcel_core::dtype::FLOAT32); - let max_val = - mlxcel_core::full_f32(&[1], self.time_step_limit.1, mlxcel_core::dtype::FLOAT32); - let dt = mlxcel_core::clip(&dt_soft, &min_val, &max_val); - - // A = -exp(A_log), cast to dt dtype (Python does .astype(dt.dtype); - // since dt is now float32, A stays in float32 for precision) - let a = mlxcel_core::negative(&mlxcel_core::exp(&self.a_log)); - let a = mlxcel_core::astype(&a, mlxcel_core::array_dtype(&dt)); - let a_reshaped = mlxcel_core::reshape(&a, &[1, 1, num_heads]); - - // dtA = dt * A - let dt_a = mlxcel_core::multiply(&dt, &a_reshaped); - - // dtx = dt * x (expand dt for broadcasting) - // Promote x to float32 before multiplication with dt (float32). - // On M5 Max (Metal GPU Family 4), mixed float32×float16 multiply - // in a lazy computation graph produces NaN via the NAx broadcast kernel. - let dt_exp = mlxcel_core::reshape(&dt, &[batch, seq_len, num_heads, 1]); - let x_f32 = mlxcel_core::astype(&x, mlxcel_core::dtype::FLOAT32); - let dtx = mlxcel_core::multiply(&dt_exp, &x_f32); - - // Promote B, C to float32 before CB matmul to prevent float16 overflow - // on M5 Max (Metal GPU Family 4) NAx GEMM kernel. The dot products - // over ssm_state_size elements can exceed float16 max (65504) for - // certain weight distributions, producing NaN that propagates through - // the entire model and causes all- output. - let b_f32 = mlxcel_core::astype(&b, mlxcel_core::dtype::FLOAT32); - let c_f32 = mlxcel_core::astype(&c, mlxcel_core::dtype::FLOAT32); - - // B: [batch, seq, groups, state_dim] -> [batch, groups, state_dim, seq] - let b_t = mlxcel_core::transpose_axes(&b_f32, &[0, 2, 3, 1]); - - // CB = C.swapaxes(1, 2) @ B (in float32 to avoid overflow) - let c_t = mlxcel_core::swap_axes(&c_f32, 1, 2); - let cb = mlxcel_core::matmul(&c_t, &b_t); - - // Repeat CB for each head in group - let cb = repeat_axis_nemotron(&cb, repeats, 1); - - // decay = exp(segsum(dtA.swapaxes(1, 2))) - let dt_a_t = mlxcel_core::swap_axes(&dt_a, 1, 2); - let seg = segsum_nemotron(&dt_a_t); - let decay = mlxcel_core::exp(&seg); - - // surrogate_attention_matrix = tril(CB * decay) - let attn_matrix = mlxcel_core::multiply(&cb, &decay); - let attn_matrix = mlxcel_core::tril(&attn_matrix, 0); - - // y = attn_matrix @ dtx.swapaxes(1, 2) - let dtx_t = mlxcel_core::swap_axes(&dtx, 1, 2); - let y = mlxcel_core::matmul(&attn_matrix, &dtx_t); - let y = mlxcel_core::swap_axes(&y, 1, 2); - - // Compute next state - // decay_last = decay[:, :, -1:, :] - let decay_shape = mlxcel_core::array_shape(&decay); - let decay_last = slice_axis(&decay, 2, decay_shape[2] - 1, decay_shape[2]); - let decay_t = mlxcel_core::transpose_axes(&decay_last, &[0, 3, 1, 2]); - - // B for state update (already float32 from promotion above) - let b_rep = repeat_axis_nemotron(&b_t, repeats, 1); - let b_sw = mlxcel_core::swap_axes(&b_rep, 2, 3); - - let dtx_decay = mlxcel_core::multiply(&dtx, &decay_t); - let dtx_decay_t = mlxcel_core::swap_axes(&dtx_decay, 1, 2); - let dtx_decay_t = mlxcel_core::swap_axes(&dtx_decay_t, 2, 3); - - let mut next_state = mlxcel_core::matmul(&dtx_decay_t, &b_sw); - - // Handle previous state if present (using Mamba2's approach) - let y = if let Some(prev_state) = state { - // exp_dtA_cumsum = exp(cumsum(dtA, axis=-2)) - // Python mx.cumsum defaults to inclusive; match mamba2.rs ssm_step - let dta_cumsum = mlxcel_core::cumsum(&dt_a, -2, false, true); - let exp_dta_cumsum = mlxcel_core::exp(&dta_cumsum); - - // Update next_state with previous state - // Python: exp_dtA_cumsum[:, -1, :, None, None] → [batch, heads, 1, 1] - let exp_shape = mlxcel_core::array_shape(&exp_dta_cumsum); - let last_exp = slice_axis(&exp_dta_cumsum, 1, exp_shape[1] - 1, exp_shape[1]); - // last_exp is [batch, 1, num_heads] — squeeze seq dim, add two trailing dims - let last_exp = mlxcel_core::reshape(&last_exp, &[batch, num_heads, 1, 1]); - - let state_contrib = mlxcel_core::multiply(&last_exp, prev_state); - next_state = mlxcel_core::add(&next_state, &state_contrib); - - // y_prev contribution (use float32 C for precision with float32 state) - let c_reshaped = - mlxcel_core::reshape(&c_f32, &[batch, seq_len, n_groups, 1, state_dim, 1]); - let state_reshaped = mlxcel_core::reshape( - prev_state, - &[batch, 1, n_groups, repeats, head_dim, state_dim], - ); - let y_prev = mlxcel_core::matmul(&state_reshaped, &c_reshaped); - let y_prev = mlxcel_core::squeeze_axis(&y_prev, -1); - let y_prev = mlxcel_core::reshape(&y_prev, &[batch, seq_len, num_heads, head_dim]); - - // Expand exp_dta_cumsum for y_prev multiplication - let mut exp_shape_y = mlxcel_core::array_shape(&exp_dta_cumsum); - exp_shape_y.push(1); - let exp_dta_exp = mlxcel_core::reshape(&exp_dta_cumsum, &exp_shape_y); - - let y_prev_scaled = mlxcel_core::multiply(&exp_dta_exp, &y_prev); - mlxcel_core::add(&y, &y_prev_scaled) - } else { - y - }; - - // Add D term: y = y + x * D (all in float32 to prevent mixed-dtype NaN - // on M5 Max NAx kernel — keep D and x in float32 matching y's dtype) - let d_f32 = mlxcel_core::astype(&self.d_param, mlxcel_core::dtype::FLOAT32); - let d_reshaped = mlxcel_core::reshape(&d_f32, &[1, 1, num_heads, 1]); - let d_contrib = mlxcel_core::multiply(&x_f32, &d_reshaped); - let y = mlxcel_core::add(&y, &d_contrib); - - // Cast y back to input dtype (dt computation was promoted to float32) - let y = mlxcel_core::astype(&y, mlxcel_core::array_dtype(hidden_states)); - - // Force evaluation of y and next_state before returning. - // On M5 Max (Metal GPU Family 4), the lazy computation graph from the - // SSM step contains mixed float32×float16 intermediate nodes. When - // this graph is fused with downstream operations (gated norm, residual - // add, MoE, attention) in a single Metal command buffer, the NAx - // kernel produces NaN. Materializing the SSM outputs here splits - // the graph at a clean float16 boundary. - mlxcel_core::eval(&y); - mlxcel_core::eval(&next_state); - - (y, next_state) + nemotron_h_ssm_attn( + &x, + &self.a_log, + &b, + &c, + &self.d_param, + dt, + &self.dt_bias, + state, + self.time_step_limit, + self.chunk_size, + ) } /// Fused SSM step using Metal kernel (single-token decode only) @@ -1034,6 +932,153 @@ fn segsum_nemotron(x: &MlxArray) -> UniquePtr { mlxcel_core::cumsum(&x_tril, -2, false, true) } +/// Chunked SSD recurrence for Nemotron-H Mamba blocks. +/// +/// The surrogate attention matrices are quadratic in their sequence dimension. +/// Processing and evaluating one configured chunk at a time bounds that transient +/// to `[batch, heads, chunk_size, chunk_size]` instead of the complete prompt. +#[allow(clippy::too_many_arguments)] +fn nemotron_h_ssm_attn( + x: &MlxArray, + a_log: &MlxArray, + b: &MlxArray, + c: &MlxArray, + d: &MlxArray, + dt: &MlxArray, + dt_bias: &MlxArray, + state: Option<&MlxArray>, + time_step_limit: (f32, f32), + chunk_size: usize, +) -> (UniquePtr, UniquePtr) { + assert!(chunk_size > 0, "chunk_size must be greater than zero"); + + let input_dtype = mlxcel_core::array_dtype(x); + let x_shape = mlxcel_core::array_shape(x); + let batch = x_shape[0]; + let seq_len = x_shape[1] as usize; + let num_heads = x_shape[2]; + let head_dim = x_shape[3]; + + let b_shape = mlxcel_core::array_shape(b); + let n_groups = b_shape[2]; + let state_dim = b_shape[3]; + let repeats = num_heads / n_groups; + + // Preserve the existing float32 precision boundary used to prevent the M5 + // mixed-dtype and CB-matmul NaNs documented in `ssm_step`. + let dt_f32 = mlxcel_core::astype(dt, mlxcel_core::dtype::FLOAT32); + let dt_bias_f32 = mlxcel_core::astype(dt_bias, mlxcel_core::dtype::FLOAT32); + let dt_biased = mlxcel_core::add(&dt_f32, &dt_bias_f32); + let dt_soft = mlxcel_core::softplus(&dt_biased); + let min_val = mlxcel_core::full_f32(&[1], time_step_limit.0, mlxcel_core::dtype::FLOAT32); + let max_val = mlxcel_core::full_f32(&[1], time_step_limit.1, mlxcel_core::dtype::FLOAT32); + let dt = mlxcel_core::clip(&dt_soft, &min_val, &max_val); + + let a = mlxcel_core::negative(&mlxcel_core::exp(a_log)); + let a = mlxcel_core::astype(&a, mlxcel_core::dtype::FLOAT32); + let a_reshaped = mlxcel_core::reshape(&a, &[1, 1, num_heads]); + let dt_a = mlxcel_core::multiply(&dt, &a_reshaped); + + let x_f32 = mlxcel_core::astype(x, mlxcel_core::dtype::FLOAT32); + let dt_exp = mlxcel_core::reshape(&dt, &[batch, seq_len as i32, num_heads, 1]); + let dtx = mlxcel_core::multiply(&dt_exp, &x_f32); + let b_f32 = mlxcel_core::astype(b, mlxcel_core::dtype::FLOAT32); + let c_f32 = mlxcel_core::astype(c, mlxcel_core::dtype::FLOAT32); + let d_f32 = mlxcel_core::astype(d, mlxcel_core::dtype::FLOAT32); + let d_reshaped = mlxcel_core::reshape(&d_f32, &[1, 1, num_heads, 1]); + + let output_capacity = seq_len / chunk_size + usize::from(!seq_len.is_multiple_of(chunk_size)); + let mut outputs = Vec::with_capacity(output_capacity); + let mut current_state = state.map(mlxcel_core::copy); + + for chunk_start in (0..seq_len).step_by(chunk_size) { + let chunk_end = chunk_start.saturating_add(chunk_size).min(seq_len); + let chunk_len = (chunk_end - chunk_start) as i32; + + let chunk_x = slice_axis(&x_f32, 1, chunk_start as i32, chunk_end as i32); + let chunk_dtx = slice_axis(&dtx, 1, chunk_start as i32, chunk_end as i32); + let chunk_dt_a = slice_axis(&dt_a, 1, chunk_start as i32, chunk_end as i32); + let chunk_b = slice_axis(&b_f32, 1, chunk_start as i32, chunk_end as i32); + let chunk_c = slice_axis(&c_f32, 1, chunk_start as i32, chunk_end as i32); + + let b_t = mlxcel_core::transpose_axes(&chunk_b, &[0, 2, 3, 1]); + let c_t = mlxcel_core::swap_axes(&chunk_c, 1, 2); + let cb = mlxcel_core::matmul(&c_t, &b_t); + let cb = repeat_axis_nemotron(&cb, repeats, 1); + + let dt_a_t = mlxcel_core::swap_axes(&chunk_dt_a, 1, 2); + let decay = mlxcel_core::exp(&segsum_nemotron(&dt_a_t)); + let attn_matrix = mlxcel_core::tril(&mlxcel_core::multiply(&cb, &decay), 0); + let dtx_t = mlxcel_core::swap_axes(&chunk_dtx, 1, 2); + let mut y_chunk = mlxcel_core::matmul(&attn_matrix, &dtx_t); + y_chunk = mlxcel_core::swap_axes(&y_chunk, 1, 2); + + let decay_shape = mlxcel_core::array_shape(&decay); + let decay_last = slice_axis(&decay, 2, decay_shape[2] - 1, decay_shape[2]); + let decay_for_state = mlxcel_core::transpose_axes(&decay_last, &[0, 3, 1, 2]); + let b_rep = repeat_axis_nemotron(&b_t, repeats, 1); + let b_sw = mlxcel_core::swap_axes(&b_rep, 2, 3); + let dtx_decay = mlxcel_core::multiply(&chunk_dtx, &decay_for_state); + let dtx_decay_t = mlxcel_core::swap_axes(&dtx_decay, 1, 2); + let dtx_decay_t = mlxcel_core::swap_axes(&dtx_decay_t, 2, 3); + let mut next_state = mlxcel_core::matmul(&dtx_decay_t, &b_sw); + + if let Some(ref previous_state) = current_state { + let dta_cumsum = mlxcel_core::cumsum(&chunk_dt_a, -2, false, true); + let exp_dta_cumsum = mlxcel_core::exp(&dta_cumsum); + + let exp_shape = mlxcel_core::array_shape(&exp_dta_cumsum); + let last_exp = slice_axis(&exp_dta_cumsum, 1, exp_shape[1] - 1, exp_shape[1]); + let last_exp = mlxcel_core::reshape(&last_exp, &[batch, num_heads, 1, 1]); + let state_contrib = mlxcel_core::multiply(&last_exp, previous_state); + next_state = mlxcel_core::add(&next_state, &state_contrib); + + let c_reshaped = + mlxcel_core::reshape(&chunk_c, &[batch, chunk_len, n_groups, 1, state_dim, 1]); + let state_reshaped = mlxcel_core::reshape( + previous_state, + &[batch, 1, n_groups, repeats, head_dim, state_dim], + ); + let y_previous = mlxcel_core::matmul(&state_reshaped, &c_reshaped); + let y_previous = mlxcel_core::squeeze_axis(&y_previous, -1); + let y_previous = + mlxcel_core::reshape(&y_previous, &[batch, chunk_len, num_heads, head_dim]); + let mut exp_shape_y = mlxcel_core::array_shape(&exp_dta_cumsum); + exp_shape_y.push(1); + let exp_dta_exp = mlxcel_core::reshape(&exp_dta_cumsum, &exp_shape_y); + let y_previous = mlxcel_core::multiply(&exp_dta_exp, &y_previous); + y_chunk = mlxcel_core::add(&y_chunk, &y_previous); + } + + let d_contrib = mlxcel_core::multiply(&chunk_x, &d_reshaped); + y_chunk = mlxcel_core::add(&y_chunk, &d_contrib); + y_chunk = mlxcel_core::astype(&y_chunk, input_dtype); + + // Evaluation inside the loop is what releases each quadratic chunk + // graph before the next one is constructed. Evaluating only the final + // concatenation would retain every chunk's intermediates. + mlxcel_core::eval(&y_chunk); + mlxcel_core::eval(&next_state); + outputs.push(y_chunk); + current_state = Some(next_state); + } + + let output = if outputs.len() == 1 { + outputs.pop().expect("one SSD output chunk") + } else { + let refs: Vec<&MlxArray> = outputs + .iter() + .map(|output| output.as_ref().expect("evaluated SSD output chunk")) + .collect(); + concatenate_many(&refs, 1) + }; + + ( + output, + current_state.expect("Nemotron-H SSM requires a non-empty sequence"), + ) +} + // Attention. struct NemotronHAttention { q_proj: UnifiedLinear, @@ -2226,6 +2271,10 @@ impl NemotronHModel { block_types.iter().filter(|t| **t == BlockType::MLP).count(), block_types.iter().filter(|t| **t == BlockType::MoE).count() ); + tracing::info!( + ssm_chunk_size = config.chunk_size, + "[NemotronH] chunked SSD prefill dispatch enabled" + ); tracing::info!("[NemotronH] Loading weights from safetensors..."); let weights = crate::models::load_text_weights(path, None)?; @@ -2358,6 +2407,7 @@ impl NemotronHModel { n_groups: config.n_groups, head_dim: config.mamba_head_dim, time_step_limit: config.effective_time_step_limit(), + chunk_size: config.chunk_size, conv_dim, conv_weight, conv_bias, diff --git a/src/models/nemotron_h_tests.rs b/src/models/nemotron_h_tests.rs index b96581978..b9ed13c91 100644 --- a/src/models/nemotron_h_tests.rs +++ b/src/models/nemotron_h_tests.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::{NemotronHConfig, NemotronLayerCache, NemotronMambaCache}; +use super::{NemotronHConfig, NemotronLayerCache, NemotronMambaCache, nemotron_h_ssm_attn}; use mlxcel_core::{dtype, generate::ModelStateSnapshot, layers::KVCache}; // --------------------------------------------------------------------------- @@ -321,6 +321,281 @@ fn post_init_no_time_step_fields_defaults_to_zero_inf() { ); } +#[test] +fn missing_chunk_size_defaults_to_128() { + let json = minimal_config_json(""); + let mut cfg: NemotronHConfig = serde_json::from_str(&json).expect("deserialize"); + assert_eq!(cfg.chunk_size, 128); + cfg.post_init().expect("default chunk_size must be valid"); +} + +#[test] +fn mamba_chunk_size_alias_is_accepted() { + let json = minimal_config_json(r#", "mamba_chunk_size": 64"#); + let mut cfg: NemotronHConfig = serde_json::from_str(&json).expect("deserialize"); + assert_eq!(cfg.chunk_size, 64); + cfg.post_init().expect("aliased chunk_size must be valid"); +} + +#[test] +fn zero_chunk_size_is_rejected() { + let json = minimal_config_json(r#", "chunk_size": 0"#); + let mut cfg: NemotronHConfig = serde_json::from_str(&json).expect("deserialize"); + let error = cfg.post_init().expect_err("zero chunk_size must fail"); + assert!(error.contains("chunk_size must be greater than zero")); +} + +#[test] +#[ignore = "requires serial MLX execution"] +fn chunked_ssd_matches_full_step_with_prior_state_and_partial_tail() { + let _runtime = crate::initialize_runtime(); + + const SEQ_LEN: i32 = 5; + const HEADS: i32 = 2; + const HEAD_DIM: i32 = 2; + const STATE_DIM: i32 = 3; + + let ramp = |len: usize, offset: f32, scale: f32| -> Vec { + (0..len) + .map(|index| offset + index as f32 * scale) + .collect() + }; + let x = mlxcel_core::from_slice_f32( + &ramp((SEQ_LEN * HEADS * HEAD_DIM) as usize, -0.2, 0.025), + &[1, SEQ_LEN, HEADS, HEAD_DIM], + ); + let b = mlxcel_core::from_slice_f32( + &ramp((SEQ_LEN * STATE_DIM) as usize, -0.1, 0.02), + &[1, SEQ_LEN, 1, STATE_DIM], + ); + let c = mlxcel_core::from_slice_f32( + &ramp((SEQ_LEN * STATE_DIM) as usize, 0.15, -0.01), + &[1, SEQ_LEN, 1, STATE_DIM], + ); + let dt = mlxcel_core::from_slice_f32( + &ramp((SEQ_LEN * HEADS) as usize, -0.4, 0.03), + &[1, SEQ_LEN, HEADS], + ); + let a_log = mlxcel_core::from_slice_f32(&[-0.4, -0.2], &[HEADS]); + let d = mlxcel_core::from_slice_f32(&[0.3, -0.2], &[HEADS]); + let dt_bias = mlxcel_core::from_slice_f32(&[-0.1, 0.2], &[HEADS]); + let prior_state = mlxcel_core::from_slice_f32( + &ramp((HEADS * HEAD_DIM * STATE_DIM) as usize, -0.05, 0.01), + &[1, HEADS, HEAD_DIM, STATE_DIM], + ); + + let (full_y, full_state) = nemotron_h_ssm_attn( + &x, + &a_log, + &b, + &c, + &d, + &dt, + &dt_bias, + Some(&prior_state), + (0.0, f32::INFINITY), + SEQ_LEN as usize, + ); + let (chunked_y, chunked_state) = nemotron_h_ssm_attn( + &x, + &a_log, + &b, + &c, + &d, + &dt, + &dt_bias, + Some(&prior_state), + (0.0, f32::INFINITY), + 2, + ); + + assert_eq!( + mlxcel_core::array_shape(&chunked_y), + vec![1, SEQ_LEN, HEADS, HEAD_DIM] + ); + let output_close = mlxcel_core::allclose(&full_y, &chunked_y, 1e-4, 1e-5); + mlxcel_core::eval(&output_close); + assert!( + mlxcel_core::item_bool(&output_close), + "chunk size 2 (2+2+1) changed SSD output" + ); + + let state_close = mlxcel_core::allclose(&full_state, &chunked_state, 1e-4, 1e-5); + mlxcel_core::eval(&state_close); + assert!( + mlxcel_core::item_bool(&state_close), + "chunk size 2 (2+2+1) changed final recurrent state" + ); +} + +fn recurrence_oracle( + x: &[f32], + a_log: &[f32], + b: &[f32], + c: &[f32], + d: &[f32], + dt: &[f32], + dt_bias: &[f32], + batch_size: usize, + seq_len: usize, + num_heads: usize, + head_dim: usize, + num_groups: usize, + state_dim: usize, +) -> (Vec, Vec) { + let heads_per_group = num_heads / num_groups; + let mut output = vec![0.0; batch_size * seq_len * num_heads * head_dim]; + let mut state = vec![0.0; batch_size * num_heads * head_dim * state_dim]; + + for batch in 0..batch_size { + for token in 0..seq_len { + for head in 0..num_heads { + let group = head / heads_per_group; + let dt_index = (batch * seq_len + token) * num_heads + head; + let delta = (dt[dt_index] + dt_bias[head]).exp().ln_1p(); + let decay = (-a_log[head].exp() * delta).exp(); + for channel in 0..head_dim { + let x_index = + ((batch * seq_len + token) * num_heads + head) * head_dim + channel; + let mut value = 0.0; + for state_channel in 0..state_dim { + let state_index = ((batch * num_heads + head) * head_dim + channel) + * state_dim + + state_channel; + let bc_index = ((batch * seq_len + token) * num_groups + group) * state_dim + + state_channel; + state[state_index] = + decay * state[state_index] + delta * x[x_index] * b[bc_index]; + value += state[state_index] * c[bc_index]; + } + output[x_index] = value + d[head] * x[x_index]; + } + } + } + } + + (output, state) +} + +#[test] +#[ignore = "requires serial MLX execution"] +fn chunked_ssd_matches_direct_recurrence_for_multigroup_f16_and_bf16() { + use mlxcel_core::utils::array_to_vec_f32; + + let _runtime = crate::initialize_runtime(); + + const BATCH: usize = 2; + const SEQ_LEN: usize = 5; + const HEADS: usize = 4; + const HEAD_DIM: usize = 2; + const GROUPS: usize = 2; + const STATE_DIM: usize = 3; + + let ramp = |len: usize, offset: f32, scale: f32| -> Vec { + (0..len) + .map(|index| offset + index as f32 * scale) + .collect() + }; + let raw_x = ramp(BATCH * SEQ_LEN * HEADS * HEAD_DIM, -0.3, 0.007); + let raw_b = ramp(BATCH * SEQ_LEN * GROUPS * STATE_DIM, -0.15, 0.004); + let raw_c = ramp(BATCH * SEQ_LEN * GROUPS * STATE_DIM, 0.2, -0.003); + let raw_dt = ramp(BATCH * SEQ_LEN * HEADS, -0.5, 0.006); + let raw_a_log = [-0.4, -0.3, -0.2, -0.1]; + let raw_d = [0.3, -0.2, 0.1, -0.05]; + let raw_dt_bias = [-0.15, -0.05, 0.05, 0.15]; + + for (dtype, tolerance) in [ + (mlxcel_core::dtype::FLOAT16, 2e-3_f32), + (mlxcel_core::dtype::BFLOAT16, 2e-2_f32), + ] { + let cast = |values: &[f32], shape: &[i32]| { + mlxcel_core::astype(&mlxcel_core::from_slice_f32(values, shape), dtype) + }; + let x = cast( + &raw_x, + &[BATCH as i32, SEQ_LEN as i32, HEADS as i32, HEAD_DIM as i32], + ); + let b = cast( + &raw_b, + &[ + BATCH as i32, + SEQ_LEN as i32, + GROUPS as i32, + STATE_DIM as i32, + ], + ); + let c = cast( + &raw_c, + &[ + BATCH as i32, + SEQ_LEN as i32, + GROUPS as i32, + STATE_DIM as i32, + ], + ); + let dt = cast(&raw_dt, &[BATCH as i32, SEQ_LEN as i32, HEADS as i32]); + let a_log = cast(&raw_a_log, &[HEADS as i32]); + let d = cast(&raw_d, &[HEADS as i32]); + let dt_bias = cast(&raw_dt_bias, &[HEADS as i32]); + + let (actual_output, actual_state) = nemotron_h_ssm_attn( + &x, + &a_log, + &b, + &c, + &d, + &dt, + &dt_bias, + None, + (0.0, f32::INFINITY), + 2, + ); + let (expected_output, expected_state) = recurrence_oracle( + &array_to_vec_f32(&x), + &array_to_vec_f32(&a_log), + &array_to_vec_f32(&b), + &array_to_vec_f32(&c), + &array_to_vec_f32(&d), + &array_to_vec_f32(&dt), + &array_to_vec_f32(&dt_bias), + BATCH, + SEQ_LEN, + HEADS, + HEAD_DIM, + GROUPS, + STATE_DIM, + ); + let expected_output = cast( + &expected_output, + &[BATCH as i32, SEQ_LEN as i32, HEADS as i32, HEAD_DIM as i32], + ); + let expected_state = mlxcel_core::from_slice_f32( + &expected_state, + &[ + BATCH as i32, + HEADS as i32, + HEAD_DIM as i32, + STATE_DIM as i32, + ], + ); + + let output_close = + mlxcel_core::allclose(&actual_output, &expected_output, 0.0, tolerance as f64); + mlxcel_core::eval(&output_close); + assert!( + mlxcel_core::item_bool(&output_close), + "chunked SSD output diverged from direct recurrence for dtype {dtype}" + ); + let state_close = + mlxcel_core::allclose(&actual_state, &expected_state, 1e-4, tolerance as f64); + mlxcel_core::eval(&state_close); + assert!( + mlxcel_core::item_bool(&state_close), + "chunked SSD state diverged from direct recurrence for dtype {dtype}" + ); + } +} + /// Config with only time_step_min: 0.001 (no time_step_limit, no time_step_max). /// After post_init, time_step_limit must be (0.0, +inf) — time_step_min is a /// softplus init bound, NOT the dt clamp lower bound (upstream PR #1026 fix).