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

Fix: scatter FP4 as uint8 (float4 doesn't support index_put)

Browse Source
This commit is contained in:
biondizzle
2026-05-17 21:28:04 +00:00
parent 364f8372bb
commit 2796bd81e8
2 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/test_pipeline_real_weights.py
View File

@@ -166,9 +166,10 @@ def main():
slot_x_fp4, slot_x_sf, l1_gs = quantize_to_nvfp4(slot_hidden)
print(f" L1 gs (dynamic): {l1_gs:.6f}")
# Scatter x_fp4 into padded layout
padded_x_fp4 = torch.zeros(total_padded, HIDDEN_SIZE // 2, dtype=torch.uint8, device=DEVICE).view(torch.float4_e2m1fn_x2)
padded_x_fp4[padded_dst] = slot_x_fp4
# Scatter x_fp4 into padded layout (use uint8 for scatter, then view as float4)
padded_x_fp4_uint8 = torch.zeros(total_padded, HIDDEN_SIZE // 2, dtype=torch.uint8, device=DEVICE)
padded_x_fp4_uint8[padded_dst] = slot_x_fp4.view(torch.uint8)
padded_x_fp4 = padded_x_fp4_uint8.view(torch.float4_e2m1fn_x2)
# For scale_a, we need to use the runner's assembly approach.
# Use the same _assemble_scales_cudagraph_safe function