perf: gather tokens per MoE expert, run large NVFP4 GEMMs on bf16 tensor cores #91

Merged

rcsheets merged 1 commit from perf/moe-gather into main 2026-09-27 18:44:58 +00:00

Conversation 1 Commits 1 Files changed 15 +536 -59

rcsheets commented 2026-09-27 18:29:59 +00:00

Owner

Copy link

A 4314-token Nemotron 3 Nano prefill on the RTX PRO 6000 goes from 9.97 s to 2.39 s, with its GPU time falling from 8.5 s to 1.0 s. The output is unchanged. The changes came out of profiling a real long prefill of the checkpoint with Nsight Systems.

What the profile showed

Per 4314-token prefill, before this PR (main after #90):

where the GPU time went	time	share
FP32 SIMT GEMM (cutlass_80_simt_sgemm), ~5,970 launches	7.16 s	84%
paged attention (decode flash kernel), 6 layers	0.57 s	6.7%
NVFP4 dequant	0.27 s	3.2%
MoE ScaleRows + Add	0.27 s	3.3%
Mamba scan	0.085 s	1.0%

Two problems stacked in that 84%:

1. Every picked expert ran over every token. Experts that no token picked were already skipped, but each picked one still ran over the whole step, weighted zero where a token didn't pick it. At 4k tokens nearly all 128 experts get picked, so the experts did about 552k token-rows of work where routing asks for about 26k (topK 6).
2. Large-batch NVFP4 ran on CUDA cores. Above nvfp4GEMVMaxRows rows, the weight was widened to f32 for an f32 GEMM, which cuBLASLt serves with a SIMT kernel.

The scan's 1% settles the open question from #90: Mamba2's SSD rewrite is not worth doing.

What changes

MoE gather (nemotronh). Each (token, picked expert) pair gets a row in packed buffers of n × topK rows, grouped by expert.

1. One Embedding gathers the inputs.
2. Each picked expert runs up → relu² → down over only its own rows.
3. One WeightedRowSum adds each token's topK outputs onto the shared expert's.

This replaces a ScaleRows and an Add per expert. Two new backend.Ops support it, with CPU, CUDA and sizing implementations:

• RowSlice(t, r0, r1) views a contiguous row range without copying, so experts share buffers allocated once per step.
• WeightedRowSum(dst, src, rows, weights) computes dst[i] += Σⱼ w[i,j] · src[rows[i,j]] in a fixed j order: one thread per output element, deterministic, no atomics.

forwardScratch gains the packed buffers' term for a pattern's MoE layers. On Nemotron 3 Nano that is the widest layer: 57,984 values per token.

bf16 NVFP4 GEMM (CUDA). Above the GEMV threshold, matMulNVFP4 now unpacks to bf16 and runs the bf16 tensor-core GEMM.

• The unpack drops the per-tensor global (gllm_dequant_nvfp4_bf16). An e2m1 code times an e4m3 block scale has at most 6 significant bits, so the unpacked weight is exact in bf16.
• The global goes into the GEMM's alpha (gllm_lt_matmul_bf16 now takes one), where a bf16 weight could only approximate it.
• The activation is rounded to bf16, the same precision choice the BF16-weight path already makes.
• The scratch weight is half the size it was.

Results

Per 4314-token prefill on the RTX PRO 6000, measured with the same Nsight Systems setup before and after:

	before	after
expert GEMMs	7.16 s (FP32 SIMT)	~0.05 s (bf16 tensor-core, nvjet_*)
NVFP4 dequant	0.27 s	0.21 s
NVFP4 GEMV (experts with ≤ 32 tokens)	—	0.19 s
paged attention	0.57 s	0.32 s
scan + conv	0.14 s	0.14 s
GPU total	8.5 s	1.0 s
wall time	9.97 s	2.39 s

• Output: unchanged (the same next token on both runs), and a chat turn still reasons and answers correctly.
• Decode: a single sequence decoded 64 tokens at about 23 ms/token.
• What's left: mostly host overhead, about 1.4 s. Every op synchronizes, and a prefill makes about 10,750 cudaMalloc/cudaFree pairs for staging and scratch. After that come prefill attention (0.32 s, the decode flash kernel on a 4k prompt) and the NVFP4 unpacks (0.21 s). The design doc records all of this.

Testing

• CPU: RowSlice (it is a view, writes reach the parent, bounds are checked) and WeightedRowSum (hand values, bounds).
• CUDA:
  • WeightedRowSum against the CPU, with repeated source rows and a pre-filled dst.
  • RowSlice: a MatMul into a slice lands in exactly those rows of the parent.
  • TestCUDAMatMulNVFP4Exact feeds bf16-exact activations through the bf16 path and holds it to the f32 tolerance, so a lossy unpack or a missing or doubled global fails.
  • The large-batch NVFP4 test now uses the bf16 GEMM's scale-aware bound, as TestCUDAMatMulBF16 does.
• Nemotron-H goldens: pass on the CPU and on both GPUs.
• Full CUDA suite: go test -tags cuda ./... passes on the RTX PRO 6000 and the RTX 3070 (24 packages each; FP8 W8A8 skips on the 3070 as designed).
• Pure-Go build: go build, go vet and go test ./... pass.

🤖 Generated with Claude Code

A 4314-token Nemotron 3 Nano prefill on the RTX PRO 6000 goes from **9.97 s to 2.39 s**, with its GPU time falling from 8.5 s to 1.0 s. The output is unchanged. The changes came out of profiling a real long prefill of the checkpoint with Nsight Systems. ## What the profile showed Per 4314-token prefill, before this PR (`main` after #90): | where the GPU time went | time | share | |---|---|---| | FP32 SIMT GEMM (`cutlass_80_simt_sgemm`), ~5,970 launches | 7.16 s | **84%** | | paged attention (decode flash kernel), 6 layers | 0.57 s | 6.7% | | NVFP4 dequant | 0.27 s | 3.2% | | MoE ScaleRows + Add | 0.27 s | 3.3% | | Mamba scan | 0.085 s | 1.0% | Two problems stacked in that 84%: 1. **Every picked expert ran over every token.** Experts that no token picked were already skipped, but each picked one still ran over the whole step, weighted zero where a token didn't pick it. At 4k tokens nearly all 128 experts get picked, so the experts did about 552k token-rows of work where routing asks for about 26k (topK 6). 2. **Large-batch NVFP4 ran on CUDA cores.** Above `nvfp4GEMVMaxRows` rows, the weight was widened to f32 for an f32 GEMM, which cuBLASLt serves with a SIMT kernel. The scan's 1% settles the open question from #90: **Mamba2's SSD rewrite is not worth doing.** ## What changes **MoE gather (`nemotronh`).** Each (token, picked expert) pair gets a row in packed buffers of `n × topK` rows, grouped by expert. 1. One `Embedding` gathers the inputs. 2. Each picked expert runs up → relu² → down over only its own rows. 3. One `WeightedRowSum` adds each token's topK outputs onto the shared expert's. This replaces a ScaleRows and an Add per expert. Two new `backend.Ops` support it, with CPU, CUDA and sizing implementations: - `RowSlice(t, r0, r1)` views a contiguous row range without copying, so experts share buffers allocated once per step. - `WeightedRowSum(dst, src, rows, weights)` computes `dst[i] += Σⱼ w[i,j] · src[rows[i,j]]` in a fixed `j` order: one thread per output element, deterministic, no atomics. `forwardScratch` gains the packed buffers' term for a pattern's MoE layers. On Nemotron 3 Nano that is the widest layer: 57,984 values per token. **bf16 NVFP4 GEMM (CUDA).** Above the GEMV threshold, `matMulNVFP4` now unpacks to bf16 and runs the bf16 tensor-core GEMM. - The unpack drops the per-tensor global (`gllm_dequant_nvfp4_bf16`). An e2m1 code times an e4m3 block scale has at most 6 significant bits, so the unpacked weight is **exact** in bf16. - The global goes into the GEMM's `alpha` (`gllm_lt_matmul_bf16` now takes one), where a bf16 weight could only approximate it. - The activation is rounded to bf16, the same precision choice the BF16-weight path already makes. - The scratch weight is half the size it was. ## Results Per 4314-token prefill on the RTX PRO 6000, measured with the same Nsight Systems setup before and after: | | before | after | |---|---|---| | expert GEMMs | 7.16 s (FP32 SIMT) | ~0.05 s (bf16 tensor-core, `nvjet_*`) | | NVFP4 dequant | 0.27 s | 0.21 s | | NVFP4 GEMV (experts with ≤ 32 tokens) | — | 0.19 s | | paged attention | 0.57 s | 0.32 s | | scan + conv | 0.14 s | 0.14 s | | **GPU total** | **8.5 s** | **1.0 s** | | **wall time** | **9.97 s** | **2.39 s** | - **Output:** unchanged (the same next token on both runs), and a chat turn still reasons and answers correctly. - **Decode:** a single sequence decoded 64 tokens at about 23 ms/token. - **What's left:** mostly host overhead, about 1.4 s. Every op synchronizes, and a prefill makes about 10,750 `cudaMalloc`/`cudaFree` pairs for staging and scratch. After that come prefill attention (0.32 s, the decode flash kernel on a 4k prompt) and the NVFP4 unpacks (0.21 s). The design doc records all of this. ## Testing - **CPU:** `RowSlice` (it is a view, writes reach the parent, bounds are checked) and `WeightedRowSum` (hand values, bounds). - **CUDA:** - `WeightedRowSum` against the CPU, with repeated source rows and a pre-filled dst. - `RowSlice`: a MatMul into a slice lands in exactly those rows of the parent. - **`TestCUDAMatMulNVFP4Exact`** feeds bf16-exact activations through the bf16 path and holds it to the f32 tolerance, so a lossy unpack or a missing or doubled global fails. - The large-batch NVFP4 test now uses the bf16 GEMM's scale-aware bound, as `TestCUDAMatMulBF16` does. - **Nemotron-H goldens:** pass on the CPU and on both GPUs. - **Full CUDA suite:** `go test -tags cuda ./...` passes on the RTX PRO 6000 and the RTX 3070 (24 packages each; FP8 W8A8 skips on the 3070 as designed). - **Pure-Go build:** `go build`, `go vet` and `go test ./...` pass. 🤖 Generated with [Claude Code](https://claude.com/claude-code)

rcsheets added 1 commit 2026-09-27 18:29:59 +00:00

perf: gather tokens per MoE expert, run large NVFP4 GEMMs on bf16 tensor cores ... 792988c8a1

A profiled 4314-token Nemotron 3 Nano prefill on the RTX PRO 6000 took
9.97 s, 8.5 s of it GPU time, and 84% of that was one kernel: cutlass's
FP32 SIMT sgemm, about 5,970 launches per prefill. Two problems stacked:

- Every picked expert ran over every token. The MoE skipped experts no
  token picked, but ran each picked one over the whole step, weighted
  zero where a token did not pick it. At 4k tokens nearly all 128
  experts are picked, so that was ~552k token-expert rows of work for
  the ~26k that routing asks for (topK 6).
- Large-batch NVFP4 ran on CUDA cores. Past nvfp4GEMVMaxRows rows it
  widened the weight to f32 for an f32 GEMM, which cuBLASLt serves with
  a SIMT kernel.

The Mamba scan was 85 ms (1%), which settles that its SSD rewrite is not
worth doing.

MoE gather (nemotronh): each (token, picked expert) pair gets a row in
packed buffers of n*topK rows grouped by expert. One Embedding gathers
the inputs, each picked expert runs up -> relu2 -> down over only its
own rows, and one WeightedRowSum adds each token's topK outputs onto
the shared expert's. That replaces a ScaleRows and an Add per expert.
Two new backend.Ops support it, with CPU, CUDA and sizing
implementations:

- RowSlice views a contiguous row range without copying, so the
  experts share buffers allocated once per step.
- WeightedRowSum sums each dst row's k weighted source rows in a fixed
  order: one thread per output element, deterministic, no atomics.

forwardScratch gains the packed buffers' term for a pattern's MoE
layers, the widest on Nemotron 3 Nano: 57,984 values per token.

bf16 NVFP4 GEMM (cuda): past the GEMV threshold, matMulNVFP4 now
unpacks to bf16 and runs the bf16 tensor-core GEMM. The unpack drops
the per-tensor global (gllm_dequant_nvfp4_bf16): an e2m1 code times an
e4m3 block scale has at most 6 significant bits, so it is exact in
bf16. The global goes into the GEMM's alpha, which gllm_lt_matmul_bf16
now takes, where a bf16 weight could only approximate it. The
activation is rounded to bf16, the same precision choice the
BF16-weight path makes. The scratch weight is half the size it was.

Result on the same prefill: 9.97 s -> 2.39 s wall, 8.5 s -> 1.0 s GPU.
Output is unchanged, and chat still reasons and answers correctly.
Single-sequence decode ran at about 23 ms/token over 64 tokens. What
remains is host overhead (about 1.4 s: a sync on every op and ~10,750
cudaMalloc/cudaFree pairs), then prefill attention (0.32 s) and the
NVFP4 unpacks (0.21 s).

Tests:
- CPU: RowSlice (a view, writes reach the parent, bounds) and
  WeightedRowSum (hand values, bounds).
- CUDA: WeightedRowSum against the CPU (repeated rows, pre-filled dst),
  and RowSlice (a MatMul into a slice lands in exactly those rows).
- TestCUDAMatMulNVFP4Exact feeds bf16-exact activations through the
  bf16 path and holds it to the f32 tolerance, so a lossy unpack or a
  missing or doubled global fails. The large-batch NVFP4 test now uses
  the bf16 GEMM's scale-aware bound, as TestCUDAMatMulBF16 does.
- The Nemotron-H goldens pass on the CPU and on both GPUs, and the full
  CUDA suite passes on the RTX PRO 6000 and an RTX 3070.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

pr-reviewer-bot commented 2026-09-27 18:29:59 +00:00

Collaborator

Copy link

Automated review by pr-reviewer v0.52.3 | Safety Check | Mistral Small | tracking id r-b960a7-61ef1e
This is an AI-generated review and may contain mistakes.

Status: ❌ Failed

---

This review couldn't be completed: that model isn't loaded on the inference service right now, and no alternate model was able to review it either. Consider splitting this PR into smaller changes. Tracking id r-b960a7-61ef1e.

Comment @pr-reviewer-bot retry to try again.

<!-- pr-reviewer:review --> *Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.52.3 | Safety Check | Mistral Small | tracking id `r-b960a7-61ef1e`* *This is an AI-generated review and may contain mistakes.* **Status:** ❌ Failed --- This review couldn't be completed: that model isn't loaded on the inference service right now, and no alternate model was able to review it either. Consider splitting this PR into smaller changes. Tracking id `r-b960a7-61ef1e`. Comment `@pr-reviewer-bot retry` to try again.

rcsheets merged commit 0b2fdf811a into main 2026-09-27 18:44:58 +00:00

rcsheets deleted branch perf/moe-gather 2026-09-27 18:44:58 +00:00

rcsheets referenced this pull request from a commit 2026-09-27 18:45:00 +00:00

Merge pull request 'perf: gather tokens per MoE expert, run large NVFP4 GEMMs on bf16 tensor cores' (#91) from perf/moe-gather into main

rcsheets referenced this pull request from a commit 2026-09-27 18:58:28 +00:00

perf: cut host overhead from long prefills (routing, matmul scratch)

rcsheets referenced this pull request 2026-09-27 18:58:49 +00:00

perf: cut host overhead from long prefills (routing, matmul scratch) #92

Sign in to join this conversation.

Reviewers

No reviewers

Labels

Clear labels

No items

No labels

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

2 participants

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference

brooktrails/gllm!91

No description provided.