nvfp4-megamoe-kernel/PERFORMANCE_AUDIT.md

# PERFORMANCE — v17 roadmap toward end-to-end NVFP4 hot path

**Verified state.** v17 has the Tier-1 indexer fixes landed (weight path,
buffer width, MQA einsum). Hot-path syncs and allocator churn from earlier
perf rounds are gone. The single_shot now genuinely runs through the
production NVFP4 kernel stack. What remains is **fusion gaps and KV-cache
dtype choices** — the difference between "uses NVFP4 kernels" and "is
NVFP4 end-to-end."

**On TurboQuant — verdict first, reasoning below.** Don't use it for DSv4.
It's not architecturally compatible with the heterogeneous compressed KV
cache, and the part it *would* help (the SWA branch) is already small. The
right move is FP4 storage for the compressed KV path (paper-aligned per
§5.2.1), not vector-quantization codebooks. Full reasoning in Section 4.

---

# PART 1 — THE NVFP4-EVERYWHERE GAP

## P0 — Fused SwiGLU exists in the library and is NEVER ENABLED

This is the biggest single-line perf bug in v17.

`dsv4/layers/moe.py:61`:
```python
self._fused_swiglu = False  # Set via set_fused_swiglu()
```

`set_fused_swiglu()` exists (`moe.py:103`), `warmup_fused_swiglu_compilation`
exists and is wired into the warmup path, the fused kernel
`run_fused_swiglu_grouped_gemm` is implemented and tested. But **searching
`single_shot_inference.py` for `set_fused_swiglu` returns zero hits.**

What this costs every layer, every token:

`moe.py:640–660` (the unfused branch that runs by default):
```python
l1_out = run_nvfp4_grouped_gemm(...)            # NVFP4 → BF16 GEMM
l1_deil = deinterleave_l1_weights(l1_out...)    # BF16 → BF16 deinterleave (extra launch)
gate = l1_deil[:, :self.intermediate_size]       # BF16 slice
up = l1_deil[:, self.intermediate_size:]         # BF16 slice
gate_silu = F.silu(gate)                         # BF16 SiLU launch
if swiglu_limit:                                  #
    gate_silu = gate_silu.clamp(...)              # BF16 clamp launch
    up = up.clamp(...)                            # BF16 clamp launch
activated = gate_silu * up                        # BF16 elementwise
slot_l2_x_fp4, slot_l2_x_sf, _ = quantize_nvfp4_gpu_fused(activated)  # back to FP4
```

That's **8 BF16-tensor-resident kernel launches** per layer per token,
moving 2× `intermediate_size × n_active_experts` BF16 elements through
HBM, between two NVFP4 GEMMs that could have been fused.

What the fused path does (`moe.py:617–625`):
- Single launch: NVFP4 GEMM + SwiGLU + clamp in kernel registers
- Output goes directly to FP4 in `deinterleave_amax_quantize_nvfp4_fused`

**For Pro (n_active=6, intermediate=3072), per token, all 30 MoE layers:**
- 30 × 6 × (3072 BF16 = 6 KB) × 2 (R+W) × 8 launches ≈ **3 MB**
  of pointless BF16 HBM traffic per token, plus 240 unfused launches.

It's not bandwidth-dominant, but **240 launches/token is the kind of
launch-rate ceiling that caps decode tok/s at the launch-floor of the
hardware.** B200 launch rate ~1–2 µs in practice. That's 240–480 µs/token
of pure launch overhead from this one missing call.

### The fix

One line in main(), in the MoE/SE setup loop:

```python
for li in range(n_layers):
    if li in moes:
        moes[li].set_fused_swiglu(True)
        moes[li].set_swiglu_limit(cfg.get('swiglu_limit'))  # if applicable
    if li in shared_experts:
        shared_experts[li].set_fused_swiglu(True)
        shared_experts[li].set_swiglu_limit(cfg.get('swiglu_limit'))
```

Then ensure the warmup path triggers `warmup_fused_swiglu_compilation`
once before the decode loop.

### Falsifiable gate

After enabling: per-MoE-layer launch count drops from ~9 to ~2 (the GEMM
+ the L2 path). Verifiable with Nsight or `cudaLaunchKernel` counter.
Numerical parity: `cos ≥ 0.9995` vs unfused, captured before the switch.

