[Bugfix] Fix benchmark_fused_collective crash on CustomOp init

[mayank-ketkar-sf]

Summary

Fix two bugs in benchmarks/kernels/benchmark_fused_collective.py:

1. Crash: VllmFusedAllreduce creates RMSNorm and QuantFP8 (CustomOp subclasses) which call dispatch_forward() → get_current_vllm_config() during __init__. Without an active VllmConfig context, the benchmark crashes immediately.

2. Incorrect dispatch: CustomOp binds its forward method (native PyTorch vs CUDA custom kernel) at init time. The original code created VllmFusedAllreduce once and reused it across different custom_ops configurations, meaning the native-vs-custom comparison was measuring the same code path.

Fix

Instantiate VllmFusedAllreduce inside each set_current_vllm_config() block so the forward dispatch matches the intended configuration.

Steps to Reproduce (crash)

# On any multi-GPU setup with vLLM installed:
torchrun --nproc_per_node=2 benchmarks/kernels/benchmark_fused_collective.py \
  --num-tokens 128 512 --hidden-dim 4096 --quant-modes none

Before fix: Crashes with AssertionError: Current vLLM config is not set
After fix: Benchmark runs, correctly dispatching native vs custom kernel per config

Test plan

• torchrun --nproc_per_node=2 benchmarks/kernels/benchmark_fused_collective.py --num-tokens 128 512 --hidden-dim 4096 --quant-modes none completes without crash
• Output includes distinct timings for _native_rms_norm vs _custom_rms_norm variants
• No changes to production code — benchmark-only fix

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request fixes a crash in the fused collective benchmark by wrapping the VllmFusedAllreduce instantiation in a VllmConfig context. The fix is correct in addressing the crash. However, I've identified a critical issue with the benchmark's logic that this change highlights. The CustomOp implementations are bound at initialization, but the benchmark reuses the same object across different configurations, leading to incorrect measurements. I've left a detailed comment explaining the issue and suggesting a path to correct the benchmark's logic.

[mayank-ketkar-sf]

Ran into this crash while benchmarking allreduce fusion on H200. Traced it to VllmFusedAllreduce being instantiated outside a VllmConfig context, and also noticed the CustomOp forward dispatch was only bound once regardless of the config switch. Updated to re-instantiate per config block.

cc @ilmarkov @mmangkad for review

[mmangkad]

Left some comments. I also encountered this.
cc @ProExpertProg

[ProExpertProg]

Thanks for the fix! @mayank-ketkar-sf the suggestions from @mmangkad sound good to me, could you apply them and then we can merge?

[mayank-ketkar-sf]

@ProExpertProg @mmangkad ofcourse. thank you for the feedback. I am in agreement, so I can update it and respond back today :) Cheers both

[mayank-ketkar-sf]

@ProExpertProg @mmangkad could you add ci tag if all looks okay to merge?

[mayank-ketkar-sf]

@ProExpertProg shall we merge this?

[mayank-ketkar-sf]

uff I am sorry about back n forth. My first PR here, should be better next time. I will rebase on latest and send back again

[mayank-ketkar-sf]

@ProExpertProg Auto-merge was disabled because I pushed new commits. Could you re-enable it when you get
a chance? Thanks!
Apologies, I am learning the ropes here