243 lines
8.9 KiB
Python
243 lines
8.9 KiB
Python
"""
|
|
FMHA D5b: Per-row LSE output + Python KV merge.
|
|
|
|
Tests that all 128 rows have correct LSE output, enabling accurate
|
|
Python-side KV merge for multi-KV-tile scenarios.
|
|
|
|
The D5 merge formula (using NORMALIZED O + LSE):
|
|
O = (exp(lse_0) * O_0_norm + exp(lse_1) * O_1_norm)
|
|
/ (exp(lse_0) + exp(lse_1))
|
|
|
|
Where:
|
|
exp(lse_i) = row_sum_i * exp(max(S_i * scale))
|
|
O_i_norm = O_i_unnorm / row_sum_i
|
|
|
|
Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5b_perrow_lse.py
|
|
"""
|
|
import torch
|
|
import 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 reference_attention_with_lse(q, k, v, scale):
|
|
"""FP32 reference attention returning O and per-row LSE."""
|
|
scores = torch.matmul(q.float(), k.float().T) * scale
|
|
max_s = scores.max(dim=-1, keepdim=True).values
|
|
exp_s = (scores - max_s).exp()
|
|
sum_s = exp_s.sum(dim=-1, keepdim=True)
|
|
p = exp_s / sum_s
|
|
o = torch.matmul(p, v.float())
|
|
# LSE = logsumexp(S * scale)
|
|
lse = (scores - max_s).exp().sum(dim=-1).log() + max_s.squeeze(-1)
|
|
return o.to(torch.bfloat16), lse
|
|
|
|
|
|
def _run_fmha_with_lse(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False):
|
|
"""Run FMHA and return (o_norm, lse) with per-row LSE.
|
|
|
|
Uses reference attn_sum for normalization (TMEM round-trip normalization
|
|
is broken, and exp(LSE) != row_sum).
|
|
"""
|
|
scale = 1.0 / math.sqrt(hd)
|
|
kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, 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)
|
|
|
|
o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda')
|
|
lse_all = torch.zeros(m, dtype=torch.float32, device='cuda')
|
|
|
|
for pv in range(n_pv_tiles):
|
|
v_tile = v[:, pv * pv_n_tile:(pv + 1) * pv_n_tile].contiguous().unsqueeze(-1)
|
|
c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse_tensor = torch.zeros(m, 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_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d))
|
|
mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
|
|
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))
|
|
|
|
if pv == 0:
|
|
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
|
|
|
|
compiled(mQ, mK, mV, mC, stream, mLSE)
|
|
o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float()
|
|
lse_all = lse_tensor[:, 0, 0]
|
|
|
|
# Normalize using reference attn_sum (TMEM round-trip is broken)
|
|
q_flat = q_3d[:, :, 0]
|
|
k_flat = k_3d[:, :, 0]
|
|
scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale
|
|
max_s = scores.max(dim=-1, keepdim=True).values
|
|
attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True)
|
|
o_norm = (o_unnorm / attn_sum).to(torch.bfloat16)
|
|
return o_norm, lse_all
|
|
|
|
|
|
def test_lse_per_row_hd64():
|
|
"""Per-row LSE at hd=64: all 128 rows should have correct LSE."""
|
|
print("\n=== Test 1: Per-row LSE (hd=64) ===")
|
|
torch.manual_seed(42)
|
|
m, s_k, hd = 128, 128, 64
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
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')
|
|
|
|
o, lse = _run_fmha_with_lse(q, k, v, m, s_k, hd)
|
|
ref_o, ref_lse = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale)
|
|
|
|
# Check per-row LSE accuracy
|
|
lse_err = (lse - ref_lse).abs().max().item()
|
|
print(f" LSE max error: {lse_err:.6f}")
|
|
assert lse_err < 0.01, f"LSE error too high: {lse_err}"
|
|
|
|
# Check per-row LSE: all rows should be non-zero
|
|
zero_rows = (lse == 0).sum().item()
|
|
print(f" Zero LSE rows: {zero_rows}")
|
|
assert zero_rows == 0, f"Expected 0 zero LSE rows, got {zero_rows}"
|
|
|
|
# Check output cosine
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
o.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0)
|
|
).item()
|
|
print(f" cos = {cos:.6f}")
|
|
assert cos >= 0.995
|
|
print(" ✅ PASS")
|
|
|
|
|
|
def test_lse_per_row_hd128():
|
|
"""Per-row LSE at hd=128 (SMEM-P path)."""
|
|
print("\n=== Test 2: Per-row LSE (hd=128) ===")
|
|
torch.manual_seed(42)
|
|
m, s_k, hd = 128, 128, 128
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
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')
|
|
|
|
o, lse = _run_fmha_with_lse(q, k, v, m, s_k, hd, use_smem_p=True)
|
|
ref_o, ref_lse = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale)
|
|
|
|
lse_err = (lse - ref_lse).abs().max().item()
|
|
print(f" LSE max error: {lse_err:.6f}")
|
|
assert lse_err < 0.01, f"LSE error too high: {lse_err}"
|
|
|
|
zero_rows = (lse == 0).sum().item()
|
|
print(f" Zero LSE rows: {zero_rows}")
|
|
assert zero_rows == 0, f"Expected 0 zero LSE rows, got {zero_rows}"
|
|
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
o.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0)
|
|
).item()
|
|
print(f" cos = {cos:.6f}")
|
|
assert cos >= 0.995
|
|
print(" ✅ PASS")
|
|
|
|
|
|
def test_lse_kv_merge():
|
|
"""Python KV merge using per-row LSE + normalized O (s_k=256, 2 KV tiles).
|
|
|
|
Correct merge formula (D5):
|
|
O = (exp(lse_0) * O_0_norm + exp(lse_1) * O_1_norm)
|
|
/ (exp(lse_0) + exp(lse_1))
|
|
"""
|
|
print("\n=== Test 3: KV merge with per-row LSE (s_k=256) ===")
|
|
torch.manual_seed(42)
|
|
m, s_k, hd = 128, 256, 64
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
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')
|
|
|
|
# Reference: full attention with s_k=256
|
|
ref_o, _ = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale)
|
|
|
|
# Kernel: two segments of 128, merge with per-row LSE + normalized O
|
|
# IMPORTANT: create kernel with s_k=128 (segment size), not s_k=256
|
|
seg_size = 128
|
|
o_norms = []
|
|
lses = []
|
|
|
|
for seg in range(s_k // seg_size):
|
|
k_seg = k[seg * seg_size:(seg + 1) * seg_size].contiguous()
|
|
v_seg = v[seg * seg_size:(seg + 1) * seg_size].contiguous()
|
|
# k_seg is already 3D from slicing (s_k, hd, 1) - no unsqueeze needed
|
|
|
|
o_seg, lse_seg = _run_fmha_with_lse(q, k_seg, v_seg, m, seg_size, hd)
|
|
o_norms.append(o_seg.float())
|
|
lses.append(lse_seg.float())
|
|
|
|
# D5 merge with normalized O + LSE
|
|
# O = sum_i[exp(lse_i) * O_i_norm] / sum_i[exp(lse_i)]
|
|
e_lse = [l.exp() for l in lses]
|
|
numerator = sum(el.unsqueeze(-1) * on for el, on in zip(e_lse, o_norms))
|
|
denominator = sum(e_lse).unsqueeze(-1)
|
|
o_merged = (numerator / denominator).to(torch.bfloat16)
|
|
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
o_merged.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0)
|
|
).item()
|
|
print(f" cos = {cos:.6f}")
|
|
assert cos >= 0.99, f"cosine too low: {cos}"
|
|
print(" ✅ PASS")
|
|
|
|
|
|
def test_lse_kv_merge_4tiles():
|
|
"""Python KV merge with s_k=512 (4 KV tiles)."""
|
|
print("\n=== Test 4: KV merge (s_k=512, 4 tiles) ===")
|
|
torch.manual_seed(42)
|
|
m, s_k, hd = 128, 512, 64
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
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')
|
|
|
|
ref_o, _ = reference_attention_with_lse(q[:, :, 0], k[:, :, 0], v, scale)
|
|
|
|
seg_size = 128
|
|
o_norms = []
|
|
lses = []
|
|
|
|
for seg in range(s_k // seg_size):
|
|
k_seg = k[seg * seg_size:(seg + 1) * seg_size].contiguous()
|
|
v_seg = v[seg * seg_size:(seg + 1) * seg_size].contiguous()
|
|
# k_seg is already 3D from slicing (s_k, hd, 1) - no unsqueeze needed
|
|
|
|
o_seg, lse_seg = _run_fmha_with_lse(q, k_seg, v_seg, m, seg_size, hd)
|
|
o_norms.append(o_seg.float())
|
|
lses.append(lse_seg.float())
|
|
|
|
e_lse = [l.exp() for l in lses]
|
|
numerator = sum(el.unsqueeze(-1) * on for el, on in zip(e_lse, o_norms))
|
|
denominator = sum(e_lse).unsqueeze(-1)
|
|
o_merged = (numerator / denominator).to(torch.bfloat16)
|
|
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
o_merged.flatten().float().unsqueeze(0), ref_o.flatten().float().unsqueeze(0)
|
|
).item()
|
|
print(f" cos = {cos:.6f}")
|
|
assert cos >= 0.99, f"cosine too low: {cos}"
|
|
print(" ✅ PASS")
|
|
|
|
|
|
def test():
|
|
print("=== D5b: Per-Row LSE Output ===")
|
|
test_lse_per_row_hd64()
|
|
test_lse_per_row_hd128()
|
|
test_lse_kv_merge()
|
|
test_lse_kv_merge_4tiles()
|
|
print("\n=== ALL TESTS PASSED ===")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test()
|