nvfp4-megamoe-kernel/PERFORMANCE_AUDIT.md

PERFORMANCE — verified hot-path audit and prioritized fixes

First: congratulations. Paris-back is the milestone. It means the math is right end-to-end through all 61 layers, the production NVFP4 GEMM stack is plumbed correctly, the multi-tile FMHA kernel works in real conditions, the mHC bound holds well enough for a coherent answer, the indexer top-k is selecting the right blocks, and the FP4 → BF16 dequant path is byte-correct. That's a real architectural validation.

Second: about the agent's "1.45s/token is slow (weight loading overhead)" line. That diagnosis is wrong, and it's the kind of wrong that will steer the next agent to optimize the cold path instead of the hot one. Weight loading happens once during Phase 1 setup, before token 0. The decode step timer (t1 = time.time() at single_shot_inference.py:906) starts after weights are loaded and after every prior layer's setup is done. 1.45s is per-token decode time, not per-token load + decode. Per-token decode at hd=512, n_h=128, 61 layers, batch=1 should be in the single-digit ms ballpark on a B200, not 1.45s. There is a ~100–300× gap, and it's not weights.

The rest of this doc identifies where it actually is.

Method. Every claim below is grounded in a line number from v15. No guessing. Per doctrine rule 3.

---

P0 — Per-call .item() D2H sync inside every NVFP4 linear

This is the biggest single contributor and almost certainly explains the order of magnitude on its own. It is also the easiest fix.

dsv4/layers/linear.py:166–168:

if getattr(self, '_use_runtime_gsa', False):
    amax = hidden_states.float().abs().max().clamp(min=1e-8).item()
    self._activation_global_scale = amax / (6.0 * 448.0)

.item() is a blocking D2H copy with full stream synchronization. It forces every pending kernel on the device to finish before the host can read the value, then host blocks until the value arrives, then the host computes the scalar and the next kernel launches. Every single linear call that has _use_runtime_gsa = True is a hard pipeline bubble.

How many times does this happen per decoded token?

Call site	Per layer	× 61 layers
make_nvfp4_linear:149 (default ON for all attention proj.)	q_a, q_b, kv, o_b → 4	244
single_shot_inference.py:693 (wo_a._use_runtime_gsa = True)	1	61
:731/:750 (router gate)	1 per non-hash layer	~58
:772 (moe runner)	1 routed expert call per layer	61
:782 (shared expert)	1	61
:810 (lm_head)	1 per decoded token	1
TOTAL D2H syncs / decoded token		~486

At conservative ~50 µs per D2H sync on a B200 with kernel queue in flight, that's ~24 ms of pure pipeline bubbles per token from this one line. That's just the syncs — the lost overlap on top of that is larger.

The fix (in priority order)

1. Stop using _use_runtime_gsa on the hot path. Compute the activation global scale during a single warmup forward at startup, store it as a tensor on device, and use it directly. The infrastructure to do this already exists at linear.py:133 (compute_activation_global_scale). One warmup token through the full pipeline before the decode loop opens, then _use_runtime_gsa = False for every Nvfp4Linear.
2. If a per-call gsa really is needed (the comment at :693 claims it's to avoid E4M3 overflow), compute it on-device, never call .item(). Keep self._activation_global_scale as a torch.Tensor of shape (1,) and pass it to the kernel as a tensor. gsa_buf.fill_() at line 188 is already designed for this — it just needs the source to be a tensor, not a Python float.
3. Profile the activation amax issue itself. The comment "to avoid E4M3 overflow" is a workaround for an upstream issue. If wo_a activations genuinely have unbounded amax, that's a numerics issue worth understanding, not a perf problem to paper over with a sync per call.

Falsifiable gate

Per-decoded-token D2H sync count (measure with Nsight or cudaMemcpyAsync counter): goes from ~486 to ≤ 5 (the 5 being: argmax+token decode + end-of-loop bookkeeping). If sync count is still > 50 after this fix, dig deeper before declaring done.

---

P1 — Multi-GPU layer pipeline is serializing every layer

single_shot_inference.py:879 (decode loop):

for li in range(n_layers):
    gpu = li % NUM_GPUS
    if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}")
    torch.cuda.set_device(gpu)
    X = forward_layer(X, ...)
X = X.to('cuda:0'); torch.cuda.set_device(0)

