Files
nvfp4-megamoe-kernel/CURRENT_BUG.md

1.7 KiB

CURRENT_BUG.md

Status: CSA/HCA kernel works. Need vLLM integration.

What We Know

  • CuTeDSL NVFP4 kernels: All pass (cosine 0.988-0.999 vs BF16)
  • Warmup gs: Irrelevant (runner recomputes per-call)
  • CSA attention kernel (cutedsl/csa_attention.py): Works with PyTorch SDPA
  • Full layer 0 forward: CuTeDSL + SDPA = cosine 0.988 vs BF16
  • Logits: std=2.98, reasonable top-5 tokens

Root Cause of vLLM Empty Output

vLLM uses two compiled CUDA kernels that DON'T work on Blackwell (SM100):

  1. torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert — fused RoPE + KV cache
  2. FlashMLA sparse attention — the actual attention computation

The model uses CSA (Compressed Sparse Attention) + HCA (Heavily Compressed Attention), NOT MLA. vLLM misnames it "MLA" in code but the architecture is CSA/HCA with mHC.

Integration Plan

Replace vLLM's broken CUDA kernels in DeepseekV4MLAAttention.forward:

  1. Replace fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert → pure PyTorch RoPE + FP8 quant + cache insert
  2. Replace FlashMLA → our CSA/HCA kernel using PyTorch SDPA
  3. Keep the compressor (it's mostly Triton which may work on SM100)
  4. Keep the indexer (it calls into sparse_attn_indexer which is also Triton)

Test Results

test_full_layer_b200.py:
  q_a_proj:     0.995 ✅  kv_proj:      0.995 ✅  q_b_proj:     0.995 ✅
  wo_b_proj:    0.995 ✅  comp.kv_proj: 0.994 ✅  comp.gate:    0.995 ✅
  shared_expert: 0.990 ✅

test_model_forward_b200.py:
  Warmup gs is IRRELEVANT (10x change → cosine 0.9993)
  CuTeDSL cosine vs BF16: 0.999

test_csa_attention_b200.py:
  Full path CuTeDSL + SDPA vs BF16: 0.988 ✅
  Logit std: 2.98 ✅