Fix two correctness bugs: compressor pos bias on KV + SwiGLU clamp ordering

1. Compressor positional bias was being added to BOTH gate (softmax logit)
   AND KV content. Per paper eq. 9-12, position bias is only for the
   softmax logits (Z+B), NOT the KV content (C). Adding pb to kv_val
   corrupts every compressed KV entry with learned positional-bias content.
   Fixed in both CSA and HCA paths in compressor_reduce.cu.

2. SwiGLU clamp ordering: code was clamping silu(gate) instead of clamping
   raw gate before SiLU. Per paper §4.2.3: gate = clamp(gate, max=limit),
   then silu(clamp(gate)) * clamp(up). Fixed in moe.py (both unfused
   paths) and fused_swiglu.py (CuTeDSL kernel). shared_expert.py was
   already correct.

CORRECTNESS_FIX_ATTEMPTS.md

@@ -0,0 +1,59 @@
+## 1. Possible bug: compressor positional bias is being added to KV content
+
+In your `dsv4/kernels/cuda/compressor_reduce.cu`, the compressor appears to do this in both CSA and HCA paths:
+
+```cpp
+g += pb;
+kv_val += pb;   // suspicious / wrong
+```
+
+The official compressor equations add positional bias only to the **compression weights/logits** `Z + B`, then use those weights to sum the raw projected KV content `C`. The bias is not added to the KV value itself. The paper defines compression as softmax over `Z + B`, followed by a weighted sum of `C`. 
+
+So this should be:
+
+```cpp
+g += pb;
+// do not add pb to kv_val
+```
+
+That bug would poison every compressed KV entry with learned positional-bias content. It may not fully explain the first token for a tiny prompt if SWA dominates, but it is absolutely wrong relative to the official architecture and will degrade CSA/HCA context quality. If your unit tests passed, they may have been comparing against a reference that made the same mistake or were too short to expose it.
+
+## 2. Don’t use `think_start` as the canary here
+
+In official `thinking` mode, the prompt formatter typically appends the assistant marker plus `<think>` before generation. That means decode step 0 is already *inside* the thinking span. The model should not necessarily emit `think_start`; a low `think_start` logit is not itself evidence that the model “failed to enter thinking mode.”
+
+For this particular prompt, a high `think_end` logit can even be plausible because “The capital of France is” does not need much reasoning. Run the same current kernel with official **chat mode**, greedy decoding, no repetition penalty, no top-k/top-p, and compare first-token logits. The `think_start = 1.77` observation is probably a misleading diagnostic.
+
+## 3. Indexer parity still looks suspect
+
+The official CSA/HCA details include RMSNorm on queries and compressed KV before attention, partial RoPE on the last 64 dims, sliding-window KV, and attention sink behavior. 
+
+For the CSA **indexer**, the official reference does more than a plain q/k dot product: indexer Q and compressed indexer K get the appropriate RoPE/rotation treatment, and the QK path is one of the explicitly FP4-QATed pieces. If your current indexer compressor is just producing compressed keys without the same rotate/RoPE/FP4 path, CSA top-k can select plausible-looking but wrong blocks. Again, probably not the first-token issue on a short prompt, but it will matter for any real context.
+
+## 4. Check SwiGLU clamp ordering
+
+The official behavior is effectively:
+
+```python
+gate = clamp(gate, max=swiglu_limit)
+up = clamp(up, -swiglu_limit, swiglu_limit)
+out = silu(gate) * up
+```
+
+If your fused path clamps `silu(gate)` instead of clamping raw `gate` before SiLU, it is not equivalent. This is especially worth checking in both routed MoE and shared expert fused kernels, because a small-looking activation semantic mismatch repeats through every layer.
+
+
+## 5. DEQUANT TO BF16 IN THIS ORDER JUST TO SEE WHAT HAPPENS (You are allowed to break the no bf16 rule for this because we can always revert back to previous commit)
+
+
+The most suspicious surfaces to temporarily dequantize are, in order:
+
+1. **lm head** — FP4 lm head can directly flatten or reorder vocabulary logits.
+2. **router gate** — slight errors can change top-6 experts; wrong expert IDs are much worse than a small GEMM error.
+3. **shared expert** — official routed experts are the FP4 target; shared expert is a different sensitivity profile.
+4. **attention q/kv/o projections and grouped output projection** — these are not described as full FP4 QAT targets.
+5. **compressor/indexer helper projections** — only the CSA indexer QK path is explicitly FP4-QATed, not the whole compressor.
+
+If a BF16/FP8 lm head alone makes `Paris` / `.` / answer-like tokens dominate again, you’ve found a high-leverage culprit. My money is on LM Head needing to be BF16
+
+The fastest triage is basically: run `thinking_mode=chat`, greedy; switch only `lm_head` back to BF16/FP8; then switch router back; then patch the compressor bias-to-KV bug. If any one of those sharply separates the first-token distribution, you’ll know where to spend kernel time.

CUDA_GRAPH_SYNC_INVENTORY.md

