A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 23:30 UTC
[Bugfix] Fix benchmark_fused_collective.py (#38082)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 26, 2026 at 14:51 UTC
2bfbdca23c60536f43850c9bec07fcc2b8f8e810
1 file changed
+19
-7
benchmarks/kernels/benchmark_fused_collective.py
+19
-7
@@ -25,6 +25,7 @@ import pandas as pd
25
import torch # type: ignore
26
import torch.distributed as dist # type: ignore
27
28
+from vllm._custom_ops import create_fp4_output_tensors
29
from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config
30
from vllm.distributed import (
31
tensor_model_parallel_all_reduce,
@@ -46,7 +47,7 @@ RMS_NORM_STATIC_FP8_QUANT_OP = torch.ops._C.rms_norm_static_fp8_quant
47
FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT_OP = (
48
torch.ops._C.fused_add_rms_norm_static_fp8_quant
49
)
49
-SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant
50
+SCALED_FP4_QUANT_OUT_OP = torch.ops._C.scaled_fp4_quant.out
51
52
logger = init_logger(__name__)
53
@@ -334,13 +335,23 @@ class VllmFusedAllreduce:
335
output_scale: torch.Tensor,
336
):
337
allreduce_out = tensor_model_parallel_all_reduce(input_tensor)
337
- rms_out = self.rms_norm(allreduce_out, residual)
338
+ rms_output = self.rms_norm(allreduce_out, residual)
339
+ if residual is None:
340
+ rms_out = rms_output
341
+ else:
342
+ rms_out, residual_out = rms_output
343
+
344
+ SCALED_FP4_QUANT_OUT_OP(
345
+ rms_out,
346
+ input_global_scale,
347
+ True,
348
+ output=quant_out,
349
+ output_scale=output_scale,
350
+ )
351
+
352
if residual is None:
339
- SCALED_FP4_QUANT_OP(quant_out, rms_out, output_scale, input_global_scale)
353
return quant_out, output_scale
354
else:
342
- rms_out, residual_out = rms_out
343
- SCALED_FP4_QUANT_OP(quant_out, rms_out, output_scale, input_global_scale)
355
return quant_out, residual_out, output_scale
356
357
@@ -362,8 +373,9 @@ def create_test_tensors(
373
scale_fp4 = torch.tensor(1.0, dtype=torch.float32)
374
quant_out_fp8 = torch.empty_like(input_tensor, dtype=FP8_DTYPE)
375
# Pre-allocate FP4 output tensors (to avoid allocation overhead in benchmarks)
365
- fp4_quant_out = torch.empty((num_tokens, hidden_dim // 2), dtype=torch.uint8)
366
- fp4_output_scale = torch.empty((128, 4), dtype=torch.int32)
376
+ fp4_quant_out, fp4_output_scale = create_fp4_output_tensors(
377
+ num_tokens, hidden_dim, input_tensor.device, True
378
+ )
379
380
return (
381
input_tensor,