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

fix lse verify

Browse Source
This commit is contained in:
biondizzle
2026-05-24 22:23:08 +00:00
parent 3577e09603
commit c47f648617
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/unit/test_d1_lse_verify.py
View File

@@ -59,13 +59,14 @@ def test_lse(hd=64):
torch.cuda.synchronize()
kernel_lse = lse_tensor.cpu()
ref_lse_cpu = ref_lse.cpu()
# Compare
lse_err = (kernel_lse - ref_lse).abs()
lse_err = (kernel_lse - ref_lse_cpu).abs()
print(f' LSE max err: {lse_err.max().item():.6f}')
print(f' LSE mean err: {lse_err.mean().item():.6f}')
print(f' Kernel LSE[:8]: {kernel_lse[:8].tolist()}')
print(f' Ref LSE[:8]: {ref_lse[:8].tolist()}')
print(f' Ref LSE[:8]: {ref_lse_cpu[:8].tolist()}')
# Check O too
ref_unnorm = attn_exp @ v.float()