biondizzle/deepseek-v4-quant
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: fix float8 randn + SF int32 packing

Browse Source
This commit is contained in:
biondizzle
2026-05-12 15:12:35 +00:00
parent 28afc2406b
commit d4c557fddc
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
patches/test_nvfp4_mega_moe.py
View File

@@ -82,9 +82,16 @@ def test_nvfp4_mega_moe():
symm_buffer.x[:num_tokens].copy_(
torch.randint(0, 256, (num_tokens, hidden // 2), dtype=torch.uint8, device=device).view(torch.int8))
# Write valid UE4M3 scales (random but non-zero)
sf_data = torch.randn(num_tokens, hidden // 64, dtype=torch.float8_e4m3fn, device=device).abs().clamp(0.1, 5.0)
symm_buffer.x_sf[:num_tokens].copy_(
sf_data.view(torch.uint8).reshape(num_tokens, hidden // 64).contiguous().view(torch.int32))
# Can't randn with float8, so generate in bf16 then cast
sf_bf16 = torch.randn(num_tokens, hidden // 64, dtype=torch.bfloat16, device=device).abs().clamp(0.1, 5.0)
# Pack as int32 (4 UE4M3 bytes per int32)
sf_u8 = sf_bf16.to(torch.float8_e4m3fn).view(torch.uint8)
# Pack 4 uint8 into int32
sf_packed = (sf_u8[..., 0::4].to(torch.int32) |
(sf_u8[..., 1::4].to(torch.int32) << 8) |
(sf_u8[..., 2::4].to(torch.int32) << 16) |
(sf_u8[..., 3::4].to(torch.int32) << 24))
symm_buffer.x_sf[:num_tokens].copy_(sf_packed)
# Write topk data directly
for i in range(num_tokens):
for j in range(top_k):