biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add blackwell_attention module and comprehensive test

Browse Source
This commit is contained in:
biondizzle
2026-05-19 15:30:29 +00:00
parent 142a4a1ad4
commit 0023fee706
2 changed files with 565 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

247
cutedsl/blackwell_attention.py Normal file
View File

@@ -0,0 +1,247 @@
"""
DeepSeek-V4 Blackwell Attention — Our own kernel.
Replaces vLLM's broken FlashMLA Blackwell path with a proper KV cache-based
attention pipeline. Does NOT depend on FlashMLA, fp8_ds_mla, or any vLLM
fused CUDA kernel.
Architecture:
- KV: (T, HD=512) single head latent, shared across all 128 Q heads
- KV Cache: fp8_e4m3 paged cache with per-token inverse scale
- RoPE: GPT-J style, applied to Q and KV before caching
- Attention: BF16 (NVFP4 is too lossy for Q×K^T, cosine 0.86)
- CSA/HCA: Compressed KV for sparse attention (compress_ratio 4 or 128)
- SWA: Sliding window attention (compress_ratio 0/1)
Pipeline:
Prefill:
1. hidden → q_a_proj → q_norm → q_b_proj → (T, NH, HD) → RoPE on Q
2. hidden → kv_proj → kv_norm → (T, HD) → RoPE → fp8 quant → write to paged cache
3. Read all cached KV → BF16 causal attention → output
Decode:
1. Same projections as prefill
2. Write new KV to cache
3. Read ALL cached KV → BF16 attention (1 query vs N KVs) → output
Output:
1. inverse RoPE on attention output
2. o_a: BMM with wo_a (BF16)
3. o_b: NVFP4 GEMM with wo_b
"""
import torch
import torch.nn.functional as F
def apply_gptj_rope(x, positions, cos_sin_cache, nope_dim, rope_dim):
"""Apply GPT-J style RoPE. Works on (T, HD) or (T, NH, HD)."""
if rope_dim == 0 or x.numel() == 0:
return x
half = rope_dim // 2
cos = cos_sin_cache[positions, :half].to(x.dtype)
sin = cos_sin_cache[positions, half:2 * half].to(x.dtype)
if x.dim() == 3:
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
x_rope = x[..., nope_dim:].clone()
even = x_rope[..., 0::2]
odd = x_rope[..., 1::2]
out = x.clone()
out[..., nope_dim:][..., 0::2] = even * cos - odd * sin
out[..., nope_dim:][..., 1::2] = even * sin + odd * cos
return out
def apply_inv_gptj_rope(x, positions, cos_sin_cache, nope_dim, rope_dim):
"""Inverse GPT-J RoPE (sin → -sin)."""
if rope_dim == 0 or x.numel() == 0:
return x
half = rope_dim // 2
cos = cos_sin_cache[positions, :half].to(x.dtype)
sin = cos_sin_cache[positions, half:2 * half].to(x.dtype)
if x.dim() == 3:
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
x_rope = x[..., nope_dim:].clone()
even = x_rope[..., 0::2]
odd = x_rope[..., 1::2]
out = x.clone()
out[..., nope_dim:][..., 0::2] = even * cos + odd * sin
out[..., nope_dim:][..., 1::2] = -even * sin + odd * cos
return out
# ── KV Cache Operations ──────────────────────────────────────────────
def kv_quantize_fp8(kv_bf16):
"""BF16 KV → fp8_e4m3 with per-token inverse scale."""
amax = kv_bf16.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
fp8_max = torch.tensor(448.0, dtype=torch.float32, device=kv_bf16.device)
scale = fp8_max / amax
kv_fp8 = (kv_bf16.float() * scale).to(torch.float8_e4m3fn)
inv_scale = (amax / fp8_max).to(torch.bfloat16)
return kv_fp8, inv_scale
def kv_dequantize_fp8(kv_fp8, inv_scale):
"""fp8 KV → BF16."""
return (kv_fp8.to(torch.bfloat16) * inv_scale).to(torch.bfloat16)
def paged_kv_write(kv_data, slot_mapping, cache, block_size):
"""Write KV into paged cache. Works for fp8 or bf16.
kv_data: (T, D) tensor to write
slot_mapping: (T,) slot indices
cache: (num_blocks, block_size, D) cache tensor
"""
for t in range(kv_data.shape[0]):
slot = slot_mapping[t].item()
block_idx = slot // block_size
offset = slot % block_size
if block_idx < cache.shape[0] and offset < cache.shape[1]:
cache[block_idx, offset] = kv_data[t]
def paged_kv_read(slot_mapping, cache, block_size, num_tokens, head_dim):
"""Read KV from paged cache."""
device = cache.device
kv = torch.zeros(num_tokens, head_dim, dtype=cache.dtype, device=device)
for t in range(num_tokens):
slot = slot_mapping[t].item()
block_idx = slot // block_size
offset = slot % block_size
if block_idx < cache.shape[0] and offset < cache.shape[1]:
kv[t] = cache[block_idx, offset]
return kv
# ── Attention ─────────────────────────────────────────────────────────
def causal_prefill_attention(q, kv, scale):
"""Full causal self-attention for prefill. q: (T, NH, HD), kv: (T, HD)."""
T, NH, HD = q.shape
q_t = q.permute(1, 0, 2) # (NH, T, HD)
kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) # (NH, T, HD)
out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=scale)
return out.permute(1, 0, 2) # (T, NH, HD)
def decode_attention(q, kv, scale):
"""Decode attention: 1 query vs N cached KVs.
q: (1, NH, HD) — single decode token
kv: (N, HD) — all cached KV (already with RoPE)
"""
NH = q.shape[1]
HD = q.shape[2]
q_t = q.permute(1, 0, 2) # (NH, 1, HD)
kv_exp = kv.unsqueeze(0).expand(NH, -1, -1) # (NH, N, HD)
out = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=False, scale=scale)
return out.permute(1, 0, 2) # (1, NH, HD)
def swa_attention(q, kv, positions, scale, window_size):
"""Sliding window attention.
q: (T, NH, HD) with RoPE
kv: (total_len, HD) — ALL cached KV with RoPE
positions: (T,) — absolute positions of the query tokens
"""
T, NH, HD = q.shape
total_len = kv.shape[0]
output = torch.zeros_like(q)
for t in range(T):
pos = positions[t].item()
window_start = max(0, pos - window_size + 1)
window_len = pos - window_start + 1
if window_len <= 0:
continue
kv_window = kv[window_start:pos + 1] # (window_len, HD)
q_t = q[t:t + 1] # (1, NH, HD)
output[t] = decode_attention(q_t, kv_window, scale).squeeze(0)
return output
# ── Full Pipeline ─────────────────────────────────────────────────────
def blackwell_attention_forward(
# Inputs
q, # (T, NH, HD) with RoPE already applied
kv, # (T, HD) kv_normed, RoPE'd — the NEW tokens' KV
positions, # (T,) absolute positions
# KV Cache
swa_kv_cache, # (num_blocks, block_size, HD) fp8 paged cache
swa_inv_scale, # (num_blocks * block_size, 1) per-token inv scale
slot_mapping, # (T,) slot indices for writing
block_size, # tokens per block
seq_lens, # (num_seqs,) total sequence lengths (prefill + history)
num_prefills, # number of prefill sequences
num_decode_tokens, # number of decode tokens
# Params
scale, # 1/sqrt(HD)
nope_dim, # 448
rope_dim, # 64
window_size, # 128
compress_ratio, # 0, 1, 4, or 128
cos_sin_cache, # (max_pos, rope_dim) for RoPE
attn_sink, # (NH,) sink weights
):
"""Full attention forward for Blackwell (SM100+).
This is what replaces vLLM's _attention_impl_blackwell.
Steps:
1. Quantize + write new KV to paged cache
2. Read ALL cached KV for each sequence
3. Attention (prefill: causal, decode: full)
4. Return attention output (T, NH, HD)
"""
T = q.shape[0]
NH = q.shape[1]
HD = q.shape[2]
device = q.device
# Step 1: Quantize new KV and write to cache
# kv already has RoPE applied (done by caller)
kv_fp8, kv_inv_scale = kv_quantize_fp8(kv)
paged_kv_write(kv_fp8, slot_mapping, swa_kv_cache, block_size)
# Write inv_scale to flat cache
for t in range(T):
slot = slot_mapping[t].item()
swa_inv_scale[slot] = kv_inv_scale[t]
# Step 2 & 3: Read cached KV and attend
# For simplicity in this initial version, we separate prefill and decode
output = torch.zeros(T, NH, HD, dtype=torch.bfloat16, device=device)
if num_decode_tokens > 0:
# Decode tokens: each needs ALL prior KV from cache
for t in range(num_decode_tokens):
pos = positions[t].item()
# Read all KV from position 0 to pos
all_slots = torch.arange(pos + 1, dtype=torch.int64, device=device)
kv_cached_fp8 = paged_kv_read(all_slots, swa_kv_cache, block_size, pos + 1, HD)
kv_inv_scales = swa_inv_scale[all_slots]
kv_cached = kv_dequantize_fp8(kv_cached_fp8, kv_inv_scales)
# Apply SWA window
window_start = max(0, pos - window_size + 1)
kv_window = kv_cached[window_start:]
q_t = q[t:t + 1] # (1, NH, HD)
output[t] = decode_attention(q_t, kv_window, scale).squeeze(0)
if num_prefills > 0:
# Prefill tokens: causal attention using the NEW kv (not from cache,
# since all KV is available from the current forward pass)
# But we DO write to cache for future decode steps
prefill_slice = slice(num_decode_tokens, T)
output[prefill_slice] = causal_prefill_attention(
q[prefill_slice], kv[prefill_slice], scale
)
return output

