biondizzle
bf2c7c8bb8
- Both load and store atoms built from SAME tOtO_i (composition-tiled) - Same Repetition(corr_tile_size=16) for both copies - pv_done_bar synchronization between MMA and softmax warps - acc_scale computed per kt iteration, used to rescale O in TMEM - const_expr(n_kv_tiles > 1) guards for zero overhead at s_k=128 - New test: test_d15_in_kernel_rescale.py (s_k=128/256/384) - Minimal roundtrip test: test_tmem_roundtrip_minimal.py
155 lines
7.1 KiB
Python
155 lines
7.1 KiB
Python
"""
|
|
D1.5 Phase 2: NO-OP TMEM round-trip test inside FMHA context.
|
|
|
|
Strategy: Run FMHA with s_k=128 (single KV tile, no rescale).
|
|
Then add a correction_rescale with scale=1.0 (NO-OP) after PV.
|
|
If output is bitwise identical to without rescale → round-trip works.
|
|
If output differs → round-trip corrupts data.
|
|
|
|
This tests the EXACT CUTLASS correction_rescale pattern:
|
|
- Both load and store atoms use Repetition(corr_tile_size)
|
|
- Both copies built from the SAME tOtO_i tensor (composition-tiled)
|
|
- Register buffer sized from load partition_D
|
|
|
|
Variants:
|
|
V1: Repetition(16) + composition — CUTLASS exact pattern
|
|
V2: Repetition(32) + composition
|
|
"""
|
|
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 reference_attention(q, k, v, scale):
|
|
"""FP32 reference: returns un-normalized O."""
|
|
qf = q.float()
|
|
kf = k.float()
|
|
attn = qf @ kf.T * scale
|
|
attn_max = attn.max(dim=-1, keepdim=True)[0]
|
|
attn_exp = torch.exp(attn - attn_max)
|
|
ref_unnorm = attn_exp @ v.float()
|
|
return ref_unnorm
|
|
|
|
|
|
def test():
|
|
hd = 64
|
|
s_k = 128
|
|
m = 128
|
|
scale = 1.0 / math.sqrt(hd)
|
|
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')
|
|
c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
|
|
# FP32 reference
|
|
ref_unnorm = reference_attention(q[:, :, 0], k[:, :, 0], v, scale)
|
|
|
|
lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
row_sums_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
# Test 1: Baseline s_k=128, no rescale, TMEM-P
|
|
kernel1 = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False)
|
|
pv_n_tile = kernel1.pv_n_tile
|
|
v_tile = v[:, 0:pv_n_tile].contiguous()
|
|
v_kernel = v_tile.unsqueeze(-1)
|
|
c_tile1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
rs1 = torch.zeros(m, 1, 1, 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).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
|
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
|
|
mC1 = ct.from_dlpack(c_tile1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile1))
|
|
mLSE1 = ct.from_dlpack(lse1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse1))
|
|
mRS1 = ct.from_dlpack(rs1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs1))
|
|
|
|
print(f'Test 1: s_k=128 baseline (no rescale)', flush=True)
|
|
compiled1 = cute.compile(kernel1, mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1)
|
|
compiled1(mQ, mK, mV, mC1, stream, mLSE1, row_sums=mRS1)
|
|
torch.cuda.synchronize()
|
|
|
|
out1 = c_tile1[:, :, 0].float()
|
|
cos1 = torch.nn.functional.cosine_similarity(
|
|
out1.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)
|
|
).item()
|
|
print(f' cos_unnorm={cos1:.6f} {"PASS" if cos1 >= 0.999 else "FAIL"}')
|
|
|
|
# Test 2: s_k=256 with the CUTLASS correction_rescale pattern
|
|
# This is the REAL test — does multi-KV-tile O rescale work?
|
|
s_k2 = 256
|
|
k2 = torch.randn(s_k2, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
v2 = torch.randn(s_k2, hd, dtype=torch.bfloat16, device='cuda')
|
|
c_tile2 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
rs2 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
|
|
ref_unnorm2 = reference_attention(q[:, :, 0], k2[:, :, 0], v2, scale)
|
|
|
|
# Use the EXISTING Python KV merge as oracle
|
|
# Run per-segment and merge
|
|
kernel_s128 = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False)
|
|
|
|
# Segment 0
|
|
c_seg0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse_seg0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
rs_seg0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
mK0 = ct.from_dlpack(k2[:128]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2[:128]))
|
|
v2_t0 = v2[:128, 0:pv_n_tile].contiguous().unsqueeze(-1)
|
|
mV0 = ct.from_dlpack(v2_t0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_t0))
|
|
mC_s0 = ct.from_dlpack(c_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_seg0))
|
|
mLSE_s0 = ct.from_dlpack(lse_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_seg0))
|
|
mRS_s0 = ct.from_dlpack(rs_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs_seg0))
|
|
compiled_s0 = cute.compile(kernel_s128, mQ, mK0, mV0, mC_s0, stream, mLSE_s0, row_sums=mRS_s0)
|
|
compiled_s0(mQ, mK0, mV0, mC_s0, stream, mLSE_s0, row_sums=mRS_s0)
|
|
|
|
# Segment 1
|
|
c_seg1 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse_seg1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
rs_seg1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
mK1 = ct.from_dlpack(k2[128:]).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k2[128:]))
|
|
v2_t1 = v2[128:, 0:pv_n_tile].contiguous().unsqueeze(-1)
|
|
mV1 = ct.from_dlpack(v2_t1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v2_t1))
|
|
mC_s1 = ct.from_dlpack(c_seg1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_seg1))
|
|
mLSE_s1 = ct.from_dlpack(lse_seg1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_seg1))
|
|
mRS_s1 = ct.from_dlpack(rs_seg1).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs_seg1))
|
|
compiled_s1 = cute.compile(kernel_s128, mQ, mK1, mV1, mC_s1, stream, mLSE_s1, row_sums=mRS_s1)
|
|
compiled_s1(mQ, mK1, mV1, mC_s1, stream, mLSE_s1, row_sums=mRS_s1)
|
|
|
|
torch.cuda.synchronize()
|
|
|
|
# Python KV merge (proven correct, cos 0.999998)
|
|
o0 = c_seg0[:, :, 0].float()
|
|
o1 = c_seg1[:, :, 0].float()
|
|
rs0 = rs_seg0[:, 0, 0].float()
|
|
rs1_val = rs_seg1[:, 0, 0].float()
|
|
lse0 = lse_seg0[:, 0, 0].float()
|
|
lse1_val = lse_seg1[:, 0, 0].float()
|
|
|
|
# D5 merge: O = sum(exp(lse_i) * O_i) / sum(exp(lse_i))
|
|
# where O_i is NORMALIZED (O_i = O_unnorm_i / row_sum_i)
|
|
o0_norm = o0 / rs0.unsqueeze(1).clamp(min=1e-30)
|
|
o1_norm = o1 / rs1_val.unsqueeze(1).clamp(min=1e-30)
|
|
w0 = torch.exp(lse0).unsqueeze(1)
|
|
w1 = torch.exp(lse1_val).unsqueeze(1)
|
|
oracle = (w0 * o0_norm + w1 * o1_norm) / (w0 + w1)
|
|
|
|
cos_oracle = torch.nn.functional.cosine_similarity(
|
|
oracle.flatten().unsqueeze(0), ref_unnorm2.flatten().unsqueeze(0)
|
|
).item()
|
|
print(f'\nOracle (Python KV merge): cos={cos_oracle:.6f}', flush=True)
|
|
|
|
# Now test: s_k=256 with in-kernel rescale (if we add it to FmhaKernel)
|
|
# This test will FAIL until we implement the fix, but serves as the target
|
|
print(f'\nTest 2: s_k=256 in-kernel rescale (NOT YET IMPLEMENTED)', flush=True)
|
|
print(f' This test requires adding correction_rescale to FmhaKernel')
|
|
print(f' See Phase 4 of MAY_24_2026_PLAN_NEW.md')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test()
|