diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py index 80b40aec..7b2b4f0b 100644 --- a/dsv4/kernels/attention/production.py +++ b/dsv4/kernels/attention/production.py @@ -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) diff --git a/tests/unit/test_production.py b/tests/unit/test_production.py index 9ac4f9d3..979053da 100644 --- a/tests/unit/test_production.py +++ b/tests/unit/test_production.py @@ -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__':