nvfp4-megamoe-kernel/CURRENT_BUG.md

Current Bug: CuTeDSLMoERunner — Status & Debug History

Current Status (May 17, 2026 13:30 UTC)

vLLM container runs, cudagraph capture succeeds, but model output is garbage (BOS token repeat).

• ✅ layertest.py — 0.988 cosine
• ✅ cudagraph_test.py — capture + replay works
• ✅ Container builds, loads weights, warmup gs computed (no L2 gs=0)
• ✅ Cudagraph capture completes (51 sizes, ~15 min)
• ✅ Server accepts requests, generates tokens
• ❌ Model output is <｜begin▁of▁sentence｜> token repeated — garbage logits

Current theory: Scale assembly layout mismatch between the fixed 128-row-per-expert approach and what the GEMM actually expects. The latest fix pads slot_hidden to num_experts * 128 rows and passes padded_expert_offsets=[0, 128, 256, ...] to the GEMM. Build is in progress on B200 to test.

---

Bugs Found & Fixed

Bug 1: Scale Assembly — Global Swizzle vs Per-Expert Swizzle

Symptom: GEMM produced all zeros even with correct global_scale.

Root cause: The original _assemble_scales_cudagraph_safe called pad_and_swizzle_single() on the ENTIRE padded buffer (all experts concatenated). But the kernel expects each expert's 128-row block to be swizzled independently (matching assemble_scales_2d_side which pads+swizzles each expert separately before concatenation).

Fix: Two-phase approach:

1. Scatter x_sf rows into 128-aligned positions in a padded buffer (GPU-only, no CPU sync)
2. Per-expert: copy 128 rows from padded buffer, pad_and_swizzle_single() each expert's block independently, then concatenate

Bug 2: searchsorted(right=False) — Wrong Expert Assignment

Symptom: Scale data in wrong positions after scatter.