This is round-robin layer-pipeline parallelism with batch=1. It does the worst thing all three vendors warn against: at batch=1 it doesn't parallelize anything, it just adds cross-GPU X.to() transfers between every layer and serializes on set_device(). Each cross-device .to() is an implicit synchronization. For 61 layers × 8 GPUs round-robin, the data hops ~61 times per token.

It also fights P0 — the next layer's compute can't even queue while the previous device is mid-.item().

The fix

For the single_shot reference: run on one GPU. B200 has 192 GB of HBM3e per GPU, which fits DSV4-Pro NVFP4 weights comfortably (1.6T params × 0.5 bytes for FP4 ≈ 800 GB for routed experts; with single-GPU you'll exceed one device's memory only if you keep all 384 experts resident — see the next point). On a single GPU you get:

• Zero cross-device transfers.
• No set_device() ping-pong.
• The CUDA stream is one continuous queue.

For the future multi-GPU case (when prefilling 100K+ tokens, or running batches): the correct shape is expert parallelism for MoE + tensor parallelism for attention, not layer pipelining at batch=1. That's a different design and belongs in the vLLM integration story, not the reference script.

What if it doesn't fit on one B200?

Two options, both better than layer-pipeline-at-batch=1:

1. Keep routed experts sharded across devices (EP=N), but keep attention, mHC, and the dense path on one GPU. Each MoE forward then has one all-to-all instead of a hop every layer.
2. Keep all per-token compute on one GPU and offload non-resident expert weights to host pinned memory, paging them in by predicted top-k. This is closer to how serving systems actually work.

For the single_shot script, option 1 is the right move. The script's job is to be a reference; "fits in one B200's HBM" is the constraint, and DSV4-Flash (284B, 13B activated) is well inside that. Run the single_shot on Flash first if Pro doesn't fit one GPU. The script being shape-agnostic between Flash and Pro is the right design.

Falsifiable gate

Single-GPU decode-step time (with P0 fixed): measure on Flash if Pro doesn't fit. Expected: decode-step time drops 5–10× from the multi-GPU layer-pipe baseline even before P0 helps, just from removing the cross-device hops.

---

P2 — Python loop in KVCache.append_swa (:272)

def append_swa(self, kv, pos):
    T = kv.shape[0]
    for i in range(T):
        idx = (self.swa_head + i) % self.ws
        self.swa[idx], self.swa_pos[idx] = kv[i], pos[i]
    ...

Per-decoded-token, T=1 so this loop runs once. But the assignment self.swa[idx], self.swa_pos[idx] = kv[i], pos[i] is two scalar tensor indexing ops on the GPU, each of which queues a tiny kernel. The single-token cost is small (~tens of µs) but it's a serialization point.

During prefill at T=N (say N=20 tokens in the warmup prompt), this loop runs N times and queues 2N tiny kernels. That's significant.

The fix

Vectorize:

def append_swa(self, kv, pos):
    T = kv.shape[0]
    idx = (self.swa_head + torch.arange(T, device=self.dev)) % self.ws
    self.swa.index_copy_(0, idx, kv)
    self.swa_pos.index_copy_(0, idx, pos)
    self.swa_head = (self.swa_head + T) % self.ws
    self.swa_len = min(self.swa_len + T, self.ws)

Two kernel launches instead of 2T. Same numerical result.

Falsifiable gate

append_swa queues exactly 2 kernels regardless of T. Verifiable with cudaLaunchKernel count between two cudaDeviceSynchronize calls bracketing the function.

---

P3 — Quadratic torch.cat growth on compressed KV (:280)

def add_compressed(self, ckv, cpos, idx_kv=None):
    if ckv is None: return
    self.comp_kv = ckv if self.comp_kv is None else torch.cat([self.comp_kv, ckv])
    ...

Each torch.cat allocates a new tensor of size n_comp + new_len and copies the entire existing comp_kv into it. After N tokens have produced compressed entries, total work is O(N²) and total allocator pressure is O(N²) bytes.

For the Paris demo with ~50 decoded tokens this is invisible. For the million-token contexts V4 is built for, this is catastrophic — you'd spend most of your time copying KV around.

The fix

Preallocate a ring or growing-power-of-2 buffer. Same pattern as swa:

# In __init__:
self.comp_kv_buf = torch.zeros(max_comp, head_dim, dtype=torch.bfloat16, device=dev)
self.comp_pos_buf = torch.zeros(max_comp, dtype=torch.long, device=dev)
self.comp_idx_buf = ...  # same
self.n_comp = 0

