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

MQA/GQA: batch Q heads into kernel batch dim, shared K/V per KV group

Browse Source
This commit is contained in:
biondizzle
2026-05-27 08:31:23 +00:00
parent 06a895ff99
commit 2412a5431b
2 changed files with 260 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

231
dsv4/kernels/attention/production.py
View File

@@ -1,16 +1,12 @@
"""DSV4 Blackwell Attention — Production kernel wrapper.
Wraps the CuTeDSL FMHA kernel with Python KV merge for multi-KV-tile.
Replaces vLLM's broken FlashMLA Blackwell implementation.
Supported attention types:
- CSA (Classical Self-Attention): 32 KV entries/block, full KV cache
- HCA (Hash-based Cross-Attention): 1 KV entry/block, compressed cache
- SWA (Sliding Window Attention): ring buffer, no KV cache
Supports MHA, MQA, and GQA attention patterns.
Limitations:
- head_dim > 256: MLIR compilation hang (known CuTeDSL issue)
- In-kernel multi-KV-tile: blocked on TMA layout matching (uses Python KV merge)
- MQA: K/V replicated across Q-head batch dim (redundant TMA loads, but parallel)
"""
import torch
import math
@@ -19,7 +15,6 @@ import cutlass.torch as ct
import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel
# Kernel cache: compiled kernels keyed by config tuple
_kernel_cache: dict = {}
@@ -27,7 +22,7 @@ def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
normalize: bool = False, apply_swa_mask: bool = False,
is_causal: bool = False, n_comp: int = 0,
apply_sink_bias: bool = False) -> tuple:
"""Get or compile a kernel for the given configuration."""
"""Get or compile a kernel for the given configuration. Cache by config."""
key = (head_dim, s_k, use_smem_p, normalize, apply_swa_mask, is_causal, n_comp, apply_sink_bias)
if key in _kernel_cache:
return _kernel_cache[key]
@@ -37,9 +32,10 @@ def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
apply_swa_mask=apply_swa_mask, is_causal=is_causal,
n_comp=n_comp, apply_sink_bias=apply_sink_bias,
)
m = 128
pv_n_tile = kernel.pv_n_tile
# Dummy tensors for compilation (single-head, single-segment)
m = 128
q = torch.randn(m, head_dim, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(s_k, head_dim, 1, dtype=torch.bfloat16, device='cuda')
v = torch.randn(s_k, pv_n_tile, dtype=torch.bfloat16, device='cuda')
@@ -48,7 +44,6 @@ def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
row_sums = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
@@ -63,55 +58,191 @@ def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
def dsv4_attention(
q: torch.Tensor, # (num_heads, seq_len, head_dim)
k: torch.Tensor, # (seq_len, head_dim) or (num_heads, seq_len, head_dim) for MQA
v: torch.Tensor, # (seq_len, head_dim) or (num_heads, seq_len, head_dim) for MQA
q: torch.Tensor, # (n_q_heads, T, head_dim)
k: torch.Tensor, # (n_kv_heads, N, head_dim) or (N, head_dim) for MQA
v: torch.Tensor, # (n_kv_heads, N, head_dim) or (N, head_dim) for MQA
scale: float = None,
swa_len: int = None,
is_causal: bool = False,
n_comp: int = 0,
sink_bias: torch.Tensor = None,
sink_bias: torch.Tensor = None, # (n_q_heads,) or scalar
) -> torch.Tensor:
"""Production DSV4 attention using CuTeDSL FMHA kernel + Python KV merge.
"""Production DSV4 attention: MHA / MQA / GQA with Python KV merge.
Args:
q: Query tensor (num_heads, seq_len, head_dim) BF16
k: Key tensor (seq_len, head_dim) or (num_heads, seq_len, head_dim) BF16
v: Value tensor (seq_len, head_dim) or (num_heads, seq_len, head_dim) BF16
scale: Softmax scale (default: 1/sqrt(head_dim))
swa_len: Sliding window length (None = no mask)
is_causal: Apply causal mask
n_comp: Compressed KV length for D5c sink bias merge
sink_bias: Per-head FP32 logit bias for sink attention
q: (n_q_heads, T, hd) BF16 — query heads
k: (n_kv_heads, N, hd) or (N, hd) BF16 — key heads
v: (n_kv_heads, N, hd) or (N, hd) BF16 — value heads
scale: 1/sqrt(hd) if None
swa_len: sliding window length
is_causal: causal mask
n_comp: compressed KV length for D5c sink bias
sink_bias: per-head FP32 logit bias
Returns:
Output tensor (num_heads, seq_len, head_dim) BF16
(n_q_heads, T, hd) BF16
"""
n_h, T, hd = q.shape
N = k.shape[-2]
n_q, T, hd = q.shape
scale = scale or (1.0 / math.sqrt(hd))
use_smem_p = hd > 64
apply_swa_mask = swa_len is not None
apply_sink_bias = sink_bias is not None
output = torch.zeros(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
# Normalize K/V to (n_kv, N, hd)
if k.dim() == 2:
k = k.unsqueeze(0) # (1, N, hd) — MQA: 1 KV head
if v.dim() == 2:
v = v.unsqueeze(0)
n_kv, N, _ = k.shape
for h in range(n_h):
q_h = q[h:h+1] # (1, T, hd)
k_h = k.unsqueeze(0) if k.dim() == 2 else k[h:h+1] # (1, N, hd)
v_h = v.unsqueeze(0) if v.dim() == 2 else v[h:h+1] # (1, N, hd)
# GQA ratio: each KV head serves (n_q // n_kv) Q heads
q_per_kv = n_q // n_kv
assert n_q % n_kv == 0, f"n_q_heads ({n_q}) must be divisible by n_kv_heads ({n_kv})"
o_h = _attention_single_head(
q_h, k_h, v_h, scale=scale,
swa_len=swa_len, is_causal=is_causal,
n_comp=n_comp, sink_bias=sink_bias[h] if sink_bias is not None else None,
use_smem_p=use_smem_p,
)
output[h] = o_h[0]
# Run attention per KV head group, batching all Q heads in that group
output = torch.zeros(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
for kv_idx in range(n_kv):
# Q heads for this KV group
q_start = kv_idx * q_per_kv
q_end = q_start + q_per_kv
q_group = q[q_start:q_end] # (q_per_kv, T, hd)
k_kv = k[kv_idx:kv_idx+1] # (1, N, hd)
v_kv = v[kv_idx:kv_idx+1] # (1, N, hd)
if q_per_kv == 1:
# Single Q head per KV — simple path
o = _attention_single_head(
q_group, k_kv, v_kv, scale=scale,
swa_len=swa_len, is_causal=is_causal, n_comp=n_comp,
sink_bias=sink_bias[q_start] if sink_bias is not None else None,
use_smem_p=use_smem_p,
)
output[q_start] = o[0]
else:
# Multiple Q heads per KV (MQA/GQA) — batch them
bias_slice = sink_bias[q_start:q_end] if sink_bias is not None else None
o = _attention_batched_mqa(
q_group, k_kv, v_kv, scale=scale,
swa_len=swa_len, is_causal=is_causal, n_comp=n_comp,
sink_bias=bias_slice, use_smem_p=use_smem_p,
)
output[q_start:q_end] = o
return output
def _attention_batched_mqa(
q: torch.Tensor, # (n_q_heads, T, hd) — multiple Q heads sharing K/V
k: torch.Tensor, # (1, N, hd) — single KV head
v: torch.Tensor, # (1, N, hd)
scale: float,
swa_len: int = None,
is_causal: bool = False,
n_comp: int = 0,
sink_bias: torch.Tensor = None, # (n_q_heads,) or None
use_smem_p: bool = False,
) -> torch.Tensor:
"""Run batched FMHA for multiple Q heads sharing one K/V (MQA/GQA).
Uses the kernel's batch dimension to run all Q heads in parallel.
K/V are expanded to (n_q_heads, N, hd) for TMA compatibility.
"""
n_q, T, hd = q.shape
N = k.shape[1]
apply_swa_mask = swa_len is not None
apply_sink_bias = sink_bias is not None
# Expand K/V across Q heads for TMA: (1, N, hd) -> (n_q, N, hd)
k_expanded = k.expand(n_q, -1, -1).contiguous() # (n_q, N, hd)
v_expanded = v.expand(n_q, -1, -1).contiguous() # (n_q, N, hd)
# Convert to 3D kernel format
q_3d = q.reshape(n_q * T, hd, 1).contiguous() # (n_q*T, hd, 1) — batched Q
k_3d_single = k[0].contiguous().unsqueeze(-1) # (N, hd, 1) — single K for TMA desc
v_2d_single = v[0].contiguous() # (N, hd) — single V for TMA desc
s_k_per_seg = 128
n_segments = (N + s_k_per_seg - 1) // s_k_per_seg
compiled, kernel = _get_or_compile_kernel(
head_dim=hd, s_k=s_k_per_seg, use_smem_p=use_smem_p,
normalize=False, apply_swa_mask=apply_swa_mask,
is_causal=is_causal, n_comp=n_comp if n_comp > 0 else None,
apply_sink_bias=apply_sink_bias,
)
pv_n_tile = kernel.pv_n_tile
n_pv_tiles = kernel.n_pv_tiles
# Per-head accumulators
o_accum = torch.zeros(n_q, T, hd, dtype=torch.float32, device='cuda')
lse_accum = torch.full((n_q, T, 1), float('-inf'), dtype=torch.float32, device='cuda')
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
for seg in range(n_segments):
k_start = seg * s_k_per_seg
k_end = min(k_start + s_k_per_seg, N)
k_seg = k_3d_single[k_start:k_end] # (s_k, hd, 1)
v_seg = v_2d_single[k_start:k_end] # (s_k, hd)
if k_end - k_start < s_k_per_seg:
pad_len = s_k_per_seg - (k_end - k_start)
k_seg = torch.cat([k_seg, torch.zeros(pad_len, hd, 1, dtype=k_seg.dtype, device='cuda')], dim=0)
v_seg = torch.cat([v_seg, torch.zeros(pad_len, hd, dtype=v_seg.dtype, device='cuda')], dim=0)
seg_o = torch.zeros(n_q, T, hd, dtype=torch.float32, device='cuda')
seg_lse = torch.zeros(n_q, T, 1, dtype=torch.float32, device='cuda')
seg_row_sums = torch.zeros(n_q, T, 1, dtype=torch.float32, device='cuda')
for nt in range(n_pv_tiles):
v_start = nt * pv_n_tile
v_end = v_start + pv_n_tile
v_tile = v_seg[:, v_start:v_end].contiguous()
v_kernel = v_tile.unsqueeze(-1) # (s_k, pv_n_tile, 1)
# Per-head output tensors
c_tile = torch.zeros(n_q * T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
lse_tensor = torch.zeros(n_q * T, 1, 1, dtype=torch.float32, device='cuda')
row_sums_tensor = torch.zeros(n_q * T, 1, 1, dtype=torch.float32, device='cuda')
mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d))
mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg))
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
torch.cuda.synchronize()
# Reshape output: (n_q*T, pv_n_tile) -> (n_q, T, pv_n_tile)
c_2d = c_tile[:, :, 0].reshape(n_q, T, pv_n_tile).float()
lse_2d = lse_tensor[:, 0, 0].reshape(n_q, T, 1).float()
rs_2d = row_sums_tensor[:, 0, 0].reshape(n_q, T, 1).float()
# Scatter PV tile output into segment accumulator
seg_o[:, :, v_start:v_end] = c_2d
if nt == 0:
seg_lse = lse_2d
seg_row_sums = rs_2d
# Normalize segment: O_norm = O_unnorm / row_sum
seg_row_sums = seg_row_sums.clamp(min=1e-30)
seg_o_norm = seg_o / seg_row_sums # (n_q, T, hd) normalized
# KV merge per head
e_old = torch.exp(lse_accum) # (n_q, T, 1)
e_new = torch.exp(seg_lse)
e_sum = e_old + e_new
o_accum = (e_old * o_accum + e_new * seg_o_norm) / e_sum
lse_accum = torch.log(e_sum)
return o_accum.to(torch.bfloat16) # (n_q, T, hd)
def _attention_single_head(
q: torch.Tensor, # (1, T, hd)
k: torch.Tensor, # (1, N, hd)
@@ -123,18 +254,16 @@ def _attention_single_head(
sink_bias: torch.Tensor = None,
use_smem_p: bool = False,
) -> torch.Tensor:
"""Run FMHA for a single head, with Python KV merge for s_k > 128."""
"""Run FMHA for a single Q head (fallback for non-MQA or single-head case)."""
_, T, hd = q.shape
N = k.shape[1]
apply_swa_mask = swa_len is not None
apply_sink_bias = sink_bias is not None
# Convert to 2D/3D shapes matching test_d1_kv_merge.py
q_3d = q[0].contiguous().unsqueeze(-1) # (T, hd, 1)
k_3d = k[0].contiguous().unsqueeze(-1) # (N, hd, 1)
v_2d = v[0].contiguous() # (N, hd)
# Segment the KV sequence into 128-entry tiles
s_k_per_seg = 128
n_segments = (N + s_k_per_seg - 1) // s_k_per_seg
@@ -147,7 +276,6 @@ def _attention_single_head(
pv_n_tile = kernel.pv_n_tile
n_pv_tiles = kernel.n_pv_tiles
# Accumulate across KV segments using log-sum-exp merge
o_accum = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
lse_accum = torch.full((T, 1), float('-inf'), dtype=torch.float32, device='cuda')
@@ -156,10 +284,9 @@ def _attention_single_head(
for seg in range(n_segments):
k_start = seg * s_k_per_seg
k_end = min(k_start + s_k_per_seg, N)
k_seg = k_3d[k_start:k_end] # (s_k, hd, 1) — same as test
v_seg = v_2d[k_start:k_end] # (s_k, hd) — same as test
k_seg = k_3d[k_start:k_end]
v_seg = v_2d[k_start:k_end]
# Pad last segment if shorter than s_k_per_seg
if k_end - k_start < s_k_per_seg:
pad_len = s_k_per_seg - (k_end - k_start)
k_seg = torch.cat([k_seg, torch.zeros(pad_len, hd, 1, dtype=k_seg.dtype, device='cuda')], dim=0)
@@ -193,21 +320,13 @@ def _attention_single_head(
seg_lse[:, 0] = lse_tensor[:, 0, 0].float()
seg_row_sums[:, 0] = row_sums_tensor[:, 0, 0].float()
# Normalize segment output: O_norm = O_unnorm / row_sum (per row)
seg_row_sums = seg_row_sums.clamp(min=1e-30)
seg_o_norm = seg_o / seg_row_sums # (T, hd) normalized
seg_o_norm = seg_o / seg_row_sums
# Merge with accumulator using CORRECT formula:
# O = sum(exp(lse_i) * O_i_norm) / sum(exp(lse_i))
# where O_i_norm = O_i_unnorm / row_sum_i
#
# o_accum is tracked in NORMALIZED form.
e_old = torch.exp(lse_accum)
e_new = torch.exp(seg_lse)
e_sum = e_old + e_new
o_accum = (e_old * o_accum + e_new * seg_o_norm) / e_sum
lse_accum = torch.log(e_sum)
output = o_accum.to(torch.bfloat16).unsqueeze(0)
return output
return o_accum.to(torch.bfloat16).unsqueeze(0)

164
tests/unit/test_production.py
View File

@@ -1,93 +1,99 @@
"""Test production DSV4 attention wrapper."""
"""Test production DSV4 attention wrapper: MHA, MQA, GQA."""
import torch
import math
from dsv4.kernels.attention.production import dsv4_attention
def test_single_head_128():
"""Single head, 1 KV segment (N=128)."""
torch.manual_seed(42)
hd = 64; n_h = 1; T = 128; N = 128
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
qf = q[0].float(); kf = k[0].float(); vf = v[0].float()
scale = 1.0 / math.sqrt(hd)
attn = qf @ kf.T * scale
ref = torch.softmax(attn, dim=-1) @ vf
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.unsqueeze(0).flatten().unsqueeze(0)).item()
print(f" hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_single_head_256():
"""Single head, 2 KV segments (N=256)."""
torch.manual_seed(42)
hd = 64; n_h = 1; T = 128; N = 256
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
qf = q[0].float(); kf = k[0].float(); vf = v[0].float()
scale = 1.0 / math.sqrt(hd)
attn = qf @ kf.T * scale
ref = torch.softmax(attn, dim=-1) @ vf
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.unsqueeze(0).flatten().unsqueeze(0)).item()
print(f" hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_single_head_512():
"""Single head, 4 KV segments (N=512)."""
torch.manual_seed(42)
hd = 64; n_h = 1; T = 128; N = 512
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
qf = q[0].float(); kf = k[0].float(); vf = v[0].float()
scale = 1.0 / math.sqrt(hd)
attn = qf @ kf.T * scale
ref = torch.softmax(attn, dim=-1) @ vf
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.unsqueeze(0).flatten().unsqueeze(0)).item()
print(f" hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_multi_head():
"""Multi-head, 2 KV segments."""
torch.manual_seed(42)
hd = 64; n_h = 4; T = 128; N = 256
q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
scale = 1.0 / math.sqrt(hd)
def _pytorch_ref(q, k, v, scale):
"""PyTorch reference attention (MHA/MQA/GQA)."""
n_q, T, hd = q.shape
n_kv, N, _ = k.shape
q_per_kv = n_q // n_kv
ref = torch.zeros_like(q)
for h in range(n_h):
qf = q[h].float(); kf = k[h].float(); vf = v[h].float()
for qi in range(n_q):
ki = qi // q_per_kv
qf = q[qi].float()
kf = k[ki].float()
vf = v[ki].float()
attn = qf @ kf.T * scale
ref[h] = (torch.softmax(attn, dim=-1) @ vf).bfloat16()
ref[qi] = (torch.softmax(attn, dim=-1) @ vf).bfloat16()
return ref
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)).item()
print(f" hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_mha():
"""Multi-head attention: n_q = n_kv (each Q head has own K/V)."""
torch.manual_seed(42)
hd = 64; T = 128; N = 256; n_q = 4; n_kv = 4
q = torch.randn(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
ref = _pytorch_ref(q, k, v, 1.0 / math.sqrt(hd))
cos = torch.nn.functional.cosine_similarity(
out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)
).item()
print(f" MHA n_q={n_q} n_kv={n_kv} N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_mqa():
"""Multi-query attention: n_q > 1, n_kv = 1 (shared K/V)."""
torch.manual_seed(42)
hd = 64; T = 128; N = 256; n_q = 8; n_kv = 1
q = torch.randn(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda') # 2D = MQA
v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
ref = _pytorch_ref(q, k.unsqueeze(0), v.unsqueeze(0), 1.0 / math.sqrt(hd))
cos = torch.nn.functional.cosine_similarity(
out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)
).item()
print(f" MQA n_q={n_q} n_kv=1 N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_gqa():
"""Grouped-query attention: n_q > n_kv > 1."""
torch.manual_seed(42)
hd = 64; T = 128; N = 256; n_q = 8; n_kv = 2
q = torch.randn(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
ref = _pytorch_ref(q, k, v, 1.0 / math.sqrt(hd))
cos = torch.nn.functional.cosine_similarity(
out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)
).item()
print(f" GQA n_q={n_q} n_kv={n_kv} N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test_mqa_2d_kv():
"""MQA with 2D K/V (no batch dim)."""
torch.manual_seed(42)
hd = 64; T = 128; N = 128; n_q = 4; n_kv = 1
q = torch.randn(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
out = dsv4_attention(q, k, v)
ref = _pytorch_ref(q, k.unsqueeze(0), v.unsqueeze(0), 1.0 / math.sqrt(hd))
cos = torch.nn.functional.cosine_similarity(
out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)
).item()
print(f" MQA-2D n_q={n_q} n_kv=1 N={N}: cos {cos:.6f} {'PASS' if cos >= 0.99 else 'FAIL'}")
def test():
print("=== Production DSV4 Attention ===\n")
test_single_head_128()
test_single_head_256()
test_single_head_512()
test_multi_head()
print("=== Production DSV4 Attention: MHA / MQA / GQA ===\n")
test_mqa_2d_kv()
test_mha()
test_mqa()
test_gqa()
if __name__ == '__main__':