Skip to content

nemotron cuda: ModelOpt checkpoints below sm_90, the lane matmul's slices reduced through a per-stream buffer - #901

Open
BobClawblaw wants to merge 1 commit into
ashhart:mainfrom
BobClawblaw:nemotron-modelopt-sm89
Open

BobClawblaw wants to merge 1 commit into
ashhart:mainfrom
BobClawblaw:nemotron-modelopt-sm89

Conversation

@BobClawblaw

Copy link
Copy Markdown
Contributor

What this changes

Nemotron's ModelOpt checkpoints (NVFP4, FP8, bf16 lane rows) now load on cards below sm_90. Before this, they were refused at load with ModeloptNeedsSm90. One commit on main.

The kernels already had what these cards need. The lane matmul decodes FP4 and FP8 to bf16 for mma.sync, and its cluster code is compiled out below sm_90 (qmmf.cu). Without clusters, qmmf.plan sums the K slices through fp32 partials and reduce_kernel. The family never passed it a partials buffer, so the load refused those cards instead.

  • cuda_state.zig: each stream's scratch gets lane_part. It holds the most partials any dense shape writes at rows short of fusing (qmmf.fused_rows - 1), computed with the lane's own plan. That is 22 MB a stream for Nemotron 3.5 Lightning, set by o_proj at split-K 8.
  • cuda_kernels.zig: Ops carries the buffer to decode rows (dense) and prompt rows (prefillDense, for the bf16 and FP8G rows the lane takes). The sm_90 refusal goes.
  • cuda_engine.zig: Engine.ops() passes the bound buffers' lane_part.

The bits. reduce_kernel adds the slices in the cluster path's order, slice 0 and then 1 to SK − 1, before the tensor scale. So a card without clusters makes the same sums as one with them. On sm_90 and later nothing changes except a buffer each stream doesn't use.

Receipt

  • Environment: RTX 4090 (24 GB, sm_89), Windows 11, driver 616.56, CUDA 13.3, Zig 0.17.0, -Dsm=89. main does not build on Windows yet, so the receipts ran this commit on Windows: the native CUDA server, and Qwen3.5-2B and Qwen3.8-27B on CUDA #836's branch (5e1c3fea); the Nemotron sources it touches are the same as main's. Checkpoint: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4, with its MTP head drafting.
  • Before: the native CUDA engine cannot load … (ModeloptNeedsSm90).
  • After, --context 16384 --parallel 2. Startup: model 19.51 GiB, two streams of 0.18 GiB, a 1.4 GiB prompt cache.
  • Exactness:
    • bench_concurrent --levels 1,2 --alone --serial: 24 of 24 alone and 12 of 12 serial equal.
    • Drafted against "draft": false on five prompts, greedy and seeded T 0.7: 10 of 10 equal.
  • Decode (bench_openai, 128 tokens): 468 code / 363 chat tok/s sampled, 553 / 414 greedy. Drafts accepted 3,977 of 6,153.
  • Prompt: bench_cache --require-cold, a cold 2,570-token prompt in 0.176 s (qualified).
  • Tests:
    • zig build test: 743 of 757 passed, 14 skipped, 0 failed.
    • test-cuda-host: 469 of 474 passed, 5 skipped.
    • tf-cuda-test quant-ref, modelopt, glue and smoke pass. quant-ref's NVFP4 [320, 2560] takes split-K 4, so on this card it runs the reduce path: within float64 tolerance, and every row count's bits are equal.
    • zig fmt --check is clean. The lean check reports 528 problems, the same as main.
  • Not run, and why: sm_90 and later, Linux, GB10, since I have no such machine; this change takes the same path there, the cluster. Ampere (sm_80/86) also takes this path but was not run.

Another checkpoint it opens. #897's Nemotron 3.5 Bolt 74B-A9B stores its Mamba in_proj and out_proj in NVFP4. With this change it serves on the 4090, text only, one stream:

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. The reduce sums the slices in the cluster path's order.
  • Prompt processing is no slower than the last release. On sm_90 and later the launches are unchanged.
  • The description says which platforms I ran: CUDA sm_89 on Windows. Not sm_90 and later, Linux or Metal.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. No new test: the reduce path already has quant-ref's coverage on these cards, and before this change the load refused them.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

…reduced through a per-stream buffer

The lane matmul already reduces its K slices through fp32 partials where a card has no clusters (qmmf.plan's
reduce, qmmf.cu's reduce_kernel), and its FP4 and FP8 decode to bf16 for mma.sync, but the family passed it no
partials buffer and refused sm_89 and below at load. Each stream's scratch now holds the most partials any dense
shape writes at rows short of fusing, Ops carries it to decode and prompt rows, and the refusal goes. The reduce
adds the slices in the cluster path's order, so the sums do not change.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant