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

Fix FP4 buffer shapes: D//2 for packed dimensions

Browse Source
This commit is contained in:
biondizzle
2026-05-17 21:26:46 +00:00
parent 5e4d674736
commit 364f8372bb
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_pipeline_real_weights.py
View File

@@ -167,7 +167,7 @@ def main():
print(f" L1 gs (dynamic): {l1_gs:.6f}")
# Scatter x_fp4 into padded layout
padded_x_fp4 = torch.zeros(total_padded, HIDDEN_SIZE, dtype=torch.uint8, device=DEVICE).view(torch.float4_e2m1fn_x2)
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
# For scale_a, we need to use the runner's assembly approach.

4
vllm/nvfp4_cutedsl.py
View File

@@ -167,13 +167,13 @@ class CuTeDSLMoERunner:
padded_max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device
),
'hidden_fp4': torch.zeros(
padded_max_slots, self.hidden_size, dtype=torch.uint8, device=self.device
padded_max_slots, self.hidden_size // 2, dtype=torch.uint8, device=self.device
).view(torch.float4_e2m1fn_x2),
'activated': torch.zeros(
padded_max_slots, self.intermediate_size, dtype=torch.bfloat16, device=self.device
),
'activated_fp4': torch.zeros(
padded_max_slots, self.intermediate_size, dtype=torch.uint8, device=self.device
padded_max_slots, self.intermediate_size // 2, dtype=torch.uint8, device=self.device
).view(torch.float4_e2m1fn_x2),
})
self._shared_bufs = CuTeDSLMoERunner._shared_padded_bufs[device_key]