biondizzle
9cbdc92744
- Split bridge.py -> ops/quantize.py, ops/layouts.py, ops/gemm_runner.py - Renamed classes: CuTeDSLNvfp4Linear -> Nvfp4Linear, etc. - Moved kernel code to dsv4/kernels/ (gemm, attention, compressor, decode, cuda) - Moved PyTorch bridges to dsv4/ops/ - Moved nn.Module layers to dsv4layers/ - Moved reference implementations to dsv4/reference/ - Moved vendored CUTLASS code to vendored/ - Archived ~190 debug tests to tests/archive/ - Kept ~15 canonical tests in tests/unit/ - Updated all import paths - Added stubs for future components (model/, cache/, loader/) - Updated pyproject.toml: dsv4-inference package name
374 lines
16 KiB
Python
374 lines
16 KiB
Python
#!/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 dsv4.layers.linear import Nvfp4Linear
|
||
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 = Nvfp4Linear(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 dsv4.layers.linear import Nvfp4Linear
|
||
|
||
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 Nvfp4Linear 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 = Nvfp4Linear(
|
||
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)
|
||
# Nvfp4Linear expects weight in (N, K//2) after permute
|
||
|
||
from dsv4.ops.quantize 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 Nvfp4Linear, 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, N_padded)
|
||
scores = scores[:, :T] # Slice to actual N=T (runner pads to 128)
|
||
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()
|