318
tests/test_blackwell_attn_b200.py Normal file
View File

@@ -0,0 +1,318 @@
#!/usr/bin/env python3
"""
DeepSeek-V4 Blackwell Attention — Full Pipeline Test
Tests the cutedsl.blackwell_attention module with real weights:
1. Prefill: process N tokens, write KV to paged cache
2. Decode: process 1 new token, read ALL cached KV, attend
3. Verify decode output matches BF16 reference
This is the core of the fix for the vLLM Blackwell garbage output bug.
Usage (on B200):
cd /root/nvfp4-megamoe-kernel
PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_blackwell_attn_b200.py
"""
import sys, os, json, torch, torch.nn.functional as F, math
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"
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 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(), freqs.cos(), freqs.sin()], dim=-1) # extra for safety
# Only use the first rope_dim cols
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 test_blackwell_attention(layer_id, compress_ratio):
"""Test the full blackwell attention pipeline for a specific layer."""
from cutedsl.blackwell_attention import (
apply_gptj_rope, apply_inv_gptj_rope,
blackwell_attention_forward,
kv_quantize_fp8, kv_dequantize_fp8,
paged_kv_write, paged_kv_read,
causal_prefill_attention, decode_attention,
)
torch.cuda.set_device(0)
torch.manual_seed(42)
torch.cuda.empty_cache()
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 = f"model.layers.{layer_id}"; a = f"{p}.self_attn"
layer_type = "SWA" if compress_ratio <= 1 else f"CSA(c={compress_ratio})"
print(f"\n{'='*70}")
print(f" Layer {layer_id} — {layer_type} — Blackwell Attention Test")
print(f"{'='*70}")
# Load weights
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")
sinks = G(f"{a}.sinks")
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")
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])
cos_sin = build_cos_sin(max_pos=4096).to(DEV)
# ── Test 1: Prefill-only attention ────────────────────────────────
print(f"\n --- Test 1: Prefill attention (8 tokens) ---")
N = 8
token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374, 2198, 643], dtype=torch.long, device=DEV)
positions = torch.arange(N, dtype=torch.int64, device=DEV)
with torch.no_grad():
hidden = emb[token_ids]
normed = rms(hidden, anorm, EPS)
qa = r_qa.run(normed)
kv = r_kv.run(normed)
qa_n = rms(qa, qn, EPS)
kv_n = rms(kv, kvn, EPS)
q = r_qb.run(qa_n).view(N, NH, HD)
q_rope = apply_gptj_rope(q, positions, cos_sin, NOPE, ROPE)
kv_rope = apply_gptj_rope(kv_n.unsqueeze(1), positions, cos_sin, NOPE, ROPE).squeeze(1)
# Causal attention
o_prefill = causal_prefill_attention(q_rope, kv_rope, SCALE)
print(f" Prefill attention output: amax={o_prefill.amax():.4f} NaN={torch.isnan(o_prefill).any()}")
# BF16 reference (same computation, different path)
q_t = q_rope.permute(1, 0, 2)
kv_exp = kv_rope.unsqueeze(0).expand(NH, -1, -1)
o_ref = F.scaled_dot_product_attention(q_t, kv_exp, kv_exp, is_causal=True, scale=SCALE).permute(1, 0, 2)
c = F.cosine_similarity(o_prefill.flatten().unsqueeze(0).float(), o_ref.flatten().unsqueeze(0).float()).item()
print(f" Prefill vs SDPA reference cosine: {c:.6f} {'✅' if c>=0.999 else '❌'}")
# ── Test 2: Decode attention with KV cache ────────────────────────
print(f"\n --- Test 2: Decode attention (1 token, 8 cached) ---")
block_size = 256
num_blocks = 64
kv_cache_fp8 = torch.zeros(num_blocks, block_size, HD, dtype=torch.float8_e4m3fn, device=DEV)
inv_scale_cache = torch.zeros(num_blocks * block_size, 1, dtype=torch.bfloat16, device=DEV)
with torch.no_grad():
# Write prefill KV to cache
kv_fp8, inv_s = kv_quantize_fp8(kv_rope)
prefill_slots = positions
paged_kv_write(kv_fp8, prefill_slots, kv_cache_fp8, block_size)
for t in range(N):
inv_scale_cache[prefill_slots[t]] = inv_s[t]
# Decode: token at position 8
decode_id = torch.tensor([991], dtype=torch.long, device=DEV)
decode_pos = torch.tensor([N], dtype=torch.int64, device=DEV)
hidden_d = emb[decode_id]
normed_d = rms(hidden_d, anorm, EPS)
qa_d = r_qa.run(normed_d)
kv_d = r_kv.run(normed_d)
qa_n_d = rms(qa_d, qn, EPS)
kv_n_d = rms(kv_d, kvn, EPS)
q_d = r_qb.run(qa_n_d).view(1, NH, HD)
q_rope_d = apply_gptj_rope(q_d, decode_pos, cos_sin, NOPE, ROPE)
kv_rope_d = apply_gptj_rope(kv_n_d.unsqueeze(1), decode_pos, cos_sin, NOPE, ROPE).squeeze(1)
# Write decode KV to cache
kv_fp8_d, inv_s_d = kv_quantize_fp8(kv_rope_d)
paged_kv_write(kv_fp8_d, decode_pos, kv_cache_fp8, block_size)
inv_scale_cache[decode_pos[0]] = inv_s_d[0]
# Read ALL 9 tokens from cache
all_slots = torch.arange(N + 1, dtype=torch.int64, device=DEV)
kv_cached_fp8 = paged_kv_read(all_slots, kv_cache_fp8, block_size, N + 1, HD)
kv_cached = kv_dequantize_fp8(kv_cached_fp8, inv_scale_cache[all_slots])
# Decode attention: 1 query vs 9 cached KVs
o_decode = decode_attention(q_rope_d, kv_cached, SCALE)
print(f" Decode attention output: amax={o_decode.amax():.4f} NaN={torch.isnan(o_decode).any()}")
# BF16 reference: process all 9 tokens at once
all_ids = torch.cat([token_ids, decode_id])
all_pos = torch.arange(N + 1, dtype=torch.int64, device=DEV)
hidden_all = emb[all_ids]
normed_all = rms(hidden_all, anorm, EPS)
qa_all = r_qa.run(normed_all)
kv_all = r_kv.run(normed_all)
qa_n_all = rms(qa_all, qn, EPS)
kv_n_all = rms(kv_all, kvn, EPS)
q_all = r_qb.run(qa_n_all).view(N + 1, NH, HD)
q_rope_all = apply_gptj_rope(q_all, all_pos, cos_sin, NOPE, ROPE)
kv_rope_all = apply_gptj_rope(kv_n_all.unsqueeze(1), all_pos, cos_sin, NOPE, ROPE).squeeze(1)
o_ref_all = causal_prefill_attention(q_rope_all, kv_rope_all, SCALE)
o_ref_decode = o_ref_all[N:] # Only the decode token
c = F.cosine_similarity(o_decode.flatten().unsqueeze(0).float(), o_ref_decode.flatten().unsqueeze(0).float()).item()
print(f" Decode vs BF16 reference cosine: {c:.6f} {'✅' if c>=0.98 else '❌'}")
# ── Test 3: Full output pipeline (inverse RoPE + o_a + o_b) ──────
print(f"\n --- Test 3: Full output pipeline ---")
with torch.no_grad():
# Using decode attention output
o_inv = apply_inv_gptj_rope(o_decode, decode_pos, cos_sin, NOPE, ROPE)
o_grouped = o_inv.view(1, OG, HPG * HD).permute(1, 0, 2)
woa_3d = woa.view(OG, OL, HPG * HD)
z_cached = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL)
attn_out_cached = r_wob.run(z_cached)
# Using BF16 reference
o_inv_ref = apply_inv_gptj_rope(o_ref_decode, decode_pos, cos_sin, NOPE, ROPE)
o_grouped_ref = o_inv_ref.view(1, OG, HPG * HD).permute(1, 0, 2)
z_ref = torch.bmm(o_grouped_ref, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(1, OG * OL)
attn_out_ref = r_wob.run(z_ref)
c_full = F.cosine_similarity(attn_out_cached.flatten().unsqueeze(0).float(), attn_out_ref.flatten().unsqueeze(0).float()).item()
print(f" Full pipeline cosine: {c_full:.6f} {'✅' if c_full>=0.98 else '❌'}")
print(f" Output amax: cached={attn_out_cached.amax():.4f} ref={attn_out_ref.amax():.4f}")
# ── Test 4: Multi-step decode (3 decode steps) ───────────────────
print(f"\n --- Test 4: Multi-step decode (3 steps) ---")
decode_ids = torch.tensor([991, 1502, 4200], dtype=torch.long, device=DEV)
with torch.no_grad():
cosines = []
for step in range(3):
pos = N + step
dpos = torch.tensor([pos], dtype=torch.int64, device=DEV)
d_id = decode_ids[step:step+1]
hidden_s = emb[d_id]
normed_s = rms(hidden_s, anorm, EPS)
qa_s = r_qa.run(normed_s)
kv_s = r_kv.run(normed_s)
qa_n_s = rms(qa_s, qn, EPS)
kv_n_s = rms(kv_s, kvn, EPS)
q_s = r_qb.run(qa_n_s).view(1, NH, HD)
q_rope_s = apply_gptj_rope(q_s, dpos, cos_sin, NOPE, ROPE)
kv_rope_s = apply_gptj_rope(kv_n_s.unsqueeze(1), dpos, cos_sin, NOPE, ROPE).squeeze(1)
# Write to cache
kv_fp8_s, inv_s_s = kv_quantize_fp8(kv_rope_s)
paged_kv_write(kv_fp8_s, dpos, kv_cache_fp8, block_size)
inv_scale_cache[dpos[0]] = inv_s_s[0]
# Read all cached KV
all_s = torch.arange(pos + 1, dtype=torch.int64, device=DEV)
kv_all_fp8 = paged_kv_read(all_s, kv_cache_fp8, block_size, pos + 1, HD)
kv_all_dequant = kv_dequantize_fp8(kv_all_fp8, inv_scale_cache[all_s])
# Decode attention
o_s = decode_attention(q_rope_s, kv_all_dequant, SCALE)
# BF16 reference
all_ids_ref = torch.cat([token_ids, decode_ids[:step+1]])
all_pos_ref = torch.arange(pos + 1, dtype=torch.int64, device=DEV)
hidden_ref = emb[all_ids_ref]
normed_ref = rms(hidden_ref, anorm, EPS)
qa_ref = r_qa.run(normed_ref)
kv_ref = r_kv.run(normed_ref)
qa_n_ref = rms(qa_ref, qn, EPS)
kv_n_ref = rms(kv_ref, kvn, EPS)
q_ref = r_qb.run(qa_n_ref).view(pos + 1, NH, HD)
q_rope_ref = apply_gptj_rope(q_ref, all_pos_ref, cos_sin, NOPE, ROPE)
kv_rope_ref = apply_gptj_rope(kv_n_ref.unsqueeze(1), all_pos_ref, cos_sin, NOPE, ROPE).squeeze(1)
o_ref_full = causal_prefill_attention(q_rope_ref, kv_rope_ref, SCALE)
o_ref_last = o_ref_full[-1:]
c = F.cosine_similarity(o_s.flatten().unsqueeze(0).float(), o_ref_last.flatten().unsqueeze(0).float()).item()
cosines.append(c)
print(f" Step {step} (pos={pos}, {pos+1} cached): cosine = {c:.6f} {'✅' if c>=0.98 else '❌'}")
# Cleanup
del r_qa, r_qb, r_kv, r_wob
torch.cuda.empty_cache()
return c_full, cosines
def main():
print("=" * 70)
print(" DeepSeek-V4 Blackwell Attention Pipeline Test")
print(" Tests cutedsl.blackwell_attention with real weights")
print("=" * 70)
# Test SWA layer (layer 60, compress_ratio=0)
c_swa, cosines_swa = test_blackwell_attention(60, 0)
print(f"\n{'='*70}")
print(f" SUMMARY")
print(f" Layer 60 (SWA):")
print(f" Full pipeline cosine: {c_swa:.6f}")
print(f" Multi-step decode: {', '.join(f'{c:.6f}' for c in cosines_swa)}")
print(f"{'='*70}")
if __name__ == "__main__":
main()