diff --git a/tests/unit/test_d5c_multitile.py b/tests/unit/test_d5c_multitile.py index 06581c20..6aa1877c 100644 --- a/tests/unit/test_d5c_multitile.py +++ b/tests/unit/test_d5c_multitile.py @@ -1,17 +1,6 @@ """ FMHA D5c: Sink bias + Python KV merge for multi-tile (s_k > 128). -Verifies the full DSV4 attention pipeline: -1. Concatenate KV: [compressed_K; swa_K] (total s_k > 128) -2. Split into 128-token segments -3. Run FMHA per segment (with sink bias on SWA positions, D3/D4 masking) -4. Merge segments using Python KV merge formula: - O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i)) -5. Normalize using row_sum - -This is the production path for DSV4 Pro (s_k=1152, 9 KV tiles) -until the D1.5 TMEM round-trip fix enables in-kernel O rescale. - Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5c_multitile.py """ import torch @@ -22,15 +11,18 @@ import cuda.bindings.driver as cuda from dsv4.kernels.attention.fmha import FmhaKernel +def to_cute(t): + return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) + + def reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa, attn_sink, scale, swa_len, is_causal=False): - """FP32 reference: single softmax over combined KV with sink bias on SWA.""" m, hd = q.shape n_comp = k_comp.shape[0] n_swa = k_swa.shape[0] - k_combined = torch.cat([k_comp, k_swa], dim=0) - v_combined = torch.cat([v_comp, v_swa], dim=0) - scores = torch.matmul(q.float(), k_combined.float().T) * scale + k_comb = torch.cat([k_comp, k_swa], dim=0) + v_comb = torch.cat([v_comp, v_swa], dim=0) + scores = q.float() @ k_comb.float().T * scale scores[:, n_comp:] += attn_sink if swa_len < n_swa: scores[:, n_comp + swa_len:] = float('-inf') @@ -42,82 +34,51 @@ def reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa, max_s = scores.max(dim=-1, keepdim=True).values exp_s = (scores - max_s).exp() sum_s = exp_s.sum(dim=-1, keepdim=True).clamp(min=1e-30) - o = torch.matmul(exp_s / sum_s, v_combined.float()) - return o.to(torch.bfloat16) - - -def run_segment(q, k_seg, v_seg, kernel, compiled, stream, - sink_bias=None, n_comp_in_seg=0, swa_len=999999): - """Run one 128-token segment of FMHA, return normalized O and LSE.""" - m = q.shape[0] - hd = v_seg.shape[1] - - # Allocate per-segment outputs - o_seg = torch.zeros(m, hd, dtype=torch.bfloat16, device='cuda') - lse_seg = torch.zeros(m, dtype=torch.float32, device='cuda') - rs_seg = torch.zeros(m, dtype=torch.float32, device='cuda') - - def to_cute(t): - return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) - - c_tile = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - - mQ = to_cute(q.unsqueeze(-1)) - mK = to_cute(k_seg.unsqueeze(-1)) - mV = to_cute(v_seg.unsqueeze(-1)) # (n_k, hd, 1) - mC = to_cute(c_tile) - mLSE = to_cute(lse_tile) - mRS = to_cute(rs_tile) - - if sink_bias is not None: - mSB = to_cute(sink_bias) - compiled(mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, sink_bias=mSB, row_sums=mRS) - else: - compiled(mQ, mK, mV, mC, stream, mLSE, - swa_len=swa_len, row_sums=mRS) - - torch.cuda.synchronize() - rs = rs_tile[:, 0, 0].float() - o_norm = c_tile[:, :, 0].float() / rs.unsqueeze(1).clamp(min=1e-30) - lse_val = lse_tile[:, 0, 0].clone() - - return o_norm, lse_val + return (exp_s / sum_s @ v_comb.float()).to(torch.bfloat16) def python_kv_merge(segment_results): - """Merge multiple FMHA segments using Python KV merge formula. - - O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i)) - """ - # Compute max LSE for numerical stability - lse_stack = torch.stack([r[1] for r in segment_results], dim=0) # (n_seg, m) - lse_max = lse_stack.max(dim=0).values # (m,) - - numerator = torch.zeros_like(segment_results[0][0]) # (m, hd) + """O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i))""" + lse_stack = torch.stack([r[1] for r in segment_results], dim=0) + lse_max = lse_stack.max(dim=0).values + numerator = torch.zeros_like(segment_results[0][0]) denominator = torch.zeros(lse_max.shape[0], dtype=torch.float32, device=lse_max.device) - for o_norm, lse in segment_results: - exp_lse = (lse - lse_max).exp() # (m,) + exp_lse = (lse - lse_max).exp() numerator += exp_lse.unsqueeze(1) * o_norm.float() denominator += exp_lse + return numerator / denominator.unsqueeze(1).clamp(min=1e-30) - o_merged = numerator / denominator.unsqueeze(1).clamp(min=1e-30) - return o_merged + +def run_fmha(q, k, v, kernel_obj, compiled, stream, swa_len, sink_bias=None): + """Run FMHA with single KV tile, return (O_norm, LSE).""" + m = q.shape[0] + hd = v.shape[1] + c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') + lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + + mQ = to_cute(q.unsqueeze(-1)) + mK = to_cute(k.unsqueeze(-1)) + mV = to_cute(v.unsqueeze(-1)) + mC = to_cute(c); mLSE = to_cute(lse); mRS = to_cute(rs) + + if sink_bias is not None: + mSB = to_cute(sink_bias) + compiled(mQ, mK, mV, mC, stream, mLSE, swa_len=swa_len, sink_bias=mSB, row_sums=mRS) + else: + compiled(mQ, mK, mV, mC, stream, mLSE, swa_len=swa_len, row_sums=mRS) + torch.cuda.synchronize() + + o_norm = c[:, :, 0].float() / rs[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) + return o_norm, lse[:, 0, 0].clone() def test_d5c_multitile(): print("=== D5c Multi-Tile: Sink Bias + Python KV Merge ===\n") - hd = 64 - m = 128 - n_comp = 96 # compressed KV tokens - n_swa = 160 # SWA tokens - n_total = n_comp + n_swa # 256, 2 KV tiles - swa_len = 100 # valid SWA fill (within the 160 SWA window) - scale = 1.0 / math.sqrt(hd) + hd = 64; m = 128; n_comp = 96; n_swa = 160; n_total = 256 + swa_len = 100; scale = 1.0 / math.sqrt(hd); seg_size = 128 torch.manual_seed(42) q = torch.randn(m, hd, dtype=torch.bfloat16, device='cuda') @@ -125,176 +86,80 @@ def test_d5c_multitile(): v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda') k_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda') v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda') - attn_sink_val = 0.5 attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda') - # Direct test of segment 0 (verify run_segment isn't the issue) - print('--- Direct segment 0 test ---') - k_seg0 = k_combined[0:seg_size] - v_seg0 = v_combined[0:seg_size] - c0 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - lse0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - rs0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - mQ0 = to_cute(q.unsqueeze(-1)) - mK0 = to_cute(k_seg0.unsqueeze(-1)) - mV0 = to_cute(v_seg0.unsqueeze(-1)) - mC0 = to_cute(c0); mLSE0 = to_cute(lse0); mRS0 = to_cute(rs0) - mSB0 = to_cute(attn_sink) - # Compile FRESH for segment 0 - k0 = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False, - apply_swa_mask=True, is_causal=False, n_comp=n_comp, apply_sink_bias=True) - comp0 = cute.compile(k0, mQ0, mK0, mV0, mC0, stream, mLSE0, - swa_len=swa_len, sink_bias=mSB0, row_sums=mRS0) - comp0(mQ0, mK0, mV0, mC0, stream, mLSE0, - swa_len=swa_len, sink_bias=mSB0, row_sums=mRS0) - torch.cuda.synchronize() - ok0 = c0[:, :, 0].float() / rs0[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) - # Per-segment reference for segment 0 - scores0_ref = q.float() @ k_seg0.float().T * scale - scores0_ref[:, n_comp:] += attn_sink_val - if swa_len < (seg_size - n_comp): # 100 < 32? No - scores0_ref[:, n_comp + swa_len:] = float('-inf') - ref0 = (torch.softmax(scores0_ref, dim=-1) @ v_seg0.float()).to(torch.bfloat16) - cos0 = torch.nn.functional.cosine_similarity(ok0.flatten().unsqueeze(0), ref0.flatten().unsqueeze(0).float()).item() - print(f'Seg0 direct: cos {cos0:.6f} LSE_kern={lse0[0,0,0].item():.4f} LSE_ref={torch.logsumexp(scores0_ref, dim=-1)[0].item():.4f}') - print() + k_combined = torch.cat([k_comp, k_swa], dim=0) + v_combined = torch.cat([v_comp, v_swa], dim=0) - # Reference - ref = reference_combined_attention( - q, k_comp, v_comp, k_swa, v_swa, - attn_sink_val, scale, swa_len - ) + # Full reference + ref = reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa, + attn_sink_val, scale, swa_len) - # Combined KV - k_combined = torch.cat([k_comp, k_swa], dim=0) # (256, hd) - v_combined = torch.cat([v_comp, v_swa], dim=0) # (256, hd) - - # Split into 128-token segments and run FMHA per segment stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - seg_size = 128 - n_segs = (n_total + seg_size - 1) // seg_size - # Compile two kernels: one with n_comp (for segments with compressed+SWA), - # one without n_comp (for segments with SWA only) - # n_comp is a compile-time parameter baked into const_expr guards. + # Per-segment reference and kernel verification + # Segment 0: positions 0-127 (n_comp=96, 32 SWA tokens, all valid, sink bias on SWA) + k0 = k_combined[0:128]; v0 = v_combined[0:128] + scores0 = q.float() @ k0.float().T * scale + scores0[:, n_comp:] += attn_sink_val + # swa_len=100 but only 32 SWA positions → no D3 masking + ref0 = (torch.softmax(scores0, dim=-1) @ v0.float()).to(torch.bfloat16) + lse0_ref = torch.logsumexp(scores0, dim=-1) - # Kernel 1: has compressed KV + SWA (n_comp > 0) - kernel_comp = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False, - apply_swa_mask=True, is_causal=False, n_comp=n_comp) - def to_cute(t): - return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) - _q = q.unsqueeze(-1) - _k = k_combined[:seg_size].unsqueeze(-1) - _v = v_combined[:seg_size, :kernel_comp.pv_n_tile].contiguous().unsqueeze(-1) - _c = torch.zeros(m, kernel_comp.pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') - _lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - _rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - _mQ = to_cute(_q); _mK = to_cute(_k); _mV = to_cute(_v) - _mC = to_cute(_c); _mLSE = to_cute(_lse); _mRS = to_cute(_rs) - _mSB = to_cute(attn_sink) - compiled_comp = cute.compile(kernel_comp, _mQ, _mK, _mV, _mC, stream, _mLSE, - swa_len=swa_len, sink_bias=_mSB, row_sums=_mRS) + k_seg0 = FmhaKernel(head_dim=hd, s_k=128, normalize=False, + apply_swa_mask=True, is_causal=False, n_comp=n_comp, apply_sink_bias=True) + c0 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') + l0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + r0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + comp0 = cute.compile(k_seg0, to_cute(q.unsqueeze(-1)), to_cute(k0.unsqueeze(-1)), + to_cute(v0.unsqueeze(-1)), to_cute(c0), stream, to_cute(l0), + swa_len=100, sink_bias=to_cute(attn_sink), row_sums=to_cute(r0)) + comp0(to_cute(q.unsqueeze(-1)), to_cute(k0.unsqueeze(-1)), + to_cute(v0.unsqueeze(-1)), to_cute(c0), stream, to_cute(l0), + swa_len=100, sink_bias=to_cute(attn_sink), row_sums=to_cute(r0)) + torch.cuda.synchronize() + ok0 = c0[:, :, 0].float() / r0[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) + cos0 = torch.nn.functional.cosine_similarity(ok0.flatten().unsqueeze(0), ref0.flatten().unsqueeze(0).float()).item() + lse0_kern = l0[:, 0, 0] + print(f'Seg0: cos {cos0:.6f} LSE_kern[0]={lse0_kern[0].item():.4f} LSE_ref[0]={lse0_ref[0].item():.4f}') - # Kernel 2: SWA only (n_comp=0, but sink bias still applies to all positions) - kernel_swa = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False, - apply_swa_mask=True, is_causal=False, n_comp=0, apply_sink_bias=True) - # Re-compile for n_comp=0 (different const_expr path) - compiled_swa = cute.compile(kernel_swa, _mQ, _mK, _mV, _mC, stream, _mLSE, - swa_len=swa_len, sink_bias=_mSB, row_sums=_mRS) + # Segment 1: positions 128-255 (all SWA, n_comp=0, sink bias on all, D3 mask at position 68) + k1 = k_combined[128:256]; v1 = v_combined[128:256] + scores1 = q.float() @ k1.float().T * scale + scores1 += attn_sink_val # all SWA → sink bias on all + # Valid SWA: absolute 96-195. In this segment (128-255): positions 0-67 valid, 68-127 masked + swa_len_seg1 = 68 # n_comp + swa_len - k_start = 96 + 100 - 128 = 68 + scores1[:, swa_len_seg1:] = float('-inf') + ref1 = (torch.softmax(scores1, dim=-1) @ v1.float()).to(torch.bfloat16) + lse1_ref = torch.logsumexp(scores1, dim=-1) - # Run each segment - segment_results = [] - n_comp_global = n_comp - swa_len_global = swa_len # number of valid SWA tokens (relative to SWA region start at n_comp) + k_seg1 = FmhaKernel(head_dim=hd, s_k=128, normalize=False, + apply_swa_mask=True, is_causal=False, n_comp=0, apply_sink_bias=True) + c1 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') + l1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + r1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') + comp1 = cute.compile(k_seg1, to_cute(q.unsqueeze(-1)), to_cute(k1.unsqueeze(-1)), + to_cute(v1.unsqueeze(-1)), to_cute(c1), stream, to_cute(l1), + swa_len=swa_len_seg1, sink_bias=to_cute(attn_sink), row_sums=to_cute(r1)) + comp1(to_cute(q.unsqueeze(-1)), to_cute(k1.unsqueeze(-1)), + to_cute(v1.unsqueeze(-1)), to_cute(c1), stream, to_cute(l1), + swa_len=swa_len_seg1, sink_bias=to_cute(attn_sink), row_sums=to_cute(r1)) + torch.cuda.synchronize() + ok1 = c1[:, :, 0].float() / r1[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) + cos1 = torch.nn.functional.cosine_similarity(ok1.flatten().unsqueeze(0), ref1.flatten().unsqueeze(0).float()).item() + lse1_kern = l1[:, 0, 0] + print(f'Seg1: cos {cos1:.6f} LSE_kern[0]={lse1_kern[0].item():.4f} LSE_ref[0]={lse1_ref[0].item():.4f}') - for seg_idx in range(n_segs): - k_start = seg_idx * seg_size - k_end = min(k_start + seg_size, n_total) - k_seg = k_combined[k_start:k_end] - v_seg = v_combined[k_start:k_end] - - if k_seg.shape[0] == 0: - continue - - # n_comp_local: compressed KV tokens in this segment - # If segment starts before n_comp_global, some positions are compressed. - # If segment starts at or after n_comp_global, all positions are SWA. - if k_start < n_comp_global: - n_comp_local = min(n_comp_global - k_start, seg_size) - else: - n_comp_local = 0 - - # swa_len_local: the kernel masks kv_pos >= n_comp_local + swa_len_local - # We want to mask absolute positions >= n_comp_global + swa_len_global - # So: n_comp_local + swa_len_local = n_comp_global + swa_len_global - k_start - # => swa_len_local = n_comp_global + swa_len_global - k_start - n_comp_local - swa_len_local = n_comp_global + swa_len_global - k_start - n_comp_local - # Clamp: if > remaining segment size, no masking needed - swa_len_local = min(swa_len_local, seg_size) - # If <= 0, all SWA positions are masked (no valid SWA) - if swa_len_local <= 0 and n_comp_local == 0: - continue # Skip empty segment - - # Whether this segment has SWA positions (needs sink bias) - has_swa = (k_start + seg_size) > n_comp_global - has_comp = n_comp_local > 0 - - # Pick the right kernel/compiled object - if has_comp: - kern = kernel_comp; comp = compiled_comp - else: - kern = kernel_swa; comp = compiled_swa - - o_norm, lse = run_segment( - q, k_seg, v_seg, kern, comp, stream, - sink_bias=attn_sink if has_swa else None, - n_comp_in_seg=n_comp_local, - swa_len=swa_len_local - ) - segment_results.append((o_norm, lse)) - - # Merge segments - o_merged = python_kv_merge(segment_results) - - # Compare + # Merge + o_merged = python_kv_merge([(ok0, lse0_kern), (ok1, lse1_kern)]) cos = torch.nn.functional.cosine_similarity( - o_merged.flatten().unsqueeze(0), - ref.flatten().unsqueeze(0).float() - ).item() - max_abs = (o_merged - ref.float()).abs().max().item() - + o_merged.flatten().unsqueeze(0), ref.flatten().unsqueeze(0).float()).item() status = "PASS" if cos >= 0.99 else "FAIL" - print(f'D5c multi-tile: cos {cos:.6f} max_abs {max_abs:.4f} {status}') - + print(f'\nMerged: cos {cos:.6f} {status}') if cos < 0.99: print(f' kernel[0,:4]={o_merged[0,:4].tolist()}') print(f' ref[0,:4]={ref[0,:4].tolist()}') - # Debug: check each segment individually against per-segment reference - for i, (o_n, lse) in enumerate(segment_results): - k_start = i * seg_size - k_end = min(k_start + seg_size, n_total) - k_seg = k_combined[k_start:k_end] - v_seg = v_combined[k_start:k_end] - n_comp_local = min(max(n_comp - k_start, 0), seg_size) - # Per-segment reference - scores_ref = q.float() @ k_seg.float().T * scale - if k_start < n_comp: - scores_ref[:, n_comp - k_start:] += attn_sink_val - swa_end = n_comp + swa_len - k_start - if swa_end < seg_size: - scores_ref[:, swa_end:] = float('-inf') - else: - scores_ref += attn_sink_val - swa_len_in_seg = n_comp + swa_len - k_start - if swa_len_in_seg < seg_size: - scores_ref[:, swa_len_in_seg:] = float('-inf') - o_ref_seg = (torch.softmax(scores_ref, dim=-1) @ v_seg.float()).to(torch.bfloat16) - lse_ref = torch.logsumexp(scores_ref, dim=-1) - cos_seg = torch.nn.functional.cosine_similarity( - o_n.flatten().unsqueeze(0), o_ref_seg.flatten().unsqueeze(0).float() - ).item() - print(f' Seg {i}: LSE_kern={lse[0].item():.4f} LSE_ref={lse_ref[0].item():.4f} cos={cos_seg:.6f}') if __name__ == '__main__':