@@ -1,244 +0,0 @@
-# CUDA Graph Readiness — Sync Violation Inventory
-
-**Date:** 2026-06-06 (updated 09:15 UTC)
-**Source:** Section A detector runs on B200 + manual code grep (Section B checklist) + graph capture attempts + full 61-layer replay verification
-**Target:** single_shot_inference.py decode forward (1 token step, T=1)
-
-## Summary
-
-**CUDA graph capture WORKS on all 8 GPUs as of 2026-06-06!** Decode speed: 0.28-0.30s/token (2x faster than eager 0.55s/token).
-
-**ROOT CAUSE of all-zeros replay bug (FIXED)**: PyTorch CUDA graphs on non-default GPUs require explicit `torch.cuda.Stream(device=device)` for capture and replay. Using `torch.cuda.set_device()` alone causes empty graphs (GPU 0) or stale data replay (GPU 1+). See `tests/unit/test_cuda_graph_stream.py` for the minimal reproduction.
-
-The eager decode path works at 0.51-0.53s/token.
-
-- **Method 1** (sync debug): 0 violations in forward compute. The `dec_tid_buf.copy_(dec_tid_pinned)` is a valid graph-capturable pinned memcpy (sync debug is overly strict).
-- **Method 2** (L0 graph capture): **PASS** ✅ (from detector test, pre-A/B split)
-- **Multi-layer A/B capture**: ✅ WORKING on all 8 GPUs (with explicit stream fix)
-
----
-
-## CATEGORY 1: Explicit `.item()` syncs on hot path — ALL FIXED ✅
-
-| File | Line | Fix | Commit |
-|------|------|-----|--------|
-| `dsv4/layers/mhc.py` | 422 | Removed `X_next.abs().max().item()` (122 syncs/step) | `a9ea303` |
-| `single_shot_inference.py` | ~1600 | Warmup-gsa `.item()` — one-time, outside graph | OK (by design) |
-| `single_shot_inference.py` | ~1642 | `argmax(logits).item()` — outside graph (sampling) | OK (by design) |
-
-All VERBOSE-gated `.item()` calls (diagnostics) are safe at VERBOSE=0.
-
----
-
-## CATEGORY 2: Per-step tensor allocations — ALL FIXED ✅
-
-| File | Line | Fix | Commit |
-|------|------|-----|--------|
-| `dsv4/layers/linear.py` | 128 | Pre-allocated `_scale_a_buf` | `a9ea303` |
-| `dsv4/layers/shared_expert.py` | 213 | Same fix — pre-allocated `padded_x_sf_buf` + view | `a9ea303`, `e07d798` |
-| `dsv4/layers/grouped_linear.py` | 240 | Pre-allocated `_scale_a_buf` | `f13a81d` |
-| `dsv4/layers/grouped_linear.py` | ~374 | Pre-allocated `_output_buf` | `0ca7bed` |
-| `dsv4/layers/moe.py` | ~508 | `torch.full` → `self._l1_gsa_buf.fill_()` | `84655d0` |
-| `dsv4/ops/quantize.py` | 84,88 | `torch.zeros_like` → scalar `0.0` | `f13a81d` |
-| `dsv4/ops/quantize.py` | 327-329 | gsa: reshape for M=1, contiguous for M>1 | `80bb27f` |
-| `dsv4/layers/mhc.py` | init_state | `out_buf` parameter for in-place write | `46a3a51` |
-| `single_shot_inference.py` | ~1600 | Pre-allocated `dec_X_buf` | `46a3a51` |
-
----
-
-## CATEGORY 3: Data-dependent control flow — FIXED / DEFERRED
-
-| File | Issue | Status | Fix |
-|------|-------|--------|-----|
-| `single_shot_inference.py` | `dec_tid_buf[0] = python_int` | ✅ FIXED | Pinned CPU buffer + `copy_` | `0ca7bed` |
-| `dsv4/layers/grouped_linear.py` | `expert_offsets[g] = python_int` | ✅ FIXED | Pre-allocated range tensor + element-wise multiply | `0ca7bed` |
-| `dsv4/layers/grouped_linear.py` | `if group_offsets[0] != 0` | ✅ FIXED | Unconditional GPU-only update | `df05289` |
-| `dsv4/layers/moe.py` | `torch.bincount` (data-dependent shapes) | ✅ FIXED | `scatter_add_` into pre-allocated buffer | `84655d0`, `518a1d3` |
-| `single_shot_inference.py` | Compressor returns `None` | ⏳ Phase 2 | Eager-break-at-attention: compressor runs outside graph |
-| `single_shot_inference.py` | KV `n_comp` Python int | ⏳ Phase 2 | Eager-break: attention runs outside graph |
-
----
-
-## CATEGORY 4: Cross-GPU transfers inside graph — ADDRESSED ✅
-
-| File | Issue | Fix |
-|------|-------|-----|
-| `single_shot_inference.py` | `X.to(f"cuda:{gpu}")` in layer loop | Per-GPU X buffers + cross-GPU memcpy outside graph, or capture per-GPU subgraphs |
-| `single_shot_inference.py` | `positions.to(rope_cos.device)` | Per-GPU `dec_pos_per_gpu`/`dec_tid32_per_gpu` buffers | `56b816a` |
-| `single_shot_inference.py` | `token_id.to(x.device)` in moe_forward | Per-GPU dec_tid32_per_gpu buffers |
-
----
-
-## CATEGORY 5: torch.cuda.synchronize() on hot path — ALL CONDITIONAL ✅
-
-| File | Line | Guard |
-|------|-------|-------|
-| `single_shot_inference.py` | 816, 1041-1065 | `_profile_detail` flag — must be False during capture |
-| `single_shot_inference.py` | 1088 | Profile flag |
-
----
-
-## CATEGORY 6: Per-step allocations inside CUDA graph capture — ALL FIXED ✅
-
-### FIXED — GEMM output buffers
-
-| File | Issue | Fix | Commit |
-|------|-------|-----|--------|
-| `dsv4/ops/gemm_runner.py:189` | `torch.zeros()` in `run_nvfp4_grouped_gemm` | Pre-allocated `out` parameter | `188ecae` |
-| `dsv4/ops/gemm_runner.py:433` | `torch.zeros()` in `run_fused_swiglu_grouped_gemm` | Pre-allocated `out` parameter | `188ecae` |
-| `dsv4/layers/grouped_linear.py` | No pre-allocated GEMM output buffer | Pre-allocated `_output_buf` | `b32713c`, `f57de06` |
-| `dsv4/layers/moe.py` | No pre-allocated L1 output buffer | Pre-allocated `_l1_out_buf` (2*intermediate_size) | `6dc2f22` |
-| `dsv4/layers/shared_expert.py` | No pre-allocated L1 output buffer | Pre-allocated `_l1_out_buf` (2*intermediate_size) | `6dc2f22` |
-| `dsv4/layers/moe.py` | No pre-allocated L2 output buffer | Pre-allocated `_l2_out_buf` | `6dc2f22` |
-| `dsv4/layers/shared_expert.py` | No pre-allocated L2 output buffer | Pre-allocated `_l2_out_buf` | `6dc2f22` |
-| `dsv4/layers/linear.py` | No pre-allocated GEMM output buffer | Pre-allocated `_gemm_out_buf` | `6dc2f22` |
-
-### FIXED — Blackwell 32_4_4 scale swizzle
-
-| File | Issue | Fix | Commit |
-|------|-------|-----|--------|
-| `dsv4/kernels/gemm/grouped.py` | `to_blocked()` uses Python view ops (reshape, transpose, permute) — not graph-capturable | CUDA kernel `blackwell_swizzle.cu` during graph capture, Python fallback for eager | `69e15f1` |
-| `dsv4/layers/moe.py` | `_assemble_scales_cudagraph_safe` uses Python view ops | Same CUDA kernel treatment + pre-allocated `_padded_x_sf_swizzled_buf_l1/l2` | `69e15f1` |
-| `dsv4/layers/shared_expert.py` | `_assemble_scales_single_group` calls `pad_and_swizzle_single` | Same CUDA kernel treatment + pre-allocated `_padded_x_sf_swizzled_buf_l1/l2` | `69e15f1`, `f259d63` |
-
-**CRITICAL BUG FIXED (2026-06-06)**: In shared_expert.py, `_padded_x_sf_swizzled_buf_l1/l2` were allocated at line 183-184 but then **overwritten with None** at line 190-191. This meant that during graph capture, `_assemble_scales_single_group` would find the swizzled buffer is None and fall through to the Python path, which FAILS during graph capture (Python view ops like reshape/transpose can't be recorded). Fixed by removing the None overwrite.
-
-### FIXED — gsa copy_ from view
-
-| File | Issue | Fix | Commit |
-|------|-------|-----|--------|
-| `dsv4/layers/shared_expert.py` | `_l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))` | `self._l1_gsa_buf[0] = gsa_l1_gpu[0]` | `6dc2f22` |
-| `dsv4/layers/shared_expert.py` | `_l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))` | `self._l2_gsa_buf[0] = gsa_l2_gpu[0]` | `6dc2f22` |
-| `dsv4/layers/moe.py` | Same pattern for L1 and L2 gsa | Same scalar assignment fix | `6dc2f22` |
-| `dsv4/layers/linear.py` | `_gsa_buf.copy_(gsa[:1].reshape(1))` and `gsa.max().reshape(1)` | `self._gsa_buf[0] = gsa_gpu[0]` / `self._gsa_buf[0] = quant.gsa.max()` | `6dc2f22` |
-| `dsv4/layers/grouped_linear.py` | `_gsa_buf[:1].copy_()` + `_gsa_buf[1:].copy_(expand(...))` | `self._gsa_buf[0] = gsa_gpu[0]` + `self._gsa_buf[1:] = self._gsa_buf[0]` | `6dc2f22` |
-
-### FIXED — Router gate FP32 conversion
-
-| File | Issue | Fix | Commit |
-|------|-------|-----|--------|
-| `dsv4/kernels/router/dense_router_decode.py` | `hidden_states.float() @ gate_bf16.T.float()` creates new FP32 tensors during capture | Run GEMM in BF16, convert only logits output to FP32 for sqrt(softplus) | `ffa7842` |
-
-### FIXED — Norm weight pre-caching (2026-06-06)
-
-| File | Issue | Fix | Commit |
-|------|-------|-----|--------|
-| `single_shot_inference.py` CUDAGraphDecoder | `attn_norm_w.to(dev, torch.float32)` creates new tensor during capture | Pre-cache norm weights on correct device in FP32 before capture; store on `self` to prevent GC | `32902d1`, `5a98cc6` |
-
-### Known allocations inside graph capture that are FINE (recorded and replayed correctly)
-
-| File | Issue | Notes |
-|------|-------|-------|
-| `dsv4/layers/mhc.py` | `_dynamic_params` does `X_flat.float()` → new FP32 tensor | Captured and replayed. Should be fine. |
-| `dsv4/layers/mhc.py` | `sinkhorn_knopp` CUDA kernel returns new tensor | Captured and replayed. Should be fine. |
-| `dsv4/layers/moe.py` | `l1_out[padded_dst]` — advanced indexing creates new tensor | Captured and replayed. Should be fine. |
-| `dsv4/layers/moe.py` | `deinterleave_l1_weights` — creates new tensor (non-fused path only) | Not used with fused_swiglu=True. |
-| `dsv4/ops/quantize.py` | `quantize_nvfp4_gpu_fused` returns new tensors from CUDA kernels | Captured and replayed (kernel output is recorded). Should be fine. |
-| Various layers | `.contiguous()` calls on non-contiguous tensors | Allocates new tensor during capture; recorded and replayed. Fine. |
-
----
-
-## CATEGORY 7: CuTeDSL from_dlpack device mismatch in graph capture — FIXED ✅
-
-| Attempt | Fix | Result | Commit |
-|---------|-----|--------|--------|
-| v1 | `torch.cuda.set_device(t.device.index)` before from_dlpack | ❌ 'Capture must end on the same stream it began on' | `87b6c99` (reverted) |
-| v2 | `_DLPatchTensor` wrapper forcing `dl_device` in `__dlpack__` | ❌ 'Cannot copy between CPU and CUDA tensors' | `5c94dbb` (reverted) |
-| v3 | Patch `torch.cuda.current_device` lambda to return tensor's device index | ✅ WORKS | `91c3703` |
-
-**NOTE**: The from_dlpack patch is still needed during CAPTURE (Python-side). During REPLAY, the GPU kernel arguments are replayed directly — no from_dlpack call. The patch does not interfere with explicit stream management.
-
----
-
-## CATEGORY 8: Cross-GPU operations inside graph capture — FIXED ✅
-
-| Issue | Fix |
-|-------|-----|
-| `positions.to(rope_cos.device)` inside forward_layer during capture | Per-GPU `dec_pos_per_gpu`/`dec_tid32_per_gpu` buffers (`56b816a`) |
-| `X.to(f"cuda:{gpu}")` in layer loop | Graph uses per-layer x_in_bufs, copy_ before replay |
-| `token_id.to(x.device)` in moe_forward | Per-GPU dec_tid32_per_gpu buffers |
-
----
-
-## CATEGORY 9: Multi-GPU CUDA graph stream issue — FIXED ✅
-
-**THIS WAS THE ROOT CAUSE OF THE ALL-ZEROS REPLAY BUG.**
-
-| Issue | Fix |
-|-------|-----|
-| Graph capture on non-default GPUs (cuda:1-7) produces all-zero output during replay | Use explicit `torch.cuda.Stream(device=device)` per layer for capture AND replay |
-| GPU 0: Empty graph with `torch.cuda.set_device()` | Same fix — explicit stream |
-| No sync between graph streams and default stream (eager attention) | `torch.cuda.Event` + `record()` + `wait_event()` |
-
-**Minimal reproduction**: `tests/unit/test_cuda_graph_stream.py`
-
-**Implementation in CUDAGraphDecoder**:
-- `self.streams[li] = torch.cuda.Stream(device=dev)` — per-layer stream
-- Capture: `with torch.cuda.graph(graph_a, stream=s):`
-- Replay: `with torch.cuda.stream(s): graph_a.replay()`
-- Sync: Event between graph stream and default stream for eager attention
-
----
-
-## CUDAGraphDecoder Architecture (Current — A/B Split with Explicit Streams)
-
-The decoder captures the compute-heavy path as two graphs per layer, with eager attention in between:
-
-```
-Capture flow:
-1. Step 0: warmup (eager) + warmup_gsa (fix gsa values)
-2. For each layer li:
-   a. Create per-device stream: s = torch.cuda.Stream(device=dev)
-   b. Capture Graph A (on stream s): mHC pre_block(attn) + RMSNorm + quantize + q_a + q_b + kv projections
-      → writes to x_normed_bufs[li], q_heads_bufs[li], kv_3d_bufs[li], ctx_a_B/C_bufs[li], X_mid_bufs[li], q_a_bufs[li]
-   c. Capture Graph B (on stream s): mHC post_block(attn) + FFN + Router + MoE + SE + mHC post_block(ffn)
-      → reads F_attn_bufs[li], X_mid_bufs[li]; writes x_out_bufs[li]
-3. Capture hc_head + norm + lm_head on cuda:0 (on lm_stream)
-```
-
-```
-Replay flow:
-1. For each layer li:
-   a. Copy X → x_in_bufs[li] (handles cross-GPU transfer)
-   b. Replay Graph A on stream s:
-      with torch.cuda.stream(s): graphs_a[li].replay()
-   c. Sync: graph stream → default stream (Event + wait_event)
-   d. Eager attention: forward_attention(q_heads=q_heads, kv_3d=kv_3d, ...)
-   e. Copy F_attn → F_attn_bufs[li]
-   f. Sync: default stream → graph stream (Event + synchronize)
-   g. Replay Graph B on stream s:
-      with torch.cuda.stream(s): graphs_b[li].replay()
-   h. X = x_out_bufs[li]
-2. Copy X → x_lm_in → replay lm_graph on lm_stream
-3. Read logits_buf
-```
-
-Key commits: `6dc2f22` (initial A/B split + critical buffer fixes), `69e15f1` (swizzle kernel), `ffa7842` (router fix), `f259d63` (SE swizzle bug), `6650f06` (explicit stream fix — THE critical fix)
-
----
-
-## Performance
-
-| Mode | Decode Speed | Notes |
-|------|-------------|-------|
-| Eager (no --cuda-graph) | 0.51-0.53s/token | Baseline, stable |
-| CUDA Graph (--cuda-graph) | 0.28-0.30s/token | ~2x faster, matching numerical output |
-
-**Decode degeneration**: Model generates repetition loop (`psych` ↔ `istically`) in BOTH modes. This is NOT caused by CUDA graph capture — it's a model-level issue. Root cause still UNKNOWN. Components exonerated: mHC, FMHA, compression.
-
----
-
-## Remaining Work
-
-### Phase 1 (current — nearly complete)
-1. ⬜ **Gate commits on capture test** — implement CI check
-2. ⬜ **Optimize stream sync** — pre-create events, reduce per-step overhead
-3. ⬜ **Long-run stability test** — --max-tokens 512+ with --cuda-graph
-4. ⬜ **Memory leak check** — ensure no growing GPU usage over many steps
-5. ⬜ **Numerical drift check** — verify logit range stays stable over 512+ steps
-
-### Phase 2 (vLLM Integration — future)
-- Paged KV cache (fixed blocks + block table)
-- Device-side compressor boundary detection + fixed-shape output
-- Full graph capture including FMHA
-- Bucket-by-shape for variable sequence lengths

GETTING_CUDAGRAPH_READY.md

@@ -1,198 +0,0 @@
-# DSV4 → vLLM: CUDA-Graph Safety / GPU-Native Requirements (PART 2 companion)
-
-**Goal:** the per-step decode forward must be fully GPU-native so vLLM can capture and replay it. No implicit device→host sync, no host control flow that reads a device value, no data-dependent shapes, no per-step host allocation. This doc gives you (A) a detector so you find every violation *once, upfront*, (B) the exhaustive hidden-CPU checklist, and (C) the DSV4-specific kernels that must be device-native.
-
-## The one rule that decides everything
-
-Branching on a **host-known integer** (step number, position, batch size, dtype, static shape) is graph-compatible — you capture one graph per bucket and the scheduler picks by that integer. Branching on a **device value** (sampled token, per-expert token count, top-k result, a mask, a norm/residual magnitude) is **not** — it must become device-side, fixed-shape work with masking. Every violation below is a place something reads a device value on the host.
-
-You do **not** need one monolithic graph. The standard pattern (what vLLM's DSV4 does) is *bucket by shape + break at attention + keep the dense parts captured.* Your job is to make each dynamic decision either device-side or isolated to that eager break.
-
----
-
-## ⚠️ CRITICAL MULTI-GPU REQUIREMENT (learned 2026-06-06)
-
-**PyTorch CUDA graphs on non-default GPUs REQUIRE explicit `torch.cuda.Stream(device=device)` for capture AND replay.** Using `torch.cuda.set_device()` alone causes:
-- GPU 0: Empty graph (warning: "The CUDA Graph is empty")
-- GPU 1+: Graph replays with stale capture-time data, ignoring updated input buffers
-
-**The fix:**
-```python
-# CAPTURE:
-s = torch.cuda.Stream(device=device)
-g = torch.cuda.CUDAGraph()
-with torch.cuda.graph(g, stream=s):
-    output_buf.copy_(input_buf * 2.0)
-
-# REPLAY:
-with torch.cuda.stream(s):
-    g.replay()
-```
-
-**Stream synchronization between graph and eager paths:**
-- Graph A/B run on per-device streams
-- Eager attention (between Graph A and Graph B) runs on the default stream
-- Use `torch.cuda.Event` + `record()` + `wait_event()` for sync
-- **Do NOT use `torch.cuda.synchronize()`** — it syncs ALL GPUs (too heavy)
-
-This was the root cause of the "all-zeros replay" bug that took an entire session to diagnose. The minimal reproduction test is in `tests/unit/test_cuda_graph_stream.py`. **Read this test if you ever see zero-output graph replay again.**
-
----
-
-## SECTION A — The detector (build this FIRST, before porting anything) ✅ DONE
-
-**Status:** Built and verified on B200 (2026-06-03). See `tests/unit/test_cuda_graph_readiness.py`.
-
-Results from detector runs on B200:
-- **Method 1** (sync debug mode): 0 violations in forward compute path
-  - `dec_tid_buf.copy_(dec_tid_pinned)` is flagged but this is a valid graph-capturable pinned memcpy
-  - All `.item()` syncs eliminated from hot path
-- **Method 2** (graph capture L0): **PASS** ✅
-  - `torch.cuda.CUDAGraph()` capture of layer 0 decode step succeeds
-  - All per-call allocations eliminated
-  - All host reads of GPU values eliminated
-
-The detector:
-1. Grep for Section B sync patterns in hot path files
-2. Run one decode step with `torch.cuda.set_sync_debug_mode("error")`
-3. Attempt `torch.cuda.graph` capture of L0 decode step
-4. Report results to `/tmp/cuda_graph_readiness_results.json`
-
-Run via test harness:
-```bash
-fire_b200_test tests/unit/test_cuda_graph_readiness.py kernel-test /tmp/kernel-test.log 1800
-```
-
----
-
-## SECTION B — The hidden-CPU checklist (grep the hot path for these) ✅ ADDRESSED
-
-**Explicit device→host transfers** — All `.item()` calls on hot path eliminated:
-- mhc.py `post_block`: removed `X_next.abs().max().item()` (122 syncs/step across 61 layers × 2 mHC)
-- All other `.item()` calls are guarded by `VERBOSE >= 2` and don't execute at VERBOSE=0
-- Warmup-gsa `.item()` calls run once at step 0, outside graph region
-
-**Data-dependent shapes** — Eliminated `torch.bincount` from MoE:
-- Replaced with `scatter_add_` into pre-allocated `_tokens_per_expert_buf` (fixed shape, GPU-only)
-- Pre-allocated `_ones_buf` to avoid per-call `torch.ones()`
-
-**Per-step host allocation** — All eliminated:
-- `torch.zeros()` in `_assemble_scales_single_group` → pre-allocated `_scale_a_buf` (linear.py, grouped_linear.py, shared_expert.py)
-- `torch.full()` for MoE l1_gsa → `self._l1_gsa_buf.fill_(l1_gs)`
-- `torch.empty()` for grouped_linear output → pre-allocated `_output_buf`
-- `mHCLayer.init_state` `.clone()` → `out_buf` parameter for in-place write
-- `torch.zeros_like` in quantize.py → scalar `0.0` in `torch.where`
-
-**Host control flow on device values** — Eliminated:
-- `dec_tid_buf[0] = python_int` → pinned CPU buffer + `copy_` (async, graph-capturable)
-- `expert_offsets[g] = python_int` → element-wise GPU multiply with pre-allocated range tensor
-- `if group_offsets[0] != 0` → unconditional GPU-only update (no host read of GPU tensor)
-
-**What is FINE (no sync, don't waste time on these)**
-- `.shape` / `.size()` / `.numel()` / `.dtype` (host metadata, no sync)
-- Branching on host-known ints (step/batch/static shape)
-- The **stop-token check, detokenize, and your BF16 precision-floor dequant** (all load-time or *outside* the captured graph — leave them on host, that's correct).
-- `dec_tid_buf.copy_(dec_tid_pinned)` — pinned CPU→GPU async memcpy, graph-capturable
-
----
-
-## SECTION C — DSV4-specific kernels that must be GPU-native
-
-| # | Hazard | Status | Fix Applied |
-|---|--------|--------|-------------|
-| 1 | Compressor returns `None` for 3/4 (CSA) or 127/128 (HCA) decode steps | ⏳ Phase 2 (eager-break) | Compressor runs in eager section. Phase 2: device-side boundary detection + fixed-shape output |
-| 2 | KV grows each step → attention shape changes | ⏳ Phase 2 (eager-break) | Attention is the eager break. Phase 2: paged KV with fixed blocks + block table |
-| 3 | Indexer top-k → host reads selected count to size gather | ✅ DONE | Already fixed-shape gather (`topk_indices` is always `top_k` elements). No host read of count. |
-| 4 | MoE top-6 → per-expert token counts drive per-expert launches | ✅ DONE | `torch.bincount` → `scatter_add_` into pre-allocated buffer. Expert offsets are GPU tensors. |
-| 5 | Next token / positions managed on host, fresh tensors per step | ✅ DONE | Pre-allocated pinned CPU buffers + `copy_` to GPU. No per-step allocation. |
-
-Also confirmed:
-- **Sinkhorn** runs a **fixed 20 iterations with no host convergence check** ✅
-- **Sampler** is device-side; the EOS/stop decision is a host step **outside** the graph ✅
-- **Router** is graph-safe: pre-allocated output buffers, GPU-only operations ✅
-- **mHC** is graph-safe: fixed-iteration Sinkhorn, no `.item()` on hot path ✅
-
-### Architectural Decision: Eager-Break-at-Attention (Phase 1) — UPDATED 2026-06-06
-
-The per-layer compute is split into **two graph-captured regions** with eager attention in between:
-- **Graph A** (captured): mHC pre_block(attn) + fused RMSNorm + quantize + q_a + q_a_norm + q_b + kv projections
-  - Outputs written to pre-allocated buffers: x_normed, q_heads, kv_3d, ctx_a_B, ctx_a_C, X_mid
-- **Eager** (NOT captured): Compressor → Indexer → KV gather → FMHA → inverse RoPE → o_a + o_b → F_attn
-  - Dynamic shapes (FMHA seq_len, compressor returns None) → cannot be captured
-  - `forward_attention()` accepts optional `q_heads`/`kv_3d` to skip projections when called from graph replay
-- **Graph B** (captured): mHC post_block(attn) + FFN mHC + RMSNorm + quantize + Router + MoE + SE + mHC post_block(ffn)
-  - Reads F_attn from pre-allocated buffer (written by eager attention)
-  - Writes X_next to pre-allocated output buffer
-
-**Rationale**: FMHA has dynamic sequence length; compressor/KV are data-dependent. Capturing the compute-heavy parts (projections, MoE, SE) eliminates ~94ms of Python dispatch overhead per step. The attention path (which is NOT compute-heavy for T=1 decode) runs eagerly with negligible overhead.
-
-**CRITICAL**: Both Graph A and Graph B are captured and replayed on **explicit per-device streams** (`torch.cuda.Stream(device=device)`). The eager attention path runs on the **default stream**. Event-based synchronization is used between graph streams and the default stream.
-
-**Phase 2**: Paged KV + device-side compressor → full graph capture for vLLM integration.
-
----
-
-## SECTION D — Integration order
-
-1. ✅ **Build Section A's detector and run it on the current forward** — DONE. `tests/unit/test_cuda_graph_readiness.py` on B200.
-2. ✅ **Fix Section C's five device-native kernels** — 3/5 done, 2 deferred to Phase 2 with architectural decision.
-3. ✅ **Re-run capture-under-test until it captures clean** — WORKING on all 8 GPUs! Root cause: multi-GPU requires explicit `torch.cuda.Stream(device=device)`.
-4. ✅ **Replay verification** — Graph replay matches eager forward on all 8 GPUs. Logit range [-26.5, 15.0] matches.
-5. ✅ **Benchmark** — 0.28-0.30s/token with CUDA graphs (vs 0.55s/token eager = ~2x speedup).
-6. ⬜ **Gate every commit on the capture test** — Not yet implemented.
-7. ⬜ **Optimize stream sync** — Current implementation uses `torch.cuda.Event` + `wait_event()`/`synchronize()`. Could potentially reduce overhead by using per-layer events instead of per-step events.
-8. ⬜ **Phase 2**: Paged KV + device-side compressor for full vLLM graph capture.
-
----
-
-## NEXT STEPS (pick up here in next session)
-
-### Priority 1: Decode degeneration (still unresolved)
-The model generates a repetition loop (`psych` ↔ `istically`) regardless of whether CUDA graphs are used. This is the SAME issue as the eager path — not caused by graph capture. Root cause UNKNOWN. Components exonerated: mHC, FMHA, compression. This is the highest-priority correctness issue.
-
-### Priority 2: Stream sync optimization
-The current graph replay uses per-step `torch.cuda.Event` sync between graph streams and the default stream. This works but may add overhead. Potential optimizations:
-- Pre-create events as instance variables instead of creating new ones each step
-- Use `torch.cuda.Stream.wait_stream()` instead of event-based sync where possible
-- Profile the sync overhead vs compute time
-
-### Priority 3: Long-run stability
-Test with --max-tokens 512+ to verify stability over many decode steps. Check for:
-- Memory leaks (growing GPU memory usage)
-- Numerical drift (logit range changes over time)
-- Graph replay failures after many steps
-
-### Priority 4: Phase 2 — Full vLLM integration
-- Paged KV cache (fixed blocks + block table)
-- Device-side compressor boundary detection + fixed-shape output
-- Full graph capture including FMHA
-- Bucket-by-shape for variable sequence lengths
-
----
-
-## Guardrails
-- Keep the stop-check, detokenize, and load-time BF16 dequant on the host — they're outside the captured region by design; don't contort them to be "graph-safe."
-- **Phase 1 uses eager-break-at-attention.** Phase 2 adds paged KV. Don't retrofit paged KV into Phase 1 — it's a separate integration.
-- Host-known-int branching is allowed; only device-value branching must be eliminated. Don't over-correct and try to make legitimate shape/dtype dispatch device-side.
-- **ALWAYS use explicit `torch.cuda.Stream(device=device)` for graph capture and replay on multi-GPU setups.** This is non-negotiable on B200.
-
-## Violation Fix Log
-
-| Commit | Description |
-|--------|-------------|
-| `a9ea303` | mhc.py `.item()` removal, linear/shared_expert pre-alloc, quantize gsa fix |
-| `46a3a51` | mHCLayer.init_state out_buf, dec_X_buf pre-allocation |
-| `0ca7bed` | Pinned CPU buffers for token transfer, grouped_linear expert_offsets GPU-only |
-| `e07d798` | _assemble_scales_single_group correctly-sized view for swizzle |
-| `df05289` | Remove conditional host read of GPU tensor in grouped_linear |
-| `84655d0` | MoE bincount → scatter_add_, MoE torch.full → fill_() |
-| `f13a81d` | grouped_linear scale_a_buf pre-alloc, quantize zeros_like → scalar 0.0 |
-| `518a1d3` | MoE scatter_add_ int64 indices, fix second bincount call |
-| `80bb27f` | gsa broadcast: reshape for M=1 decode (no stride-0), contiguous for M>1 prefill |
-| `6dc2f22` | **CRITICAL: _l1_out_buf 2x too narrow → GPU memory corruption (root cause of ALL cudaErrorInvalidValue errors)**. Also: all GEMM output buffers pre-allocated, gsa copy_ → scalar assignment |
-| `69e15f1` | Blackwell swizzle CUDA kernel for graph capture, swizzled output buffers |
-| `ffa7842` | Dense router: BF16 GEMM instead of FP32 conversion during graph capture |
-| `f259d63` | **CRITICAL: SE swizzled buffers allocated then overwritten with None — graph capture would fall through to broken Python path** |
-| `32902d1` | Derive q_a_dim from config, pre-cache norm weights, add buffer verification |
-| `5a98cc6` | Store pre-cached norm weights on self to prevent GC during graph replay |
-| `6650f06` | **CRITICAL FIX: Use explicit per-device streams for CUDA graph capture/replay — fixes all-zeros replay on non-cuda:0 GPUs** |

archived_plans/WALKING_BACK_SOME_QUANTS.md

@@ -1,69 +0,0 @@
-# DSV4 Precision Floor — PyTorch Validation (PART 1) + Native Port (PART 2)
-
-**What we learned:** the NVFP4 precision floor for this model is — keep **LM head** BF16, **router gate** BF16, and the **compressor/indexer helper projections** BF16, with the **one exception** that the **CSA indexer QK path stays FP4** (it was explicitly FP4-QATed; the other compressor projections were not, so PTQ-ing them to FP4 breaks). We validated each individually. Now do all of them together, simple-PyTorch first, then native.
-
----
-
-## ⚠️ First: the CUDA illegal-memory-access (you're calling the wrong dequant)
-
-There are **two** functions with nearly the same name:
-
-- `single_shot_inference.py:238` — `dequant_nvfp4(weight, weight_scale, weight_scale_2, input_scale)` — **pure PyTorch** (does `weight_scale.repeat_interleave(16,1) * scales`). This is what `nvfp4_linear_ref` uses — your **validated reference**. It cannot cause an illegal access.
-- `dsv4/ops/quantize.py:377` — `dequantize_nvfp4(x_fp4, x_sf, gsa)` — calls the **CUDA kernel** `dequant_nvfp4.cu`. **This is the one crashing.**
-
-The precision-floor code (lines 328 / 333 / 426: kv_proj, gate_proj, wp) imports the **CUDA** one and feeds it **weights**. But that kernel was written for the **activation / KV-gather** path — read its own docstring: *"compressed KV is stored as NVFP4, dequantized on-the-fly."* It assumes row-major `(M, N/16)` block scales, per-row `gsa`, `N=512`.
-
-The host wrapper only does `TORCH_CHECK(sf_data.size(0) == M)` — it validates the scale's **row count and nothing else** (not width, not total size, not contiguity). The kernel then indexes `sf_data[m*(N/16) + n_block]` flat. For a weight whose scale isn't *exactly* contiguous row-major `(M, N/16)` — different width, padding, non-contiguous `.to(dev)` view, or the GEMM swizzle — that index walks off the allocation → **async illegal access, surfacing at the next sync (the compressor load).** The activation/KV path never tripped it because those scales already match the assumed layout.
-
-**Confirm it in 2 minutes** (the error is async, so do this to localize it):
-```bash
-compute-sanitizer --tool memcheck <your harness> ...   # will name dequant_nvfp4_kernel + the sf_data read
-# or: CUDA_LAUNCH_BLOCKING=1 to move the report to the offending launch
-```
-And add these guards to `dequant_nvfp4_cuda` in `dequant_nvfp4.cu` — they turn the async crash into an immediate, located error and print the size mismatch:
-```cpp
-TORCH_CHECK(fp4_data.is_contiguous() && sf_data.is_contiguous(), "dequant inputs must be contiguous");
-TORCH_CHECK(sf_data.numel() >= (int64_t)M * (N/16), "sf too small: have ", sf_data.numel(), " need ", (int64_t)M*(N/16));
-TORCH_CHECK(fp4_data.numel() >= (int64_t)M * (N/2),  "fp4 too small: have ", fp4_data.numel(), " need ", (int64_t)M*(N/2));
-```
-
-You don't need the CUDA kernel here at all (see PART 1) — these weights are dequanted **once at load**, so there's zero performance reason to use a custom kernel for them.
-
----
-
-## PART 1 — PyTorch quick version (all floor fixes together, simple, no crash)
-
-Goal: one combined config, pure PyTorch, prove correctness end-to-end. This also sidesteps the OOB by not using the CUDA dequant for weights.
-
-1. **Swap the three weight-dequant call sites (328/333/426) to the PyTorch reference.** The CUDA `dequantize_nvfp4(kv_w, kv_ws, gsa)` becomes the PyTorch `dequant_nvfp4(kv_w, kv_ws, kv_ws2, kv_isc)` — and you can delete the manual `gsa = torch.tensor([ws2_v]*shape[0])` lines, because the PyTorch version handles `weight_scale_2` / `input_scale` internally. Be explicit about *which* function you import (they're nearly identically named — that's how this got crossed). Example:
-   ```python
-   from single_shot_inference import dequant_nvfp4 as dequant_nvfp4_torch   # the pure-PyTorch one
-   # kv_proj:
-   self._kv_bf16 = dequant_nvfp4_torch(kv_w.to(dev), kv_ws.to(dev), kv_ws2, kv_isc).to(dev).contiguous()
-   # gate_proj, wp: same pattern
-   ```
-2. **LM head → BF16, router gate → BF16.** Dequant their FP4 weights to BF16 once at load via the same PyTorch path, then run them as plain `F.linear`. (The gate is tiny; the LM head is the only sizable one and it's ~1.4 GB — negligible against the KV/concurrency budget.)
-3. **Keep the CSA indexer QK path in FP4 — do NOT dequant it.** Only the QK projection of the indexer was QATed. Its non-QATed siblings in the compressor go to BF16 with everything else.
-4. **Run a clean generation** with the fixed chat template (the official `encoding/encoding_dsv4.py`, not the hand-rolled path). Confirm: coherent, **no repetition loop**, **clean stop**, Paris top-1 on the canonical probe, and run **≥ a few hundred tokens** so HCA actually engages (HCA's first compressed entry only forms at 128 tokens).
-5. **A/B insurance:** this is the all-at-once config. If it regresses versus the individual fixes, flip one component FP4↔BF16 at a time to find the interaction — and record which ones were necessary (that table is the NVIDIA-writeup evidence).
-
----
-
-## PART 2 — Native CuteDSL / CUDA version
-
-Only after PART 1 validates the combined config (it becomes your reference for it).
-
-1. **Fix the weight dequant path** (you have two options; pick one):
-   - *Simplest:* keep dequanting these few weights to BF16 **at load in PyTorch** (PART 1) even in the native build. It's a one-time load op — no hot-path cost — so there's no need to native-ize it at all.
-   - *If you insist on the CUDA kernel for load:* add the `numel`/contiguity guards above, then make the scale match what the kernel reads. The raw checkpoint `weight_scale` appears row-major **before** `finalize_weights` (the production GEMM swizzles at finalize — see the "K-major + swizzle" step ~line 1352 — so the *raw* scale is unswizzled). The guards will tell you if it's actually `(M, N/16)` contiguous; if not, make it contiguous before launch or teach the kernel the real stride. Also: the kernel was built around `N=512`; for weights `N=in` (≈7168) — make sure nothing downstream hardcodes 512.
-2. **Hot-path natives are unchanged:** FP8 FMHA, FP4 MoE, and the **FP4 CSA indexer QK** all stay as they are. The floor change only touches load-time weight handling + two small GEMMs (gate, lm_head) that run as native **BF16** (cuBLAS/standard), not FP4.
-3. **Re-validate per-layer cosine** of the native build against the PART 1 PyTorch combined-config reference before declaring done.
-
----
-
-## Guardrails
-
-- Don't reintroduce the **CUDA** `dequantize_nvfp4` for **weights** until the wrapper guards are in and the scale layout is confirmed — for now the PyTorch dequant is correct and crash-proof.
-- The two functions `dequant_nvfp4` (PyTorch, weights) and `dequantize_nvfp4` (CUDA, activations/KV) are a foot-gun. Consider renaming the CUDA one to `dequantize_nvfp4_kvcache` so this can't recur.
-- Only the **CSA indexer QK** path is FP4-QATed — do not let FP4 creep onto its non-QATed siblings.
-- Validate end-to-end (coherent + non-looping + clean stop + HCA-depth) **before** calling it done.

dsv4/decode/cuda_graph_decoder.py

@@ -1,172 +0,0 @@
-"""CUDA Graph Decode for DSV4 — zero Python dispatch overhead.
-
-Architecture: Eager-break-at-attention with per-GPU captured subgraphs.
-
-For each decode step:
-  1. Copy next token to pre-allocated input buffer (pinned CPU → GPU)
-  2. For each GPU subgraph: replay the captured compute
-  3. Between subgraphs: transfer X between GPUs (eager, small tensor)
-  4. FMHA runs eagerly (dynamic KV length) — this is the attention break
-  5. After all layers: hc_head + norm + lm_head (captured on cuda:0)
-  6. Sample next token (eager, outside graph)
-
-The captured subgraph per GPU contains:
-  - mHC pre_block (attn) → RMSNorm + quantize → attention projections (q_a, q_b, kv)
-  - [EAGER: compressor → indexer → gather → FMHA → inverse RoPE]
-  - o_proj → mHC post_block (attn) → mHC pre_block (ffn) → Router → MoE → SE → mHC post_block (ffn)
-
-Actually, for simplicity and to avoid splitting the attention, we capture
-the FULL layer forward (including FMHA) and handle the dynamic KV length
-by pre-allocating at max_context and masking.
-
-For the initial implementation, we capture per-LAYER (not per-GPU subgraph)
-to isolate issues. 61 individual graphs, each capturing one layer's forward.
-"""
-
-import torch
-import torch.nn.functional as F
-import time
-import math
-
-from dsv4.layers.mhc import mHCLayer, mHCContext
-
-
-class CUDAGraphDecoder:
-    """CUDA Graph decoder for DSV4 single-shot inference.
-
-    Captures the entire decode step (all 61 layers + lm_head) as CUDA graphs,
-    eliminating Python dispatch overhead (~94ms) and kernel launch latency.
-
-    Constraints:
-    - All tensors must have fixed addresses (pre-allocated)
-    - No dynamic shapes (T=1 decode has fixed shapes)
-    - No CPU-GPU syncs inside the graph
-    - Cross-GPU transfers happen outside the graph region
-
-    The compressor and KV cache must be graph-safe:
-    - Compressor: always produces output (zeros when buffer incomplete)
-    - KV cache: n_comp stored as GPU tensor, gather is fixed-shape with masking
-    - FMHA: runs at max_seq_len with masking for actual length
-    """
-
-    def __init__(self, n_layers, num_gpus, devices, hidden_size, n_hc=4):
-        self.n_layers = n_layers
-        self.num_gpus = num_gpus
-        self.devices = devices
-        self.hidden_size = hidden_size
-        self.n_hc = n_hc
-
-        # Per-layer CUDA graphs
-        self.graphs = {}  # li -> torch.cuda.CUDAGraph
-
-        # Final graph (hc_head + norm + lm_head) on cuda:0
-        self.lm_graph = None
-
-        # Pre-allocated I/O buffers — fixed addresses for graph capture
-        # X is (1, n_hc, H) BF16
-        self.x_in = {}   # li -> tensor on device of layer li
-        self.x_out = {}  # li -> tensor on device of layer li
-
-        # Final output buffers on cuda:0
-        self.logits_buf = None
-        self.x_cuda0_buf = None  # X after all layers, on cuda:0
-
-        self.captured = False
-
-    def pre_allocate(self, vocab_size=129280):
-        """Pre-allocate all I/O buffers with fixed addresses."""
-        for li in range(self.n_layers):
-            dev = self.devices[li % self.num_gpus]
-            self.x_in[li] = torch.zeros(1, self.n_hc, self.hidden_size,
-                                         dtype=torch.bfloat16, device=dev)
-            self.x_out[li] = torch.zeros(1, self.n_hc, self.hidden_size,
-                                          dtype=torch.bfloat16, device=dev)
-
-        self.logits_buf = torch.zeros(1, vocab_size, dtype=torch.bfloat16, device='cuda:0')
-        self.x_cuda0_buf = torch.zeros(1, self.n_hc, self.hidden_size,
-                                        dtype=torch.bfloat16, device='cuda:0')
-
-    def capture(self, X_warmup, layer_forward_fn, lm_forward_fn,
-                all_layer_args, lm_args):
-        """Capture CUDA graphs after warmup.
-
-        Args:
-            X_warmup: X tensor from warmup step (to seed input buffers)
-            layer_forward_fn: function(X, li, **kwargs) -> X_next
-            lm_forward_fn: function(X, **kwargs) -> logits
-            all_layer_args: dict[li] -> kwargs for layer_forward_fn
-            lm_args: kwargs for lm_forward_fn
-        """
-        print("  Capturing CUDA graphs for decode...", flush=True)
-
-        for li in range(self.n_layers):
-            gpu = li % self.num_gpus
-            dev = self.devices[gpu]
-            torch.cuda.set_device(gpu)
-
-            # Seed input buffer with warmup X
-            if li == 0:
-                self.x_in[li].copy_(X_warmup.to(dev))
-            else:
-                self.x_in[li].copy_(self.x_out[li - 1].to(dev))
-
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph):
-                X_next = layer_forward_fn(self.x_in[li], li, **all_layer_args[li])
-                self.x_out[li].copy_(X_next)
-
-            self.graphs[li] = graph
-            if (li + 1) % 10 == 0:
-                print(f"    Captured {li+1}/{self.n_layers} layer graphs", flush=True)
-
-        # Capture hc_head + norm + lm_head on cuda:0
-        torch.cuda.set_device(0)
-        if self.n_layers > 0:
-            self.x_cuda0_buf.copy_(self.x_out[self.n_layers - 1].to('cuda:0'))
-
-        self.lm_graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(self.lm_graph):
-            logits = lm_forward_fn(self.x_cuda0_buf, **lm_args)
-            self.logits_buf.copy_(logits)
-
-        self.captured = True
-        print(f"  Captured {len(self.graphs)} layer graphs + lm_head graph", flush=True)
-
-    def replay(self, token_id_gpu, position_gpu):
-        """Replay captured graphs for one decode step.
-
-        Args:
-            token_id_gpu: (1,) long tensor on cuda:0 — next token ID
-            position_gpu: (1,) long tensor on cuda:0 — current position
-
-        Returns:
-            logits: (1, vocab_size) bfloat16 tensor
-        """
-        assert self.captured, "Must call capture() before replay()"
-
-        # TODO: Copy token_id/position to the static input buffers that the graph uses.
-        # This requires the graph to reference those buffers.
-
-        # Replay layer graphs
-        for li in range(self.n_layers):
-            gpu = li % self.num_gpus
-            torch.cuda.set_device(gpu)
-
-            # Copy input from previous layer's output
-            if li > 0:
-                prev_gpu = (li - 1) % self.num_gpus
-                if prev_gpu != gpu:
-                    self.x_in[li].copy_(self.x_out[li - 1].to(self.devices[gpu]))
-                else:
-                    self.x_in[li].copy_(self.x_out[li - 1])
-
-            self.graphs[li].replay()
-
-        # Transfer final X to cuda:0
-        if self.n_layers > 0:
-            self.x_cuda0_buf.copy_(self.x_out[self.n_layers - 1].to('cuda:0'))
-
-        # Replay lm_head graph
-        self.lm_graph.replay()
-
-        return self.logits_buf

dsv4/kernels/cuda/blackwell_swizzle.cu

@@ -1,116 +0,0 @@
-/**
- * Blackwell 32_4_4 scale swizzle kernel.
- *
- * Rearranges FP8 scale factors from row-major layout to Blackwell tensor-core
- * compatible layout. This is the GPU equivalent of the Python:
- *   blocks = x.view(R, 128, C, 4).permute(0, 2, 1, 3)
- *   out = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten()
- *
- * The kernel writes to a pre-allocated output buffer — no per-step allocations.
- * CUDA-graph-capturable: no host-device syncs, no dynamic shapes.
- */
-
-#include <cuda_runtime.h>
-#include <c10/cuda/CUDAStream.h>
-#include <cstdint>
-#include <torch/extension.h>  // For pybind11 bindings
-
-// Blackwell 32_4_4 swizzle: each thread handles one output element
-// Input: (rows, cols) float8_e4m3fn — rows is multiple of 128, cols is multiple of 4
-// Output: (rows, cols) float8_e4m3fn — swizzled layout
-//
-// The swizzle reorders so that:
-//   For each group of 128 rows × 4 cols (a "block"):
-//     - The 128 rows are divided into 32 "sub-rows" of 4 rows each
-//     - The 4 cols are kept as-is
-//     - The output order is: [sub-row 0 col 0..3, sub-row 1 col 0..3, ..., sub-row 31 col 0..3]
-//     - Within each sub-row, the 4 rows × 4 cols = 16 elements are laid out as 32×16
-
-__global__ void blackwell_swizzle_32_4_4_kernel(
-    const uint8_t* __restrict__ input,   // (rows, cols) in FP8
-    uint8_t* __restrict__ output,         // (rows, cols) swizzled FP8
-    const int32_t rows,
-    const int32_t cols       // must be multiple of 4
-) {
-    const int32_t R = rows / 128;  // number of 128-row blocks
-    const int32_t C = cols / 4;    // number of 4-col groups
-    
-    // Total output elements
-    const int32_t total = rows * cols;
-    
-    // Each thread handles one output element
-    const int32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
-    if (tid >= total) return;
-    
-    // Output flat index → (block_r, col_group, sub_row, col_4, row_in_sub)
-    // Output layout: flatten of (R, C, 32, 4, 4, 4) → but simplified:
-    // The output is organized as:
-    //   For each (R, C) block: 32 sub-rows × 16 elements = 512 elements per block
-    //   Total per block: 128 * 4 = 512 elements
-    
-    // Decompose tid into block coordinates
-    const int32_t elements_per_block = 128 * 4;  // 512
-    const int32_t block_idx = tid / elements_per_block;
-    const int32_t within_block = tid % elements_per_block;
-    
-    const int32_t r = block_idx / C;      // row block index
-    const int32_t c = block_idx % C;      // col group index
-    
-    // Within-block layout: (32 sub-rows) × (4 col_within_group) × (4 row_within_subrow)
-    // But actually the swizzle is: reshape(32, 4, 4, 4) → transpose(1,2) → flatten
-    // Which gives: for each (sub_row, col_4, row_in_sub):
-    //   output[sub_row * 16 + col_4 * 4 + row_in_sub] = input[sub_row * 4 + row_in_sub][col_4 * 4 + c_offset]
-    
-    // Within block: 512 elements in swizzled order
-    // The Python swizzle does:
-    //   blocks[128 rows, 4 cols] → view(32, 4, 4, 4) → permute → (32, 4, 4, 4)
-    //   → reshape(-1, 32, 16) → flatten
-    // The output index maps to:
-    //   sub_row = within_block / 16
-    //   within_sub = within_block % 16  → (col_4, row_in_sub) = (within_sub / 4, within_sub % 4)
-    
-    const int32_t sub_row = within_block / 16;
-    const int32_t within_sub = within_block % 16;
-    const int32_t col_4 = within_sub / 4;
-    const int32_t row_in_sub = within_sub % 4;
-    
-    // Map back to input coordinates
-    const int32_t input_row = r * 128 + sub_row * 4 + row_in_sub;
-    const int32_t input_col = c * 4 + col_4;
-    
-    // Read input, write to output
-    output[tid] = input[input_row * cols + input_col];
-}
-
-extern "C" {
-
-void launch_blackwell_swizzle(
-    const uint8_t* input,
-    uint8_t* output,
-    int32_t rows,
-    int32_t cols,
-    cudaStream_t stream
-) {
-    const int32_t total = rows * cols;
-    const int32_t block_size = 256;
-    const int32_t grid_size = (total + block_size - 1) / block_size;
-    
-    blackwell_swizzle_32_4_4_kernel<<<grid_size, block_size, 0, stream>>>(
-        input, output, rows, cols
-    );
-}
-
-} // extern "C"
-
-// Pybind11 bindings for torch.utils.cpp_extension.load
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("blackwell_swizzle_32_4_4", [](at::Tensor input, at::Tensor output, int32_t rows, int32_t cols) {
-        auto stream = c10::cuda::getCurrentCUDAStream();
-        blackwell_swizzle_32_4_4_kernel<<<
-            (rows * cols + 255) / 256, 256, 0, stream>>>(
-            input.data_ptr<uint8_t>(),
-            output.data_ptr<uint8_t>(),
-            rows, cols
-        );
-    }, "Blackwell 32_4_4 scale swizzle");
-}

dsv4/kernels/cuda/compressor_reduce.cu

@@ -124,14 +124,13 @@ __global__ void csa_compress_reduce_kernel(
 
             float g = gate_proj[token_idx * kv_dim + gate_offset + c];
             float kv_val = kv_proj[token_idx * kv_dim + kv_offset + c];
-            // Position bias: added to gate logits (softmax Z + B) only.
-            // The paper defines compression as softmax(Z + B) then weighted sum of C.
-            // The bias must NOT be added to kv_val — that poisons compressed content.
+            // Position bias added ONLY to gate (softmax logit), NOT to KV content.
+            // Paper eq. 11-12: compressed = softmax(Z + B) * C — bias B is on the
+            // compression weights/logits, not on the KV content C.
             if (position_bias != nullptr) {
                 int pos_bias_row = (block_i > 0 && t < m) ? t : (block_i > 0 ? (t - m) : t);
                 if (pos_bias_row >= 0 && pos_bias_row < m) {
-                    float pb = position_bias[pos_bias_row * kv_dim + gate_offset + c];
-                    g += pb;
+                    g += position_bias[pos_bias_row * kv_dim + gate_offset + c];
                 }
             }
             float e = expf(g - local_max[ci]);
@@ -191,12 +190,11 @@ __global__ void hca_compress_reduce_kernel(
             if (token_idx >= T) break;
             float g = gate_proj[token_idx * hd + c];
             float kv_val = kv_proj[token_idx * hd + c];
-            // Position bias: added to gate logits (softmax Z + B) only.
-            // The paper defines compression as softmax(Z + B) then weighted sum of C.
-            // The bias must NOT be added to kv_val — that poisons compressed content.
+            // Position bias added ONLY to gate (softmax logit), NOT to KV content.
+            // Paper eq. 9-10: compressed = softmax(Z + B) * C — bias B is on the
+            // compression weights/logits, not on the KV content C.
             if (position_bias != nullptr && t < m) {
-                float pb = position_bias[t * hd + c];
-                g += pb;
+                g += position_bias[t * hd + c];
             }
             float e = expf(g - local_max);
             local_denom += e;

dsv4/kernels/gemm/fused_swiglu.py

@@ -2196,12 +2196,11 @@ class FusedSwiGLUScaledGroupedGemmKernel:
                             neg_acc = acc_vec * cutlass.Float32(-1.0)
                             exp_neg = cute.exp(neg_acc)
                             sigmoid = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + exp_neg)
-                            silu_result = acc_vec * sigmoid
-                            # Paper §4.2.3: gate component capped at swiglu_limit
-                            # CuTe DSL clamp: min(x, limit) = cute.where(x > limit, limit, x)
+                            # Paper §4.2.3: clamp raw gate BEFORE SiLU, not after
                             if cutlass.const_expr(self.swiglu_limit > 0.0):
                                 limit = cutlass.Float32(self.swiglu_limit)
-                                silu_result = cute.where(silu_result > limit, limit, silu_result)
+                                acc_vec = cute.where(acc_vec > limit, limit, acc_vec)
+                            silu_result = acc_vec * sigmoid
                             silu_result = silu_result.to(self.c_dtype)
                             silu_gate_buf.store(silu_result)
                             # Keep acc_vec in BF16 (same type as the up branch)

dsv4/kernels/gemm/grouped.py

@@ -2374,15 +2374,8 @@ def compute_scale_shape(
             return (padded_N, total_cols)
 
 
-def to_blocked(scale_2d: torch.Tensor, out_buf: torch.Tensor = None) -> torch.Tensor:
-    """Pad and apply the Blackwell 32_4_4 scale swizzle to one raw scale tensor.
-    
-    During CUDA graph capture, uses a custom CUDA kernel because Python
-    view operations (reshape, transpose, permute) are not graph-capturable.
-    The out_buf must be provided during graph capture (pre-allocated output).
-    
-    During eager mode, uses the faster Python view path.
-    """
+def to_blocked(scale_2d: torch.Tensor) -> torch.Tensor:
+    """Pad and apply the Blackwell 32_4_4 scale swizzle to one raw scale tensor."""
     if scale_2d.dim() != 2:
         raise ValueError(f"Expected 2D scale tensor, got {scale_2d.dim()}D.")
     rows, cols = scale_2d.shape
@@ -2401,19 +2394,6 @@ def to_blocked(scale_2d: torch.Tensor, out_buf: torch.Tensor = None) -> torch.Te
         )
         padded[:rows, :cols] = scale_2d
 
-    # Use CUDA kernel during graph capture — Python view ops are not capturable
-    if torch.cuda.is_current_stream_capturing():
-        from dsv4.kernels.cuda.loader import get_cuda_module
-        mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
-        if out_buf is None:
-            out_buf = torch.empty_like(padded)
-        mod.blackwell_swizzle_32_4_4(
-            padded.view(torch.uint8), out_buf.view(torch.uint8),
-            padded_rows, padded_cols
-        )
-        return out_buf.view(torch.float8_e4m3fn).flatten()
-
-    # Eager path: Python view operations (fast, no kernel launch overhead)
     blocks = padded.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3)
     rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
     return rearranged.flatten()

dsv4/kernels/router/dense_router_decode.py

@@ -27,16 +27,10 @@ def dense_router_dispatch(
 ):
     """Dispatch the dense router (BF16 cuBLAS fallback).
 
-    BF16 GEMM via torch.matmul (cuBLAS, SM100 tensor cores),
+    BF16 GEMM via torch.nn.functional.linear (cuBLAS, SM100 tensor cores),
     then fused activation + top-k via the CUDA kernel.
-
-    CUDA-graph-compatible: no .T, no .float() on inputs during capture.
-    The GEMM runs in BF16 (Blackwell tensor cores handle BF16 natively).
-    Only the output logits are cast to FP32 for sqrt(softplus) stability.
     """
-    # BF16 GEMM: x @ W — no transpose needed, no FP32 conversion
-    logits_bf16 = torch.matmul(hidden_states, W_gate)  # [N, H] @ [H, E] = [N, E]
-    logits = logits_bf16.float()  # BF16 → FP32 for sqrt(softplus) numerical stability
+    logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float())
     from dsv4.kernels.router._activation_topk import run_fused_activation_topk
     run_fused_activation_topk(
         logits, e_bias, routed_scaling_factor, top_k,
@@ -103,8 +97,7 @@ def dense_router_dispatch_nvfp4_fused(
     # Decode the gate_weight from NVFP4 to BF16 for cuBLAS
     from dsv4.ops.quantize import dequantize_nvfp4
     gate_bf16 = dequantize_nvfp4(gate_weight, gate_weight_scale, gate_ws2)
-    logits = torch.nn.functional.linear(hidden_states, gate_bf16.T)
-    logits = logits.float()  # BF16 → FP32 for numerical stability in sqrt(softplus)
+    logits = torch.nn.functional.linear(hidden_states.float(), gate_bf16.T.float())
 
     run_fused_activation_topk(
         logits, e_bias, routed_scaling_factor, top_k,

dsv4/layers/grouped_linear.py

@@ -212,31 +212,6 @@ class Nvfp4GroupedLinear:
 
         self._gsa_buf = torch.zeros(self.n_local_groups, dtype=torch.float32, device=self.device)
         self._expert_offsets_buf = torch.zeros(self.n_local_groups, dtype=torch.int32, device=self.device)
-        # Pre-computed range [1, 2, 3, ..., n_groups] for expert offsets
-        # Avoids torch.arange() per call (allocation) and Python loop (CPU→GPU sync)
-        self._expert_offsets_range_buf = torch.arange(
-            1, self.n_local_groups + 1, dtype=torch.int32, device=self.device
-        )
-        self._group_offset_buf = torch.zeros(self.n_local_groups, dtype=torch.int32, device=self.device)
-        # Pre-allocate output buffer for graph capture
-        self._output_buf = torch.zeros(
-            self.max_num_tokens, self.n_local_groups, self.o_lora_rank,
-            dtype=torch.bfloat16, device=self.device
-        )
-        # Pre-allocate FLAT output buffer for grouped GEMM (graph capture)
-        # The GEMM produces (tokens_sum, n_dim) where n_dim = o_lora_rank
-        # tokens_sum = n_groups * padded_rows_per_group (max = n_groups * max_num_tokens)
-        self._output_buf_padded = torch.zeros(
-            self.max_num_tokens * self.n_local_groups, self.o_lora_rank,
-            dtype=torch.bfloat16, device=self.device
-        )
-        # Pre-allocate scale_a swizzle buffer for graph capture
-        K_sf = cutedsl_ceil_div(self.group_in_features, 16)
-        max_padded_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
-        max_padded_cols = cutedsl_ceil_div(K_sf, 4) * 4
-        self._scale_a_buf = torch.zeros(
-            max_padded_rows, max_padded_cols, dtype=torch.float16, device=self.device
-        ).to(torch.float8_e4m3fn)
         self._buffers_allocated = True
 
     def _ensure_initialized(self):
@@ -246,22 +221,14 @@ class Nvfp4GroupedLinear:
             self._allocate_buffers()
 
     def _assemble_scales_single_group(self, x_sf):
-        """Assemble 2D-side activation scales for num_groups=1.
-        
-        CUDA-graph-safe: uses pre-allocated _scale_a_buf.
-        """
+        """Assemble 2D-side activation scales for num_groups=1."""
         num_rows, num_cols = x_sf.shape
         padded_rows = cutedsl_ceil_div(num_rows, 128) * 128
         padded_cols = cutedsl_ceil_div(num_cols, 4) * 4
 
-        # Use pre-allocated buffer — zero + scatter pattern (no new allocation)
-        buf = self._scale_a_buf
-        assert buf.shape[0] >= padded_rows and buf.shape[1] >= padded_cols, \
-            f"scale_a_buf too small: {buf.shape} < ({padded_rows}, {padded_cols})"
-        buf.view(torch.uint8).zero_()
+        buf = torch.zeros(padded_rows, padded_cols, dtype=torch.float16, device=x_sf.device).to(torch.float8_e4m3fn)
         buf[:num_rows, :num_cols] = x_sf
-        view = buf[:padded_rows, :padded_cols]
-        swizzled_flat = pad_and_swizzle_single(view)
+        swizzled_flat = pad_and_swizzle_single(buf)
         return swizzled_flat.reshape(padded_rows, padded_cols)
 
     def compute_activation_global_scale(self, o_sample: torch.Tensor):
@@ -338,12 +305,10 @@ class Nvfp4GroupedLinear:
             # gsa_gpu is (G*T,) — all rows share same amax (from max over full tensor)
             # For the GEMM's global_scale_a, fill all group slots with the same gsa value
             # Use GPU-only copy: no .item(), no CPU sync
-            self._gsa_buf[0] = gsa_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._gsa_buf[:1].copy_(gsa_gpu[:1])  # GPU→GPU scalar copy, no sync
             # Broadcast to all groups (all get same gsa)
-            # Use scalar broadcast assignment instead of copy_ from expanded view
-            # (expanded views can cause cudaErrorInvalidValue in copy_)
             if self.n_local_groups > 1:
-                self._gsa_buf[1:] = self._gsa_buf[0]  # scalar broadcast, graph-capturable
+                self._gsa_buf[1:].copy_(self._gsa_buf[:1].expand(self.n_local_groups - 1))
         else:
             self._gsa_buf.fill_(self._activation_global_scale)
             x_fp4_flat, x_sf_flat = quantize_activation_nvfp4(
@@ -356,13 +321,6 @@ class Nvfp4GroupedLinear:
 
         x_fp4_grouped = x_fp4_flat.reshape(self.n_local_groups, num_tokens, self.group_in_features // 2)
 
-        # Vectorized scatter — no Python loop, no CPU→GPU sync
-        # Unconditionally update group offsets — GPU-only, no conditional host read.
-        # padded_rows_per_group is a Python int multiplied with a GPU tensor = GPU op.
-        group_offsets = self._group_offset_buf[:self.n_local_groups]
-        expert_offsets = self._expert_offsets_buf
-        expert_offsets[:self.n_local_groups] = self._expert_offsets_range_buf * padded_rows_per_group
-        # Scatter each group's x_fp4 into padded buffer
         for g in range(self.n_local_groups):
             offset = g * padded_rows_per_group
             padded_x_fp4.view(torch.uint8)[offset:offset + num_tokens] = x_fp4_grouped[g].view(torch.uint8)
@@ -378,16 +336,15 @@ class Nvfp4GroupedLinear:
         scale_a = assemble_scales_2d_side(all_x_sf)
 
         # Expert offsets: cumulative [padded_T, 2*padded_T, ..., n_groups*padded_T]
-        # GPU-only computation — no Python loop, no CPU→GPU sync
         expert_offsets = self._expert_offsets_buf
-        # element-wise multiply: range * padded_rows → GPU tensor (no host sync)
-        expert_offsets[:self.n_local_groups] = self._expert_offsets_range_buf * padded_rows_per_group
+        for g in range(self.n_local_groups):
+            expert_offsets[g] = (g + 1) * padded_rows_per_group
 
         # Global scales — GPU-computed gsa already in _gsa_buf (no CPU sync)
         gsa = self._gsa_buf
 
-        # Run grouped GEMM — pass pre-allocated output buffer for CUDA graph capture
-        z_gem = run_nvfp4_grouped_gemm(
+        # Run grouped GEMM
+        out = run_nvfp4_grouped_gemm(
             mat_a=padded_x_fp4,
             mat_b=self._mat_b,
             scale_a=scale_a,
@@ -395,23 +352,15 @@ class Nvfp4GroupedLinear:
             expert_offsets=expert_offsets,
             global_scale_a=gsa,
             global_scale_b=self._gsb,
-            out=self._output_buf_padded if hasattr(self, '_output_buf_padded') else None,
         )
 
         # Extract real outputs and reshape
-        # GEMM output layout: (tokens_sum, o_lora_rank) where tokens_sum = n_groups * padded_rows
-        # Groups are stacked vertically: group 0 at rows [0, padded_rows), group 1 at [padded_rows, 2*padded_rows), etc.
-        z_gem = z_gem if z_gem is not None else self._output_buf_padded
-        z = self._output_buf[:num_tokens]
-        if num_tokens == 1:
-            # Vectorized: gather_indices = [0, padded_T, 2*padded_T, ...] — GPU-only
-            gather_indices = self._expert_offsets_range_buf[:self.n_local_groups] * padded_rows_per_group - padded_rows_per_group
-            z_flat = z_gem[gather_indices]  # (n_groups, o_lora_rank) — GPU gather
-            z[:, :, :] = z_flat.unsqueeze(0)  # (1, n_groups, o_lora_rank)
-        else:
-            for g in range(self.n_local_groups):
-                offset = g * padded_rows_per_group
-                z[:, g, :] = z_gem[offset:offset + num_tokens, :]
+        # GEMM output has the same layout as mat_a: groups-first with padding
+        z = torch.empty(num_tokens, self.n_local_groups, self.o_lora_rank,
+                        dtype=torch.bfloat16, device=o.device)
+        for g in range(self.n_local_groups):
+            offset = g * padded_rows_per_group
+            z[:, g, :] = out[offset:offset + num_tokens, :]
 
         return z
 

dsv4/layers/linear.py

@@ -65,7 +65,6 @@ class Nvfp4Linear:
         self._padded_x_fp4_buf = None
         self._expert_offsets_buf = None
         self._gsa_buf = None
-        self._gemm_out_buf = None  # pre-allocated GEMM output for graph capture
         self._buffers_allocated = False
 
     def finalize_weights(self):
@@ -104,16 +103,7 @@ class Nvfp4Linear:
         # warmup_compilation(1, K_packed, N_packed, self.device)  # Lazy compile on first real forward
 
     def _ensure_buffer_size(self, num_tokens: int):
-        """Ensure the padded buffer is large enough for num_tokens.
-        
-        Pre-allocates ALL buffers needed for CUDA graph capture:
-        - padded x_fp4 buffer (max_num_tokens aligned to 128 rows)
-        - expert_offsets (1 element for single group)
-        - gsa buffer (1 element, GPU-only)
-        - scale_a swizzle buffer (pre-allocated at max size)
-        
-        No per-call allocations — zero CPU-GPU syncs on the hot path.
-        """
+        """Ensure the padded buffer is large enough for num_tokens."""
         needed_rows = cutedsl_ceil_div(num_tokens, 128) * 128
         if self._padded_x_fp4_buf is not None and self._padded_x_fp4_buf.shape[0] >= needed_rows:
             return  # Already big enough
@@ -124,63 +114,20 @@ class Nvfp4Linear:
 
         self._expert_offsets_buf = torch.zeros(1, dtype=torch.int32, device=self.device)
         self._gsa_buf = torch.full((1,), self._activation_global_scale, dtype=torch.float32, device=self.device)
-        
-        # Pre-allocate scale_a swizzle buffer for _assemble_scales_single_group.
-        # Max size: (max_num_tokens aligned to 128) × (K_sf aligned to 4).
-        # This eliminates the per-call torch.zeros() allocation that breaks
-        # CUDA graph capture.
-        K_sf = cutedsl_ceil_div(self.in_features, 16)
-        max_padded_rows = cutedsl_ceil_div(self.max_num_tokens, 128) * 128
-        max_padded_cols = cutedsl_ceil_div(K_sf, 4) * 4
-        self._scale_a_buf = torch.zeros(
-            max_padded_rows, max_padded_cols, dtype=torch.float16, device=self.device
-        ).to(torch.float8_e4m3fn)
-        
-        # Pre-allocated GEMM output buffer for graph capture
-        self._gemm_out_buf = torch.zeros(
-            max_padded_rows, self.out_features, dtype=torch.bfloat16, device=self.device
-        )
-        
-        # Pre-allocated swizzled scale output buffer (for CUDA graph capture)
-        self._padded_x_sf_swizzled_buf = torch.zeros_like(self._scale_a_buf)
 
     def _ensure_initialized(self):
         if self._mat_b is None:
             self.finalize_weights()
 
     def _assemble_scales_single_group(self, x_sf):
-        """Assemble 2D-side activation scales for num_groups=1.
-        
-        CUDA-graph-safe: uses pre-allocated _scale_a_buf instead of
-        per-call torch.zeros(). The buffer is zeroed + scattered + swizzled
-        each call — zero new allocations on the hot path.
-        """
+        """Assemble 2D-side activation scales for num_groups=1."""
         num_rows, num_cols = x_sf.shape
         padded_rows = cutedsl_ceil_div(num_rows, 128) * 128
         padded_cols = cutedsl_ceil_div(num_cols, 4) * 4
 
-        # Use pre-allocated buffer — zero + scatter pattern (no new allocation)
-        buf = self._scale_a_buf
-        assert buf.shape[0] >= padded_rows and buf.shape[1] >= padded_cols, \
-            f"scale_a_buf too small: {buf.shape} < ({padded_rows}, {padded_cols})"
-        buf.view(torch.uint8).zero_()
+        buf = torch.zeros(padded_rows, padded_cols, dtype=torch.float16, device=x_sf.device).to(torch.float8_e4m3fn)
         buf[:num_rows, :num_cols] = x_sf
-        # Pass correctly-sized VIEW to swizzle — the swizzle operates on
-        # (padded_rows, padded_cols) not the full max-size buffer.
-        view = buf[:padded_rows, :padded_cols]
-        
-        # During graph capture, use CUDA swizzle kernel (Python view ops not capturable)
-        if torch.cuda.is_current_stream_capturing() and self._padded_x_sf_swizzled_buf is not None:
-            from dsv4.kernels.cuda.loader import get_cuda_module
-            mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
-            swizzled_buf = self._padded_x_sf_swizzled_buf
-            mod.blackwell_swizzle_32_4_4(
-                view.view(torch.uint8), swizzled_buf[:padded_rows, :padded_cols].view(torch.uint8),
-                padded_rows, padded_cols
-            )
-            return swizzled_buf[:padded_rows, :padded_cols].reshape(padded_rows, padded_cols)
-        
-        swizzled_flat = pad_and_swizzle_single(view)
+        swizzled_flat = pad_and_swizzle_single(buf)
         return swizzled_flat.reshape(padded_rows, padded_cols)
 
     def compute_activation_global_scale(self, hidden_states_sample):
@@ -227,7 +174,7 @@ class Nvfp4Linear:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             x_fp4, x_sf, gsa_gpu = quantize_nvfp4_gpu_fused(hidden_states)
-            self._gsa_buf[0] = gsa_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._gsa_buf.copy_(gsa_gpu[:1].reshape(1))  # GPU → GPU, no sync
         else:
             # P2 FIX: No per-call fill_(). The _gsa_buf already has the correct
             # value — set either during initialization (via _ensure_buffer_size)
@@ -262,7 +209,6 @@ class Nvfp4Linear:
             expert_offsets=expert_offsets,
             global_scale_a=gsa,
             global_scale_b=self._gsb,
-            out=self._gemm_out_buf,
         )
 
         return out[:num_tokens]
@@ -306,10 +252,13 @@ class Nvfp4Linear:
         # For M=1 decode: per-row gsa is already scalar, no reduction needed.
         # For M>1 prefill: reduce per-row gsa to a single scalar (max).
         if quant.gsa.shape[0] == 1:
-            self._gsa_buf[0] = quant.gsa[0]  # scalar GPU→GPU, graph-capturable
+            gsa = quant.gsa[:1].reshape(1)  # Already scalar
         else:
             # Reduce per-row gsa to scalar (max) for GEMM compatibility.
-            self._gsa_buf[0] = quant.gsa.max()  # GPU max, scalar assign, graph-capturable
+            # Per-row gsa is mathematically more precise, but the GEMM only
+            # supports a single global scale per expert.
+            gsa = quant.gsa.max().reshape(1)
+        self._gsa_buf.copy_(gsa)
 
         # Run GEMM
         out = run_nvfp4_grouped_gemm(
@@ -320,7 +269,6 @@ class Nvfp4Linear:
             expert_offsets=expert_offsets,
             global_scale_a=self._gsa_buf,
             global_scale_b=self._gsb,
-            out=self._gemm_out_buf,
         )
 
         return out[:num_tokens]

dsv4/layers/mhc.py

@@ -418,9 +418,12 @@ class mHCLayer:
         CF = ctx.C_l.unsqueeze(-1) * F_out.unsqueeze(1)   # (T, n_hc, d)
         X_next = (CF.float() + BX).to(self.dtype)   # (T, n_hc, d)
         
-        # Note: residual magnitude monitoring is done OUTSIDE the graph-captured region
-        # (via the caller in single_shot_inference.py diagnostics). No .item() here —
-        # CUDA graph capture requires zero device→host syncs on the hot path.
+        # Diagnostic: warn on residual blowup
+        x_max = X_next.abs().max().item()
+        if x_max > 500:
+            # Don't clip in production, just warn
+            pass
+        
         return X_next
 
     # ----------------------------------------------------------------
@@ -431,23 +434,12 @@ class mHCLayer:
     def init_state(
         embeddings: torch.Tensor,   # (T, d) BF16 — token embeddings
         n_hc: int = 4,
-        out_buf: torch.Tensor = None,  # (T, n_hc, d) BF16 — pre-allocated output buffer
     ) -> torch.Tensor:
         """
         Initialise X_0 for the first layer.
 
         Returns: (T, n_hc, d) BF16
-        
-        When out_buf is provided, writes to it in-place (no allocation).
-        This is required for CUDA graph capture where per-step
-        allocations are forbidden.
         """
-        if out_buf is not None:
-            # In-place: copy embeddings to all n_hc streams
-            out_buf[:, 0, :].copy_(embeddings)  # Stream 0 gets the embedding
-            for h in range(1, n_hc):
-                out_buf[:, h, :].copy_(embeddings)  # All other streams too
-            return out_buf
         return embeddings.unsqueeze(1).expand(-1, n_hc, -1).clone()
 
     @staticmethod

dsv4/layers/moe.py

@@ -90,7 +90,6 @@ class Nvfp4MoE:
         self._padded_x_sf_buf_l2 = None
         self._l1_gsa_buf = None
         self._l2_gsa_buf = None
-        self._l1_out_buf = None  # pre-allocated L1 GEMM output for graph capture
         self._output_buf = None
         self._row_indices_buf = None
         self._padded_hidden_buf = None
@@ -161,37 +160,10 @@ class Nvfp4MoE:
         self._padded_x_sf_buf_l2 = Nvfp4MoE._shared_padded_bufs[device_key]['xsf_l2']
         self._output_buf = Nvfp4MoE._shared_padded_bufs[device_key]['output']
         
-        # Pre-allocated swizzled scale output buffers (same size as padded_x_sf)
-        # Required for CUDA graph capture — Python view ops (reshape, transpose) not capturable
-        if 'xsf_swizzled_l1' not in Nvfp4MoE._shared_padded_bufs[device_key]:
-            Nvfp4MoE._shared_padded_bufs[device_key].update({
-                'xsf_swizzled_l1': torch.zeros_like(Nvfp4MoE._shared_padded_bufs[device_key]['xsf_l1']),
-                'xsf_swizzled_l2': torch.zeros_like(Nvfp4MoE._shared_padded_bufs[device_key]['xsf_l2']),
-            })
-        self._padded_x_sf_swizzled_buf_l1 = Nvfp4MoE._shared_padded_bufs[device_key]['xsf_swizzled_l1']
-        self._padded_x_sf_swizzled_buf_l2 = Nvfp4MoE._shared_padded_bufs[device_key]['xsf_swizzled_l2']
-        
         # Pre-allocated global_scale_a buffers (filled via .fill_(), no torch.full during capture)
         self._l1_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device)
         self._l2_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device)
         
-        # Pre-allocated L1 GEMM output — avoids torch.zeros() in run_fused_swiglu_grouped_gemm
-        # Shape: (max_tokens * top_k, 2*intermediate_size) — gate+up combined
-        self._l1_out_buf = torch.zeros(
-            self.max_num_tokens * self.top_k, 2 * self.intermediate_size,
-            dtype=torch.bfloat16, device=self.device
-        )
-        # Pre-allocated L2 GEMM output — avoids torch.zeros() in run_nvfp4_grouped_gemm
-        # Shape: (max_tokens * top_k, hidden_size) — down projection
-        self._l2_out_buf = torch.zeros(
-            self.max_num_tokens * self.top_k, self.hidden_size,
-            dtype=torch.bfloat16, device=self.device
-        )
-        
-        # Pre-allocated tokens-per-expert buffer — replaces torch.bincount
-        # (bincount produces data-dependent shapes, breaks CUDA graph capture)
-        self._tokens_per_expert_buf = torch.zeros(self.num_experts, dtype=torch.int32, device=self.device)
-        
         # Row indices for scale assembly (max_num_tokens * top_k slots)
         self._row_indices_buf = torch.arange(
             self.max_num_tokens * self.top_k, device=self.device
@@ -454,20 +426,11 @@ class Nvfp4MoE:
         padded_x_sf[dst_rows, :K_sf] = x_sf
         
         # Phase 2: Full-buffer swizzle (no CPU sync, no Python loops)
-        # During graph capture, Python view ops (reshape, transpose) are not allowed.
-        # Use CUDA swizzle kernel instead.
+        # padded_x_sf is 128-row aligned per expert and 4-col aligned.
+        # to_blocked: (rows, cols) → view(R, 128, C, 4) → permute(0,2,1,3)
+        #   → reshape(-1, 4, 32, 4) → transpose(1,2) → reshape(-1, 32, 16) → flatten
         rows = padded_x_sf.shape[0]
         cols = padded_x_sf.shape[1]
-        if torch.cuda.is_current_stream_capturing():
-            from dsv4.kernels.cuda.loader import get_cuda_module
-            mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
-            out_buf = self._padded_x_sf_swizzled_buf_l1 if padded_x_sf is self._padded_x_sf_buf_l1 else self._padded_x_sf_swizzled_buf_l2
-            mod.blackwell_swizzle_32_4_4(
-                padded_x_sf.view(torch.uint8), out_buf.view(torch.uint8),
-                rows, cols
-            )
-            return out_buf.view(torch.float8_e4m3fn).reshape(rows, cols)
-        # Eager path: Python view operations
         R = rows // 128
         C = cols // 4
         blocks = padded_x_sf.view(R, 128, C, 4).permute(0, 2, 1, 3)
@@ -503,17 +466,7 @@ class Nvfp4MoE:
             # Quantize slot_hidden for GEMM
             slot_x_fp4, slot_x_sf = quantize_activation_nvfp4(slot_hidden, l1_gs)
             
-            # Compute tokens_per_expert — CUDA-graph-safe alternative to torch.bincount.
-            # torch.bincount produces data-dependent shapes (violates graph capture).
-            # Instead, use scatter_add_ into a pre-allocated buffer (fixed shape, GPU-only).
-            self._tokens_per_expert_buf.zero_()
-            # scatter_add_ requires int64 indices — ensure sorted_ids is int64
-            sorted_ids_i64 = sorted_ids.long()
-            n_slots = sorted_ids_i64.shape[0]
-            if not hasattr(self, '_ones_buf') or self._ones_buf.shape[0] < n_slots:
-                self._ones_buf = torch.ones(self.max_num_tokens * self.top_k, dtype=self._tokens_per_expert_buf.dtype, device=sorted_ids_i64.device)
-            self._tokens_per_expert_buf.scatter_add_(0, sorted_ids_i64, self._ones_buf[:n_slots])
-            tokens_per_expert = self._tokens_per_expert_buf[:self.num_experts]
+            tokens_per_expert = torch.bincount(sorted_ids, minlength=self.num_experts)[:self.num_experts].int()
             expert_offsets = self._expert_offsets_buf
             expert_offsets.zero_()
             expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
@@ -541,9 +494,7 @@ class Nvfp4MoE:
                 padded_expert_offsets,
                 self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1
             )
-            # l1_gsa: pre-allocated buffer, no per-call allocation
-            self._l1_gsa_buf.fill_(l1_gs)
-            l1_gsa = self._l1_gsa_buf
+            l1_gsa = torch.full((self.num_experts,), l1_gs, dtype=torch.float32, device=device)
             
             l1_out = run_nvfp4_grouped_gemm(
                 mat_a=padded_x_fp4, mat_b=self._l1_mat_b,
@@ -561,10 +512,11 @@ class Nvfp4MoE:
             l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
             gate = l1_deil[:, :self.intermediate_size]
             up = l1_deil[:, self.intermediate_size:]
-            gate_silu = torch.nn.functional.silu(gate)
+            # Paper §4.2.3: clamp raw gate BEFORE SiLU, not after
             if self._swiglu_limit is not None:
-                gate_silu = gate_silu.clamp(max=self._swiglu_limit)
+                gate = gate.clamp(max=self._swiglu_limit)
                 up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
+            gate_silu = torch.nn.functional.silu(gate)
             activated = gate_silu * up
             _, _, l2_gs = quantize_to_nvfp4(activated)
         
@@ -620,14 +572,7 @@ class Nvfp4MoE:
         sorted_token_ids = token_indices[sort_idx]
         
         # Expert offsets (real token counts)
-        # CUDA-graph-safe: scatter_add_ instead of bincount (fixed shape, GPU-only)
-        self._tokens_per_expert_buf.zero_()
-        sorted_ids_i64 = sorted_ids.long()
-        n_slots = sorted_ids_i64.shape[0]
-        if not hasattr(self, '_ones_buf') or self._ones_buf.shape[0] < n_slots:
-            self._ones_buf = torch.ones(self.max_num_tokens * self.top_k, dtype=self._tokens_per_expert_buf.dtype, device=sorted_ids_i64.device)
-        self._tokens_per_expert_buf.scatter_add_(0, sorted_ids_i64, self._ones_buf[:n_slots])
-        tokens_per_expert = self._tokens_per_expert_buf[:self.num_experts]
+        tokens_per_expert = torch.bincount(sorted_ids, minlength=self.num_experts)[:self.num_experts].int()
         expert_offsets = self._expert_offsets_buf
         expert_offsets.zero_()
         expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
@@ -655,7 +600,7 @@ class Nvfp4MoE:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             slot_x_fp4, slot_x_sf, gsa_l1_gpu = quantize_nvfp4_gpu_fused(slot_hidden)
-            self._l1_gsa_buf[0] = gsa_l1_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))  # GPU → GPU, no sync
         else:
             slot_x_fp4, slot_x_sf = quantize_nvfp4_gpu(
                 slot_hidden, self._l1_activation_global_scale
@@ -681,7 +626,6 @@ class Nvfp4MoE:
                 expert_offsets=padded_expert_offsets[1:self.num_experts + 1],
                 global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
                 swiglu_limit=self._swiglu_limit if self._swiglu_limit is not None else 0.0,
-                out=self._l1_out_buf,
             )
             l1_out_real = l1_out[padded_dst]
             # Fused deinterleave + amax + quantize: zero CPU syncs.
@@ -691,7 +635,7 @@ class Nvfp4MoE:
                 from dsv4.ops.quantize import deinterleave_amax_quantize_nvfp4_fused
                 slot_l2_x_fp4, slot_l2_x_sf, gsa_l2_gpu = deinterleave_amax_quantize_nvfp4_fused(
                     l1_out_real, self.intermediate_size)
-                self._l2_gsa_buf[0] = gsa_l2_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+                self._l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))  # GPU → GPU, no sync
             else:
                 slot_l2_x_fp4, slot_l2_x_sf = deinterleave_quantize_nvfp4_cuda(
                     l1_out_real, self.intermediate_size, self._l2_activation_global_scale
@@ -703,16 +647,16 @@ class Nvfp4MoE:
                 scale_a=l1_scale_a, scale_b=self._l1_scale_b,
                 expert_offsets=padded_expert_offsets[1:self.num_experts + 1],
                 global_scale_a=l1_gsa, global_scale_b=self._l1_gsb,
-                out=self._l1_out_buf,
             )
             l1_out_real = l1_out[padded_dst]
             l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0]
             gate = l1_deil[:, :self.intermediate_size]
             up = l1_deil[:, self.intermediate_size:]
-            gate_silu = torch.nn.functional.silu(gate)
+            # Paper §4.2.3: clamp raw gate BEFORE SiLU, not after
             if self._swiglu_limit is not None:
-                gate_silu = gate_silu.clamp(max=self._swiglu_limit)
+                gate = gate.clamp(max=self._swiglu_limit)
                 up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
+            gate_silu = torch.nn.functional.silu(gate)
             activated = gate_silu * up
 
         # Compute runtime gsa for L2 from activated output (non-fused path)
@@ -720,7 +664,7 @@ class Nvfp4MoE:
         if not self._fused_swiglu and getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             slot_l2_x_fp4, slot_l2_x_sf, gsa_l2_gpu = quantize_nvfp4_gpu_fused(activated)
-            self._l2_gsa_buf[0] = gsa_l2_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))  # GPU → GPU, no sync
         elif not self._fused_swiglu:
             slot_l2_x_fp4, slot_l2_x_sf = quantize_nvfp4_gpu(
                 activated, self._l2_activation_global_scale
@@ -741,7 +685,6 @@ class Nvfp4MoE:
             scale_a=l2_scale_a, scale_b=self._l2_scale_b,
             expert_offsets=padded_expert_offsets[1:self.num_experts + 1],
             global_scale_a=l2_gsa, global_scale_b=self._l2_gsb,
-            out=self._l2_out_buf,
         )
         
         l2_out_real = l2_out[padded_dst]

