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`:

```python
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):

```python
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`)

```python
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:

```python
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`)

```python
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`:

```python
# 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`)

```python
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:

```python
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:

```python
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`)

```python
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):

```python
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`)

```python
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`:

```python
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:

```python
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

```python
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.