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 8, 01:15 UTC
[Bugfix] Fix benchmark_fused_collective crash on CustomOp init (#34665)
Signed-off-by: Mayank Ketkar <mketkar@zoox.com> Signed-off-by: Mayank Ketkar <mayket04@gmail.com> Co-authored-by: Mayank Ketkar <mketkar@zoox.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Mayank Ketkar committed
Feb 19, 2026 at 16:01 UTC
648951a9c3ab7d8ade25b80edb55eb4018acfd58
1 file changed
+11
-6
benchmarks/kernels/benchmark_fused_collective.py
+11
-6
@@ -408,18 +408,18 @@ def run_benchmarks(
408
409
rms_eps = 1e-6
410
results = {}
411
- vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
411
use_oneshot_options = [False] if no_oneshot else [True, False]
412
414
- # Create RMSNorm and QuantFP8 layers once for native benchmarks
415
-
413
if "none" in quant_modes:
414
# Standard AllReduce + RMSNorm
415
+ # Re-create VllmFusedAllreduce per config so CustomOp binds the
416
+ # correct forward method (native vs custom kernel).
417
for custom_op in ["-rms_norm", "+rms_norm"]:
418
with set_current_vllm_config(
419
VllmConfig(compilation_config=CompilationConfig(custom_ops=[custom_op]))
420
):
421
try:
422
+ vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
423
suffix = (
424
"_custom_rms_norm" if "+" in custom_op else "_native_rms_norm"
425
)
@@ -438,6 +438,7 @@ def run_benchmarks(
438
VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
439
):
440
try:
441
+ vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
442
standard_allreduce_rmsnorm_native_compiled = torch.compile(
443
vllm_fused_allreduce.allreduce_rmsnorm,
444
fullgraph=True,
@@ -482,7 +483,7 @@ def run_benchmarks(
483
"_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm"
484
)
485
for quant_fp8_custom_op in ["-quant_fp8", "+quant_fp8"]:
485
- suffix += (
486
+ op_suffix = suffix + (
487
"_custom_quant_fp8"
488
if "+" in quant_fp8_custom_op
489
else "_native_quant_fp8"
@@ -495,16 +496,17 @@ def run_benchmarks(
496
)
497
):
498
try:
499
+ vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
500
time_ms = benchmark_operation(
501
vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
502
input_tensor,
503
residual=residual,
504
scale_factor=scale_fp8,
505
)
504
- results[f"standard_allreduce{suffix}"] = time_ms
506
+ results[f"standard_allreduce{op_suffix}"] = time_ms
507
except Exception as e:
508
logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e)
507
- results[f"standard_allreduce{suffix}"] = float("inf")
509
+ results[f"standard_allreduce{op_suffix}"] = float("inf")
510
511
# Standard AllReduce + RMSNorm + FP8 Quant Native Compiled
512
with set_current_vllm_config(
@@ -515,6 +517,7 @@ def run_benchmarks(
517
)
518
):
519
try:
520
+ vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
521
standard_allreduce_rmsnorm_fp8_quant_native_compiled = torch.compile(
522
vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
523
fullgraph=True,
@@ -580,6 +583,7 @@ def run_benchmarks(
583
)
584
):
585
try:
586
+ vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
587
time_ms = benchmark_operation(
588
vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
589
input_tensor,
@@ -598,6 +602,7 @@ def run_benchmarks(
602
VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
603
):
604
try:
605
+ vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
606
standard_allreduce_rmsnorm_fp4_quant_native_compiled = torch.compile(
607
vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
608
fullgraph=True,