Files
nvfp4-megamoe-kernel/CURRENT_BUG.md

7.8 KiB
Raw Blame History

Current Bug: CuTeDSLMoERunner — Status & Debug History

Current Status (May 17, 2026 16:52 UTC)

vLLM runs, cudagraph capture succeeds, but model output is empty/invisible tokens (garbage logits). Going back to layer tests to debug GEMM output quality.

  • layertest.py — 0.988 cosine (with dynamic gs reference)
  • cudagraph_test.py — capture + replay works
  • Container builds, loads weights, warmup gs computed (no L2 gs=0)
  • With --gpu_memory_utilization=0.9 and max_cudagraph_capture_size=8, container starts and serves
  • Model output is empty content (30 tokens of invisible/BS token) — MoE GEMM output is wrong

Next step: Debug WHY the runner produces 0.988 cosine in layertest but garbage in vLLM. Likely issues:

  1. The warmup gs values (computed from random data) don't match real runtime activation magnitudes
  2. The scale assembly layout is subtly wrong for 48 experts vs 3 experts in the test
  3. The padded buffer scatter (clamped_local) is dropping real token data

Current vLLM launch config:

--gpu_memory_utilization=0.9
--compilation-config='{"cudagraph_mode": "FULL_DECODE_ONLY", "custom_ops": ["all"], "cudagraph_capture_sizes": [1, 2, 4, 8], "max_cudagraph_capture_size": 8}'

Bugs Found & Fixed (121)

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

Fix: Two-phase scatter + per-expert swizzle.

Bug 2: searchsorted(right=False)

Fix: Changed to right=True.

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

Symptom: _token_indices all zeros after JIT. Root cause: cute.compile corrupts GPU memory. Fix: _fill_token_indices() builds on CPU, copies to GPU. _needs_token_refill flag.

Bug 4: expert_offsets With Leading 0

Fix: Pass expert_offsets[1:] to GEMM.

Bug 5: Checkpoint input_scale Wrong for Runtime gs

Root cause: Calibration value, too-small gs → block scale overflow. Fix: compute_activation_global_scales() warmup method.

Bug 6: L1/L2 Need Separate gs

Fix: Compute L2 gs from L1 output after SiLU*up.

Bug 7: L1/L2 Need Separate Scale Buffers

Fix: Separate _padded_x_sf_buf_l1/_l2, separate per-expert bufs.

Bug 8: Global→Local Expert ID Mismatch — CUDA_ERROR_ASSERT

Root cause: topk_ids contains global IDs (0-255), runner treated as local. Fix: experts_start_idx, remap global→local, mask non-local tokens.

Bug 8b: .cpu() Sync Breaking Cudagraph

Fix: _token_indices on GPU, _fill_token_indices() CPU→GPU copy.

Bug 9: padded_x_sf Buffer Too Small

Fix: Iterative — see Bugs 14, 16.

Bug 10: Wrong top_k/max_num_tokens Defaults

Fix: Pass from deepseek_v4.py.

Bug 11: Full-Buffer Swizzle Wrong for GEMM

Fix: Per-expert swizzle.

Bug 12: torch.full() During Cudagraph Capture

Symptom: cudaErrorStreamCaptureUnsupported. Fix: Pre-allocated buffers, .fill_() instead of torch.full().

Bug 13: Warmup Passed Global Expert IDs

Symptom: L2 gs=0.0 on EP5/EP7. Fix: Pass local IDs.

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

Symptom: BOS token repeat (garbage logits). Root cause: Scale assembly at e*128 offsets, GEMM reads by real expert_offsets. Expert with 500 tokens → GEMM reads 500 scale rows but only 128 have data. Fix: Fixed-layout padding: each expert gets max_chunks * 128 rows. Pad slot_hidden. Pass padded_expert_offsets to GEMM. Extract via l1_out[padded_dst].

Bug 15: OOM — Per-Layer Padded Buffers (4.3 GB)

Root cause: padded_hidden_buf + padded_activated_buf at 72 MB × 60 layers. Fix: Shared buffers (Bug 21).

