diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py index f0de728b..e790ce81 100644 --- a/tests/test_pipeline_real_weights.py +++ b/tests/test_pipeline_real_weights.py @@ -160,10 +160,16 @@ def main(): l1_gsb = torch.tensor(l1_gs_list, dtype=torch.float32, device=DEVICE) # Quantize activation (dynamic gs, not warmup) + # KEY FIX: quantize slot_hidden (sorted tokens), NOT padded_hidden. + # padded_hidden has zeros in padding rows; quantizing it gives wrong x_sf layout. print("\n--- L1 GEMM (dynamic gs) ---") - x_fp4, x_sf, l1_gs = quantize_to_nvfp4(padded_hidden) + 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, 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. # Use the same _assemble_scales_cudagraph_safe function from vllm.nvfp4_cutedsl import CuTeDSLMoERunner @@ -192,13 +198,13 @@ def main(): # Just use the runner's scale assembly l1_gsa = torch.full((NUM_EXPERTS,), l1_gs, dtype=torch.float32, device=DEVICE) l1_scale_a = runner._assemble_scales_cudagraph_safe( - x_sf[:num_slots], expert_offsets[:NUM_EXPERTS+1], + slot_x_sf, expert_offsets[:NUM_EXPERTS+1], padded_expert_offsets, runner._padded_x_sf_buf_l1, runner._per_expert_scale_bufs_l1 ) l1_out = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=l1_mat_b, + mat_a=padded_x_fp4, mat_b=l1_mat_b, scale_a=l1_scale_a, scale_b=l1_scale_b, expert_offsets=padded_expert_offsets[1:NUM_EXPERTS+1], global_scale_a=l1_gsa, global_scale_b=l1_gsb,