biondizzle
13be3ad443
FMHA kernel (fmha_6warp_tma_multirow_multitile.cuh): - Added sink_bias field to FmhaTmaMultiRowMultiTileParams - After KV tile loop, sink logit is included in online softmax rescale: new_max = max(running_max, sink_bias * scale) rescale existing O_unnorm and running_sum running_sum += exp(sink_bias * scale - new_max) No PV contribution from sink (D5c: single softmax) - C API: fmha_multitile_decode_launch now takes sink_bias_ptr - Python: fmha_multitile_decode_raw accepts attn_sink tensor single_shot_inference.py: - Full rewrite to use production kernel stack - mHC: uses dsv4.layers.mhc.mHCLayer (proper Sinkhorn-Knopp) - Projections: uses Nvfp4Linear (CuTeDSL GEMM) for q_a, q_b, kv, o_b - FMHA: 6-warp TMA multi-tile with sink bias (no SDPA fallback) - MoE: Nvfp4MoE + Nvfp4SharedExpert (no reference fallback) - Router: production dense/hash dispatch - Compressor/Indexer: reference dequant (not yet on tensor cores) - NO try/except fallbacks on production paths
89 lines
3.2 KiB
Python
89 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Test FMHA kernel with attention sink bias.
|
|
|
|
Validates that the kernel's sink bias correction matches PyTorch reference:
|
|
softmax([QK^T * scale, sink_bias])[:N] @ V
|
|
|
|
Tests HD=64,128,256,512 with and without sinks.
|
|
"""
|
|
import torch
|
|
import math
|
|
import sys
|
|
|
|
def reference_fmha_with_sink(q, k, v, scale, sink_bias=None):
|
|
"""PyTorch reference: softmax([QK^T * scale, sink_bias]) @ V.
|
|
|
|
q: (n_h, T, hd), k: (1, N, hd), v: (1, N, hd)
|
|
sink_bias: (n_h,) FP32 or None
|
|
Returns: (n_h, T, hd) BF16
|
|
"""
|
|
n_h, T, hd = q.shape
|
|
N = k.shape[1]
|
|
# QK^T: (n_h, T, N)
|
|
scores = torch.matmul(q, k.transpose(-1, -2)) * scale # (n_h, T, N)
|
|
|
|
if sink_bias is not None:
|
|
# Concatenate sink as extra column: (n_h, T, N+1)
|
|
sb = sink_bias.reshape(n_h, 1, 1).expand(-1, T, 1)
|
|
combined = torch.cat([scores, sb], dim=-1)
|
|
attn = torch.softmax(combined.float(), dim=-1)[:, :, :N] # drop sink column
|
|
else:
|
|
attn = torch.softmax(scores.float(), dim=-1)
|
|
|
|
out = torch.matmul(attn.bfloat16(), v) # (n_h, T, hd)
|
|
return out
|
|
|
|
def test_fmha_sink():
|
|
from dsv4.kernels.attention.production import dsv4_attention
|
|
|
|
torch.manual_seed(42)
|
|
device = 'cuda'
|
|
passed = 0
|
|
failed = 0
|
|
|
|
for hd in [64, 128, 256, 512]:
|
|
for N in [9, 32, 128, 256]:
|
|
for use_sink in [False, True]:
|
|
n_h = 4 # small for speed
|
|
T = 1
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device=device)
|
|
k = torch.randn(1, N, hd, dtype=torch.bfloat16, device=device)
|
|
v = torch.randn(1, N, hd, dtype=torch.bfloat16, device=device)
|
|
sink = torch.randn(n_h, dtype=torch.float32, device=device) * 2 if use_sink else None
|
|
|
|
# Production kernel
|
|
try:
|
|
o_kernel = dsv4_attention(q, k, v, scale=scale, sink_bias=sink)
|
|
except Exception as e:
|
|
print(f" FAIL hd={hd} N={N} sink={use_sink}: kernel error: {e}")
|
|
failed += 1
|
|
continue
|
|
|
|
# PyTorch reference
|
|
o_ref = reference_fmha_with_sink(q, k, v, scale, sink)
|
|
|
|
# Compare
|
|
o_kf = o_kernel.float()
|
|
o_rf = o_ref.float()
|
|
cos = torch.nn.functional.cosine_similarity(o_kf.flatten().unsqueeze(0),
|
|
o_rf.flatten().unsqueeze(0)).item()
|
|
max_diff = (o_kf - o_rf).abs().max().item()
|
|
|
|
status = "PASS" if cos > 0.999 else "FAIL"
|
|
if status == "PASS":
|
|
passed += 1
|
|
else:
|
|
failed += 1
|
|
print(f" {status} hd={hd} N={N} sink={use_sink} cos={cos:.6f} max_diff={max_diff:.6f}")
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Results: {passed} PASSED, {failed} FAILED")
|
|
print(f"{'='*60}")
|
|
return failed == 0
|
|
|
|
if __name__ == "__main__":
|
|
success = test_fmha_sink()
|
|
sys.exit(0 if success else 1)
|