Bug 16: padded_max_slots Mismatch

Root cause: Sized for max_tokens*top_k but needed num_experts*max_chunks*128. Fix: Size correctly.

Bug 17: Shape Mismatch (49152 vs 3072)

Root cause: Cap max_num_tokens to 512 made buffers too small for 8192-token warmup. Fix: Reverted cap, use shared buffers (Bug 21).

Bug 1820: Cudagraph Capture Failures

Root cause: Dynamic tensor allocation (torch.zeros), variable-trip loops, GPU scalars in Python control flow. Fix: Pre-allocate everything, fixed loop counts, Python constants for offsets.

Bug 21: OOM (correct fix) — Shared Padded Buffers

Root cause: Per-layer allocation of padded buffers. Fix: Class-level shared buffers dict keyed by device. Layers execute sequentially → safe to share. Also shared padded_x_sf_buf and output_buf. Total ~150 MB instead of ~4.3 GB.


Current Architecture: Fixed-Layout Padding

Each expert gets max_chunks * 128 rows at offset (e * max_chunks * 128).

Scatter: padded_dst = expert_assign * max_rows_per_expert + clamped_local_row
GEMM input: padded_hidden (total = num_experts * max_chunks * 128 rows)
GEMM offsets: [0, max_rows, 2*max_rows, ...] (fixed, pre-computed)
GEMM output: same total rows
Extract: l1_out[padded_dst] → only real token rows

Scale assembly:
  Phase 1: Scatter x_sf into padded_x_sf at same fixed offsets
  Phase 2: Per-expert, per-chunk swizzle (fixed loop: max_chunks iterations)
  No dynamic tensor allocation, no GPU→CPU syncs

Shared buffers (class-level):
  padded_hidden, padded_activated, padded_xsf_l1, padded_xsf_l2, output
  ~150 MB total (not per-layer)

Cudagraph Constraints (All Resolved)

  • No .item(), .cpu(), .tolist()
  • No torch.zeros/ones/full/empty/arange() during capture
  • No dynamic Python control flow from GPU values
  • Per-expert Python loops OK (fixed num_experts)
  • Shared buffers OK (layers sequential during capture and replay)

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

  • 256 total experts, top_k=6
  • EP=8 → 48 local experts per rank
  • experts_start_idx = rank × 32
  • max_num_tokens = 8192
  • max_chunks_per_expert = ceil(8192 × 6 / (48 × 128)) = 8

Outstanding Issue: Garbage Model Output

Symptom: Model generates 30 tokens of empty/invisible content (BOS or thinking token). Not meaningful text.

What works: layertest gives 0.988 cosine with 3 experts, 8 tokens, top_k=8.

What doesn't: vLLM with 48 experts, variable tokens, top_k=6 produces garbage.

Hypotheses to investigate:

  1. Warmup gs from random data ≠ real activation magnitudes. The warmup uses torch.randn (amax ~3) but real activations have amax ~8-10. The gs values would be wrong, causing quantization errors.
  2. Scale assembly with 48 experts × 8 chunks. With max_chunks=8 and 48 experts, there are 384 swizzle blocks. The fixed-layout scatter with clamped_local may be dropping tokens that overflow the expert's max_rows section.
  3. clamped_local = local_row.clamp(max=max_rows_per_expert - 1). If an expert has more than max_chunks*128 real tokens, overflow tokens all map to the same row, overwriting each other. This silently drops data.
  4. The _needs_token_refill path. After GEMM JIT, _token_indices may get corrupted. The refill happens AFTER the first run, but the first run already used corrupted indices.

Test Files

File Purpose
tests/layertest.py Reference vs runner, 3 experts. Must pass ≥0.98 cosine.
tests/cudagraph_test.py Cudagraph capture + replay. Must pass.
tests/test_runner_vs_pipeline.py Runner vs pipeline comparison.
tests/test_scale_assembly.py Scale assembly comparison.
tests/test_warmup_gs.py Warmup gs computation.

Run order after any code change:

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

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