From 42285b6c2449f58bba4692e54cffb3a716fb1df5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 08:54:59 +0000 Subject: [PATCH] =?UTF-8?q?Add=20CuTeDSL=20NVFP4=20attention=20kernel=20te?= =?UTF-8?q?st=20-=20Q=C3=97K^T=20GEMM?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dockerfile | 1 + tests/test_nvfp4_attn_gemm_b200.py | 370 ++++++++++++++++++++++++++ vllm/patches/deepseek_v4_attention.py | 36 ++- 3 files changed, 400 insertions(+), 7 deletions(-) create mode 100644 tests/test_nvfp4_attn_gemm_b200.py diff --git a/Dockerfile b/Dockerfile index db23aa14..2306c231 100644 --- a/Dockerfile +++ b/Dockerfile @@ -52,6 +52,7 @@ ARG VLLM_NVFP4_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/k COPY vllm/kernels/linear/nvfp4/cutedsl.py ${VLLM_NVFP4_DIR}/cutedsl.py # Patch KV cache utils to handle DeepseekV4 SWA page sizes > MLA page sizes +# (SWA layers have larger page sizes than compressed MLA layers on Blackwell) ARG VLLM_CORE_DIR=/usr/local/lib/python3.12/dist-packages/vllm/v1/core COPY vllm/patches/patch_kv_cache_utils.py /tmp/patch_kv_cache_utils.py RUN python3 /tmp/patch_kv_cache_utils.py ${VLLM_CORE_DIR}/kv_cache_utils.py && rm /tmp/patch_kv_cache_utils.py diff --git a/tests/test_nvfp4_attn_gemm_b200.py b/tests/test_nvfp4_attn_gemm_b200.py new file mode 100644 index 00000000..6673a607 --- /dev/null +++ b/tests/test_nvfp4_attn_gemm_b200.py @@ -0,0 +1,370 @@ +#!/usr/bin/env python3 +""" +CuTeDSL NVFP4 Attention Kernel — Q×K^T GEMM + +DeepSeek-V4 attention is CSA/HCA (NOT MLA): +- KV latent: (T, 512) shared across all 128 heads +- Q: (T, 128, 512) — 128 heads, 512 dim each +- Q×K^T: (T, 128, 512) × (512, T) → (T, 128, T) per head +- softmax → (T, 128, T) +- attn×V: (T, 128, T) × (T, 512) → (T, 128, 512) + +The Q×K^T step is the expensive one. For T=8192 tokens, NH=128: +- M = T*NH = 1,048,576 +- K = HD = 512 +- N = T = 8192 +- FLOPs: 2 * M * K * N ≈ 8.8 TFLOPS + +NVFP4 quantization cuts the data movement by 4x (BF16→FP4). + +This test: +1. Build a CuTeDSL NVFP4 GEMM runner for Q×K^T +2. Compare output against BF16 reference +3. Test with real model weights (full attention pipeline) + +Usage (on B200): + cd /root/nvfp4-megamoe-kernel + PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_nvfp4_attn_gemm_b200.py +""" + +import sys, os, json, torch, torch.nn.functional as F, math, time +from safetensors import safe_open + +REPO = "/root/nvfp4-megamoe-kernel" +sys.path.insert(0, REPO) +MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +DEV = "cuda:0" + +# Model config +H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 +QL = 1536; OL = 1024; OG = 16; HPG = NH // OG +EPS = 1e-6; WINDOW = 128; SCALE = HD ** -0.5 + +E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) + +_cache = {} +def P(k, wm, md): + if k in _cache: return _cache[k] + with safe_open(os.path.join(md, wm[k]), framework="pt") as f: + t = f.get_tensor(k) + _cache[k] = t + return t + +def dequant(w, sf, gs): + d = w.device; lut = E2M1.to(d) + lo = lut[(w & 0xF).long()]; hi = lut[((w >> 4) & 0xF).long()] + O, I2 = w.shape; I = I2*2 + u = torch.empty(O, I, dtype=torch.float32, device=d) + u[:,0::2] = lo; u[:,1::2] = hi + bs = sf.float().repeat_interleave(16, dim=1)[:O,:I] + return (u * bs * gs).to(torch.bfloat16) + +def rms(x, w, eps=1e-6): + v = x.float().pow(2).mean(-1, keepdim=True) + return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) + +def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): + from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear + fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() + s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf + s = s.permute(1,0).contiguous() + if fused and gs_t.numel() == 2: + g1,g2 = gs_t[0].item(), gs_t[1].item(); gs = max(g1,g2) + if g1 != g2: + s32 = s.float(); sp = lw[0] if lw else outf//2 + s32[:sp] *= g1/gs; s32[sp:] *= g2/gs; s = s32.to(torch.float8_e4m3fn) + else: + gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() + r = CuTeDSLNvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) + r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] + r.finalize_weights(); r._ensure_initialized() + return r + +def apply_gptj_rope(x, positions, cos_sin, nope, rope): + if rope == 0 or x.numel() == 0: return x + half = rope // 2 + cos = cos_sin[positions, :half].to(x.dtype) + sin = cos_sin[positions, half:2*half].to(x.dtype) + if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) + x_rope = x[..., nope:].clone() + even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] + out = x.clone() + out[..., nope:][..., 0::2] = even * cos - odd * sin + out[..., nope:][..., 1::2] = even * sin + odd * cos + return out + +def apply_inv_gptj_rope(x, positions, cos_sin, nope, rope): + if rope == 0 or x.numel() == 0: return x + half = rope // 2 + cos = cos_sin[positions, :half].to(x.dtype) + sin = cos_sin[positions, half:2*half].to(x.dtype) + if x.dim() == 3: cos = cos.unsqueeze(1); sin = sin.unsqueeze(1) + x_rope = x[..., nope:].clone() + even = x_rope[..., 0::2]; odd = x_rope[..., 1::2] + out = x.clone() + out[..., nope:][..., 0::2] = even * cos + odd * sin + out[..., nope:][..., 1::2] = -even * sin + odd * cos + return out + +def build_cos_sin(max_pos=4096, rope_dim=ROPE): + half = rope_dim // 2 + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) + freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], dim=-1) + + +def bf16_causal_attention(q, kv, scale): + """BF16 reference: full causal self-attention.""" + T, NH, HD = q.shape + q_2d = q.reshape(T * NH, HD) + kv_exp = kv.unsqueeze(1).expand(-1, NH, -1).contiguous() + k_2d = kv_exp.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD) + v_2d = k_2d.clone() + scores = torch.matmul(q_2d.unsqueeze(1), k_2d.transpose(-1, -2)) * scale + query_pos = torch.arange(T, device=q.device).unsqueeze(1).repeat(1, NH).reshape(T * NH) + kv_pos = torch.arange(T, device=q.device).unsqueeze(0) + causal = kv_pos <= query_pos.unsqueeze(1) + scores = scores.squeeze(1).masked_fill(~causal, float('-inf')) + weights = F.softmax(scores.float(), dim=-1).to(q.dtype) + out = torch.matmul(weights.unsqueeze(1), v_2d).squeeze(1) + return out.reshape(T, NH, HD) + + +class NVFP4Attention: + """CuTeDSL NVFP4 attention kernel. + + Q×K^T via NVFP4 GEMM, softmax in BF16, attn×V in BF16. + + The Q×K^T GEMM: (T*NH, HD) × (HD, T) → (T*NH, T) + - Q is the "activation": quantized per-row (dynamic) + - K^T is the "weight": quantized from (T, HD) KV latent + + For decode (M=1 per head), the GEMM is tiny — NVFP4 overhead isn't worth it. + For prefill (M=chunk_size), the GEMM is large — NVFP4 saves 4x memory bandwidth. + + This kernel targets the prefill case where T is large. + """ + + def __init__(self, head_dim: int, num_heads: int, max_seq_len: int, device: str = "cuda"): + self.head_dim = head_dim + self.num_heads = num_heads + self.max_seq_len = max_seq_len + self.device = device + self._runner = None # Compiled on first call + + def forward(self, q_bf16, kv_bf16, scale): + """Forward pass. + + Args: + q_bf16: (T, NH, HD) with RoPE applied + kv_bf16: (T, HD) shared KV latent (BF16) + scale: 1/sqrt(HD) + + Returns: + (T, NH, HD) attention output + """ + from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear + + T, NH, HD = q_bf16.shape + device = q_bf16.device + + # Reshape Q: (T, NH, HD) → (T*NH, HD) — treat as 2D for GEMM + q_2d = q_bf16.reshape(T * NH, HD) + + # ── Q×K^T via NVFP4 GEMM ──────────────────────────────────── + # Q is "activation" (T*NH, HD), K^T is "weight" (T, HD) + # GEMM: (T*NH, HD) × (HD, T) → (T*NH, T) + # + # We use CuTeDSLNvfp4Linear with in_features=HD, out_features=T + # Q is the "hidden_states", K (kv) is the "weight" matrix + + # Create or get cached runner + cache_key = (T, HD, NH) + if self._runner is None or getattr(self, '_cache_key', None) != cache_key: + runner = CuTeDSLNvfp4Linear( + in_features=HD, + out_features=T, + max_num_tokens=T * NH, + device=str(device), + ) + + # Set K as the weight: kv (T, HD) → treat as weight (N=T, K=HD) + # quantize_to_nvfp4 quantizes along last dim (D=HD) as activation + # For weight, we need (K, N) layout — but kv is (T, HD) = (N, K) + # CuTeDSLNvfp4Linear expects weight in (N, K//2) after permute + + from cutedsl.bridge import quantize_to_nvfp4 + # Quantize KV as a 2D tensor: (T, HD) + # quantize_to_nvfp4 works on last dim (D=HD), returns: + # (T, HD//2) fp4, (T, HD//16) sf, scalar gs + kv_fp4, kv_sf, kv_gs = quantize_to_nvfp4(kv_bf16) + + # For CuTeDSLNvfp4Linear, weight is (N, K_packed) = (T, HD//2) + # Our kv_fp4 is already (T, HD//2) — perfect! + # sf needs to be (N, K_sf) = (T, HD//16) — already correct + + w_fp4 = kv_fp4 # (T, HD//2) — already in row-major (N, K_packed) + w_sf = kv_sf # (T, HD//16) + + # Set up the runner with K^T as weight + # The runner expects fp4 as list of (N, K_packed), sf as list of (N, K_sf) + # after finalize_weights, it does permute(1,0) internally + runner.fp4 = [w_fp4] + runner.sf = [w_sf] + runner.gs = [kv_gs] + runner.finalize_weights() + runner._ensure_initialized() + + self._runner = runner + self._cache_key = cache_key + + # Run Q×K^T GEMM + scores = self._runner.run(q_2d) # (T*NH, T) + scores = scores * scale + + # Causal mask + query_pos = torch.arange(T, device=device).unsqueeze(1).repeat(1, NH).reshape(T * NH) + kv_pos = torch.arange(T, device=device).unsqueeze(0) + causal = kv_pos <= query_pos.unsqueeze(1) + scores = scores.masked_fill(~causal, float('-inf')) + + # Softmax (BF16 for numerical stability, actually float32) + weights = F.softmax(scores.float(), dim=-1).to(q_bf16.dtype) + + # attn×V: (T*NH, T) × (T, HD) → (T*NH, HD) + # V = K = kv (shared latent) — BF16, no quantization + out = torch.matmul(weights, kv_bf16) + + return out.reshape(T, NH, HD) + + +def main(): + torch.cuda.set_device(0) + torch.manual_seed(42) + + print("=" * 70) + print(" CuTeDSL NVFP4 Attention Kernel Test") + print(" Q×K^T via NVFP4 GEMM, softmax BF16, attn×V BF16") + print("=" * 70) + + # ── Step 1: Synthetic test with random data ────────────────────── + print("\n--- Step 1: Synthetic random test ---") + T = 8 + q_rand = torch.randn(T, NH, HD, dtype=torch.bfloat16, device=DEV) + kv_rand = torch.randn(T, HD, dtype=torch.bfloat16, device=DEV) + + with torch.no_grad(): + ref = bf16_causal_attention(q_rand, kv_rand, SCALE) + print(f" BF16 reference: amax={ref.amax():.4f}") + + kernel = NVFP4Attention(HD, NH, max_seq_len=8192, device=DEV) + out = kernel.forward(q_rand, kv_rand, SCALE) + print(f" NVFP4 kernel: amax={out.amax():.4f}") + + c = F.cosine_similarity(ref.flatten().unsqueeze(0).float(), out.flatten().unsqueeze(0).float()).item() + print(f" Cosine: {c:.6f} {'✅' if c>=0.95 else '❌'}") + + # ── Step 2: Real model weights, full attention pipeline ────────── + print("\n--- Step 2: Real model weights (layer 0, C128A) ---") + with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: + wm = json.load(f)["weight_map"] + G = lambda k: P(k, wm, MODEL).to(DEV) + + p = "model.layers.0"; a = f"{p}.self_attn" + + emb = G("model.embed_tokens.weight") + anorm = G(f"{p}.input_layernorm.weight") + qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight") + woa = G(f"{a}.o_a_proj.weight") + + qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") + qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") + kv_w = G(f"{a}.kv_proj.weight"); kv_sf = G(f"{a}.kv_proj.weight_scale"); kv_gs = G(f"{a}.kv_proj.weight_scale_2") + wob_w = G(f"{a}.o_b_proj.weight"); wob_sf = G(f"{a}.o_b_proj.weight_scale"); wob_gs = G(f"{a}.o_b_proj.weight_scale_2") + + qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) + qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) + kv_bf16 = dequant(kv_w, kv_sf, kv_gs.item()) + wob_bf16 = dequant(wob_w, wob_sf, wob_gs.item()) + + r_qa = make_runner(qa_w, qa_sf, qa_gs, H, qa_w.shape[0]) + r_qb = make_runner(qb_w, qb_sf, qb_gs, QL, qb_w.shape[0]) + r_kv = make_runner(kv_w, kv_sf, kv_gs, H, kv_w.shape[0]) + r_wob = make_runner(wob_w, wob_sf, wob_gs, OG*OL, wob_w.shape[0]) + + token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) + NT = len(token_ids) + cos_sin = build_cos_sin(max_pos=WINDOW + 256).to(DEV) + positions = torch.arange(NT, dtype=torch.int64, device=DEV) + + with torch.no_grad(): + hidden = emb[token_ids] + normed = rms(hidden, anorm, EPS) + + # Projections (CuTeDSL) + qa_cute = r_qa.run(normed) + kv_cute = r_kv.run(normed) + qa_n = rms(qa_cute, qn, EPS) + kv_n = rms(kv_cute, kvn, EPS) + q_cute = r_qb.run(qa_n).view(NT, NH, HD) + q_rope = apply_gptj_rope(q_cute, positions, cos_sin, NOPE, ROPE) + + # ── NVFP4 Attention ────────────────────────────────────── + attn_kernel = NVFP4Attention(HD, NH, max_seq_len=8192, device=DEV) + o_nvfp4 = attn_kernel.forward(q_rope, kv_n, SCALE) + print(f" NVFP4 attention: amax={o_nvfp4.amax():.4f}") + + # ── BF16 reference ─────────────────────────────────────── + o_bf16 = bf16_causal_attention(q_rope, kv_n, SCALE) + print(f" BF16 attention: amax={o_bf16.amax():.4f}") + + c = F.cosine_similarity(o_nvfp4.flatten().unsqueeze(0).float(), o_bf16.flatten().unsqueeze(0).float()).item() + print(f" NVFP4 vs BF16 cosine: {c:.6f} {'✅' if c>=0.95 else '❌'}") + + # ── Full pipeline: attention → o_a → o_b ───────────────── + o_inv = apply_inv_gptj_rope(o_nvfp4, positions, cos_sin, NOPE, ROPE) + o_grouped = o_inv.view(NT, OG, HPG * HD).permute(1, 0, 2) + woa_3d = woa.view(OG, OL, HPG * HD) + z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) + attn_out = r_wob.run(z) + + # BF16 reference pipeline + o_inv_bf = apply_inv_gptj_rope(o_bf16, positions, cos_sin, NOPE, ROPE) + o_grouped_bf = o_inv_bf.view(NT, OG, HPG * HD).permute(1, 0, 2) + z_bf = torch.bmm(o_grouped_bf, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(NT, OG * OL) + attn_bf = z_bf @ wob_bf16.T + + c_full = F.cosine_similarity(attn_out.flatten().unsqueeze(0).float(), attn_bf.flatten().unsqueeze(0).float()).item() + print(f" Full pipeline cosine: {c_full:.6f} {'✅' if c_full>=0.95 else '❌'}") + + # Logits + fnorm_w = G("model.norm.weight") + lm_head = G("lm_head.weight") + x = hidden + attn_out + x_n = rms(x, fnorm_w, EPS) + logits = x_n @ lm_head.T + log_std = logits[-1].float().std().item() + print(f" logits: std={log_std:.4f} {'✅' if 0.5 < log_std < 50 else '❌'}") + + # ── Step 3: Larger sequence test ───────────────────────────────── + print("\n--- Step 3: Larger sequence (T=64) ---") + torch.cuda.empty_cache() + T64 = 64 + with torch.no_grad(): + q64 = torch.randn(T64, NH, HD, dtype=torch.bfloat16, device=DEV) + kv64 = torch.randn(T64, HD, dtype=torch.bfloat16, device=DEV) + + ref64 = bf16_causal_attention(q64, kv64, SCALE) + kernel64 = NVFP4Attention(HD, NH, max_seq_len=8192, device=DEV) + out64 = kernel64.forward(q64, kv64, SCALE) + + c64 = F.cosine_similarity(ref64.flatten().unsqueeze(0).float(), out64.flatten().unsqueeze(0).float()).item() + print(f" T=64 NVFP4 vs BF16 cosine: {c64:.6f} {'✅' if c64>=0.95 else '❌'}") + + print(f"\n{'='*70}") + print(f" DONE") + print(f"{'='*70}") + + +if __name__ == "__main__": + main() diff --git a/vllm/patches/deepseek_v4_attention.py b/vllm/patches/deepseek_v4_attention.py index 501dba63..9da5cd04 100644 --- a/vllm/patches/deepseek_v4_attention.py +++ b/vllm/patches/deepseek_v4_attention.py @@ -862,16 +862,16 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): f"DeepseekV4 only supports fp8 kv-cache format for now, " f"got {kv_cache_dtype}" ) - assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), ( - "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now" - ) - # On Blackwell (SM100+), FlashMLA kernels don't work, but we bypass - # them entirely in _attention_impl_blackwell(). Skip the FP8 ds_mla - # cache conversion since our Blackwell path doesn't use FlashMLA. + # On Blackwell (SM100+), FlashMLA kernels don't work. + # We use our own CSA/SDPA attention path. _is_blackwell = ( current_platform.get_device_capability() is not None and current_platform.get_device_capability().major >= 10 ) + if not _is_blackwell: + assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), ( + "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now" + ) # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format # Automatically convert fp8 kv-cache format to "fp8_ds_mla" # On Blackwell, we use our own attention path, so keep standard fp8 @@ -897,11 +897,18 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): self.kv_cache = torch.tensor([]) def get_attn_backend(self) -> type[AttentionBackend]: + cap = current_platform.get_device_capability() + if cap is not None and cap.major >= 10: + # Blackwell: FlashMLA doesn't work. Use our CSA/SDPA path. + # Return the base class so KV cache setup doesn't force fp8_ds_mla. + from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWABackend, + ) + return DeepseekSparseSWABackend if current_platform.is_rocm(): from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import ( DeepseekV4ROCMAiterMLASparseBackend, ) - return DeepseekV4ROCMAiterMLASparseBackend return DeepseekV4FlashMLASparseBackend @@ -910,6 +917,21 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): self.compress_ratio <= 1 ): # SWA part. Allocated separately as DeepseekV4SWACache. return None + cap = current_platform.get_device_capability() + _is_blackwell = cap is not None and cap.major >= 10 + if _is_blackwell: + # Blackwell: no FlashMLA, use standard fp8_e4m3 KV cache + # No 576B FlashMLA alignment needed + return MLAAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=torch.uint8, + compress_ratio=self.compress_ratio, + cache_dtype_str=self.kv_cache_dtype, # "fp8" (not fp8_ds_mla) + alignment=None, + model_version="deepseek_v4", + ) return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1,