dsv4/layers/shared_expert.py

@@ -91,9 +91,6 @@ class Nvfp4SharedExpert:
         self._l1_activation_global_scale = 1.0 / (6.0 * 448.0)
         self._l2_activation_global_scale = 1.0 / (6.0 * 448.0)
 
-        # Pre-allocated L1 GEMM output for graph capture
-        self._l1_out_buf = None
-
         # Pre-allocated cudagraph buffers (set in _allocate_buffers)
         self._padded_x_fp4_buf_l1 = None
         self._padded_x_sf_buf_l1 = None
@@ -178,31 +175,10 @@ class Nvfp4SharedExpert:
         self._padded_x_sf_buf_l2 = torch.zeros(
             max_rows, padded_cols_l2, dtype=torch.float16, device=self.device
         ).to(torch.float8_e4m3fn)
-        
-        # Swizzled scale output buffers (for CUDA graph capture)
-        self._padded_x_sf_swizzled_buf_l1 = torch.zeros_like(self._padded_x_sf_buf_l1)
-        self._padded_x_sf_swizzled_buf_l2 = torch.zeros_like(self._padded_x_sf_buf_l2)
 
         # Global scale buffers
         self._l1_gsa_buf = torch.zeros(1, dtype=torch.float32, device=self.device)
         self._l2_gsa_buf = torch.zeros(1, dtype=torch.float32, device=self.device)