Root cause: torch.searchsorted([4, 8, 8], 4, right=False) returns 0, assigning row 4 (expert 1's first token) to expert 0.

Fix: Changed to right=True:

expert_assign = torch.searchsorted(expert_offsets[1:], row_indices, right=True)

Bug 3: CuTeDSL cute.compile GPU Memory Corruption — CRITICAL

Symptom: _token_indices was all zeros, making every token map to token 0.

Root cause: CuTeDSL's cute.compile (JIT compilation) corrupts GPU memory. Tensors allocated on GPU before or during JIT compilation get zeroed.

Fix: Allocate _token_indices with _fill_token_indices() which builds on CPU and copies to GPU. Added _needs_token_refill flag to handle GEMM JIT corruption on first call.

Bug 4: expert_offsets With Leading 0

Symptom: GEMM produced wrong output with correct scale data.

Root cause: The runner passed expert_offsets[:num_experts + 1] = [0, 4, 8, 8] (4 elements with leading 0) but the kernel expects [4, 8, 8] (cumulative sum without leading 0).

Fix: Pass expert_offsets[1:num_experts + 1] to the GEMM.

Bug 5: Checkpoint input_scale Is Wrong for Activation Global Scale

Symptom: Block scales all saturate at float8 max (448), producing garbage quantization.

Root cause: The checkpoint's input_scale (~0.000286) is a calibration value computed from a different input magnitude (amax ≈ 0.77) than what runtime produces (amax ≈ 8.17). Too-small gs → block scale overflow → garbage.

Fix: compute_activation_global_scales() warmup method that runs quantize_to_nvfp4 (dynamic gs with .max()) before cudagraph capture.

Bug 6: L1 and L2 Need Separate Activation Global Scales

Symptom: L2 output was garbage even with correct L1 gs.

Root cause: After SiLU(gate)*up, the activation has amax ~286. The L1 gs is 30x too small for L2.

Fix: compute_activation_global_scales() computes L2 gs from the actual L1 output (after SiLU*up).

Bug 7: L1 and L2 Need Separate Padded Scale Buffers

Symptom: IndexError when quantizing L2 activation — K_sf differs between L1 (448) and L2 (192).

Fix: Separate _padded_x_sf_buf_l1 and _padded_x_sf_buf_l2, plus separate _per_expert_scale_bufs_l1 and _per_expert_scale_bufs_l2.

Bug 8: Global→Local Expert ID Mismatch — CUDA_ERROR_ASSERT

Symptom: IndexKernel.cu:111 assertion failed, cascading into CUDA_ERROR_ASSERT (710) across all workers.

Root cause: With EP=8, topk_ids contains global expert IDs (0-255), but CuTeDSLMoERunner treated them as local IDs (0-31/48).

Fix: Added experts_start_idx param; in run(), remap global→local and mask non-local tokens.

Bug 8b: .cpu() Sync Breaking Cudagraph Compatibility

Fix: Moved _token_indices to GPU, added _fill_token_indices() (CPU→GPU copy), _needs_token_refill for GEMM JIT.

Bug 9: padded_x_sf Buffer Too Small — Index Out of Bounds

Symptom: IndexKernel.cu:111 OOB in scale assembly scatter. dst_rows exceeded buffer size.

Root cause: Buffer was sized for num_experts * 128 rows, but scatter positions were computed from actual token distribution (not fixed 128 per expert). With 8192 tokens and top_k=6, dst_rows could exceed 6144.

Fix (attempted): Sized buffer for max_num_tokens * top_k rows. Later reverted to num_experts * 128 with fixed 128-row-per-expert scatter layout.

Bug 10: Wrong top_k and max_num_tokens Defaults

Symptom: _token_indices max=6143 instead of 8191 (built with top_k=8, actual top_k=6).

Root cause: CuTeDSLMoERunner defaulted to max_num_tokens=8192, top_k=8, but vLLM uses top_k=6. deepseek_v4.py didn't pass these values.

Fix: Pass max_num_tokens and top_k from deepseek_v4.py to the runner constructor.

Bug 11: Full-Buffer Swizzle Produced Wrong GEMM Input

Symptom: L2 gs=0.0 on EP5/EP7 during warmup. Model produced BOS token.

Root cause: Applied the Blackwell 32_4_4 swizzle to the entire padded buffer at once, but the GEMM expects per-expert swizzled blocks. The combined swizzle layout doesn't match expert_offsets indexing.

Fix (in progress): Reverted to per-expert swizzle with fixed 128-row slots.

Bug 12: torch.full() During Cudagraph Capture

Symptom: cudaErrorStreamCaptureUnsupported on all 8 workers during cudagraph capture.

Root cause: torch.full() in run() allocates a new tensor during stream capture, which CUDA doesn't allow.

Fix: Pre-allocated _l1_gsa_buf and _l2_gsa_buf, use .fill_() instead of torch.full(). Also pre-allocated _output_buf, _row_indices_buf.

Bug 13: Warmup Passed Global Expert IDs Instead of Local

Symptom: L2 gs=0.0 on EP5/EP7 (all ranks except EP0).

Root cause: _warmup_activation_global_scales() passed global IDs (e.g. 336+) to compute_activation_global_scales(), which matches against expert_id_range (0..47). No tokens matched → zero L1 GEMM output → L2 gs=0.

Fix: Pass local expert IDs (0..num_experts-1) in warmup.

Bug 14 (CURRENT): GEMM Scale Layout Mismatch — 128-Row Fixed vs Variable

Symptom: Model generates BOS token repeatedly. Tokens are produced but logits are garbage.

Root cause: Scale assembly places data at fixed e*128 offsets (128 rows per expert). But the GEMM reads scale_a according to expert_offsets (real token counts, e.g. expert 0 = 500 rows). When expert 0 has 500 tokens, GEMM reads scale_a[0:500] but only rows 0-127 have valid scale data. Rows 128-499 are zeros → GEMM produces zeros for those tokens → garbage output.

Fix (in progress): Pad slot_hidden to num_experts * 128 rows (128 per expert) and pass padded_expert_offsets=[0, 128, 256, ...] to the GEMM. The GEMM processes exactly 128 tokens per expert. Padding tokens' output is discarded by scatter_add. Pre-allocated _padded_hidden_buf, _padded_activated_buf, _padded_expert_offsets_buf.

---

vLLM Integration Status

Component	Status	Notes
Weight loading	✅	Direct NVFP4 path, no BF16 round-trip
Weight stacking	✅	make_b_k_major + assemble_scales_3d_side
Global→local ID remap	✅	experts_start_idx, mask non-local tokens
Warmup gs computation	✅	Per-layer, local expert IDs, L1+L2 gs
Scale assembly	⚠️	128-row fixed layout, pending GEMM alignment fix
Cudagraph capture	✅	No dynamic allocations, no CPU syncs
Model output	❌	Garbage (BOS repeat) — scale/GEMM layout mismatch

---

Test Files

File	Purpose
tests/layertest.py	Reference: moe_pipeline with dynamic gs, 3 experts, layer 0. Must pass (≥0.98 cosine).
tests/cudagraph_test.py	CuTeDSLMoERunner cudagraph capture + replay. Must pass.
tests/test_warmup_gs.py	Warmup gs computation, safety margin sweep.
tests/test_runner_vs_pipeline.py	Compare runner.run() vs moe_pipeline.
tests/test_scale_assembly.py	Compare cudagraph-safe vs reference scale assembly.

Run order after any code change:

1. python3 tests/layertest.py — must pass
2. python3 tests/cudagraph_test.py — must pass

---

Key Architecture: CuTeDSL NVFP4 MoE

Data Flow

hidden_states (BF16) ──→ global→local remap ──→ sort by expert
    │
    ├── L1 (gate+up)
    │   quantize_activation_nvfp4 → x_fp4, x_sf
    │   _assemble_scales_cudagraph_safe → scale_a (swizzled)
    │   run_nvfp4_grouped_gemm → l1_out (BF16)
    │
    ├── SiLU(gate) * up → activated
    │
    ├── L2 (down)
    │   quantize_activation_nvfp4 → l2_x_fp4, l2_x_sf
    │   _assemble_scales_cudagraph_safe → scale_a (swizzled)
    │   run_nvfp4_grouped_gemm → l2_out (BF16)
    │
    └── scatter_add → y (BF16)

Cudagraph Constraints

• No .item(), .cpu(), .tolist() — zero CPU-GPU syncs
• No torch.zeros/ones/full/empty/arange during capture — pre-allocate everything
• No dynamic shapes — num_tokens equals the captured budget
• Per-expert Python loops are OK (fixed num_experts, unrolled at capture time)
• pad_and_swizzle_single is OK on pre-padded 128×4-aligned buffers (no internal allocation)

EP Configuration (DeepSeek-V4-Pro on 8×B200)

• 256 total experts, top_k=6
• EP=8 → 32 local experts per rank (in practice 48 based on logs)
• experts_start_idx = rank * 32 (0, 32, 64, ..., 224)
• max_num_tokens from scheduler_config.max_num_batched_tokens

---

Repo Info

• Kernel: sweetapi.com/biondizzle/nvfp4-megamoe-kernel (master)
• Local: ~/dev/nvfp4-megamoe-kernel/
• B200: /root/nvfp4-megamoe-kernel/
• Model: /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4 (read-only)
• Never edit on B200 directly — edit locally → commit → push → pull on B200