def add_compressed(self, ckv, cpos, idx_kv=None):
    if ckv is None: return
    T = ckv.shape[0]
    end = self.n_comp + T
    self.comp_kv_buf[self.n_comp:end] = ckv
    self.comp_pos_buf[self.n_comp:end] = cpos
    if idx_kv is not None: self.comp_idx_buf[self.n_comp:end] = idx_kv
    self.n_comp = end

comp_kv getters return comp_kv_buf[:n_comp] (a view, no copy).

max_comp for 1M context with m=4: 250K entries × 512 × 2 bytes = 256 MB. For 1M context with m=128 (HCA): ~16K entries × 512 × 2 = 16 MB. Both fit.

Falsifiable gate

Memory growth across 1000 decode steps stays flat (within 100 MB of steady-state). Decode-step time stays flat instead of growing.

---

P4 — v = k.clone() in _run_production_fmha (:318)

def _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx):
    ...
    q = q_heads.permute(1, 0, 2).contiguous()
    k = all_kv.unsqueeze(0).contiguous()
    v = k.clone()  # <-- this
    ...
    attn_out = dsv4_attention(q=q, k=k, v=v, ...)

DSV4 (like MLA-derivatives) uses shared KV — k and v are the same tensor. The clone() allocates and copies the entire KV buffer per call. At context=1M, this is GB-per-token of bandwidth wasted on a copy you don't need.

This is the LLM equivalent of putting the same coffee in two cups so each person can hold one.

The fix

Pass the same tensor:

attn_out = dsv4_attention(q=q, k=k, v=k, ...)

Verify dsv4_attention doesn't mutate v. From the kernel source: it doesn't — both are read-only inputs to the FMHA TMA loads.

If the kernel API insists on distinct strides or you're worried about future mutation: pass as a view, not a clone:

attn_out = dsv4_attention(q=q, k=k, v=k.view_as(k), ...)

Falsifiable gate

Memory peak per decode step drops by n_h × N × hd × 2 bytes (the V buffer). At context 8K, hd=512, n_h=128 that's 2 GB per layer per token. Not optional.

---

P5 — RoPE allocates and clones the whole tensor (:65)

def _apply_rope(x, pos, cos, sin, rope_dim, inverse=False):
    T, nh, hd = x.shape; nope = hd - rope_dim
    if pos.device != cos.device: pos = pos.to(cos.device)
    c, s = cos[pos].unsqueeze(1), sin[pos].unsqueeze(1)
    xr = x[:, :, nope:].float(); ev, od = xr[..., 0::2], xr[..., 1::2]
    if inverse: rev, rod = ev*c + od*s, -ev*s + od*c
    else: rev, rod = ev*c - od*s, ev*s + od*c
    out = x.clone(); ro = torch.empty_like(xr)
    ro[..., 0::2], ro[..., 1::2] = rev, rod
    out[:, :, nope:] = ro.bfloat16(); return out