## P1 — Shared expert has the same fused-path gap

The shared expert (`shared_expert.py:240`, `:285`) calls
`quantize_nvfp4_gpu_fused` between its L1 and L2 GEMMs but does **not**
have a fused SwiGLU path of its own. Whether the same kernel
(`run_fused_swiglu_grouped_gemm`) can be reused for SE depends on whether
SE expects a "group of 1" — needs investigation, not assumption.

### Action (read, don't guess)

Print the shapes and dtypes of SE's L1 GEMM input/output and compare to
what `run_fused_swiglu_grouped_gemm` expects. If they match (modulo
groups=1), wire it. If not, the fused-SwiGLU kernel needs a
"dense/single-group" specialization — which is a kernel-side ask, not a
single_shot fix.

### Falsifiable gate

Either SE uses the same fused kernel as MoE (same launch-count savings),
or there's a documented `.md` paper trail explaining why it can't and
what the production path is.

## P2 — Linear `.run()` per-call FP32 scale uploads still exist

`dsv4/layers/linear.py:188`:
```python
gsa = self._gsa_buf.fill_(self._activation_global_scale)
```

After the earlier P0 fix (`_use_runtime_gsa = False`), this no longer
syncs via `.item()`. But it still does a CPU→GPU scalar fill per call.
For Pro, 4 Nvfp4Linears in attention × 61 layers = 244 `fill_()` calls
per token. At ~5 µs each that's ~1.2 ms/token of CPU→GPU dispatch.

### The fix

Make `_activation_global_scale` a 1-element `torch.Tensor` on device, set
once at warmup. The fill becomes redundant — pass `self._gsa_buf` directly
to the kernel, no per-call fill needed.

```python
# In Nvfp4Linear.__init__:
self._gsa_buf = torch.full((1,), 1.0 / (6.0 * 448.0), dtype=torch.float32, device=device)

# After compute_activation_global_scale (runs once at warmup):
self._gsa_buf.fill_(gs)   # ONE TIME, not per call

# In run():
self.kernel(..., global_scale_a=self._gsa_buf)  # no fill
```

### Falsifiable gate

Zero CPU→GPU scalar fills on the hot path. Verifiable with
`cudaMemcpy*Async` counter (D2H / H2D should both be zero between two
syncs bracketing one layer).

## P3 — In-kernel RoPE fusion (still on the table, deferred from prior audit)

P5 from the v15 audit: in-place RoPE eliminated the clone problem, but
RoPE is still 3 separate launches per attention block × 61 layers ≈ 183
launches per token. Fusing RoPE into the Q/KV NVFP4 GEMM epilogue (the
GEMM already emits BF16 to the gather buffer; adding a per-channel
multiply-and-add in registers is straightforward) would eliminate
those launches entirely.

**This is a kernel-side change**, not a single_shot fix. Production target,
not single_shot target. Track it but don't gate the perf rollup on it.

### Falsifiable gate (when kernel work lands)

RoPE launch count: 183/token → 0/token. End-to-end cos ≥ 0.999998 vs
unfused.

---

# PART 2 — KV CACHE: WHAT'S ALREADY FP4-COMPATIBLE, WHAT ISN'T

DSv4's three KV streams have very different characteristics. Treating them
uniformly is the trap.

| Stream | Stored width | At 1M ctx | Per-access pattern | Quantizable? |
|---|---|---|---|---|
| **CSA main compressed** | hd=512 BF16 | 256 MB × 30 = ~7.5 GB | Random access via top-k (~1024 entries / query) | **Yes — FP4 strongly indicated** |
| **CSA indexer keys** | c_I=128 BF16 | 64 MB × 30 = ~2 GB | Streamed full-cache for top-k scoring | **Yes — FP4 paper-specified §5.2.1** |
| **HCA compressed** | hd=512 BF16 | 8 MB × 30 = 240 MB | Full sequential read every layer | **Yes — FP4 indicated** |
| **SWA** | hd=512 BF16 | 128 KB × 61 = 8 MB | Sequential ring buffer, recent 128 tokens | **No — too small to matter** |

