biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: test_d5c use float for attn_sink in reference functions

Browse Source
This commit is contained in:
biondizzle
2026-05-26 15:01:31 +00:00
parent 9d64434954
commit 04b66e0f9c
1 changed files with 8 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/unit/test_d5c_fused.py
View File

@@ -91,7 +91,7 @@ def reference_sink_merge(q, k_comp, v_comp, k_swa, v_swa,
lse_swa = torch.logsumexp(attn_swa, dim=-1, keepdim=True)
# Sink merge (normalized formula)
exp_sink = attn_sink.exp()
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
@@ -117,7 +117,8 @@ def test_d5c_combined():
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 = torch.tensor([0.5], dtype=torch.float32, device='cuda')
attn_sink_val = 0.5
attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
# === FP32 References ===
qf = q[:, :, 0]
@@ -125,13 +126,13 @@ def test_d5c_combined():
# 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[0].item(), scale, swa_len
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[0].item(), scale, swa_len
attn_sink_val, scale, swa_len
)
# Verify the two references agree
@@ -222,12 +223,13 @@ def test_d5c_with_causal():
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 = torch.tensor([0.3], dtype=torch.float32, 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[0].item(), scale, swa_len, is_causal=True
attn_sink_val, scale, swa_len, is_causal=True
)
# Concatenate KV