Called 3× per attention block (Q, KV, inverse) × 61 layers = 183 RoPE calls per token. Each call does: cos[pos] gather, FP32 cast of 64 dims, multiply-add, x.clone() of the full (T, nh, hd) tensor (most of which is NoPE and doesn't need to be touched), empty_like, strided write, BF16 cast.

For T=1, hd=512, nope=448, n_h=128 per call: cloning 128×512 BF16 = 128 KB per call × 183 = 23 MB of pointless memcpy per token. Negligible bandwidth-wise on a B200, but it's 183 kernel launches that contribute to the launch-rate ceiling.

The fix

In-place RoPE for the last 64 dims, no full clone, no FP32 round-trip on the NoPE half (which is already BF16 in-place):

def _apply_rope_inplace(x, pos, cos, sin, rope_dim, inverse=False):
    nope = x.shape[-1] - rope_dim
    c = cos[pos]  # (T, rope_dim/2)
    s = sin[pos]
    # Only touch the last rope_dim dims
    xr = x[..., nope:]  # view, not copy
    ev = xr[..., 0::2].clone()  # need the original ev for the mix
    od = xr[..., 1::2]          # view; will write back below
    if inverse:
        xr[..., 0::2] = ev * c[..., None, :] + od * s[..., None, :]
        xr[..., 1::2] = -ev * s[..., None, :] + od.clone() * c[..., None, :]
    else:
        ...
    return x  # mutated in place

If the caller relies on x being unmodified, document that and have one explicit copy outside the rope, not three implicit clones inside.

Even better: fuse RoPE into the Q/KV projection kernel. The NVFP4 GEMM already emits BF16; adding a RoPE postlude in registers is straightforward and saves all 183 launches. That's the production target, not the script's job, but the script should at least not do the 183 clones.

Falsifiable gate

RoPE kernel launch count per decoded token drops from 183 (or more — each call queues multiple sub-kernels) to ≤ 3 (Q, KV, inverse, each a single launch). When fused into GEMM: 0.

---

P6 — Indexer scoring is eager einsum + topk on FP32 (:257)

scores = torch.einsum('tnd,cnd->tnc', q_idx.float(), k_idx.float())
scores = F.relu(scores)
total = (scores * w_h.unsqueeze(-1).float()).sum(1)
tk = min(self.top_k, n_comp)
_, idx = total.topk(tk, -1)

This is the lightning indexer — the paper §5.2.1 specifies this runs in FP4 on tensor cores for "99.7% recall." Right now it's running in FP32 on CUDA cores via einsum. The LUT fix made it correct; it's not fast.

Cost grows linearly with n_comp. At 1M context with m=4 → n_comp ~ 250K per CSA layer. The einsum is T × n_ih × ihd × n_comp = 1 × 64 × 128 × 250K = ~2B fmas. At FP32 on CUDA cores at ~30 TFLOPS = ~70 ms. Per CSA layer. That alone is a wall at long context.

The fix

This is roadmap item E7 ("Stage F: Lightning indexer FP4 tensor-core scoring"). For the single_shot today, FP32 einsum is correct and acceptable at Paris-scale (n_comp ≤ ~30 for 50-token decode). Do not optimize this now. Just know it's the next wall after context exceeds ~5K tokens.

The right next step when E7 lands: replace this block with a kernel call that does tcgen05.mma on FP4 keys + FP4 queries + warp-level top-k. The infrastructure is already there in dsv4/kernels/indexer/.

Falsifiable gate

Defer this gate to E7's gate (recall ≥ 99.7%). For perf today: just don't profile-test at long context until E7 lands.

---

P7 — Compressor re-runs full pipeline every step including when it won't produce output

forward_attention:357:

if compressor is not None and compressor.ratio > 0:
    comp_kv, comp_pos, block_bias = compressor.forward(x_normed, positions)
    ...

Compressor.forward at single_shot_inference.py:193 does:

def forward(self, hidden_states, positions):
    if self.ratio == 0 or self.kv_lin is None: return None, None, None
    T = hidden_states.shape[0]; r = self.ratio; dev = hidden_states.device
    n_complete = T // r
    if n_complete == 0: return None, None, None   # <-- the early return
    # ... but only AFTER computing T // r
    kv = self.kv_lin(hidden_states).float()    # NVFP4 GEMM
    gate = self.gate_lin(hidden_states).float()
    ...

At T=1 in decode, n_complete = 1 // r = 0, so we return None early — good.

But on prefill at T=20 with r=128 (HCA), n_complete = 0 and the NVFP4 GEMMs ran for nothing. Move the n_complete == 0 check above the GEMMs.

For CSA (r=4) at decode T=1, we get to do this check ~every 4 decoded tokens. The compressor still runs an NVFP4 GEMM for the kv/gate projections even though only the 4th token produces output. Either buffer hidden_states across 4 decode steps and run the GEMM once on the batched (4, H) input, or skip the projection on steps 1-3 and run only on step 4. Pre-buffering is cleaner and matches the compressor's actual semantics (windowed reduction over m consecutive tokens).

The fix

def forward(self, hidden_states, positions):
    if self.ratio == 0 or self.kv_lin is None: return None, None, None
    T = hidden_states.shape[0]; r = self.ratio
    if T // r == 0: return None, None, None  # <-- moved up
    ...

Plus the decode-time buffering: hold hidden_states from the last (r-1) steps, run the compressor only on the step where you have a complete block. This is a cache state change, not a perf trick per se — it makes the work match the math.

Falsifiable gate

For r=128 (HCA), compressor.forward does zero GEMM work for the first 127 decode steps. Measurable via Nsight kernel timeline.

---

P8 — Layer-level fusion candidates (the production future)

Below are not doctrine-acceptable shortcuts; they are the structural performance targets the production stack should reach once E5+ from the Stage E roadmap lands. Listed so they're on the radar.

1. NVFP4-1.2: Fuse FP4 quant into FMHA output → wo_a (roadmap E6). The kernel already has the epilogue_op hook. Today the FMHA writes BF16 to GMEM, then wo_a re-quantizes it via quantize_activation_nvfp4 inside its run. One full BF16 R/W pass of (T, n_h × hd) memory saved per layer. For T=1 n_h=128 hd=512 that's 128 KB per layer × 61 = 7.8 MB/token. Small absolute, but each pass is also a kernel launch and a memory round-trip in the wrong direction.

2. Fuse RMSNorm + Q/KV projection. Q/KV input is RMSNormed; the norm is a separate kernel (rmsnorm at :93). Either fuse it into the next NVFP4 linear's input quantization (the quantizer already does an amax-based normalization, ride alongside), or write a rmsnorm_quantize fused kernel.

3. Fuse RoPE into Q/KV GEMM epilogue (as in P5 above). The CuTeDSL epilogue path supports it; CUTLASS examples ship with rope-fused GEMMs.

4. mHC pre_block + RMSNorm fusion. forward_layer:469: x_in, ctx_a = attn_mhc.pre_block(X_l); x_normed = rmsnorm(x_in, attn_norm_w). That's three kernels (pre_block matmul, RMS, scale). Fusable. The mHC layer already exposes _project_and_rms per STATUS.md — wire it through.

5. CUDA graph capture (roadmap E9). Single biggest single-token win after P0/P1. The single_shot is structurally almost graph-capturable already (Nvfp4MoE says "CUDA-graph-compatible: all buffers pre-allocated, no CPU-GPU syncs, no dynamic shapes"). The blockers are:

   • P0 (.item() calls)
   • P3 (torch.cat for compressed KV — dynamic shape)
   • The 9 explicit torch.cuda.synchronize calls (grep -nE "torch\.cuda\.synchronize" single_shot_inference.py returns 9)
   • The Python if/for control flow on tensor data (e.g. topk_ids.max().item() at :432)

   Fix P0–P3 and the syncs, then graph capture is straightforward.

---

Priority order

#	Item	Effort	Win	Where
P0	Kill .item() in _use_runtime_gsa; warmup-once gsa	S	Huge (~24 ms/token)	linear.py:166–168
P1	Stop layer-pipeline at batch=1; run on 1 GPU	S	Huge (5-10×)	single_shot:879,913
P2	Vectorize KVCache.append_swa	XS	Small/medium (prefill)	single_shot:272
P3	Preallocate comp_kv, kill torch.cat	S	Critical at long ctx	single_shot:280
P4	v = k instead of v = k.clone()	XS	Big (memory + BW)	single_shot:318
P5	In-place / fused RoPE	S	Medium (-180 launches)	single_shot:65
P6	(Defer — E7 covers it)	—	—	single_shot:257
P7	Move n_complete == 0 check above GEMMs; buffer hidden for CSA	S	Medium	single_shot:193
P8	Production fusion targets — see Stage E roadmap	L	Where the real wins live	—

Do P0 and P1 first. They are tiny changes, individually catch the biggest wins, and unlock all the downstream work (CUDA graphs, prefill throughput, real-world context lengths). The script will probably go from 1.45 s/token to well under 100 ms/token from those two alone — at which point the rest of this list becomes worth measuring against, instead of swamped by syncs.

---

DOCTRINE — what to refuse during this perf pass

1. DSL wall → raw CUDA C++, not Python. Doesn't apply here — we're removing Python perf bugs, not adding kernels. But: if an agent says "I'll cache the amax in Python state," that's still Python on the hot path. The right cache lives in a torch.Tensor on device.

2. Raw CUDA ≠ scalar math. Relevant for P5/P8: when someone reaches for "let's just write a scalar fused RoPE kernel," remind them the production target is tensor-core throughput in the NVFP4 GEMM epilogue. Don't ship a scalar fused kernel as "fast enough."

3. Print, don't guess. Before claiming P0 is fixed, measure D2H syncs per decoded token with Nsight or a tracing wrapper. The "we removed .item()" claim is not verified until the sync count drops.

4. Integration over exploration. Do not write linear_v2.py with "perf improvements." Edit linear.py. The four _use_runtime_gsa = True flags in single_shot_inference.py are the test surface: flip them, run, compare.

5. Falsifiable gates. Every priority above has a measured number. "It feels faster" does not close the gate.

6. Do not optimize cold paths. Weight loading is cold. mHC weight conversion is cold. Anything that runs once during main() setup is cold. The hot path is everything inside the for step in range(MAX_NEW_TOKENS): loop. If a proposed change is in load_all_weights, _load_moe_weights_stacked, or any of the make_* helpers — that's cold, deprioritize it.