Fix padded FP4 scaling, and enable its usage in flashinfer_scaled_fp4_mm

Signed-off-by: Roi Koren <roik@nvidia.com>

Author: roikoren755

vllm/_custom_ops.py

@@ -1128,7 +1128,7 @@ def scaled_fp4_quant(
         f'input.dtype needs to be fp16 or bf16 but got {input.dtype}.')
 
     # Two fp4 values will be packed into an uint8.
-    output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
+    output = torch.zeros((m, n // 2), device=device, dtype=torch.uint8)
 
     # We use the rounded values to store the swizzled values. Due to the
     # requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
@@ -1139,7 +1139,7 @@ def scaled_fp4_quant(
     rounded_m = round_up(m, 128)
     scale_n = n // block_size
     rounded_n = round_up(scale_n, 4)
-    output_scale = torch.empty((rounded_m, rounded_n // 4),
+    output_scale = torch.zeros((rounded_m, rounded_n // 4),
                                device=device,
                                dtype=torch.int32)
 

vllm/utils/flashinfer.py

@@ -368,8 +368,6 @@ def flashinfer_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,
     assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2
     assert a.stride(-1) == 1 and b.stride(-1) == 1
     assert a.shape[1] == b.shape[1]
-    assert block_scale_a.shape[1] == a.shape[1] // 8
-    assert block_scale_b.shape[1] == b.shape[1] // 8
 
     if backend == "cutlass":
         block_scale_a = block_scale_a.view(torch.uint8)