From 364f8372bb3832d6facc6c9bd4c4127304cd2e67 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 21:26:46 +0000 Subject: [PATCH] Fix FP4 buffer shapes: D//2 for packed dimensions --- tests/test_pipeline_real_weights.py | 2 +- vllm/nvfp4_cutedsl.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py index e790ce81..f1760c4a 100644 --- a/tests/test_pipeline_real_weights.py +++ b/tests/test_pipeline_real_weights.py @@ -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. diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index bc98606e..8594208c 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -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]