diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 1580525f..1aaa5b51 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -465,7 +465,7 @@ class FmhaKernel: ) cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) for k in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[k] = tTMrO_i[k] * Float32(1.0) # DEBUG: NO-OP round-trip test + tTMrO_i[k] = tTMrO_i[k] * acc_scale cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) cute.arch.fence_view_async_tmem_store() @@ -505,14 +505,15 @@ class FmhaKernel: # Compute LSE: lse = ln(row_sum) + row_max * ln(2) # Only when emitting un-normalized output (D5a path). # When normalize=True, LSE is not needed (in-kernel normalization). + # Each thread writes its row's LSE. With 128 softmax threads and 128 rows, + # each thread (sfw_idx) owns exactly one row. if const_expr(not self.normalize): _row_max_safe = row_max if row_max == -cutlass.Float32.inf: _row_max_safe = Float32(0.0) - if sfw_idx == 0: - _ln2 = Float32(0.6931471805599453) # ln(2) - lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2 - mLSE[0] = lse_val + _ln2 = Float32(0.6931471805599453) # ln(2) + lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2 + mLSE[sfw_idx] = lse_val tmem.relinquish_alloc_permit() tmem.free(tmem_ptr) diff --git a/tests/unit/test_d1_kv_merge_v2.py b/tests/unit/test_d1_kv_merge_v2.py new file mode 100644 index 00000000..49ed50dc --- /dev/null +++ b/tests/unit/test_d1_kv_merge_v2.py @@ -0,0 +1,122 @@ +""" +D1: Multi-KV-tile merge using per-row LSE and log-sum-exp. + +Strategy: Run s_k=128 kernel per KV segment, get per-row O and LSE. +Merge using the D5 formula: + O = (exp(lse_0) * O_0 + exp(lse_1) * O_1) / (exp(lse_0) + exp(lse_1)) + +This avoids the broken TMEM round-trip O rescale. +""" +import torch, math +import cutlass.cute as cute +import cutlass.torch as ct +import cuda.bindings.driver as cuda +from dsv4.kernels.attention.fmha import FmhaKernel + + +def test_multi_kv_merge(hd=64, s_k=256): + m = 128 + n_kv_segments = s_k // 128 + torch.manual_seed(42) + + q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') + k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') + v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') + + # FP32 reference (full attention) + qf = q[:, :, 0].float() + kf = k[:, :, 0].float() + scale = 1.0 / math.sqrt(hd) + attn_scores = qf @ kf.T * scale + attn_max = attn_scores.max(dim=-1, keepdim=True)[0] + attn_exp = torch.exp(attn_scores - attn_max) + attn_sum = attn_exp.sum(dim=-1, keepdim=True) + ref_norm = (attn_exp / attn_sum) @ v.float() + + # Run s_k=128 kernel per KV segment + kernel = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False) + pv_n_tile = kernel.pv_n_tile + n_pv_tiles = kernel.n_pv_tiles + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # Compile once + k_seg0 = k[:128] + v_tile0 = v[:128, 0:pv_n_tile].contiguous() + v_kernel0 = v_tile0.unsqueeze(-1) + c_tile0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') + # Per-row LSE: shape (m,) + lse_tensor = torch.zeros(m, dtype=torch.float32, device='cuda') + + mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) + mK = ct.from_dlpack(k_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg0)) + mV = ct.from_dlpack(v_kernel0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel0)) + mC = ct.from_dlpack(c_tile0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile0)) + mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) + + print(f' Compiling (hd={hd}, s_k=128, {n_kv_segments} segments)...', flush=True) + compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) + + # Accumulate across KV segments using log-sum-exp merge + o_accum = None # Will be (m, hd) FP32 + lse_accum = None # Will be (m,) FP32 + + for seg in range(n_kv_segments): + k_start = seg * 128 + k_end = k_start + 128 + k_seg = k[k_start:k_end] + v_seg = v[k_start:k_end] + + seg_o = torch.zeros(m, hd, 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) + c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') + lse_tensor.zero_() + + mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) + 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)) + + compiled(mQ, mK, mV, mC, stream, mLSE) + torch.cuda.synchronize() + + seg_o[:, v_start:v_end] = c_tile[:, :, 0].float() + + seg_lse = lse_tensor.clone() # (m,) per-row LSE + + # Log-sum-exp merge with accumulator + if o_accum is None: + o_accum = seg_o + lse_accum = seg_lse + else: + e_old = torch.exp(lse_accum) # (m,) + e_new = torch.exp(seg_lse) # (m,) + e_sum = e_old + e_new # (m,) + + o_accum = (e_old.unsqueeze(-1) * o_accum + e_new.unsqueeze(-1) * seg_o) / e_sum.unsqueeze(-1) + lse_accum = torch.log(e_sum) + + cos = torch.nn.functional.cosine_similarity( + o_accum.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) + ).item() + print(f' hd={hd}, s_k={s_k} ({n_kv_segments} segments): cos_norm {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') + return cos + + +def test(): + print("=== D1: Multi-KV Merge via Per-Row Log-Sum-Exp ===\n") + + test_multi_kv_merge(64, 256) + test_multi_kv_merge(64, 384) + test_multi_kv_merge(64, 512) + test_multi_kv_merge(64, 1024) + + +if __name__ == '__main__': + test()