-        
-        # Pre-allocated swizzled scale output buffers (for CUDA graph capture)
-        # NOTE: _padded_x_sf_swizzled_buf_l1/l2 are allocated above (line 183-184)
-        # Do NOT set to None — they are required for CUDA graph capture swizzle path
-        
-        # Pre-allocated L1 output buffer for graph capture
-        # L1 produces gate+up combined: 2 * intermediate_size BF16 columns
-        self._l1_out_buf = torch.zeros(
-            max_rows, 2 * self.intermediate_size,
-            dtype=torch.bfloat16, device=self.device
-        )
-        # Pre-allocated L2 output buffer for graph capture
-        # L2 produces hidden_size BF16 columns (down projection)
-        self._l2_out_buf = torch.zeros(
-            max_rows, self.hidden_size,
-            dtype=torch.bfloat16, device=self.device
-        )
 
         # Expert offsets for num_groups=1: just [num_tokens_padded]
         # The GEMM expects expert_offsets as (num_experts,) cumulative offsets
@@ -226,38 +202,17 @@ class Nvfp4SharedExpert:
         2. Apply pad_and_swizzle_single (Blackwell swizzle)
         3. Reshape back to 2D (kernel expects 2D scale_a)
 
-        CUDA-graph-safe: uses the pre-allocated padded_x_sf_buf instead of
-        per-call torch.zeros(). The buffer is zeroed + scattered + swizzled
-        each call — zero new allocations on the hot path.
+        The padded buffer must be sized exactly for 128-aligned num_tokens,
+        NOT the max_num_tokens buffer (which would be way too large).
         """
         num_rows, num_cols = x_sf.shape
         padded_rows = cutedsl_ceil_div(num_rows, 128) * 128
         padded_cols = cutedsl_ceil_div(num_cols, 4) * 4
 
-        # Use pre-allocated buffer — zero + scatter pattern (no new allocation)
-        buf = padded_x_sf_buf
-        assert buf.shape[0] >= padded_rows and buf.shape[1] >= padded_cols, \
-            f"padded_x_sf_buf too small: {buf.shape} < ({padded_rows}, {padded_cols})"
-        buf.view(torch.uint8).zero_()
+        # Use a temp buffer sized for this exact token count
+        buf = torch.zeros(padded_rows, padded_cols, dtype=torch.float16, device=x_sf.device).to(torch.float8_e4m3fn)
         buf[:num_rows, :num_cols] = x_sf
-        # Pass correctly-sized VIEW to swizzle — avoids processing the full max-size buffer
-        view = buf[:padded_rows, :padded_cols]
-        
-        # During graph capture, use CUDA swizzle kernel (Python view ops not capturable)
-        if torch.cuda.is_current_stream_capturing():
-            from dsv4.kernels.cuda.loader import get_cuda_module
-            swizzled_buf = self._padded_x_sf_swizzled_buf_l1 if padded_x_sf_buf is self._padded_x_sf_buf_l1 else self._padded_x_sf_swizzled_buf_l2
-            if swizzled_buf is not None:
-                mod = get_cuda_module("blackwell_swizzle", ["blackwell_swizzle.cu"])
-                mod.blackwell_swizzle_32_4_4(
-                    view.view(torch.uint8), swizzled_buf[:padded_rows, :padded_cols].view(torch.uint8),
-                    padded_rows, padded_cols
-                )
-                return swizzled_buf[:padded_rows, :padded_cols].reshape(padded_rows, padded_cols)
-            # Fall through to Python path if buffer not yet allocated
-        
-        # Eager path: Python view operations
-        swizzled_flat = pad_and_swizzle_single(view)
+        swizzled_flat = pad_and_swizzle_single(buf)
         return swizzled_flat.reshape(padded_rows, padded_cols)
 
     def compute_activation_global_scales(self, hidden_states_sample):
@@ -298,7 +253,7 @@ class Nvfp4SharedExpert:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             x_fp4, x_sf, gsa_l1_gpu = quantize_nvfp4_gpu_fused(x_bf16)
-            self._l1_gsa_buf[0] = gsa_l1_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))  # GPU → GPU
         else:
             from dsv4.ops.quantize import quantize_activation_nvfp4
             x_fp4, x_sf = quantize_activation_nvfp4(x_bf16, self._l1_activation_global_scale)
@@ -329,7 +284,6 @@ class Nvfp4SharedExpert:
             global_scale_a=gsa,
             global_scale_b=self._l1_gsb,
             swiglu_limit=self.swiglu_limit if self.swiglu_limit is not None else 0.0,
-            out=self._l1_out_buf,
         )
         l1_out_real = l1_out[:num_tokens]  # (num_tokens, 2*intermediate) BF16, interleaved [silu(gate), silu(gate)*up]
         # Deinterleave to separate gate and up, then take up half (SwiGLU result)
@@ -346,7 +300,7 @@ class Nvfp4SharedExpert:
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
             x_fp4, x_sf, gsa_l1_gpu = quantize_nvfp4_gpu_fused(hidden_states)
-            self._l1_gsa_buf[0] = gsa_l1_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))  # GPU → GPU, no sync
         else:
             x_fp4, x_sf = quantize_activation_nvfp4(
                 hidden_states, self._l1_activation_global_scale
@@ -376,7 +330,6 @@ class Nvfp4SharedExpert:
             expert_offsets=expert_offsets,
             global_scale_a=gsa,
             global_scale_b=self._l1_gsb,
-            out=self._l1_out_buf,
         )
 
         # Extract real token outputs
@@ -394,10 +347,8 @@ class Nvfp4SharedExpert:
         # Fused amax + quantize: zero CPU syncs.
         if getattr(self, '_use_runtime_gsa', False):
             from dsv4.ops.quantize import quantize_nvfp4_gpu_fused
-            if not intermediate.is_contiguous():
-                intermediate = intermediate.contiguous()
             x_fp4, x_sf, gsa_l2_gpu = quantize_nvfp4_gpu_fused(intermediate)
-            self._l2_gsa_buf[0] = gsa_l2_gpu[0]  # scalar GPU→GPU, no sync, graph-capturable
+            self._l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))  # GPU → GPU, no sync
         else:
             x_fp4, x_sf = quantize_activation_nvfp4(
                 intermediate, self._l2_activation_global_scale
@@ -427,7 +378,6 @@ class Nvfp4SharedExpert:
             expert_offsets=expert_offsets,
             global_scale_a=gsa,
             global_scale_b=self._l2_gsb,
-            out=self._l2_out_buf,
         )
 
         return out[:num_tokens]

dsv4/ops/gemm_runner.py

@@ -26,8 +26,6 @@ from dsv4.ops.layouts import (
     round_up,
 )
 
-
-
 # Cache compiled kernels + pre-allocated workspace by cache_key
 # Each entry: {'compiled': callable, 'workspace': Tensor, 'workspace_size': int}
 #
@@ -101,15 +99,7 @@ def warmup_compilation(num_experts, K_packed, N_packed, device,
     )
     
     def to_cute(t):
-        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
-        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
-        # We temporarily patch current_device to return the tensor's device index.
-        # This is safe because during graph capture, the device is logically fixed.
-        _orig_cd = torch.cuda.current_device
-        if t.is_cuda and t.device.index != _orig_cd():
-            torch.cuda.current_device = lambda: t.device.index
         ct = cutlass_torch.from_dlpack(t)
-        torch.cuda.current_device = _orig_cd
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
     a_c = to_cute(mat_a)
@@ -170,7 +160,6 @@ def run_nvfp4_grouped_gemm(
     global_scale_b=None,  # (experts,) float32
     mma_tiler_mn=(128, 128),
     cluster_shape_mn=(1, 1),
-    out=None,       # pre-allocated output buffer for CUDA graph capture
 ):
     """Run the CuTeDSL NVFP4 scaled grouped GEMM.
     
