359 lines
14 KiB
Python
359 lines
14 KiB
Python
"""
|
|
FMHA D5c: Fused sparse + SWA attention via combined KV + sink bias.
|
|
|
|
Mathematical insight: the sink merge is equivalent to a single attention
|
|
pass over the concatenated KV with a logit bias (attn_sink) applied to
|
|
the SWA portion. No two passes needed, no merge epilogue.
|
|
|
|
S = [S_comp, S_swa + attn_sink]
|
|
O = softmax(S) @ [V_comp; V_swa]
|
|
|
|
This is identical to:
|
|
O = (exp(lse_sparse)*O_sparse + exp(sink)*exp(lse_swa)*O_swa)
|
|
/ (exp(lse_sparse) + exp(sink)*exp(lse_swa))
|
|
|
|
The kernel changes are minimal:
|
|
1. K = [compressed_K; swa_K], V = [compressed_V; swa_V]
|
|
2. n_comp = length of compressed KV (sink bias applies to positions >= n_comp)
|
|
3. attn_sink = per-head logit bias for SWA positions
|
|
4. D3/D4 masking applies to SWA region (positions >= n_comp)
|
|
|
|
Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5c_fused.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_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]
|
|
|
|
# Concatenate KV
|
|
k_combined = torch.cat([k_comp, k_swa], dim=0) # (n_comp + n_swa, hd)
|
|
v_combined = torch.cat([v_comp, v_swa], dim=0)
|
|
|
|
# Compute combined logits
|
|
scores = torch.matmul(q.float(), k_combined.float().T) * scale # (m, n_comp + n_swa)
|
|
|
|
# Add sink bias to SWA positions
|
|
scores[:, n_comp:] += attn_sink
|
|
|
|
# D3: SWA length mask (only SWA region)
|
|
if swa_len < n_swa:
|
|
scores[:, n_comp + swa_len:] = float('-inf')
|
|
|
|
# D4: causal mask (only SWA region)
|
|
if is_causal:
|
|
# Within SWA region: mask k_coord > m_coord
|
|
for i in range(m):
|
|
for j in range(n_swa):
|
|
if j > i: # k_coord > m_coord
|
|
scores[i, n_comp + j] = float('-inf')
|
|
|
|
# Softmax + PV
|
|
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)
|
|
p = exp_s / sum_s
|
|
o = torch.matmul(p, v_combined.float())
|
|
return o.to(torch.bfloat16)
|
|
|
|
|
|
def reference_sink_merge(q, k_comp, v_comp, k_swa, v_swa,
|
|
attn_sink, scale, swa_len, is_causal=False):
|
|
"""FP32 reference: separate attention + sink merge (original D5b formula)."""
|
|
m, hd = q.shape
|
|
n_comp = k_comp.shape[0]
|
|
n_swa = k_swa.shape[0]
|
|
|
|
# Compressed KV attention (no mask)
|
|
attn_comp = torch.matmul(q.float(), k_comp.float().T) * scale
|
|
o_norm_comp = torch.softmax(attn_comp, dim=-1) @ v_comp.float()
|
|
lse_comp = torch.logsumexp(attn_comp, dim=-1, keepdim=True) # (m, 1)
|
|
|
|
# SWA KV attention (with swa_len mask)
|
|
attn_swa = torch.matmul(q.float(), k_swa.float().T) * scale
|
|
if swa_len < n_swa:
|
|
attn_swa[:, swa_len:] = float('-inf')
|
|
if is_causal:
|
|
for i in range(m):
|
|
for j in range(n_swa):
|
|
if j > i:
|
|
attn_swa[i, j] = float('-inf')
|
|
o_norm_swa = torch.softmax(attn_swa, dim=-1) @ v_swa.float()
|
|
lse_swa = torch.logsumexp(attn_swa, dim=-1, keepdim=True)
|
|
|
|
# Sink merge (normalized formula)
|
|
exp_sink = math.exp(attn_sink)
|
|
numerator = lse_comp.exp() * o_norm_comp + exp_sink * lse_swa.exp() * o_norm_swa
|
|
denominator = (lse_comp.exp() + exp_sink * lse_swa.exp()).clamp(min=1e-30)
|
|
o = numerator / denominator
|
|
return o.to(torch.bfloat16)
|
|
|
|
|
|
def test_d5c_combined():
|
|
print("=== Stage D5c: Fused Sparse+SWA via Combined KV + Sink Bias ===\n")
|
|
|
|
hd = 64
|
|
m = 128 # query rows
|
|
n_comp = 128 # compressed KV length
|
|
n_swa = 128 # SWA window length
|
|
n_total = n_comp + n_swa # combined KV length = 256
|
|
swa_len = 64 # actual SWA fill
|
|
scale = 1.0 / math.sqrt(hd)
|
|
torch.manual_seed(42)
|
|
|
|
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
k_comp = torch.randn(n_comp, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda')
|
|
k_swa = torch.randn(n_swa, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda')
|
|
|
|
# Sink weight (in log domain, per-head). n_h=1 so it's a scalar.
|
|
attn_sink_val = 0.5
|
|
attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
|
|
|
|
# === FP32 References ===
|
|
qf = q[:, :, 0]
|
|
|
|
# Reference 1: Combined softmax with sink bias (our kernel's approach)
|
|
ref_combined = reference_combined_attention(
|
|
qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa,
|
|
attn_sink_val, scale, swa_len
|
|
)
|
|
|
|
# Reference 2: Separate attention + sink merge (original D5b formula)
|
|
ref_merge = reference_sink_merge(
|
|
qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa,
|
|
attn_sink_val, scale, swa_len
|
|
)
|
|
|
|
# Verify the two references agree
|
|
cos_ref = torch.nn.functional.cosine_similarity(
|
|
ref_combined.flatten().unsqueeze(0).float(),
|
|
ref_merge.flatten().unsqueeze(0).float()
|
|
).item()
|
|
print(f"Reference: combined softmax vs sink merge cos = {cos_ref:.6f}")
|
|
assert cos_ref > 0.999, f"References don't match: cos={cos_ref}"
|
|
|
|
# === Kernel ===
|
|
# Concatenate KV for the kernel
|
|
k_combined = torch.cat([k_comp, k_swa], dim=0) # (n_total, hd, 1)
|
|
v_combined = torch.cat([v_comp, v_swa], dim=0) # (n_total, hd)
|
|
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
kernel = FmhaKernel(
|
|
head_dim=hd,
|
|
s_k=n_total, # combined KV length
|
|
normalize=False, # D5a: emit un-normalized O + LSE
|
|
apply_swa_mask=True, # D3: mask SWA positions
|
|
is_causal=False, # D4: no causal mask for this test
|
|
n_comp=n_comp, # D5c: compressed KV length (sink bias starts here)
|
|
)
|
|
|
|
# Allocate output
|
|
c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
|
|
# Prepare CuTe tensors
|
|
def to_cute(t):
|
|
return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
|
|
|
|
mQ = to_cute(q)
|
|
mK = to_cute(k_combined)
|
|
mV = to_cute(v_combined.unsqueeze(-1).contiguous())
|
|
mC = to_cute(c_out)
|
|
mLSE = to_cute(lse_out)
|
|
mRowSums = to_cute(row_sum_out)
|
|
|
|
# Compile
|
|
print('Compiling D5c kernel (combined KV + sink bias)...', flush=True)
|
|
mSinkBias = to_cute(attn_sink)
|
|
compiled = cute.compile(
|
|
kernel, mQ, mK, mV, mC, stream, mLSE,
|
|
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
|
|
)
|
|
|
|
# Run
|
|
print('Running D5c kernel...', flush=True)
|
|
compiled(
|
|
mQ, mK, mV, mC, stream, mLSE,
|
|
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# Check results
|
|
# Kernel outputs UN-NORMALIZED O (normalize=False). Normalize using per-row row_sum.
|
|
# O_norm[i] = O_unnorm[i] / row_sum[i]
|
|
o_kernel_unnorm = c_out[:, :, 0].float() # (m, hd)
|
|
row_sums = row_sum_out[:, 0, 0].float() # (m,)
|
|
|
|
# Normalize each row by its row_sum
|
|
o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30)
|
|
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
o_kernel.flatten().unsqueeze(0),
|
|
ref_combined.flatten().unsqueeze(0).float()
|
|
).item()
|
|
max_abs = (o_kernel - ref_combined.float()).abs().max().item()
|
|
|
|
status = "PASS" if cos >= 0.99 else "FAIL"
|
|
print(f'\nD5c result: cos {cos:.6f} max_abs {max_abs:.4f} {status}')
|
|
|
|
if cos < 0.99:
|
|
print(f' kernel[0,:4]={o_kernel[0,:4].tolist()}')
|
|
print(f' ref[0,:4]={ref_combined[0,:4].tolist()}')
|
|
print(f' row_sum range: {row_sums.min().item():.4f} to {row_sums.max().item():.4f}')
|
|
print(f' LSE range: {lse_out[:,0,0].min().item():.4f} to {lse_out[:,0,0].max().item():.4f}')
|
|
|
|
|
|
def test_d5c_with_causal():
|
|
"""D5c with causal mask on SWA branch."""
|
|
print("\n=== Stage D5c: Fused Sparse+SWA with Causal Mask ===\n")
|
|
|
|
hd = 64
|
|
m = 128
|
|
n_comp = 64
|
|
n_swa = 128
|
|
n_total = n_comp + n_swa
|
|
swa_len = 96 # partially filled SWA
|
|
scale = 1.0 / math.sqrt(hd)
|
|
torch.manual_seed(123)
|
|
|
|
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
k_comp = torch.randn(n_comp, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda')
|
|
k_swa = torch.randn(n_swa, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda')
|
|
|
|
attn_sink_val = 0.3
|
|
attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
|
|
qf = q[:, :, 0]
|
|
|
|
ref = reference_combined_attention(
|
|
qf, k_comp[:, :, 0], v_comp, k_swa[:, :, 0], v_swa,
|
|
attn_sink_val, scale, swa_len, is_causal=True
|
|
)
|
|
|
|
# Concatenate KV
|
|
k_combined = torch.cat([k_comp, k_swa], dim=0)
|
|
v_combined = torch.cat([v_comp, v_swa], dim=0)
|
|
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
kernel = FmhaKernel(
|
|
head_dim=hd,
|
|
s_k=n_total,
|
|
normalize=False,
|
|
apply_swa_mask=True,
|
|
is_causal=True,
|
|
n_comp=n_comp,
|
|
)
|
|
|
|
c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
|
|
def to_cute(t):
|
|
return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
|
|
|
|
mQ = to_cute(q)
|
|
mK = to_cute(k_combined)
|
|
mV = to_cute(v_combined.unsqueeze(-1).contiguous())
|
|
mC = to_cute(c_out)
|
|
mLSE = to_cute(lse_out)
|
|
mRowSums = to_cute(row_sum_out)
|
|
mSinkBias = to_cute(attn_sink)
|
|
compiled = cute.compile(
|
|
kernel, mQ, mK, mV, mC, stream, mLSE,
|
|
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
|
|
)
|
|
compiled(
|
|
mQ, mK, mV, mC, stream, mLSE,
|
|
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
o_kernel_unnorm = c_out[:, :, 0].float()
|
|
row_sums = row_sum_out[:, 0, 0].float()
|
|
o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30)
|
|
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
o_kernel.flatten().unsqueeze(0),
|
|
ref.flatten().unsqueeze(0).float()
|
|
).item()
|
|
max_abs = (o_kernel - ref.float()).abs().max().item()
|
|
|
|
status = "PASS" if cos >= 0.99 else "FAIL"
|
|
print(f'D5c causal result: cos {cos:.6f} max_abs {max_abs:.4f} {status}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# Test 0: baseline with s_k=128 (single KV tile, no O rescale)
|
|
print('=== Baseline: single-tile D3 mask (s_k=128, no D5c) ===\n')
|
|
_hd = 64; _m = 128; _s_k = 128; _swa_len = 64
|
|
_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')
|
|
_qf = _q[:, :, 0].float(); _kf = _k[:, :, 0].float(); _vf = _v.float()
|
|
_scores = _qf @ _kf.T * _scale
|
|
_scores[:, _swa_len:] = float('-inf')
|
|
_ref = (torch.softmax(_scores, dim=-1) @ _vf).to(torch.bfloat16)
|
|
_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
_kernel = FmhaKernel(head_dim=_hd, s_k=_s_k, normalize=False, apply_swa_mask=True, is_causal=False, n_comp=0)
|
|
_c_out = 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')
|
|
def _tc(t): return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
|
|
_mQ = _tc(_q); _mK = _tc(_k); _mV = _tc(_v.unsqueeze(-1).contiguous())
|
|
_mC = _tc(_c_out); _mLSE = _tc(_lse); _mRS = _tc(_rs)
|
|
_compiled = cute.compile(_kernel, _mQ, _mK, _mV, _mC, _stream, _mLSE, swa_len=_swa_len, row_sums=_mRS)
|
|
_compiled(_mQ, _mK, _mV, _mC, _stream, _mLSE, swa_len=_swa_len, row_sums=_mRS)
|
|
torch.cuda.synchronize()
|
|
_o_k = _c_out[:, :, 0].float() / _rs[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
|
|
_cos = torch.nn.functional.cosine_similarity(_o_k.flatten().unsqueeze(0), _ref.flatten().unsqueeze(0).float()).item()
|
|
print(f'Baseline (s_k=128, D3 mask): cos {_cos:.6f} {"PASS" if _cos > 0.99 else "FAIL"}\n')
|
|
|
|
# Test 1: D5c with single KV tile (n_comp=64, n_swa=64, s_k=128)
|
|
print('=== D5c: single-tile combined KV + sink bias ===\n')
|
|
_n_comp = 64; _n_swa = 64; _n_total = 128; _swa_len2 = 40
|
|
_attn_sink = torch.tensor([0.5], dtype=torch.float32, device='cuda')
|
|
_k_comp = torch.randn(_n_comp, _hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
_v_comp = torch.randn(_n_comp, _hd, dtype=torch.bfloat16, device='cuda')
|
|
_k_swa = torch.randn(_n_swa, _hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
_v_swa = torch.randn(_n_swa, _hd, dtype=torch.bfloat16, device='cuda')
|
|
_k_comb = torch.cat([_k_comp, _k_swa], dim=0)
|
|
_v_comb = torch.cat([_v_comp, _v_swa], dim=0)
|
|
_ref2 = reference_combined_attention(_qf, _k_comp[:,:,0], _v_comp, _k_swa[:,:,0], _v_swa, 0.5, _scale, _swa_len2)
|
|
_kernel2 = FmhaKernel(head_dim=_hd, s_k=_n_total, normalize=False, apply_swa_mask=True, is_causal=False, n_comp=_n_comp)
|
|
_c2 = torch.zeros(_m, _hd, 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')
|
|
_mK2 = _tc(_k_comb); _mV2 = _tc(_v_comb.unsqueeze(-1).contiguous())
|
|
_mC2 = _tc(_c2); _mLSE2 = _tc(_lse2); _mRS2 = _tc(_rs2)
|
|
_mSB2 = _tc(_attn_sink)
|
|
_comp2 = cute.compile(_kernel2, _mQ, _mK2, _mV2, _mC2, _stream, _mLSE2, swa_len=_swa_len2, sink_bias=_mSB2, row_sums=_mRS2)
|
|
_comp2(_mQ, _mK2, _mV2, _mC2, _stream, _mLSE2, swa_len=_swa_len2, sink_bias=_mSB2, row_sums=_mRS2)
|
|
torch.cuda.synchronize()
|
|
_ok2 = _c2[:, :, 0].float() / _rs2[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
|
|
_cos2 = torch.nn.functional.cosine_similarity(_ok2.flatten().unsqueeze(0), _ref2.flatten().unsqueeze(0).float()).item()
|
|
print(f'D5c single-tile: cos {_cos2:.6f} {"PASS" if _cos2 > 0.99 else "FAIL"}')
|
|
if _cos2 < 0.99:
|
|
print(f' kernel[0,:4]={_ok2[0,:4].tolist()}')
|
|
print(f' ref[0,:4]={_ref2[0,:4].tolist()}')
|
|
print(f' row_sum range: {_rs2[:,0,0].min().item():.4f} to {_rs2[:,0,0].max().item():.4f}')
|
|
|
|
test_d5c_combined()
|
|
test_d5c_with_causal()
|