nvfp4-megamoe-kernel/CURRENT_BUG.md

CURRENT_BUG.md

Status: Container starts, model generates tokens, but output is GARBAGE (empty/NaN)

THE FUNDAMENTAL PROBLEM

Mike was right — we need our own kernels. Not just for the NVFP4 GEMMs, but for the ENTIRE attention pipeline. The current approach of patching individual vLLM functions is a house of cards.

Here's what happened: we spent hours patching vLLM to "work" on Blackwell. We patched:

1. VLLM_NVFP4_GEMM_BACKEND=cutedsl → invalid, removed env var
2. KV cache page size assertion → patched kv_cache_utils.py
3. 91 missing compressor cache layers → patched alignment in 3 cache specs
4. softmax_scale AttributeError → fixed to self.scale
5. NaN from missing RoPE on KV → added _apply_rope_kv()
6. Shape mismatch in apply_gptj_rope → rewrote as inline RoPE

After ALL of that, the container starts and runs. But the model output is GARBAGE — empty strings, NaN logprobs, zero meaningful text. Because the attention pipeline is fundamentally broken:

• The KV cache is never written to (the C++ kernel is FlashMLA-only)
• The attention reads from raw projection output, not the cache
• The compressor/indexer runs but the Blackwell path doesn't integrate with it
• Everything is deeply coupled — patch one thing, three more break

THE ONLY FIX: Build CuTeDSL kernels for the ENTIRE attention pipeline.

Do NOT try to patch vLLM's FlashMLA code. Do NOT use pure PyTorch as a workaround. Do NOT skip the KV cache write. BUILD THE KERNELS.

Container Crash History (May 19)

Each crash was "fixed" with a patch. Each patch led to the next crash. This is the house of cards:

1. VLLM_NVFP4_GEMM_BACKEND=cutedsl — invalid choice in envs.py → removed env var
2. assert max(sm_page_sizes) <= max(all_page_sizes) — KV cache page size mismatch → patched kv_cache_utils.py
3. Some layers are not correctly initialized — 91 missing compressor cache layers (alignment=576 wrong on Blackwell) → patched SWA, indexer, compressor cache specs
4. AttributeError: softmax_scale — wrapper uses self.scale not self.softmax_scale → fixed
5. 200 GiB KV cache for 512 tokens → reduced max_model_len to 256, patched cache specs to remove FlashMLA alignment
6. NaN output (logprobs) → KV wasn't getting RoPE → added _apply_rope_kv()
7. Shape mismatch in apply_gptj_rope → rewrote as inline 2D RoPE
8. Garbage/empty output — the attention pipeline is fundamentally broken

What Actually Works (standalone B200 venv tests)

Every single kernel works when tested individually. The problem is ONLY in the vLLM integration.

Kernel	Test File	Result
CuTeDSL NVFP4 Linear	test_full_layer_b200.py	cosine 0.994+ ✅
CuTeDSL NVFP4 MoE	layertest.py	cosine 0.988 ✅
FP8 KV quantize/dequant	test_kv_cache_b200.py	cosine 0.9997 ✅
NVFP4 KV quantize/dequant	test_kv_cache_b200.py	cosine 0.9943 ✅
Paged KV cache read/write	test_kv_cache_b200.py	cosine 1.0 ✅
FP8 KV → full attention	test_kv_cache_b200.py	cosine 0.9997 ✅
CSA sparse attention (cr=4)	test_sparse_attn_b200.py	works, no NaN ✅
HCA sparse attention (cr=128)	test_sparse_attn_b200.py	works, no NaN ✅
Merged CSA+SWA attention	test_sparse_attn_b200.py	works, no NaN ✅
Full pipeline (all layer types)	test_v4_attention_b200.py	cosine 0.981-0.995 ✅
NVFP4 Q×K^T GEMM	test_nvfp4_attn_gemm_b200.py	cosine 0.86 ❌ (too lossy)

Key Lessons (READ THESE OR REPEAT THE SAME MISTAKES)

1. NVFP4 is NOT suitable for attention Q×K^T. The per-element dot products are too sensitive. Cosine 0.86. Keep attention in BF16, use NVFP4 only for weight GEMMs.