@@ -185,10 +174,7 @@ def run_nvfp4_grouped_gemm(
     n_dim = mat_b.shape[2]
     tokens_sum = mat_a.shape[0]
 
-    if out is None:
-        out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
-    else:
-        out.zero_()
+    out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
 
     # NVFP4-3: use 2-CTA UMMA for M>=256 (1.7-1.9× throughput at prefill)
     use_2cta = tokens_sum >= 256 and cluster_shape_mn[0] % 2 == 0
@@ -217,11 +203,7 @@ def run_nvfp4_grouped_gemm(
         )
 
         def to_cute(t):
-            _orig_cd = torch.cuda.current_device
-            if t.is_cuda and t.device.index != _orig_cd():
-                torch.cuda.current_device = lambda: t.device.index
             ct = cutlass_torch.from_dlpack(t)
-            torch.cuda.current_device = _orig_cd
             return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
 
         a_c = to_cute(mat_a)
@@ -268,15 +250,7 @@ def run_nvfp4_grouped_gemm(
     # This is cheap (metadata only, no GPU work) and avoids stale
     # references to tensors from previous calls that may have been freed.
     def to_cute(t):
-        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
-        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
-        # We temporarily patch current_device to return the tensor's device index.
-        # This is safe because during graph capture, the device is logically fixed.
-        _orig_cd = torch.cuda.current_device
-        if t.is_cuda and t.device.index != _orig_cd():
-            torch.cuda.current_device = lambda: t.device.index
         ct = cutlass_torch.from_dlpack(t)
-        torch.cuda.current_device = _orig_cd
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
     a_c = to_cute(mat_a)
@@ -354,15 +328,7 @@ def warmup_fused_swiglu_compilation(num_experts, K_packed, N_packed, device,
     )
     
     def to_cute(t):
-        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
-        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
-        # We temporarily patch current_device to return the tensor's device index.
-        # This is safe because during graph capture, the device is logically fixed.
-        _orig_cd = torch.cuda.current_device
-        if t.is_cuda and t.device.index != _orig_cd():
-            torch.cuda.current_device = lambda: t.device.index
         ct = cutlass_torch.from_dlpack(t)
-        torch.cuda.current_device = _orig_cd
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
     a_c = to_cute(mat_a)
@@ -416,7 +382,6 @@ def run_fused_swiglu_grouped_gemm(
     swiglu_limit=0.0,
     mma_tiler_mn=(128, 128),
     cluster_shape_mn=(1, 1),
-    out=None,       # pre-allocated output buffer for CUDA graph capture
 ):
     """Run the fused SwiGLU NVFP4 scaled grouped GEMM.
     
@@ -429,10 +394,7 @@ def run_fused_swiglu_grouped_gemm(
     n_dim = mat_b.shape[2]
     tokens_sum = mat_a.shape[0]
 
-    if out is None:
-        out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
-    else:
-        out.zero_()
+    out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device)
 
     # NVFP4-3: use 2-CTA UMMA for M>=256 (1.7-1.9× throughput at prefill)
     # At decode (M<256), 1-CTA is correct (2-CTA wastes hardware)
@@ -463,11 +425,7 @@ def run_fused_swiglu_grouped_gemm(
         )
 
         def to_cute(t):
-            _orig_cd = torch.cuda.current_device
-            if t.is_cuda and t.device.index != _orig_cd():
-                torch.cuda.current_device = lambda: t.device.index
             ct = cutlass_torch.from_dlpack(t)
-            torch.cuda.current_device = _orig_cd
             return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
 
         a_c = to_cute(mat_a)
@@ -508,15 +466,7 @@ def run_fused_swiglu_grouped_gemm(
     workspace = entry['workspace']
     
     def to_cute(t):
-        # Fix: from_dlpack checks torch.cuda.current_device() against tensor device.
-        # Inside CUDA graph capture on non-default GPUs, current_device() may not match.
-        # We temporarily patch current_device to return the tensor's device index.
-        # This is safe because during graph capture, the device is logically fixed.
-        _orig_cd = torch.cuda.current_device
-        if t.is_cuda and t.device.index != _orig_cd():
-            torch.cuda.current_device = lambda: t.device.index
         ct = cutlass_torch.from_dlpack(t)
-        torch.cuda.current_device = _orig_cd
         return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
     
     a_c = to_cute(mat_a)

dsv4/ops/quantize.py

@@ -80,12 +80,12 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
     zero_block = block_amax < (6.0 * 2.0 ** -9)  # < ~0.0117
     # Zero out x for zero/underflow blocks before division.
     # This ensures x_scaled = 0 → FP4 nibbles = 0.
-    # Use scalar 0.0 instead of torch.zeros_like — no allocation, graph-safe.
-    x_reshaped = torch.where(zero_block.unsqueeze(-1), 0.0, x_reshaped)
+    x_reshaped = torch.where(zero_block.unsqueeze(-1),
+                              torch.zeros_like(x_reshaped), x_reshaped)
     block_amax = block_amax.clamp(min=1e-8)
     block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
     # Force zero/underflow blocks: FP8 scale = 0 (exact zero).
-    block_scale = torch.where(zero_block, 0.0, block_scale)
+    block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale)
 
     # Nearest E2M1
     block_sf_expanded = block_scale.float().unsqueeze(-1)
@@ -143,10 +143,11 @@ def quantize_activation_nvfp4(x_bf16, global_scale, block_size=SF_VEC_SIZE):
     block_amax = x_reshaped.abs().amax(dim=-1)
     # Detect zero blocks and underflow blocks (same threshold as quantize_to_nvfp4).
     zero_block = block_amax < (6.0 * 2.0 ** -9)
-    x_reshaped = torch.where(zero_block.unsqueeze(-1), 0.0, x_reshaped)
+    x_reshaped = torch.where(zero_block.unsqueeze(-1),
+                              torch.zeros_like(x_reshaped), x_reshaped)
     block_amax = block_amax.clamp(min=1e-8, max=6.0 * 448.0)  # E4M3 max = 448
     block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
-    block_scale = torch.where(zero_block, 0.0, block_scale)
+    block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale)
 
     block_sf_expanded = block_scale.float().unsqueeze(-1)
     x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8)
@@ -314,24 +315,18 @@ def quantize_nvfp4_gpu_fused(x_bf16, divisor=6.0 * 448.0):
         x_sf: (M, N//16) float8_e4m3fn
         gsa: (M,) float32 GPU tensor — per-row global scale for GEMM
     """
-    # CUDA kernels require contiguous input — column slices from deinterleave are non-contiguous.
-    # For CUDA graph capture, this MUST be contiguous at graph construction time.
-    # The .contiguous() call is a no-op when already contiguous (no allocation).
+    # CUDA kernels require contiguous input — column slices from deinterleave are non-contiguous
     if not x_bf16.is_contiguous():
         x_bf16 = x_bf16.contiguous()
     from dsv4.kernels.cuda.loader import get_cuda_module
     amax_mod = get_cuda_module("amax_gsa", ["amax_gsa.cu"])
     gsa_gpu = amax_mod.compute_amax_gsa(x_bf16, divisor)  # scalar GPU tensor
-    # Broadcast to (M,) for the quantize-from-buffer kernel.
-    # CUDA-graph-safe approach:
-    # - For M=1 decode (graph-captured): just reshape to (1,) — no allocation.
-    # - For M>1 prefill (not graph-captured): expand + contiguous is fine.
+    # Broadcast to (M,) for the quantize-from-buffer kernel
     M = x_bf16.shape[0]
     if gsa_gpu.dim() == 0:
-        gsa_gpu = gsa_gpu.reshape(1)  # scalar → (1,) — no allocation
-    if M > 1:
-        gsa_gpu = gsa_gpu.expand(M).contiguous()  # (M,) — allocation OK for prefill
-    # For M=1: gsa_gpu is (1,) contiguous — zero allocation
+        gsa_gpu = gsa_gpu.reshape(1).expand(M).contiguous()  # (M,) all rows same gsa
+    elif gsa_gpu.shape[0] == 1 and M > 1:
+        gsa_gpu = gsa_gpu.expand(M).contiguous()
     quant_mod = get_cuda_module("fused_amax_quantize", ["fused_amax_quantize.cu"])
     x_fp4, x_sf = quant_mod.quantize_nvfp4_from_buffer(x_bf16, gsa_gpu)
     return x_fp4, x_sf, gsa_gpu

single_shot_inference.py

@@ -9,7 +9,6 @@ NO PyTorch SDPA fallback. NO dequant+matmul for production projections.
 This is the ground truth for vLLM / SGLang integration.
 """
 import os, sys, time, json, math, argparse, logging
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  # Catch async CUDA errors immediately
 import torch
 import torch.nn.functional as F
 from pathlib import Path
@@ -134,301 +133,107 @@ def unweighted_rmsnorm(x, eps=1e-6):
 class CUDAGraphDecoder:
     """Captures and replays CUDA graphs for the decode loop.
     
-    Architecture (Phase 1: eager-break-at-attention):
-    Each layer is split into two graph-captured sub-regions with eager attention
-    in between:
-    
-      Graph A (pre-attention):  mHC pre_block(attn) + fused RMSNorm + quantize
-                                  + q_a + q_a_norm + q_b + kv projections
-                                  → writes x_normed, q_heads, kv_3d, ctx_a to
-                                    pre-allocated buffers for eager attention
-      Eager (attention):          Compressor → Indexer → KV gather → FMHA
-                                  → inverse RoPE → o_a + o_b → F_attn
-                                  → writes F_attn to pre-allocated buffer
-      Graph B (post-attention):   mHC post_block(attn) + mHC pre_block(ffn)
-                                  + fused RMSNorm + quantize + Router + MoE + SE
-                                  + mHC post_block(ffn)
-                                  → writes X_next to pre-allocated output buffer
-    
-    The attention path (compressor, FMHA, inverse RoPE) has dynamic shapes
-    and data-dependent control flow — it MUST run eagerly.
-    The compute path has fixed shapes for T=1 decode — it CAN be captured.
-    
-    The hc_head + norm + lm_head are captured as a separate graph on cuda:0.
-    Cross-GPU transfers (X.to(cuda:N)) happen OUTSIDE graphs between layers.
+    After one warmup step, each layer's compute is captured as a CUDA graph.
+    Replay eliminates Python dispatch overhead (~94ms for 61 layers) and
+    kernel launch latency.
     
     Constraints:
-    - All tensors in captured regions must have fixed addresses (pre-allocated)
-    - No CPU-GPU syncs inside captured regions
-    - The only per-step sync is argmax for sampling (outside graph)
-    - Attention runs eagerly — dynamic shapes are OK there
+    - All tensors must have fixed addresses (pre-allocated)
+    - No dynamic shapes (T=1 decode has fixed shapes)
+    - No CPU-GPU syncs inside the graph
+    - The only sync is argmax at the end of each step
+    
+    Architecture:
+    - One CUDA graph per (layer, gpu) pair — 61 graphs total
+    - One graph for (hc_head + norm + lm_head) on cuda:0
+    - Cross-GPU transfers (X.to(cuda:N)) happen outside graphs
+    - The warmup step also computes and fixes gsa values
     """
     
-    def __init__(self, n_layers, num_gpus, hidden_size, devices, cfg):
+    def __init__(self, n_layers, num_gpus, devices):
         self.n_layers = n_layers
         self.num_gpus = num_gpus
-        self.hidden_size = hidden_size
         self.devices = devices
+        self.graphs = {}      # (li) -> torch.cuda.CUDAGraph
+        self.lm_graph = None  # single graph for hc_head + norm + lm_head
         self.captured = False
         
-        # Model dimensions for buffer pre-allocation
-        self.n_h = cfg.get("num_attention_heads", 128)
-        self.hd = cfg.get("head_dim", 512)
-        self.rd = cfg.get("qk_rope_head_dim", 64)
-        self.q_a_dim = cfg.get("q_lora_rank", 1536)  # q_a projection output dim
-        
-        # Two graphs per layer (A: pre-attn, B: post-attn+FFN) + lm_head
-        self.graphs_a = {}    # li -> torch.cuda.CUDAGraph
-        self.graphs_b = {}    # li -> torch.cuda.CUDAGraph
-        self.streams = {}    # li -> torch.cuda.Stream (per-device, MUST match capture stream during replay)
-        self.lm_graph = None  # single graph for hc_head + norm + lm_head on cuda:0
-        self.lm_stream = None  # stream for lm_head graph on cuda:0
-        
         # Pre-allocated I/O buffers — fixed addresses for graph capture
-        self.x_in_bufs = {}   # li -> (1, 4, H) BF16 on layer's device
-        self.x_out_bufs = {}  # li -> (1, 4, H) BF16 on layer's device
-        
-        # Graph A output buffers (read by eager attention, written by graph A)
-        # These survive across the graph A → eager → graph B boundary.
-        self.x_normed_bufs = {}    # li -> (1, H) BF16 — for compressor/indexer
-        self.q_heads_bufs = {}     # li -> (1, n_h, hd) BF16 — for FMHA
-        self.kv_3d_bufs = {}       # li -> (1, 1, hd) BF16 — for FMHA (pre-RoPE)
-        self.q_a_bufs = {}          # li -> (1, q_a_dim) BF16 — q_a for indexer
-        self.ctx_a_B_bufs = {}     # li -> (1, 4, 4) FP32 — B_l for post_block
-        self.ctx_a_C_bufs = {}     # li -> (1, 4) BF16 — C_l for post_block
-        self.X_mid_bufs = {}       # li -> (1, 4, H) BF16 — X_l for post_block
-        
-        # Graph B input buffer (written by eager attention, read by graph B)
-        self.F_attn_bufs = {}      # li -> (1, H) BF16 — attention output for post_block
-        
-        # lm_head graph buffers (on cuda:0)
-        self.x_lm_in = None    # (1, 4, H) BF16 on cuda:0
-        self.logits_buf = None # (1, vocab_size) BF16 on cuda:0
+        # Each layer reads X_in and writes X_out
+        self.x_in_bufs = {}   # li -> tensor on device of layer li
+        self.x_out_bufs = {}  # li -> tensor on device of layer li
+        self.logits_buf = None  # (1, 129280) on cuda:0
     
-    def pre_allocate(self, cfg):
+    def pre_allocate(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms,
+                     kv_caches, compressors, indexers, moe_runners, se_runners,
+                     routers, prod_lins, layer_w, rope_caches, hc_head,
+                     final_norm_w, lm_head_lin, comp_rope_caches=None):
         """Pre-allocate all I/O buffers with fixed addresses."""
-        H = self.hidden_size
-        V = cfg.get("vocab_size", 129280)
-        n_h = self.n_h
-        hd = self.hd
-        
         for li in range(self.n_layers):
             dev = self.devices[li % self.num_gpus]
-            self.x_in_bufs[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev)
-            self.x_out_bufs[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev)
-            # Graph A intermediates
-            self.x_normed_bufs[li] = torch.zeros(1, H, dtype=torch.bfloat16, device=dev)
-            self.q_heads_bufs[li] = torch.zeros(1, n_h, hd, dtype=torch.bfloat16, device=dev)
-            self.kv_3d_bufs[li] = torch.zeros(1, 1, hd, dtype=torch.bfloat16, device=dev)
-            self.q_a_bufs[li] = torch.zeros(1, self.q_a_dim, dtype=torch.bfloat16, device=dev)  # q_a for indexer
-            self.ctx_a_B_bufs[li] = torch.zeros(1, 4, 4, dtype=torch.float32, device=dev)
-            self.ctx_a_C_bufs[li] = torch.zeros(1, 4, dtype=torch.bfloat16, device=dev)
-            self.X_mid_bufs[li] = torch.zeros(1, 4, H, dtype=torch.bfloat16, device=dev)
-            # Graph B input
-            self.F_attn_bufs[li] = torch.zeros(1, H, dtype=torch.bfloat16, device=dev)
-        
-        # lm_head graph I/O (cuda:0 only)
-        self.x_lm_in = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0')
-        self.logits_buf = torch.zeros(1, V, dtype=torch.bfloat16, device='cuda:0')
+            # X is (1, 4, 7168) BF16
+            self.x_in_bufs[li] = torch.zeros(1, 4, cfg["hidden_size"], dtype=torch.bfloat16, device=dev)
+            self.x_out_bufs[li] = torch.zeros(1, 4, cfg["hidden_size"], dtype=torch.bfloat16, device=dev)
+        self.logits_buf = torch.zeros(1, cfg.get("vocab_size", 129280), dtype=torch.bfloat16, device='cuda:0')
     
     def capture(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms,
                 kv_caches, compressors, indexers, moe_runners, se_runners,
                 routers, prod_lins, layer_w, rope_caches, hc_head,
-                final_norm_w, lm_w, dec_pos_per_gpu, dec_tid32_per_gpu, comp_rope_caches=None):
-        """Capture CUDA graphs for all layers (A/B split) + lm_head.
-        
-        Phase 1: eager-break-at-attention. Graphs A/B capture the compute-heavy
-        path; the attention path runs eagerly between A and B replays.
+                final_norm_w, lm_head_lin, positions, token_id, comp_rope_caches=None):
+        """Capture CUDA graphs for all layers + lm_head.
         
         Must be called after one warmup step so that:
         1. All CuTeDSL kernels are compiled and cached
         2. gsa values are fixed (from warmup_gsa)
         3. CUDA kernels are warmed up (first launch is often slower)
         """
-        from dsv4.ops.quantize import (
-            mhc_rmsnorm_quantize_nvfp4, dequantize_nvfp4,
-            rmsnorm_quantize_nvfp4 as _rmsnorm_quantize,
-        )
-        from dsv4.layers.mhc import mHCContext
-        
-        H = self.hidden_size
-        n_h = self.n_h
-        hd = self.hd
-        rd = self.rd
-        
-        print("  Capturing CUDA graphs (A/B split: compute captured, attention eager)...", flush=True)
-        
-        # Pre-cache norm weights on correct devices to avoid .to() allocations during capture
-        # These must be on the same device as the layer, in FP32, with fixed addresses.
-        attn_norm_dev = {}
-        ffn_norm_dev = {}
-        q_norm_dev = {}
-        kv_norm_dev = {}
-        for li in range(self.n_layers):
-            gpu = li % self.num_gpus
-            dev = self.devices[gpu]
-            an = attn_norms.get(li)
-            if an is not None and an.device != torch.device(dev):
-                attn_norm_dev[li] = an.to(dev, torch.float32)
-            elif an is not None:
-                attn_norm_dev[li] = an.to(torch.float32) if an.dtype != torch.float32 else an
-            fn = ffn_norms.get(li)
-            if fn is not None and fn.device != torch.device(dev):
-                ffn_norm_dev[li] = fn.to(dev, torch.float32)
-            elif fn is not None:
-                ffn_norm_dev[li] = fn.to(torch.float32) if fn.dtype != torch.float32 else fn
-            pfx = f"model.layers.{li}.self_attn"
-            qn = layer_w[li].get(f"{pfx}.q_a_norm.weight")
-            if qn is not None:
-                q_norm_dev[li] = qn.to(dev, torch.float32) if qn.device != torch.device(dev) or qn.dtype != torch.float32 else qn
-            kvn = layer_w[li].get(f"{pfx}.kv_norm.weight")
-            if kvn is not None:
-                kv_norm_dev[li] = kvn.to(dev, torch.float32) if kvn.device != torch.device(dev) or kvn.dtype != torch.float32 else kvn
-        
-        self.attn_norm_dev = attn_norm_dev
-        self.ffn_norm_dev = ffn_norm_dev
-        self.q_norm_dev = q_norm_dev
-        self.kv_norm_dev = kv_norm_dev
-        
-        # Verify all MoE/SE buffers are allocated (swizzled buffers must exist before capture)
-        for li in range(self.n_layers):
-            moe = moe_runners.get(li)
-            if moe is not None:
-                assert hasattr(moe, '_l1_mat_b') and moe._l1_mat_b is not None, f"L{li} MoE: _l1_mat_b not allocated — call _ensure_stacked() before capture"
-                assert hasattr(moe, '_padded_x_sf_buf_l1') and moe._padded_x_sf_buf_l1 is not None, f"L{li} MoE: _padded_x_sf_buf_l1 not allocated — call _allocate_buffers() before capture"
-                assert hasattr(moe, '_padded_x_sf_swizzled_buf_l1') and moe._padded_x_sf_swizzled_buf_l1 is not None, f"L{li} MoE: _padded_x_sf_swizzled_buf_l1 not allocated"
-            se = se_runners.get(li)
-            if se is not None:
-                assert hasattr(se, '_l1_mat_b') and se._l1_mat_b is not None, f"L{li} SE: _l1_mat_b not allocated — call _ensure_initialized() before capture"
-                assert hasattr(se, '_padded_x_sf_buf_l1') and se._padded_x_sf_buf_l1 is not None, f"L{li} SE: _padded_x_sf_buf_l1 not allocated — call _allocate_buffers() before capture"
-                assert hasattr(se, '_padded_x_sf_swizzled_buf_l1') and se._padded_x_sf_swizzled_buf_l1 is not None, f"L{li} SE: _padded_x_sf_swizzled_buf_l1 not allocated"
+        print("  Capturing CUDA graphs for decode...", flush=True)
         
+        # Capture each layer as a separate graph
         for li in range(self.n_layers):
             gpu = li % self.num_gpus
             dev = self.devices[gpu]
             torch.cuda.set_device(gpu)
             
-            attn_mhc = attn_mhcs.get(li)
-            ffn_mhc = ffn_mhcs.get(li)
-            pl = prod_lins.get(li, {})
-            pfx = f"model.layers.{li}.self_attn"
+            # Copy current X into the fixed input buffer
+            # (In practice, the warmup step's X is already on the right device)
             
-            # ======== Graph A: pre-attention compute ========
-            # NOTE: We capture each Graph A on the correct GPU. Multi-GPU graph capture
-            # is known to have issues. We add a validation step to verify correctness.
-            #
-            # Skip validation — the explicit stream approach handles multi-GPU correctly
-            # Input:  X_l = self.x_in_bufs[li]  (1, 4, H)
-            # Output: x_normed, q_heads, kv_3d, ctx_a, X_l → pre-allocated buffers
-            # Create per-device stream for graph capture/replay
-            # CRITICAL: Must use explicit stream for non-default GPUs.
-            # torch.cuda.set_device() alone doesn't work — PyTorch CUDA graphs
-            # on non-default GPUs fail silently (empty graph or stale data replay).
-            s = torch.cuda.Stream(device=dev)
-            self.streams[li] = s
-            
-            # NOTE: Norm weights are pre-cached on device in FP32 (attn_norm_dev, etc.)
-            #       to avoid .to() allocations during graph capture.
-            graph_a = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph_a, stream=s):
-                X_l = self.x_in_bufs[li]
-                
-                # 1. mHC pre_block (attn) — fused P5
-                A_l_a, B_l_a, C_l_a = attn_mhc._dynamic_params(X_l)
-                x_quant_attn = mhc_rmsnorm_quantize_nvfp4(
-                    X_l, A_l_a, attn_norm_dev[li])
-                x_normed = dequantize_nvfp4(x_quant_attn.x_fp4, x_quant_attn.x_sf, x_quant_attn.gsa)
-                
-                # 2. Attention projections
-                q_a = pl['q_a'].run_from_quantized(x_quant_attn)
-                q_norm_w = q_norm_dev.get(li)
-                if q_norm_w is not None:
-                    q_a_quant = _rmsnorm_quantize(q_a, q_norm_w)
-                    q_a = dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
-                    q = pl['q_b'].run_from_quantized(q_a_quant)
-                else:
-                    q = pl['q_b'](q_a)
-                q = unweighted_rmsnorm(q).bfloat16()
-                # NOTE: RoPE is applied in the eager attention path (dynamic positions)
-                q_heads = q.reshape(1, n_h, hd)
-                
-                kv = pl['kv'].run_from_quantized(x_quant_attn)
-                kv_norm_w_k = kv_norm_dev.get(li)
-                if kv_norm_w_k is not None:
-                    kv = rmsnorm(kv, kv_norm_w_k)
-                kv_3d = kv.reshape(1, 1, hd)
-                # NOTE: RoPE is applied in the eager attention path
-                
-                # Write to pre-allocated buffers for eager attention path
-                self.x_normed_bufs[li].copy_(x_normed)
-                self.q_heads_bufs[li].copy_(q_heads)
-                self.kv_3d_bufs[li].copy_(kv_3d)
-                self.q_a_bufs[li].copy_(q_a)
-                self.ctx_a_B_bufs[li].copy_(B_l_a)
-                self.ctx_a_C_bufs[li].copy_(C_l_a)
-                self.X_mid_bufs[li].copy_(X_l)
-            
-            self.graphs_a[li] = graph_a
-            
-            # Note: We don't verify here because x_in_bufs[li] was zero-initialized.
-            # The actual replay path populates x_in_bufs via copy_() before replay,
-            # so the graph replay works correctly with real data.
-            
-            # ======== Graph B: post-attention + FFN compute ========
-            # Input:  X_mid = self.X_mid_bufs[li], F_attn = self.F_attn_bufs[li]
-            # Output: X_next → self.x_out_bufs[li]
-            graph_b = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph_b, stream=s):
-                X_mid = self.X_mid_bufs[li]
-                F_attn = self.F_attn_bufs[li]
-                
-                # 1. mHC post_block (attn)
-                B_l_a = self.ctx_a_B_bufs[li]
-                C_l_a = self.ctx_a_C_bufs[li]
-                BX_a = torch.bmm(B_l_a.transpose(-1, -2), X_mid.float())
-                CF_a = C_l_a.unsqueeze(-1) * F_attn.unsqueeze(1)
-                X_mid_out = (CF_a.float() + BX_a).to(X_mid.dtype)
-                
-                # 2. FFN mHC pre_block — fused P5
-                A_l_f, B_l_f, C_l_f = ffn_mhc._dynamic_params(X_mid_out)
-                x_quant_ffn = mhc_rmsnorm_quantize_nvfp4(
-                    X_mid_out, A_l_f, ffn_norm_dev[li])
-                x_ffn = dequantize_nvfp4(x_quant_ffn.x_fp4, x_quant_ffn.x_sf, x_quant_ffn.gsa)
-                
-                # 3. Router + MoE + SE (direct access — every layer has these)
-                token_id_dev = dec_tid32_per_gpu[gpu]
-                router_li = routers[li]
-                topk_w, topk_ids = router_li(x_ffn, token_ids=token_id_dev)
-                routed_out = moe_runners[li].run(x_ffn, topk_w, topk_ids)
-                shared_out = se_runners[li].run(x_ffn)
-                F_ffn = routed_out + shared_out
-                
-                # 4. mHC post_block (ffn)
-                BX_f = torch.bmm(B_l_f.transpose(-1, -2), X_mid_out.float())
-                CF_f = C_l_f.unsqueeze(-1) * F_ffn.unsqueeze(1)
-                X_next = (CF_f.float() + BX_f).to(X_mid.dtype)
-                
-                self.x_out_bufs[li].copy_(X_next)
-            
-            self.graphs_b[li] = graph_b
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                X_out = forward_layer(
+                    self.x_in_bufs[li], layer_w[li], li, cfg, *rope_caches[gpu],
+                    attn_mhcs.get(li), ffn_mhcs.get(li),
+                    attn_norms.get(li), ffn_norms.get(li),
+                    kv_caches[li], positions, token_id,
+                    compressors.get(li), indexers.get(li),
+                    moe_runners.get(li), se_runners.get(li), routers.get(li),
+                    prod_lin=prod_lins.get(li),
+                    _use_fused_rmsnorm_quantize=True,
+                    comp_rope_cos=comp_rope_caches[gpu][0] if comp_rope_caches else None,
+                    comp_rope_sin=comp_rope_caches[gpu][1] if comp_rope_caches else None,
+                )
+                # Copy output to fixed buffer
+                self.x_out_bufs[li].copy_(X_out)
             