Total BF16: ~10 GB at 1M context. Per the prior audit rewrite, this fits
comfortably on 8×B200. So **KV quantization is a throughput question, not
a memory question.**

## Why FP4 storage is the right answer for the compressed streams

Three reasons, in priority order:

1. **Paper-aligned.** §5.2.1 explicitly specifies the indexer QK path
   runs entirely in FP4. The main compressed KV cache being FP4 is
   consistent with the rest of the NVFP4 model — the cache is, after all,
   just stored projections of NVFP4 weights × BF16 hidden states.

2. **Bandwidth.** Decode is KV-read-bound at long context. Reading
   FP4 instead of BF16 quarters the bytes-per-token loaded by FMHA.
   At top_k=1024, hd=512, 30 CSA layers: that's `30 × 1024 × 512 × 1.5 bytes
   saved = 23 MB/token saved`. Across batch=8 and millions of decode
   steps, real money.

3. **Kernel-native on Blackwell.** Loading FP4 → tcgen05.mma is a
   first-class path with TMA + UMMA + the `mxf4nvf4` MMA kind. The
   in-kernel dequant happens for free during the MMA. **The infrastructure
   exists in the production FMHA kernel already** (per the prior
   `epilogue_op` work and the `ENABLE_FP4_EPILOGUE` template param).

## What this looks like in code

The compressed KV write path currently lands BF16 in `comp_kv_buf`. The
production sequence should be:

1. Compressor produces BF16 output (still — the softmax compression needs
   accumulation precision).
2. Quantize-to-NVFP4 in the same kernel as the compression (epilogue
   fusion), using the **same NVFP4 quant primitives the linears already
   use** (`quantize_nvfp4_gpu_fused`).
3. Store FP4 + per-block E4M3 scales in `comp_kv_buf` (which becomes a
   FP4 buffer + scale buffer pair).
4. FMHA reads FP4, dequants in-kernel via TMA + tcgen05's native FP4
   path. No `__constant__` LUT needed — the hardware decodes E2M1.

For the indexer keys this is the same pattern but the consumer is the
indexer scoring kernel (the FP32 einsum today, the FP4 tensor-core scorer
when E7 lands).

### Falsifiable gate (per stream)

- **CSA main + HCA + indexer:** end-to-end output cos ≥ 0.999 with FP4
  storage vs BF16. KV cache memory at 8K context drops by ~3.5× (8 → 2.3
  GB). FMHA-bound decode latency at 8K context drops measurably.
- **Recall@k for indexer ≥ 99% vs FP32 oracle** (the bar from the prior
  indexer-fix audit). Critical — FP4 must not corrupt top-k ranking.

---

# PART 3 — OTHER FUSION WINS, RANKED BY EFFORT/IMPACT

## P4 — Fuse RMSNorm into the next NVFP4 quantize

Q/KV projection input is RMSNormed; RMSNorm is a separate launch. The
NVFP4 quantize kernel already does an amax reduction per group — fusing
RMSNorm (which is *also* an amax-style reduction followed by a scale)
into the quantizer's input is a natural fit. Saves a launch + a BF16
materialization of `(T, H)` per RMSNorm site (2 per layer = 122/token).

**Effort:** S (kernel-side, but the quantizer already has the right shape).
**Impact:** Medium. 122 launches/token, ~0.7 ms/token from launch overhead alone.

## P5 — Fuse mHC pre_block + RMSNorm into a single op

Same logic as P4 but for mHC. `attn_mhc.pre_block(X_l)` → `rmsnorm` is 3
kernels back-to-back. Fusable. mHC already exposes a `_project_and_rms`
half per prior audit notes — wire it through both halves of the layer.

**Effort:** S. **Impact:** Medium. ~120 launches/token.

## P6 — CUDA graph capture (the big one, last)

Single biggest single-token win after everything above. Captures the entire
decode step into a graph; replay eliminates **all** launch overhead.
Probably worth 2–3× speedup at batch=1.

Blockers in v17:
1. `set_device()` boundaries in the layer pipeline (the `cuda.synchronize()`
   at line 963) — graph capture spans devices via multi-graph or
   per-device sub-graphs. Manageable but not free.
2. Dynamic shape in `KVCache.add_compressed` — `self.n_comp` grows.
   Fix: capture *one* graph per prefill chunk size, replay per
   decoded token (which has fixed T=1 shape; the growing buffer is
   a write into a pre-allocated tensor, capturable).
3. Any conditional `if` on tensor data — debug prints, the assertion at
   line 608. Strip from the capture path with a flag.

**Effort:** L. **Impact:** Huge (the biggest remaining single win).
**Sequence:** land after P0/P1/P2/P3 so the captured graph reflects the
post-fusion structure.

---

# PART 4 — TURBOQUANT: ARCHITECTURAL VERDICT

Reading `turboquant/`: this is an **ICLR 2026 paper implementation** of
vector-quantization KV compression. Two algorithms:
- MSE-quantize keys/values via codebook (3 bit by default)
- Inner-product-aware quantize keys (preserves dot products) via Algorithm 2
- Per-vector L2-norm preserved separately, plus QJL sign sketch for
  residual recovery

Operational shape:
- Operates on **standard MHA/GQA shape** `(..., n_heads, head_dim)`,
  head_dim typically 128.
- Requires a `head_dim × head_dim` rotation matrix per layer (precomputed
  from random seed, shared across heads).
- Has a Triton fused-decode kernel that computes attention scores directly
  from packed codebook indices.
- vLLM integration via `turboquant/vllm_attn_backend.py`.

## Why it doesn't fit DSv4

Three structural mismatches, in order of severity:

### 1. The DSv4 KV cache is already a learned compression

DSv4 doesn't store per-token KV. The CSA compressor's whole job is to
reduce m=4 tokens into 1 compressed entry via a softmax-weighted mix.
That entry is what gets cached. TurboQuant quantizes the *post-projection
per-token KV* of standard attention — exactly the thing DSv4 has
already replaced with a learned compressor. **You'd be applying a lossy
compression on top of an already-lossy compression**, which (a) compounds
loss in an uncontrolled way and (b) attacks the wrong dimension. The
compressed entries are already 4× (CSA) or 128× (HCA) reduced in the
sequence dimension; further reducing the *head dimension* via codebook
gives little additional savings (you're already attending over very few
entries per query) at high quality cost.

### 2. Wrong shape, wrong primitive

TurboQuant operates on `(..., n_heads, head_dim=128)` per-token vectors
and uses a `128×128` random rotation. DSv4's compressed cache is shape
`(n_comp, head_dim=512)` — no head dimension. The whole "rotate the head
dim" abstraction needs to be reworked, and once you do, you're writing
new code that isn't TurboQuant anymore.

For the indexer keys, the storage *is* per-block 128-dim, which is closer
to TurboQuant's natural shape. But the indexer's scoring math is
`ReLU(q·k) · w_h` summed across heads — TurboQuant's "preserve inner
products" guarantee from Algorithm 2 doesn't compose with the ReLU
nonlinearity. The quantization error becomes worst-case at the threshold,
which is where top-k decisions get made. **Bad fit precisely where it
matters most.**

### 3. NVFP4 hardware exists; TurboQuant is software-only

TurboQuant runs as bit-packed uint8 + Triton kernels. It can't use
tcgen05 FP4 tensor cores because its values aren't FP4 — they're
codebook *indices*. So you'd be paying CPU/GPU cycles to dequant via
gathers and per-token rotation matrix-vector multiplies, when the same
storage cost (4 bits/value) is available natively as FP4 with hardware
dequant during MMA.

The TurboQuant benchmark numbers (+3–5% throughput at 3-bit) are
real, but they're against `bf16_kv` baselines on architectures that
don't have FP4 tensor cores. On Blackwell with NVFP4, the comparison
should be FP4 storage + FP4 MMA — which is strictly better in every
axis (bandwidth, capacity, dequant cost).

## Where TurboQuant *would* help, and the verdict on whether it's worth it

The only DSv4 stream where TurboQuant's shape is a natural fit is the
**SWA branch** — uncompressed per-token KV in the sliding window, 128
tokens × `n_layers` × `hd=512` = 8 MB at 1M context.

**It's 8 MB.** Not worth a new dependency, a paper-grade extra failure
mode, or the rotation overhead. The SWA branch fits in L2 cache on B200.

### Verdict

Don't use TurboQuant. The right move for DSv4's KV cache is **FP4 storage
+ FP4 MMA on the compressed streams**, fully Blackwell-native, paper-
aligned (§5.2.1), with no codebook lookup overhead. The infrastructure to
do this is already in your kernel library (the `ENABLE_FP4_EPILOGUE`
template, the FP4 MMA path).

If you want a paper to cite for "what's the state-of-the-art KV
compression in 2026," TurboQuant is one. If you want the highest-perf
production-grade DSv4 implementation, native FP4 is the answer.

---

# PRIORITY ORDER

| # | Item | Effort | Win | Type |
|---|---|---|---|---|
| **P0** | Call `set_fused_swiglu(True)` on all MoEs | **XS** | **240–480 µs/token** | one-line script fix |
| **P1** | Same for shared expert (after print-and-confirm) | S | ~120 µs/token | likely script fix |
| **P2** | Drop per-call `fill_()` in Nvfp4Linear | S | ~1.2 ms/token | library fix |
| **KV-1** | FP4 storage for CSA main compressed KV | M | Huge at long context | kernel + script |
| **KV-2** | FP4 storage for HCA compressed KV | M | Same pattern as KV-1 | reuses KV-1 work |
| **KV-3** | FP4 storage for indexer keys (pair with E7) | M | Throughput + paper compliance | kernel work |
| **P3** | RoPE fused into Q/KV GEMM epilogue | M | 183 launches/token | kernel work |
| **P4** | RMSNorm fused into next quantize | S | 122 launches/token | kernel work |
| **P5** | mHC pre_block + RMSNorm fused | S | ~120 launches/token | kernel work |
| **P6** | CUDA graph capture | L | **2–3× total** | after everything above |

**P0 first.** It's a one-line edit that unlocks the fused kernel that
already exists. It is the most embarrassingly easy and most embarrassingly
overlooked perf bug in v17. The kernel author already did the hard work;
the script just isn't asking for it.

After P0/P1/P2 land, the linear hot path is genuinely tight and the
remaining wins are kernel-side fusion (P3/P4/P5) and the KV cache dtype
question (KV-1/KV-2/KV-3). Land all of those before attempting CUDA
graphs — the captured graph should reflect the final fused structure, not
the pre-fusion one.

---

# DOCTRINE

1. **DSL wall → raw CUDA C++, not Python.** Applies to P3/P4/P5 (kernel-
   side fusion work). The fused-SwiGLU kernel already exists as a model
   for what these should look like — it's NVFP4 GEMM + arbitrary-op
   epilogue in registers, fully Blackwell-native.

2. **Raw CUDA ≠ scalar math.** Applies to KV-1/KV-2/KV-3. The FP4
   storage path on the read side uses `tcgen05.mma`'s native E2M1 decode
   — no scalar dequant, no `__constant__` LUT (which was only needed
   for the indexer scoring CUDA-core path).

3. **Print, don't guess.** Applies in particular to P1 (verify SE
   shapes can use the MoE fused kernel) and KV-1/KV-2 (print the actual
   compressor output before deciding the FP4 quant boundary — same
   pattern that found the indexer bug). Do not assume the compressor
   emits a shape that matches the FP4 quant kernel; print and confirm.

4. **Integration over exploration.** Do not write `Nvfp4MoE_v2`. Do not
   write `KVCache_fp4_v2`. Edit the existing classes. P0 is one line in
   `main()`. KV-1/KV-2 are 2-tensor type changes plus the kernel-side
   read path.

5. **Falsifiable gates.** Already listed per priority. Meta-gate: after
   P0–P5 land, decode latency at 8K context should be **single-digit
   ms**, not three-digit. If it isn't, something is still on the hot
   path that shouldn't be, and the answer is "profile, don't guess
   next."

6. **Don't optimize for problems you don't have.** TurboQuant is the
   cautionary tale here. The KV cache at 1M is 10 GB on 8 × B200 — that
   is not a problem that needs solving with a new dependency. The
   problem is throughput, and the right answer is FP4 storage + FP4 MMA,
   which is hardware-native and doesn't require codebook lookups.