2. DeepSeek-V4 is NOT MLA. It uses CSA (Compressed Sparse Attention) + HCA (Heavily Compressed Attention). vLLM misnames everything "MLA" internally — don't be confused by class names like DeepseekV4MLAAttention.

3. The fp8_ds_mla format is FlashMLA-specific. 584 bytes per token (448 NoPE FP8 + 128 RoPE FP8 + 8 scale). This is NOT a standard fp8 tensor. You can't just view() it as [slot, 512] uint8.

4. The SWA cache, indexer cache, and compressor cache all use alignment=576 for FlashMLA. On Blackwell, this must be None (no FlashMLA). There are 4 separate classes that set this, and you must patch ALL of them.

5. DeepseekV4MultiHeadLatentAttentionWrapper registers ITSELF (not the inner MLA attention) in static_forward_context. The custom op deepseek_v4_attention looks up the wrapper. So attention_impl must be on the WRAPPER, and it must use self.scale (not self.softmax_scale).

6. The Triton compressor and indexer DO work on Blackwell. They're not the problem. The problem is that the Blackwell attention path doesn't integrate with them.

THE PLAN: Build CuTeDSL Attention Backend

STOP. Do NOT touch the vLLM container. Build and test kernels on the B200 venv first.

Step 1: KV Cache Write Kernel

• BF16 KV → apply RoPE → fp8 quantize → write to paged cache
• Test in tests/test_kv_cache_write_b200.py:
  • Write KV for N tokens, read it back, compare against BF16 reference
  • Must handle: slot mapping, block_size, fp8 per-token scale

Step 2: KV Cache Read Kernel

• Paged cache → fp8 dequant → BF16 KV with RoPE
• Test: write then read, cosine >= 0.99

Step 3: BF16 Attention Kernel

• Q (with RoPE) × K^T → causal mask → softmax → attn × V
• Keep in BF16 (NVFP4 too lossy for attention scores)
• Handle CSA sparse (gather top-k positions from compressed cache)
• Handle HCA sparse (gather from 1/128 positions)
• Handle SWA (sliding window, full causal within window)
• Test: compare against PyTorch SDPA, cosine >= 0.99

Step 4: Full Pipeline Integration

• KV cache read → attention → inverse RoPE → o_a BMM → o_b NVFP4 projection
• Wire CSA/HCA/SWA with sink weight merge
• Test: compare full pipeline against BF16 reference, cosine >= 0.98
• Test: run through ALL 61 layers, verify logits are reasonable (std between 0.5 and 50)

Step 5: vLLM Attention Backend

• Create a proper AttentionBackend subclass (e.g., CuTeDSLBlackwellBackend)
• Override DeepseekSparseSWABackend on Blackwell
• Handle metadata, slot mapping, cache format properly
• ONLY THEN test in the container

Step 6: Test in Container

• Build container with the new backend
• Test with real prompts
• If output is garbage, DO NOT declare success. Fix it.

Architecture: CSA + HCA + mHC (NOT MLA!)

• CSA (Compress Ratio 4): Compressed Sparse Attention — KV compressed 4x with overlap (coff=2). Indexer finds per-layer top-k.
• HCA (Compress Ratio 128): Heavily Compressed Attention — KV compressed 128x.
• mHC: Manifold-Constrained Hyper-Connections — replaces standard residual connections.
• SWA: Sliding Window Attention — local window (compress_ratio=0, last layer only)
• KV latent: (T, 512) shared across all 128 heads. After kv_norm, it's a single 512-dim vector per token.
• Sink weights: Merge sparse + SWA attention outputs. sigmoid(sinks) controls the mix.

B200 Server Details

Server: root@45.76.247.107 (password: <B200_PASSWORD>)
Repo: /root/nvfp4-megamoe-kernel (branch: proper-nvfp4-integration)
Model: /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4
Venv: /root/nvfp4-megamoe-kernel/tests/venv/
Test command: PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/<test>.py
Container: docker compose up (takes ~14 min to build, ~10 min to warm up)

PLEASE LISTEN

If you're the next session reading this: BUILD THE KERNELS. Test them standalone. Don't waste time patching vLLM. The container will still be there when you're ready. The kernels won't build themselves.

Mike said it best: "The only way to do this is to do our own kernels."

Just make the fucking kernel.