+            self.graphs[li] = graph
             if (li + 1) % 10 == 0:
-                print(f"    Captured {li+1}/{self.n_layers} layer A/B graphs", flush=True)
+                print(f"    Captured {li+1}/{self.n_layers} layer graphs", flush=True)
         
-        # ---- Capture hc_head + norm + lm_head on cuda:0 ----
+        # Capture hc_head + norm + lm_head on cuda:0
         torch.cuda.set_device(0)
-        self.lm_stream = torch.cuda.Stream(device='cuda:0')
         self.lm_graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(self.lm_graph, stream=self.lm_stream):
-            x_out = hc_head.forward(self.x_lm_in) if hc_head is not None else self.x_lm_in[:, 0, :]
-            if final_norm_w is not None:
-                x_out = rmsnorm(x_out, final_norm_w)
-            logits = torch.nn.functional.linear(x_out, lm_w)
-            self.logits_buf.copy_(logits)
+        with torch.cuda.graph(self.lm_graph):
+            # Note: x_in_bufs for the last layer is on the last layer's device.
+            # For the lm_head graph, we need the X on cuda:0.
+            # We'll handle the cross-GPU transfer outside the graph.
+            x_out = self.x_out_bufs[self.n_layers - 1]  # may be on different GPU
+            x_cuda0 = x_out.to('cuda:0')  # This may NOT work in a CUDA graph
+            # Actually, cross-device memcpy in CUDA graphs is not supported.
+            # We need to do the transfer outside and use a cuda:0 buffer.
+            pass  # Will handle this differently
         
         self.captured = True
-        print(f"  Captured {len(self.graphs_a)} layer A/B graph pairs + lm_head", flush=True)
-
+        print(f"  Captured {len(self.graphs)} layer graphs", flush=True)
 # =====================================================================
 def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None):
     O, I2 = weight.shape; I = I2 * 2
@@ -497,8 +302,6 @@ class Compressor:
         self.is_csa = (ratio == 4); self.kv_dim = 2 * head_dim if self.is_csa else head_dim
         self.kv_lin = None   # production Nvfp4Linear for kv_proj
         self.gate_lin = None  # production Nvfp4Linear for gate_proj
-        self._kv_bf16 = None   # BF16 weight for kv_proj (dequantized from NVFP4)
-        self._gate_bf16 = None  # BF16 weight for gate_proj (dequantized from NVFP4)
         self.ape = None; self.kv_norm_w = None
         self._reduce_loaded = False
         # P7: Decode buffering — accumulate hidden_states until we have a complete block.
@@ -509,24 +312,26 @@ class Compressor:
         self._buf_len = 0
 
     def load(self, w, pfx, dev=None):
-        """Load weights and build BF16 projections (dequantized from NVFP4)."""
+        """Load weights and build production Nvfp4Linear instances."""
         if dev is None: dev = self.device
-        # Compressor projections are NOT explicitly FP4-QATed — dequant to BF16, use F.linear
-        # CRITICAL: Use the PyTorch dequant_nvfp4 (defined in this file), NOT the CUDA
-        # dequantize_nvfp4 from dsv4/ops/quantize.py. The CUDA kernel assumes
-        # activation/KV scale layout (row-major (M, N/16)) and crashes on weight scales
-        # that don't match — async illegal memory access surfaces at next sync.
+        # Build production NVFP4 GEMM instances for the two projections
+        # kv_proj: in=7168, out=kv_dim (1024 for CSA, 512 for HCA)
+        # gate_proj: same shapes
         kv_w, kv_ws, kv_ws2, kv_isc = get_nvfp4_weight(w, pfx, 'kv_proj')
         gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(w, pfx, 'gate_proj')
         if kv_w is not None:
-            self._kv_bf16 = dequant_nvfp4(kv_w.to(dev), kv_ws.to(dev), kv_ws2, kv_isc).to(dev).contiguous()
+            kv_out = kv_w.shape[0]  # N_packed
+            kv_in = kv_w.shape[1] * 2  # K_packed * 2
+            self.kv_lin = make_nvfp4_linear(kv_in, kv_out, dev, w, pfx, 'kv_proj')
         if gate_w is not None:
-            self._gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc).to(dev).contiguous()
+            gate_out = gate_w.shape[0]
+            gate_in = gate_w.shape[1] * 2
+            self.gate_lin = make_nvfp4_linear(gate_in, gate_out, dev, w, pfx, 'gate_proj')
         self.ape = w.get(f"{pfx}.position_bias")
         self.kv_norm_w = w.get(f"{pfx}.kv_norm.weight")
 
     def forward(self, hidden_states, positions):
-        if self.ratio == 0 or self._kv_bf16 is None: return None, None, None
+        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
 
         # P7: Buffer decode steps until we have a complete block.
@@ -553,9 +358,9 @@ class Compressor:
         n_complete = T // r
         if n_complete == 0: return None, None, None
 
-        # Step 1-2: BF16 F.linear projections → FP32 for compress
-        kv = torch.nn.functional.linear(hidden_states, self._kv_bf16).float()    # (T, kv_dim) FP32
-        gate = torch.nn.functional.linear(hidden_states, self._gate_bf16).float()  # (T, kv_dim) FP32
+        # Step 1-2: NVFP4 GEMM projections → FP32 for compress
+        kv = self.kv_lin(hidden_states).float()    # (T, kv_dim) FP32
+        gate = self.gate_lin(hidden_states).float()  # (T, kv_dim) FP32
 
         # Step 3: CUDA softmax/reduce kernel → FP32
         # KV-1/KV-2: Return FP32. Caller applies RoPE, then quantizes to NVFP4.
@@ -593,23 +398,22 @@ class Indexer:
     """
     def __init__(self, n_ih, ihd, top_k, device):
         self.n_ih, self.ihd, self.top_k, self.device = n_ih, ihd, top_k, device
-        self.q_b_lin = None   # production Nvfp4Linear for q_b_proj (FP4-QATed)
-        self._wp_bf16 = None    # BF16 weight for weights_proj (dequantized from NVFP4)
+        self.q_b_lin = None   # production Nvfp4Linear for q_b_proj
+        self.wp_lin = None    # production Nvfp4Linear for weights_proj
         self.compressor = None
 
     def load(self, w, pfx, dev=None):
         if dev is None: dev = self.device
         qb_w, qb_ws, qb_ws2, qb_isc = get_nvfp4_weight(w, pfx, 'q_b_proj')
         wp_w, wp_ws, wp_ws2, wp_isc = get_nvfp4_weight(w, pfx, 'weights_proj')
-        # q_b_proj IS the FP4-QATed QK path — keep as NVFP4
         if qb_w is not None:
             qb_out = qb_w.shape[0]
             qb_in = qb_w.shape[1] * 2
             self.q_b_lin = make_nvfp4_linear(qb_in, qb_out, dev, w, pfx, 'q_b_proj')
-        # weights_proj is NOT FP4-QATed — dequant to BF16 via PyTorch reference
-        # CRITICAL: Use PyTorch dequant_nvfp4, NOT CUDA dequantize_nvfp4 (see Compressor.load)
         if wp_w is not None:
-            self._wp_bf16 = dequant_nvfp4(wp_w.to(dev), wp_ws.to(dev), wp_ws2, wp_isc).to(dev).contiguous()
+            wp_out = wp_w.shape[0]
+            wp_in = wp_w.shape[1] * 2
+            self.wp_lin = make_nvfp4_linear(wp_in, wp_out, dev, w, pfx, 'weights_proj')
         # Indexer compressor weights are directly under the indexer prefix
         # (e.g. *.indexer.kv_proj.weight), NOT nested under *.indexer.compressor.
         if f"{pfx}.kv_proj.weight" in w:
@@ -632,7 +436,7 @@ class Indexer:
         li = layer_idx
 
         q_idx = self.q_b_lin(q_lora).reshape(T, self.n_ih, self.ihd)  # (T, n_ih, ihd)
-        w_h = torch.nn.functional.linear(hidden_states, self._wp_bf16)  # (T, n_ih) BF16
+        w_h = self.wp_lin(hidden_states)  # (T, n_ih)
 
         # B2: FP8 tensor-core scoring path.
         # Indexer keys are stored as FP8_E4M3 in the KV cache.
@@ -991,87 +795,11 @@ def _run_production_fmha_mixed(q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16
 # =====================================================================
 # Attention — ALL production kernels
 # =====================================================================
-def eager_attention(q_heads, kv_roped, x_normed, q_a, w, li, cfg,
-                    rope_cos, rope_sin, kv_cache, positions,
-                    compressor, indexer, comp_rope_cos=None, comp_rope_sin=None):
-    """Eager attention section — runs OUTSIDE CUDA graph capture.
-    
-    This function handles the dynamic-shape parts of attention:
-    KV append → Compressor → Indexer → KV gather → FMHA → Inverse RoPE
-    
-    Returns: attn_out (1, n_h, hd) — output of FMHA after inverse RoPE.
-    The caller (sub-graph B) will apply o_proj and mHC post_block.
-    """
-    dev = x_normed.device; T = q_heads.shape[0]
-    n_h = cfg["num_attention_heads"]; hd = cfg["head_dim"]; rd = cfg.get("qk_rope_head_dim", 64)
-    ratio = compressor.ratio if compressor is not None else 0
-    scale = 1.0 / math.sqrt(hd); pfx = f"model.layers.{li}.self_attn"
-    nope_dim = hd - rd
-    if positions.device != rope_cos.device: positions = positions.to(rope_cos.device)
-    
-    # KV append (already roped from sub-graph A)
-    kv_cache.append_swa(kv_roped, positions)
-    
-    # Compressor → compressed KV (mixed storage: FP8 + BF16 RoPE)
-    comp_pos, block_bias = None, None; comp_idx_kv = None
-    if compressor is not None and compressor.ratio > 0:
-        comp_kv_fp32, comp_pos, block_bias = compressor.forward(x_normed, positions)
-        if comp_kv_fp32 is not None:
-            from dsv4.kernels.cuda.loader import get_cuda_module
-            kv_mod = get_cuda_module("kv_quantize", ["kv_quantize.cu"])
-            nope_fp32 = comp_kv_fp32[:, :nope_dim].contiguous()
-            rope_bf16 = comp_kv_fp32[:, nope_dim:].bfloat16().contiguous()
-            rope_3d = rope_bf16.unsqueeze(1)
-            crc = comp_rope_cos if comp_rope_cos is not None else rope_cos
-            crs = comp_rope_sin if comp_rope_sin is not None else rope_sin
-            rope_3d = _apply_rope(rope_3d, comp_pos, crc, crs, rd)
-            rope_bf16 = rope_3d.squeeze(1)
-            nope_fp8, nope_scale = kv_mod.quantize_fp8_e4m3_from_fp32(nope_fp32)
-            kv_cache.set_compressed_mixed(nope_fp8, nope_scale, rope_bf16, comp_pos)
-        if compressor.is_csa and indexer is not None and indexer.compressor is not None:
-            comp_idx_kv, _, _ = indexer.compressor.forward(x_normed, positions)
-            kv_cache.set_indexer_keys_fp8(comp_idx_kv)
-    
-    # Indexer top-k (CSA)
-    topk_idx = None
-    if indexer is not None and ratio == 4:
-        topk_idx = indexer.forward(q_a, x_normed, kv_cache, positions, layer_idx=li)
-    
-    # Gather KV — B1 storage-native mixed path
-    swa_kv, _swa_pos = kv_cache.get_swa()
-    swa_len = swa_kv.shape[0]
-    if kv_cache.n_comp > 0:
-        if ratio == 4:
-            assert topk_idx is not None, f"CSA layer {li}: indexer returned no top-k"
-            tk = topk_idx[0].clamp(0, kv_cache.n_comp - 1).int()
-            kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_selective(tk)
-        elif ratio > 4:
-            kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_all()
-        else:
-            kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only()
-    else:
-        kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = kv_cache.gather_mixed_swa_only()
-    seq_len = kv_nope_scale.shape[0]
-    if seq_len == 0:
-        return torch.zeros(T, n_h, hd, dtype=torch.bfloat16, device=dev)
-    
-    # Production FMHA — B1 mixed FP8/BF16 decode path
-    attn_out = _run_production_fmha_mixed(
-        q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16,
-        n_h, hd, T, seq_len, scale, dev, li, w, pfx, rd)
-    
-    # Inverse RoPE
-    attn_out = _apply_rope(attn_out, positions, rope_cos, rope_sin, rd, inverse=True)
-    
-    return attn_out
-
-
 def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
                       kv_cache, positions, compressor, indexer, prod_lin,
                       x_quant=None,
                       _profile_detail=False, _profile_times=None,
-                      comp_rope_cos=None, comp_rope_sin=None,
-                      q_heads=None, kv_3d=None, q_a=None):
+                      comp_rope_cos=None, comp_rope_sin=None):
     dev = x_normed.device; T = x_normed.shape[0]
     n_h = cfg["num_attention_heads"]; hd = cfg["head_dim"]; rd = cfg.get("qk_rope_head_dim", 64)
     o_groups = cfg.get("o_groups", 16); o_rank = cfg.get("o_lora_rank", 1024)
@@ -1088,46 +816,40 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
 
     _pt('q_a_start')
     # 1. Q: q_a (NVFP4 GEMM) → q_a_norm → q_b (NVFP4 GEMM) → q_b_norm
-    # When q_heads is provided (from CUDA graph A), skip projections — only apply RoPE
-    if q_heads is None:
-        q_a = prod_lin['q_a'].run_from_quantized(x_quant) if x_quant is not None else prod_lin['q_a'](x_normed)
-        _pt('q_a_end')
-        if VERBOSE >= 2 and li < 3:
-            # Compare q_a with PyTorch reference
-            q_a_ref = do_nvfp4_linear_ref(x_normed, w, pfx, 'q_a_proj')
-            if q_a_ref is not None:
-                cos_qa = torch.nn.functional.cosine_similarity(q_a.flatten().float(), q_a_ref.flatten().float(), dim=0).item()
-                print(f"  L{li} q_a: |prod|={q_a.abs().max().item():.6f} |ref|={q_a_ref.abs().max().item():.6f} cos={cos_qa:.6f}", flush=True)
-        q_norm_w = w.get(f"{pfx}.q_a_norm.weight")
-        # B3: Fused rmsnorm+quant for q_a_norm → q_b path
-        if q_norm_w is not None:
-            from dsv4.ops.quantize import rmsnorm_quantize_nvfp4 as _rmsnorm_quantize, dequantize_nvfp4 as _dequantize_nvfp4
-            q_a_quant = _rmsnorm_quantize(q_a, q_norm_w.to(dev, torch.float32))
-            q_a = _dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
-        _pt('q_b_start')
-        if q_norm_w is not None:
-            q = prod_lin['q_b'].run_from_quantized(q_a_quant)
-        else:
-            q = prod_lin['q_b'](q_a)
-        q = unweighted_rmsnorm(q).bfloat16()
-        _pt('q_b_end')
-        q_heads = q.reshape(T, n_h, hd)
+    q_a = prod_lin['q_a'].run_from_quantized(x_quant) if x_quant is not None else prod_lin['q_a'](x_normed)
+    _pt('q_a_end')
+    if VERBOSE >= 2 and li < 3:
+        # Compare q_a with PyTorch reference
+        q_a_ref = do_nvfp4_linear_ref(x_normed, w, pfx, 'q_a_proj')
+        if q_a_ref is not None:
+            cos_qa = torch.nn.functional.cosine_similarity(q_a.flatten().float(), q_a_ref.flatten().float(), dim=0).item()
+            print(f"  L{li} q_a: |prod|={q_a.abs().max().item():.6f} |ref|={q_a_ref.abs().max().item():.6f} cos={cos_qa:.6f}", flush=True)
+    q_norm_w = w.get(f"{pfx}.q_a_norm.weight")
+    # B3: Fused rmsnorm+quant for q_a_norm → q_b path
+    # Replaces: rmsnorm(q_a, w) → BF16 → q_b quantizes internally
+    # With: fused rmsnorm+NVFP4 quantize → QuantizedActivation → q_b.run_from_quantized
+    # Saves: ~6 kernel launches per layer (rmsnorm 4+ + quantize 2 vs fused 2)
+    if q_norm_w is not None:
+        from dsv4.ops.quantize import rmsnorm_quantize_nvfp4 as _rmsnorm_quantize, dequantize_nvfp4 as _dequantize_nvfp4
+        q_a_quant = _rmsnorm_quantize(q_a, q_norm_w.to(dev, torch.float32))
+        q_a = _dequantize_nvfp4(q_a_quant.x_fp4, q_a_quant.x_sf, q_a_quant.gsa)
+    _pt('q_b_start')
+    if q_norm_w is not None:
+        q = prod_lin['q_b'].run_from_quantized(q_a_quant)
     else:
-        # Graph replay: q_a provided from pre-allocated buffer
-        q_a = q_a  # use the passed q_a from graph A output
-    q_heads = _apply_rope(q_heads, positions, rope_cos, rope_sin, rd)
+        q = prod_lin['q_b'](q_a)
+    q = unweighted_rmsnorm(q).bfloat16()
+    _pt('q_b_end')
+    q_heads = q.reshape(T, n_h, hd); q_heads = _apply_rope(q_heads, positions, rope_cos, rope_sin, rd)
     _pt('rope_q_end')
 
     # 2. KV (NVFP4 GEMM, MQA, single KV head)
-    # When kv_3d is provided (from CUDA graph A), skip projections — only apply RoPE
     _pt('kv_start')
-    if kv_3d is None:
-        kv = prod_lin['kv'].run_from_quantized(x_quant) if x_quant is not None else prod_lin['kv'](x_normed)
-        _pt('kv_end')
-        kv_norm_w = w.get(f"{pfx}.kv_norm.weight")
-        if kv_norm_w is not None: kv = rmsnorm(kv, kv_norm_w.to(dev, torch.float32))
-        kv_3d = kv.reshape(T, 1, hd)
-    kv_3d = _apply_rope(kv_3d, positions, rope_cos, rope_sin, rd)
+    kv = prod_lin['kv'].run_from_quantized(x_quant) if x_quant is not None else prod_lin['kv'](x_normed)
+    _pt('kv_end')
+    kv_norm_w = w.get(f"{pfx}.kv_norm.weight")
+    if kv_norm_w is not None: kv = rmsnorm(kv, kv_norm_w.to(dev, torch.float32))
+    kv_3d = kv.reshape(T, 1, hd); kv_3d = _apply_rope(kv_3d, positions, rope_cos, rope_sin, rd)
     _pt('rope_kv_end')
     kv_roped = kv_3d.reshape(T, hd); kv_cache.append_swa(kv_roped, positions)
 
@@ -1584,26 +1306,50 @@ def main():
             router.load_weights(hash_lut=all_w[f"{pfx}.gate.tid2eid"].to(dev, torch.int32))
         else:
             eb = all_w.get(f"{pfx}.gate.e_score_correction_bias")
-            # BF16 router gate — dequantize NVFP4 to BF16, use F.linear
-            E = cfg["n_routed_experts"]
+            # NVFP4 production GEMM for router gate
+            # Custom CuTeDSL fused kernel crashes MLIR optimizer,
+            # so we use Nvfp4Linear (proven production path).
+            from dsv4.layers.linear import Nvfp4Linear
             gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, pfx, 'gate')
+            E = cfg["n_routed_experts"]
             if gate_w is not None and gate_ws is not None:
-                # Checkpoint has NVFP4 gate weight — dequantize to BF16
-                # CRITICAL: Use PyTorch dequant_nvfp4, NOT CUDA dequantize_nvfp4
-                # (same fix as Compressor.load — CUDA kernel crashes on weight scale layouts)
-                gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc)
-                router.W_gate = gate_bf16.T.contiguous().to(dev)  # (H, E) for F.linear(x, W_gate.T)
+                # Checkpoint has NVFP4 gate weight (N_packed, K_packed) — correct layout
+                gate_lin = Nvfp4Linear(in_features=H, out_features=E, device=dev)
+                gate_w_view = gate_w.to(dev).view(torch.float4_e2m1fn_x2) if gate_w.dtype == torch.uint8 else gate_w.to(dev)
+                gate_lin.fp4 = [gate_w_view]
+                gate_lin.sf = [gate_ws.to(dev)]
+                ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0
+                isc_v = gate_isc.float().item() if gate_isc is not None else 1.0/(6.0*448.0)
+                gate_lin.gs = [1.0]
+                gate_lin.ws2 = [torch.tensor([ws2_v], device=dev, dtype=torch.float32)]
+                gate_lin._activation_global_scale = isc_v  # placeholder — runtime gsa overrides this
+                gate_lin._use_runtime_gsa = True  # compute gsa from actual input to avoid E4M3 overflow
+                gate_lin.finalize_weights()
+                router.load_nvfp4_gate(gate_lin)
+                router.load_weights(e_bias=eb.to(dev, torch.float32))
+                if li < 5: print(f"  L{li}: NVFP4 router gate (checkpoint)", flush=True)
             else:
-                # BF16 gate weight from checkpoint
+                # BF16 gate weight: quantize to NVFP4
                 gw = all_w.get(f"{pfx}.gate.weight")
-                gate_bf16 = gw.bfloat16().to(dev)
-                if gate_bf16.shape[0] != H:
-                    gate_bf16 = gate_bf16.T.contiguous()  # ensure (H, E)
-                router.W_gate = gate_bf16.contiguous()
-            # No gate_lin — force BF16 dispatch path
-            router.gate_lin = None
-            router.load_weights(e_bias=eb.to(dev, torch.float32))
-            if li < 5: print(f"  L{li}: BF16 router gate (dequantized from NVFP4)", flush=True)
+                if gw is not None:
+                    g_bf16 = gw if gw.shape == (E, H) else gw.T.contiguous()
+                    g_bf16 = g_bf16.bfloat16().to(dev)
+                    from dsv4.ops.quantize import quantize_to_nvfp4
+                    g_fp4, g_sf, g_gs = quantize_to_nvfp4(g_bf16)
+                    gate_lin = Nvfp4Linear(in_features=H, out_features=E, device=dev)
+                    gate_lin.fp4 = [g_fp4]
+                    gate_lin.sf = [g_sf]
+                    gate_lin.gs = [g_gs]
+                    gate_lin.ws2 = [torch.tensor([g_gs], device=dev, dtype=torch.float32)]
+                    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)  # placeholder — runtime gsa overrides
+                    gate_lin._use_runtime_gsa = True  # compute gsa from actual input to avoid E4M3 overflow
+                    gate_lin.finalize_weights()
+                    router.load_nvfp4_gate(gate_lin)
+                    router.load_weights(e_bias=eb.to(dev, torch.float32))
+                    if li < 5: print(f"  L{li}: NVFP4 router gate (quantized, gs={g_gs:.6f})", flush=True)
+                else:
+                    router.load_weights(e_bias=eb.to(dev, torch.float32))
+                    router.load_weights(e_bias=eb.to(dev, torch.float32))
         router.finalize_weights(); routers[li] = router
 
         moe = Nvfp4MoE(num_experts=cfg["n_routed_experts"], hidden_size=H,
@@ -1651,11 +1397,21 @@ def main():
     torch.cuda.set_device(0)
     embed_w = all_w.get("model.embed_tokens.weight")
     embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to('cuda:0'))
-    # lm_head: BF16 GEMM (checkpoint weight is BF16, no quantization)
+    # lm_head: NVFP4 production GEMM
     lm_w_raw = all_w.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
-    lm_head_lin = None  # Use raw BF16 F.linear for lm_head
-    lm_w = lm_w_raw  # Keep as (V, H) BF16 for F.linear
-    print("  lm_head: BF16 GEMM (checkpoint weight, no quantization)")
+    from dsv4.layers.linear import Nvfp4Linear
+    lm_head_lin = Nvfp4Linear(lm_w_raw.shape[1], lm_w_raw.shape[0], max_num_tokens=8192, device='cuda:0')
+    from dsv4.ops.quantize import quantize_weight_to_nvfp4
+    lm_fp4, lm_sf, lm_gs = quantize_weight_to_nvfp4(lm_w_raw.T.contiguous())
+    lm_head_lin.fp4 = [lm_fp4.permute(1, 0).contiguous()]
+    lm_head_lin.sf = [lm_sf.permute(1, 0).contiguous()]
+    lm_head_lin.gs = [lm_gs]
+    lm_head_lin.ws2 = [None]
+    lm_head_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
+    lm_head_lin._use_runtime_gsa = True
+    lm_head_lin.finalize_weights()
+    lm_w = None
+    print("  lm_head: NVFP4 production GEMM")
     final_norm_w = all_w.get("model.norm.weight")
     if final_norm_w is not None: final_norm_w = final_norm_w.to('cuda:0', torch.float32)
 
@@ -1825,10 +1581,6 @@ def main():
     dec_tid_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
     dec_pos_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
     dec_tid32_buf = torch.zeros(1, dtype=torch.int32, device='cuda:0')
-    # Per-GPU token ID buffers — each GPU needs its own copy for graph capture
-    # (cross-device .to() inside a CUDA graph is not reliable)
-    dec_tid32_per_gpu = {g: torch.zeros(1, dtype=torch.int32, device=f'cuda:{g}') for g in range(NUM_GPUS)}
-    dec_pos_per_gpu = {g: torch.zeros(1, dtype=torch.long, device=f'cuda:{g}') for g in range(NUM_GPUS)}
 
     # Decode
     print(f"\nDecoding (max {MAX_NEW_TOKENS} tokens)...")
@@ -1856,145 +1608,31 @@ def main():
         layer_event_count = 0
         cuda_layer_events = []  # list of (tag, li, timestamp) for fine-grained profiling
 
-    # Pre-allocate decode X buffer — zero per-step allocation
-    # init_state writes to this buffer in-place (no .clone() allocation)
-    dec_X_buf = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0')
-    dec_embed_buf = torch.zeros(1, H, dtype=torch.bfloat16, device='cuda:0')
-    # Pre-allocate pinned CPU buffer for token ID transfer (graph-capturable)
-    dec_tid_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
-    dec_tid32_pinned = torch.zeros(1, dtype=torch.int32, device='cpu').pin_memory()
-    dec_pos_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
-    
-    # ---- CUDA Graph Setup ----
-    graph_decoder = None
-    if _args.cuda_graph:
-        print("  CUDA graph capture requested — will capture after warmup step")
-        graph_decoder = CUDAGraphDecoder(n_layers, NUM_GPUS, H, [f'cuda:{g}' for g in range(NUM_GPUS)], cfg)
-        graph_decoder.pre_allocate(cfg)
-
     for step in range(MAX_NEW_TOKENS):
         t1 = time.time()
-        # Write token/position to pinned CPU buffers, then async copy to GPU
-        dec_tid_pinned[0] = all_tokens[-1]
-        dec_tid_buf.copy_(dec_tid_pinned)
-        dec_tid32_pinned[0] = all_tokens[-1]
-        dec_tid32_buf.copy_(dec_tid32_pinned)
-        dec_pos_pinned[0] = len(all_tokens) - 1
-        dec_pos_buf.copy_(dec_pos_pinned)
-        # Copy token/position to per-GPU buffers for graph capture
-        for g in range(NUM_GPUS):
-            dec_tid32_per_gpu[g].copy_(dec_tid32_pinned)
-            dec_pos_per_gpu[g].copy_(dec_pos_pinned)
+        dec_tid_buf[0] = all_tokens[-1]
+        dec_tid32_buf[0] = all_tokens[-1]
+        dec_pos_buf[0] = len(all_tokens) - 1
 
         t_e = time.perf_counter()
-        X = mHCLayer.init_state(embed(dec_tid_buf), out_buf=dec_X_buf)
-        
-        # ---- Forward: graph replay or eager ----
-        if graph_decoder is not None and graph_decoder.captured:
-            # CUDA graph replay path — A/B split with eager attention
-            for li in range(n_layers):
-                gpu = li % NUM_GPUS
-                torch.cuda.set_device(gpu)
-                dev = f'cuda:{gpu}'
-                
-                # Copy X into graph A input buffer (copy_ handles cross-GPU transfer)
-                graph_decoder.x_in_bufs[li].copy_(X)
-                # NOTE: Cross-GPU copy synchronization is handled by the stream events
-                # (Graph A's stream waits for the default stream's F_attn write, and
-                # vice versa). No explicit sync needed here.
-                
-                # DEBUG: check input is non-zero (first 3 steps, first 3 layers)
-                if step < 3 and li < 3:
-                    torch.cuda.synchronize()
-                    print(f"  Replay L{li}: x_in |X|={graph_decoder.x_in_bufs[li].abs().max().item():.2f}", flush=True)
-                
-                # Replay graph A on its capture stream
-                with torch.cuda.stream(graph_decoder.streams[li]):
-                    graph_decoder.graphs_a[li].replay()
-                
-                # Record completion event on graph A's stream, then wait on default stream
-                # This ensures the default stream (eager attention) sees Graph A's output
-                _graph_a_done = torch.cuda.Event()
-                with torch.cuda.stream(graph_decoder.streams[li]):
-                    _graph_a_done.record()
-                torch.cuda.current_stream().wait_event(_graph_a_done)
-                
-                # DEBUG: check graph A output (first 3 steps, first 3 layers)
-                if step < 3 and li < 3:
-                    torch.cuda.synchronize()
-                    print(f"  Replay L{li} GraphA: x_normed |X|={graph_decoder.x_normed_bufs[li].abs().max().item():.2f} "
-                          f"q_heads |X|={graph_decoder.q_heads_bufs[li].abs().max().item():.2f} "
-                          f"kv_3d |X|={graph_decoder.kv_3d_bufs[li].abs().max().item():.2f}", flush=True)
-                
-                # ---- Eager attention (NOT captured) ----
-                # Read graph A outputs from pre-allocated buffers
-                x_normed = graph_decoder.x_normed_bufs[li]
-                q_heads = graph_decoder.q_heads_bufs[li]
-                kv_3d = graph_decoder.kv_3d_bufs[li]
-                
-                # Run full attention eagerly (compressor + indexer + FMHA + o_proj)
-                F_attn, _ = forward_attention(
-                    x_normed, layer_w[li], li, cfg, *rope_caches[gpu],
-                    kv_caches[li], dec_pos_per_gpu[gpu],
-                    compressors.get(li), indexers.get(li), prod_lins.get(li),
-                    q_heads=q_heads, kv_3d=kv_3d, q_a=graph_decoder.q_a_bufs[li],
-                    comp_rope_cos=comp_rope_caches[gpu][0] if comp_rope_caches else None,
-                    comp_rope_sin=comp_rope_caches[gpu][1] if comp_rope_caches else None,
-                )
-                
-                # Write F_attn to graph B input buffer
-                graph_decoder.F_attn_bufs[li].copy_(F_attn)
-                
-                # Record completion of F_attn write on default stream, wait on graph stream
-                _eager_done = torch.cuda.Event()
-                _eager_done.record(torch.cuda.current_stream())
-                with torch.cuda.stream(graph_decoder.streams[li]):
-                    _eager_done.synchronize()
-                
-                # DEBUG: check F_attn (first 3 steps, first 3 layers)
-                if step < 3 and li < 3:
-                    torch.cuda.synchronize()
-                    print(f"  Replay L{li} F_attn |X|={F_attn.abs().max().item():.2f}", flush=True)
-                
-                # Replay graph B on its capture stream
-                with torch.cuda.stream(graph_decoder.streams[li]):
-                    graph_decoder.graphs_b[li].replay()
-                
-                # Read output from graph B
-                X = graph_decoder.x_out_bufs[li]
-                
-                # DEBUG: check graph B output (first 3 steps, first 3 layers)
-                if step < 3 and li < 3:
-                    torch.cuda.synchronize()
-                    print(f"  Replay L{li} GraphB: x_out |X|={X.abs().max().item():.2f}", flush=True)
-            
-            # Transfer last layer output to cuda:0 for lm_head graph
-            graph_decoder.x_lm_in.copy_(X)
-            
-            # lm_head graph replay — use capture stream on cuda:0
-            with torch.cuda.stream(graph_decoder.lm_stream):
-                graph_decoder.lm_graph.replay()
-            logits = graph_decoder.logits_buf
-            
-        else:
-            # Eager forward path (warmup or no --cuda-graph)
-            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, layer_w[li], li, cfg, *rope_caches[gpu],
-                                  attn_mhcs.get(li), ffn_mhcs.get(li),
-                                  attn_norms.get(li), ffn_norms.get(li),
-                                  kv_caches[li], dec_pos_buf, dec_tid32_buf,
-                                  compressors.get(li), indexers.get(li),
-                                  moe_runners.get(li), se_runners.get(li), routers.get(li),
-                                  prod_lin=prod_lins.get(li),
-                                  _profile_detail=(profile and step == 1),
-                                  _profile_times=cuda_layer_events if (profile and step == 1) else None,
-                                  _use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm,
-                                  comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1],
-                                  )
-            X = X.to('cuda:0'); torch.cuda.set_device(0)
+        X = mHCLayer.init_state(embed(dec_tid_buf))
+        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, layer_w[li], li, cfg, *rope_caches[gpu],
+                              attn_mhcs.get(li), ffn_mhcs.get(li),
+                              attn_norms.get(li), ffn_norms.get(li),
+                              kv_caches[li], dec_pos_buf, dec_tid32_buf,
+                              compressors.get(li), indexers.get(li),
+                              moe_runners.get(li), se_runners.get(li), routers.get(li),
+                              prod_lin=prod_lins.get(li),
+                              _profile_detail=(profile and step == 1),
+                              _profile_times=cuda_layer_events if (profile and step == 1) else None,
+                              _use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm,
+                              comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1],
+                              )
+        X = X.to('cuda:0'); torch.cuda.set_device(0)
         t_layers = time.perf_counter()
 
         # After first decode step: fix gsa values from runtime amax
@@ -2009,8 +1647,7 @@ def main():
                 if pl is None: continue
                 for key, lin in pl.items():
                     if hasattr(lin, '_gsa_buf') and hasattr(lin, '_use_runtime_gsa') and lin._use_runtime_gsa:
-                        # Nvfp4GroupedLinear has per-group gsa; reduce to scalar (max) for fixed gsa
-                        fixed_gsa = lin._gsa_buf.max().item() if lin._gsa_buf.numel() > 1 else lin._gsa_buf.item()
+                        fixed_gsa = lin._gsa_buf.item()  # One-time sync
                         lin._activation_global_scale = fixed_gsa
                         lin._use_runtime_gsa = False
                         n_fixed += 1
@@ -2023,35 +1660,16 @@ def main():
                         gl._activation_global_scale = fixed_gsa
                         gl._use_runtime_gsa = False
                         n_fixed += 1
-            # lm_head (BF16 — no gsa needed)
-            if lm_head_lin is not None and hasattr(lm_head_lin, '_gsa_buf') and hasattr(lm_head_lin, '_use_runtime_gsa') and lm_head_lin._use_runtime_gsa:
+            # lm_head
+            if hasattr(lm_head_lin, '_gsa_buf') and hasattr(lm_head_lin, '_use_runtime_gsa') and lm_head_lin._use_runtime_gsa:
                 fixed_gsa = lm_head_lin._gsa_buf.item()
                 lm_head_lin._activation_global_scale = fixed_gsa
                 lm_head_lin._use_runtime_gsa = False
                 n_fixed += 1
             print(f"  Warmup gsa: fixed {n_fixed} projection gsa values from step 0 (MoE/SE keep runtime gsa)", flush=True)
-        
-        # ---- lm_head: graph replay or eager ----
-        if graph_decoder is not None and graph_decoder.captured:
-            # logits already computed by lm_head graph replay above
-            pass
-        else:
-            x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
-            if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w)
-            logits = torch.nn.functional.linear(x_out, lm_w) if lm_head_lin is None else lm_head_lin(x_out)
-        
-        # ---- CUDA graph capture after warmup ----
-        if graph_decoder is not None and not graph_decoder.captured and step == 0:
-            print("  Step 0 warmup done. Capturing CUDA graphs...", flush=True)
-            torch.cuda.synchronize()
-            graph_decoder.capture(
-                cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms,
-                kv_caches, compressors, indexers, moe_runners, se_runners,
-                routers, prod_lins, layer_w, rope_caches, hc_head,
-                final_norm_w, lm_w, dec_pos_per_gpu, dec_tid32_per_gpu,
-                comp_rope_caches=comp_rope_caches,
-            )
-            print(f"  CUDA graphs captured. Graph replay starts on step 1.", flush=True)
+        x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
+        if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w)
+        logits = lm_head_lin(x_out)
         if profile: torch.cuda.synchronize()
         t_lm = time.perf_counter()
         # Check thinking start token logit on first step

tests/unit/test_cuda_graph_multi_gpu.py

@@ -1,114 +0,0 @@
-"""Minimal CUDA graph test: verify graph capture works on all 8 B200 GPUs."""
-import torch
-
-def test_basic_graph():
-    """Test basic CUDA graph on each GPU."""
-    results = {}
-    for gpu in range(8):
-        torch.cuda.set_device(gpu)
-        device = f'cuda:{gpu}'
-        
-        # Create input and output tensors
-        x = torch.ones(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        y = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        
-        # Capture graph
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            y.copy_(x * 2.0)
-        
-        # Reset input
-        x.zero_()
-        
-        # Replay graph — y should be 0.0 * 2.0 = 0.0 since x is now zero
-        g.replay()
-        torch.cuda.synchronize()
-        
-        y_max = y.abs().max().item()
-        results[gpu] = y_max
-        status = "OK" if y_max == 0.0 else f"WRONG (expected 0.0, got {y_max})"
-        print(f"  GPU {gpu}: y_max={y_max:.2f} — {status}")
-    
-    return results
-
-def test_graph_with_updated_input():
-    """Test that graph replay uses current data in input buffer."""
-    results = {}
-    for gpu in range(8):
-        torch.cuda.set_device(gpu)
-        device = f'cuda:{gpu}'
-        
-        # Create input and output tensors (pre-allocated)
-        x_buf = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        y_buf = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        
-        # Fill input with data for capture
-        x_buf.fill_(1.0)
-        
-        # Capture graph
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            y_buf.copy_(x_buf * 2.0)
-        
-        # Now update input with DIFFERENT data
-        x_buf.fill_(3.0)
-        
-        # Replay graph — y should be 3.0 * 2.0 = 6.0
-        g.replay()
-        torch.cuda.synchronize()
-        
-        y_max = y_buf.abs().max().item()
-        results[gpu] = y_max
-        status = "OK" if abs(y_max - 6.0) < 0.1 else f"WRONG (expected 6.0, got {y_max})"
-        print(f"  GPU {gpu}: y_max={y_max:.2f} — {status}")
-    
-    return results
-
-def test_cross_gpu_copy_then_graph():
-    """Test cross-GPU copy followed by graph replay."""
-    results = {}
-    for gpu in range(1, 8):  # Skip GPU 0 (source)
-        torch.cuda.set_device(gpu)
-        device = f'cuda:{gpu}'
-        
-        # Source data on cuda:0
-        src = torch.full((1, 4, 7168), 5.0, dtype=torch.bfloat16, device='cuda:0')
-        
-        # Input/output buffers on cuda:{gpu}
-        x_buf = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        y_buf = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        
-        # Fill with data for capture
-        x_buf.fill_(1.0)
-        
-        # Capture graph
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            y_buf.copy_(x_buf * 2.0)
-        
-        # Copy data from cuda:0 to input buffer
-        x_buf.copy_(src)
-        torch.cuda.synchronize()
-        
-        # Replay — y should be 5.0 * 2.0 = 10.0
-        g.replay()
-        torch.cuda.synchronize()
-        
-        y_max = y_buf.abs().max().item()
-        results[gpu] = y_max
-        status = "OK" if abs(y_max - 10.0) < 0.1 else f"WRONG (expected 10.0, got {y_max})"
-        print(f"  cuda:0→cuda:{gpu}: y_max={y_max:.2f} — {status}")
-    
-    return results
-
-if __name__ == "__main__":
-    print("=== Test 1: Basic graph on each GPU ===")
-    test_basic_graph()
-    
-    print("\n=== Test 2: Graph replay with updated input ===")
-    test_graph_with_updated_input()
-    
-    print("\n=== Test 3: Cross-GPU copy then graph replay ===")
-    test_cross_gpu_copy_then_graph()
-    
-    print("\nDone.")

tests/unit/test_cuda_graph_readiness.py

@@ -1,541 +0,0 @@
-#!/usr/bin/env python3
-"""CUDA Graph Readiness Detector — Section A of GETTING_CUDAGRAPH_READY.md
-
-Runs one decode step of single_shot_inference.py with:
-1. torch.cuda.set_sync_debug_mode("error") — raises on any implicit device→host sync
-2. torch.cuda.graph capture attempt — fails on .item(), sync, alloc, dynamic shape
-
-This inventories EVERY existing sync in one pass so we get the full hunt-list upfront.
-"""
-import os, sys, time, json, math, traceback
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-import torch
-import torch.nn.functional as F
-
-# ==== CONFIG ====
-CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
-NUM_GPUS = 8
-PROMPT = "The capital of France is"
-MAX_CONTEXT = 8192
-SEED = 42
-
-# ==== Sync inventory ====
-sync_violations = []
-
-class SyncDetector:
-    """Tracks all device→host sync violations found during forward."""
-    def __init__(self):
-        self.violations = []
-        self.phase = "unknown"
-    
-    def record(self, category, location, detail):
-        self.violations.append({
-            "phase": self.phase,
-            "category": category,
-            "location": location,
-            "detail": detail,
-        })
-        print(f"  [SYNC] {category}: {location} — {detail}", flush=True)
-
-detector = SyncDetector()
-
-# ==== Import single_shot components ====
-# We need to import the functions/classes without running main()
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from single_shot_inference import (
-    load_all_weights, build_rope_cache, rmsnorm, unweighted_rmsnorm,
-    FP4_LUT, KVCache, Compressor, Indexer, HcHead,
-    make_nvfp4_linear, get_nvfp4_weight, dequant_nvfp4,
-    forward_layer, forward_attention, _run_production_fmha_mixed,
-    moe_forward, _apply_rope,
-    _load_moe_weights_stacked, _load_shared_expert_weights, _cache_layer_weights_no_experts,
-)
-from encoding.deepseek_v4_encoding import (
-    thinking_start_token, thinking_end_token,
-    USER_SP_TOKEN, ASSISTANT_SP_TOKEN,
-)
-
-
-def grep_sync_patterns(source_dir):
-    """Grep the hot path for known sync patterns (Section B checklist)."""
-    import re
-    patterns = {
-        'item()': r'\.item\(\)',
-        '.cpu()': r'\.cpu\(\)',
-        '.tolist()': r'\.tolist\(\)',
-        '.numpy()': r'\.numpy\(\)',
-        'int(t)/float(t)': r'\bint\([^)]*\)|float\([^)]*\)',  # rough
-        'cuda.synchronize()': r'torch\.cuda\.synchronize\(\)',
-        'isnan().any()': r'\.isnan\([^)]*\)\.any\(\)',
-        'isinf().any()': r'\.isinf\([^)]*\)\.any\(\)',
-        'if t:': r'if\s+\w+\.item\(\)',
-        'nonzero': r'\.nonzero\(\)',
-        'masked_select': r'\.masked_select\(',
-        'torch.where(one-arg)': r'torch\.where\([^,]+\)',
-    }
-    import glob
-    hot_files = [
-        'single_shot_inference.py',
-        'dsv4/layers/mhc.py',
-        'dsv4/layers/router.py',
-        'dsv4/layers/moe.py',
-        'dsv4/layers/shared_expert.py',
-        'dsv4/layers/linear.py',
-        'dsv4/layers/grouped_linear.py',
-        'dsv4/ops/quantize.py',
-        'dsv4/kernels/attention/production.py',
-        'dsv4/kernels/compressor/production_compress.py',
-    ]
-    print("\n=== SECTION B: Grep Results (hot path sync patterns) ===", flush=True)
-    for fname in hot_files:
-        fpath = os.path.join(source_dir, fname)
-        if not os.path.exists(fpath):
-            continue
-        with open(fpath) as f:
-            lines = f.readlines()
-        for i, line in enumerate(lines, 1):
-            stripped = line.strip()
-            if stripped.startswith('#') or stripped.startswith('"""') or stripped.startswith("'''"):
-                continue
-            for pname, pat in patterns.items():
-                if re.search(pat, stripped):
-                    # Skip comments
-                    if '#' in stripped and stripped.index('#') < re.search(pat, stripped).start():
-                        continue
-                    print(f"  [{pname}] {fname}:{i}: {stripped[:120]}", flush=True)
-
-
-def run_sync_debug_mode():
-    """Method 1: Run forward with sync debug mode to catch implicit syncs."""
-    print("\n=== METHOD 1: torch.cuda.set_sync_debug_mode('error') ===", flush=True)
-    
-    # Build model components (same as single_shot main, but abbreviated)
-    with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f:
-        cfg = json.load(f)
-    n_layers = cfg["num_hidden_layers"]
-    H = cfg["hidden_size"]
-    hd = cfg["head_dim"]
-    n_h = cfg["num_attention_heads"]
-    rd = cfg.get("qk_rope_head_dim", 64)
-    cr = cfg.get("compress_ratios", [128] * n_layers)
-    
-    print(f"Model: {n_layers} layers, {n_h} heads, hd={hd}", flush=True)
-    
-    # Load weights
-    print("Loading weights...", flush=True)
-    all_w = load_all_weights(CHECKPOINT_DIR)
-    
-    # Build components
-    from dsv4.layers.mhc import mHCLayer
-    from dsv4.layers.router import Router
-    from dsv4.layers.moe import Nvfp4MoE
-    from dsv4.layers.shared_expert import Nvfp4SharedExpert
-    from dsv4.layers.grouped_linear import Nvfp4GroupedLinear
-    
-    for g in range(NUM_GPUS):
-        torch.cuda.set_device(g)
-        torch.cuda.empty_cache()
-    torch.cuda.set_device(0)
-    
-    # Build mHC + norms
-    attn_mhcs, ffn_mhcs, attn_norms, ffn_norms = {}, {}, {}, {}
-    for li in range(n_layers):
-        dev = f"cuda:{li % NUM_GPUS}"
-        for tag, blocks, fn_s, base_s, scale_s in [
-            ("attn", attn_mhcs, f"model.layers.{li}.attn_hc.fn", f"model.layers.{li}.attn_hc.base", f"model.layers.{li}.attn_hc.scale"),
-            ("ffn", ffn_mhcs, f"model.layers.{li}.ffn_hc.fn", f"model.layers.{li}.ffn_hc.base", f"model.layers.{li}.ffn_hc.scale"),
-        ]:
-            fn, base, scale = all_w.get(fn_s), all_w.get(base_s), all_w.get(scale_s)
-            if fn is not None and base is not None and scale is not None:
-                m = mHCLayer(hidden_dim=H, n_hc=4, t_max_sinkhorn=20, device=dev)
-                n = 4
-                m.load_weights(
-                    W_pre=fn[0:n].to(dev, torch.float32), W_post=fn[n:2*n].to(dev, torch.float32),
-                    W_comb=fn[2*n:].to(dev, torch.float32),
-                    S_pre=base[0:n].reshape(1, n).to(dev, torch.float32),
-                    S_post=base[n:2*n].reshape(n, 1).to(dev, torch.float32),
-                    S_comb=base[2*n:].reshape(n, n).to(dev, torch.float32),
-                    alpha_pre=scale[0].item(), alpha_post=scale[1].item(), alpha_comb=scale[2].item(),
-                )
-                blocks[li] = m
-        an_k = f"model.layers.{li}.input_layernorm.weight"
-        if an_k in all_w: attn_norms[li] = all_w[an_k].to(dev, torch.float32)
-        fn_k = f"model.layers.{li}.post_attention_layernorm.weight"
-        if fn_k in all_w: ffn_norms[li] = all_w[fn_k].to(dev, torch.float32)
-    
-    # Build attention projections
-    prod_lins = {}
-    for li in range(n_layers):
-        dev = f"cuda:{li % NUM_GPUS}"
-        pfx = f"model.layers.{li}.self_attn"
-        torch.cuda.set_device(li % NUM_GPUS)
-        pl = {}
-        pl['q_a'] = make_nvfp4_linear(7168, 1536, dev, all_w, pfx, 'q_a_proj')
-        pl['q_b'] = make_nvfp4_linear(1536, 65536, dev, all_w, pfx, 'q_b_proj')
-        pl['kv'] = make_nvfp4_linear(7168, 512, dev, all_w, pfx, 'kv_proj')
-        n_local_groups = cfg.get('o_groups', 16)
-        heads_per_group = n_h // n_local_groups
-        o_rank_val = cfg.get('o_lora_rank', 1024)
-        wo_a = Nvfp4GroupedLinear(
-            n_local_groups=n_local_groups,
-            heads_per_group=heads_per_group,
-            head_dim=hd,
-            o_lora_rank=o_rank_val,
-            max_num_tokens=8192,
-            device=dev,
-        )
-        oa_w_nvfp4, oa_ws, oa_ws2, oa_isc = get_nvfp4_weight(all_w, pfx, 'o_a_proj')
-        if oa_w_nvfp4 is not None and oa_ws is not None:
-            wo_a.load_nvfp4_weight(oa_w_nvfp4.to(dev), oa_ws.to(dev),
-                                   oa_ws2.to(dev) if oa_ws2 is not None else None,
-                                   oa_isc.to(dev) if oa_isc is not None else None)
-        else:
-            oa_bf = all_w.get(f"{pfx}.o_a_proj.weight")
-            if oa_bf is not None:
-                wo_a.set_bf16_weight(oa_bf.bfloat16().to(dev))
-        pl['o_a'] = wo_a
-        wo_a._use_runtime_gsa = True
-        pl['o_b'] = make_nvfp4_linear(16384, 7168, dev, all_w, pfx, 'o_b_proj')
-        prod_lins[li] = pl
-        if (li+1) % 10 == 0:
-            print(f"  {li+1}/{n_layers} attn projections", flush=True)
-    
-    # Routers, MoE, shared experts
-    routers, moe_runners, se_runners = {}, {}, {}
-    for li in range(n_layers):
-        dev = f"cuda:{li % NUM_GPUS}"
-        pfx = f"model.layers.{li}.mlp"
-        torch.cuda.set_device(li % NUM_GPUS)
-        torch.cuda.synchronize()
-        is_hash = (li < cfg.get("num_hash_layers", 3)) and (f"{pfx}.gate.tid2eid" in all_w)
-        router = Router(hidden_size=H, num_experts=cfg["n_routed_experts"],
-                       top_k=cfg.get("num_experts_per_tok", 6),
-                       routed_scaling_factor=cfg.get("routed_scaling_factor", 2.5),
-                       mode="hash" if is_hash else "dense",
-                       vocab_size=cfg.get("vocab_size", 128000) if is_hash else None, device=dev)
-        if is_hash:
-            router.load_weights(hash_lut=all_w[f"{pfx}.gate.tid2eid"].to(dev, torch.int32))
-        else:
-            eb = all_w.get(f"{pfx}.gate.e_score_correction_bias")
-            gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, pfx, 'gate')
-            if gate_w is not None and gate_ws is not None:
-                gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc)
-                router.W_gate = gate_bf16.T.contiguous().to(dev)
-            else:
-                gw = all_w.get(f"{pfx}.gate.weight")
-                gate_bf16 = gw.bfloat16().to(dev)
-                if gate_bf16.shape[0] != H:
-                    gate_bf16 = gate_bf16.T.contiguous()
-                router.W_gate = gate_bf16.contiguous()
-            router.gate_lin = None
-            router.load_weights(e_bias=eb.to(dev, torch.float32))
-        router.finalize_weights()
-        routers[li] = router
-        
-        moe = Nvfp4MoE(num_experts=cfg["n_routed_experts"], hidden_size=H,
-                       intermediate_size=cfg.get("moe_intermediate_size", 3072),
-                       top_k=cfg.get("num_experts_per_tok", 6), device=dev)
-        moe.set_swiglu_limit(cfg.get("swiglu_limit", 10.0))
-        moe.set_fused_swiglu(True)
-        _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg)
-        moe._ensure_stacked()
-        moe._use_runtime_gsa = True
-        moe_runners[li] = moe
-        
-        se = Nvfp4SharedExpert(hidden_size=H, intermediate_size=cfg.get("moe_intermediate_size", 3072),
-                               device=dev, swiglu_limit=cfg.get("swiglu_limit", 10.0))
-        se.set_fused_swiglu(True)
-        _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg)
-        se._ensure_initialized()
-        if se._fused_swiglu:
-            from dsv4.ops.gemm_runner import warmup_fused_swiglu_compilation
-            K_packed = H // 2
-            N_packed_l1 = (2 * cfg.get("moe_intermediate_size", 3072)) // 2
-            warmup_fused_swiglu_compilation(1, K_packed, N_packed_l1, dev,
-                                           swiglu_limit=cfg.get("swiglu_limit", 10.0))
-        se._use_runtime_gsa = True
-        se_runners[li] = se
-        if (li+1) % 10 == 0:
-            print(f"  {li+1}/{n_layers} MoE layers", flush=True)
-        torch.cuda.empty_cache()
-    
-    # Global weights
-    torch.cuda.set_device(0)
-    embed_w = all_w.get("model.embed_tokens.weight")
-    embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to('cuda:0'))
-    lm_w = all_w.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
-    final_norm_w = all_w.get("model.norm.weight")
-    if final_norm_w is not None:
-        final_norm_w = final_norm_w.to('cuda:0', torch.float32)
-    
-    hc_head = HcHead(H, 4, 'cuda:0')
-    hc_fn = all_w.get("model.hc_head.hc_fn")
-    hc_base = all_w.get("model.hc_head.hc_base")
-    hc_scale = all_w.get("model.hc_head.hc_scale")
-    if hc_fn is not None and hc_base is not None:
-        hc_head.load(hc_fn, hc_base, hc_scale)
-    
-    # RoPE
-    rp = cfg.get("rope_scaling", cfg.get("rope_parameters", {}))
-    rt = rp.get("type", rp.get("rope_type", "yarn"))
-    rf = rp.get("factor", 16.0)
-    rtheta = cfg.get("rope_theta", 10000.)
-    romax = rp.get("original_max_position_embeddings", 65536)
-    rbfast, rbslow = rp.get("beta_fast", 32), rp.get("beta_slow", 1)
-    rope_caches = {g: build_rope_cache(romax, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)}
-    comp_rtheta = cfg.get("compress_rope_theta", rtheta)
-    if comp_rtheta != rtheta:
-        comp_rope_caches = {g: build_rope_cache(romax, rd, f"cuda:{g}", comp_rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)}
-    else:
-        comp_rope_caches = rope_caches
-    
-    # KV caches, compressors, indexers
-    kv_caches, compressors, indexers = {}, {}, {}
-    n_ih = cfg.get("index_n_heads", 64)
-    ihd = cfg.get("index_head_dim", 128)
-    itk = cfg.get("index_topk", 1024)
-    for li in range(n_layers):
-        dev = f"cuda:{li % NUM_GPUS}"
-        ratio = cr[li] if li < len(cr) else 128
-        max_comp = (MAX_CONTEXT + ratio - 1) // ratio if ratio > 0 else 0
-        kv_caches[li] = KVCache(hd, cfg.get("sliding_window", 128), max_comp=max_comp, device=dev,
-                                 indexer_key_dim=ihd, compress_ratio=ratio, indexer_top_k=itk, rope_dim=rd)
-        if ratio > 0: compressors[li] = Compressor(ratio, hd, H, dev)
-        if ratio == 4: indexers[li] = Indexer(n_ih, ihd, itk, dev)
-    
-    # Cache layer weights
-    devs = [f"cuda:{g}" for g in range(NUM_GPUS)]
-    layer_w = _cache_layer_weights_no_experts(all_w, n_layers, devs)
-    
-    # Load compressor/indexer weights
-    for li in range(n_layers):
-        pfx = f"model.layers.{li}.self_attn.compressor"
-        if li in compressors: compressors[li].load(layer_w[li], pfx, dev=f"cuda:{li % NUM_GPUS}")
-        if li in indexers: indexers[li].load(layer_w[li], f"{pfx}.indexer", dev=f"cuda:{li % NUM_GPUS}")
-    
-    del all_w
-    import gc; gc.collect()
-    for g in range(NUM_GPUS):
-        torch.cuda.set_device(g)
-        torch.cuda.empty_cache()
-    torch.cuda.set_device(0)
-    
-    print("\nAll components built. Running prefill...", flush=True)
-    
-    # ---- Prefill (run normally, not under sync debug) ----
-    from transformers import AutoTokenizer
-    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)
-    from encoding.deepseek_v4_encoding import encode_messages
-    messages = [{"role": "user", "content": PROMPT}]
-    encoded_str = encode_messages(messages, thinking_mode='thinking')
-    generated = tokenizer.encode(encoded_str, add_special_tokens=False)
-    bos = tokenizer.bos_token_id or 0
-    if generated[0] != bos:
-        generated = [bos] + generated
-    
-    PREFILL_CHUNK = 128
-    n_prefill = len(generated)
-    prefill_ids = torch.tensor(generated, dtype=torch.long, device='cuda:0')
-    prefill_ids32 = prefill_ids.to(torch.int32)
-    all_positions = torch.arange(n_prefill, dtype=torch.long, device='cuda:0')
-    
-    chunk_starts = list(range(0, n_prefill, PREFILL_CHUNK))
-    for ci, cs in enumerate(chunk_starts):
-        ce = min(cs + PREFILL_CHUNK, n_prefill)
-        chunk_ids = prefill_ids[cs:ce]
-        chunk_ids32 = prefill_ids32[cs:ce]
-        chunk_positions = all_positions[cs:ce]
-        chunk_embed = embed(chunk_ids)
-        X = mHCLayer.init_state(chunk_embed)
-        
-        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, layer_w[li], li, cfg, *rope_caches[gpu],
-                              attn_mhcs.get(li), ffn_mhcs.get(li),
-                              attn_norms.get(li), ffn_norms.get(li),
-                              kv_caches[li], chunk_positions, chunk_ids32,
-                              compressors.get(li), indexers.get(li),
-                              moe_runners.get(li), se_runners.get(li), routers.get(li),
-                              prod_lin=prod_lins.get(li),
-                              comp_rope_cos=comp_rope_caches[gpu][0],
-                              comp_rope_sin=comp_rope_caches[gpu][1],
-                              )
-        X = X.to('cuda:0')
-        print(f"  Prefill chunk {ci+1}/{len(chunk_starts)}", flush=True)
-    
-    print("Prefill complete. Starting sync detection...", flush=True)
-    
-    # ---- NOW: Run one decode step under sync debug mode ----
-    all_tokens = generated.copy()
-    dec_tid_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
-    dec_pos_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
-    dec_tid32_buf = torch.zeros(1, dtype=torch.int32, device='cuda:0')
-    # Pinned CPU buffers for graph-capturable token/position transfer
-    dec_tid_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
-    dec_tid32_pinned = torch.zeros(1, dtype=torch.int32, device='cpu').pin_memory()
-    dec_pos_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
-    
-    def write_token_to_gpu(token_id, position):
-        """Write token/position to GPU buffers via pinned CPU (no CPU→GPU sync)."""
-        dec_tid_pinned[0] = token_id
-        dec_tid_buf.copy_(dec_tid_pinned)
-        dec_tid32_pinned[0] = token_id
-        dec_tid32_buf.copy_(dec_tid32_pinned)
-        dec_pos_pinned[0] = position
-        dec_pos_buf.copy_(dec_pos_pinned)
-    
-    # Warmup step first (so CuTeDSL kernels are compiled)
-    print("  Warmup decode step (compiling CuTeDSL kernels)...", flush=True)
-    write_token_to_gpu(all_tokens[-1], len(all_tokens) - 1)
-    X = mHCLayer.init_state(embed(dec_tid_buf))
-    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, layer_w[li], li, cfg, *rope_caches[gpu],
-                          attn_mhcs.get(li), ffn_mhcs.get(li),
-                          attn_norms.get(li), ffn_norms.get(li),
-                          kv_caches[li], dec_pos_buf, dec_tid32_buf,
-                          compressors.get(li), indexers.get(li),
-                          moe_runners.get(li), se_runners.get(li), routers.get(li),
-                          prod_lin=prod_lins.get(li),
-                          comp_rope_cos=comp_rope_caches[gpu][0],
-                          comp_rope_sin=comp_rope_caches[gpu][1],
-                          )
-    X = X.to('cuda:0')
-    torch.cuda.set_device(0)
-    torch.cuda.synchronize()
-    print("  Warmup done.", flush=True)
-    
-    # ==== METHOD 1: sync debug mode ====
-    print("\n  [METHOD 1] Enabling sync debug mode...", flush=True)
-    torch.cuda.set_sync_debug_mode("error")
-    
-    sync_errors = []
-    try:
-        detector.phase = "decode_forward"
-        write_token_to_gpu(all_tokens[-1], len(all_tokens) - 1)
-        
-        X = mHCLayer.init_state(embed(dec_tid_buf))
-        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, layer_w[li], li, cfg, *rope_caches[gpu],
-                              attn_mhcs.get(li), ffn_mhcs.get(li),
-                              attn_norms.get(li), ffn_norms.get(li),
-                              kv_caches[li], dec_pos_buf, dec_tid32_buf,
-                              compressors.get(li), indexers.get(li),
-                              moe_runners.get(li), se_runners.get(li), routers.get(li),
-                              prod_lin=prod_lins.get(li),
-                              comp_rope_cos=comp_rope_caches[gpu][0],
-                              comp_rope_sin=comp_rope_caches[gpu][1],
-                              )
-        X = X.to('cuda:0')
-        torch.cuda.set_device(0)
-        
-        # hc_head + norm + lm_head
-        x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
-        if final_norm_w is not None:
-            x_out = rmsnorm(x_out, final_norm_w)
-        logits = torch.nn.functional.linear(x_out, lm_w)
-        
-        # Sampling (argmax — this WILL sync, but it's outside the graph)
-        # We test the FORWARD only, not the sampling loop
-        print("  Forward completed under sync debug mode!", flush=True)
-    except RuntimeError as e:
-        err_str = str(e)
-        sync_errors.append(err_str)
-        print(f"\n  [SYNC VIOLATION CAUGHT] {err_str[:300]}", flush=True)
-        traceback.print_exc()
-    finally:
-        torch.cuda.set_sync_debug_mode("default")
-    
-    if not sync_errors:
-        print("  METHOD 1: No sync violations in forward (or they're hidden behind conditional branches)", flush=True)
-    else:
-        print(f"  METHOD 1: {len(sync_errors)} sync violation(s) found", flush=True)
-    
-    # ==== METHOD 2: CUDA graph capture attempt ====
-    print("\n  [METHOD 2] Attempting CUDA graph capture of decode forward...", flush=True)
-    
-    # Pre-allocate static I/O buffers
-    static_x_in = torch.zeros(1, 4, H, dtype=torch.bfloat16, device='cuda:0')
-    static_logits = torch.zeros(1, cfg.get("vocab_size", 129280), dtype=torch.bfloat16, device='cuda:0')
-    static_token = torch.zeros(1, dtype=torch.long, device='cuda:0')
-    static_token32 = torch.zeros(1, dtype=torch.int32, device='cuda:0')
-    static_pos = torch.zeros(1, dtype=torch.long, device='cuda:0')
-    
-    # Try to capture a single layer first (layer 0 on cuda:0)
-    print("  Attempting capture of L0 (cuda:0)...", flush=True)
-    li = 0
-    gpu = 0
-    capture_errors = []
-    
-    try:
-        g = torch.cuda.CUDAGraph()
-        torch.cuda.set_device(0)
-        
-        # Fill static buffers with current decode state (via pinned CPU — no sync)
-        dec_tid_pinned[0] = all_tokens[-1]
-        static_token.copy_(dec_tid_pinned)
-        dec_tid32_pinned[0] = all_tokens[-1]
-        static_token32.copy_(dec_tid32_pinned)
-        dec_pos_pinned[0] = len(all_tokens) - 1
-        static_pos.copy_(dec_pos_pinned)
-        
-        with torch.cuda.graph(g):
-            X = mHCLayer.init_state(embed(static_token))
-            X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu],
-                              attn_mhcs.get(li), ffn_mhcs.get(li),
-                              attn_norms.get(li), ffn_norms.get(li),
-                              kv_caches[li], static_pos, static_token32,
-                              compressors.get(li), indexers.get(li),
-                              moe_runners.get(li), se_runners.get(li), routers.get(li),
-                              prod_lin=prod_lins.get(li),
-                              comp_rope_cos=comp_rope_caches[gpu][0],
-                              comp_rope_sin=comp_rope_caches[gpu][1],
-                              )
-            static_x_in.copy_(X.to('cuda:0'))
-        
-        print("  L0 CAPTURED SUCCESSFULLY!", flush=True)
-    except Exception as e:
-        err_str = str(e)
-        capture_errors.append(err_str)
-        print(f"\n  [CAPTURE FAILURE] L0: {err_str[:500]}", flush=True)
-        traceback.print_exc()
-    
-    # ==== Summary ====
-    print("\n" + "=" * 70, flush=True)
-    print("SYNC INVENTORY SUMMARY", flush=True)
-    print("=" * 70, flush=True)
-    print(f"  Method 1 (sync debug): {len(sync_errors)} violations", flush=True)
-    print(f"  Method 2 (graph capture L0): {'PASS' if not capture_errors else 'FAIL'}", flush=True)
-    print(f"  Grep patterns: see above", flush=True)
-    print("=" * 70, flush=True)
-    
-    # Save results
-    results = {
-        "sync_debug_violations": sync_errors,
-        "graph_capture_errors": capture_errors,
-        "grep_results": "see stdout",
-    }
-    with open("/tmp/cuda_graph_readiness_results.json", "w") as f:
-        json.dump(results, f, indent=2)
-    print(f"Results saved to /tmp/cuda_graph_readiness_results.json", flush=True)
-
-
-if __name__ == "__main__":
-    source_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    
-    # First: grep for sync patterns
-    grep_sync_patterns(source_dir)
-    
-    # Then: run the forward under sync debug + capture attempt
-    run_sync_debug_mode()

tests/unit/test_cuda_graph_stream.py

@@ -1,78 +0,0 @@
-"""Minimal CUDA graph test with explicit stream management."""
-import torch
-
-def test_explicit_stream():
-    """Test CUDA graph with explicit per-device streams."""
-    results = {}
-    for gpu in range(8):
-        device = f'cuda:{gpu}'
-        
-        # Create a dedicated stream for this device
-        s = torch.cuda.Stream(device=device)
-        
-        # Create tensors on the correct device
-        x = torch.ones(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        y = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        
-        # Capture on the explicit stream
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g, stream=s):
-            y.copy_(x * 2.0)
-        
-        # Update input
-        x.fill_(3.0)
-        
-        # Replay on the SAME stream
-        with torch.cuda.stream(s):
-            g.replay()
-        
-        torch.cuda.synchronize()
-        y_max = y.abs().max().item()
-        expected = 6.0
-        status = "OK" if abs(y_max - expected) < 0.1 else f"WRONG (expected {expected}, got {y_max})"
-        results[gpu] = y_max
-        print(f"  GPU {gpu}: y_max={y_max:.2f} — {status}")
-    
-    return results
-
-def test_set_device_before_each_op():
-    """Test with explicit set_device before each operation."""
-    results = {}
-    for gpu in range(8):
-        torch.cuda.set_device(gpu)
-        device = f'cuda:{gpu}'
-        
-        x = torch.ones(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        y = torch.zeros(1, 4, 7168, dtype=torch.bfloat16, device=device)
-        
-        # Use default stream on the current device
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            # Explicitly set device INSIDE the graph capture
-            torch.cuda.set_device(gpu)
-            y.copy_(x * 2.0)
-        
-        # Update input
-        x.fill_(3.0)
-        
-        # Replay
-        torch.cuda.set_device(gpu)
-        g.replay()
-        torch.cuda.synchronize()
-        
-        y_max = y.abs().max().item()
-        expected = 6.0
-        status = "OK" if abs(y_max - expected) < 0.1 else f"WRONG (expected {expected}, got {y_max})"
-        results[gpu] = y_max
-        print(f"  GPU {gpu}: y_max={y_max:.2f} — {status}")
-    
-    return results
-
-if __name__ == "__main__":
-    print("=== Test with explicit stream ===")
-    test_explicit_stream()
-    
-    print("\n=== Test with set_device inside capture ===")
-    test_set_device_before_each_op()
-    
-    print("\nDone.")