Related: #32358 Status: Under Discussion Written with: Claude

Summary

This RFC discusses how batch invariance should be handled within the vLLM IR dispatching model. Currently toggled via VLLM_BATCH_INVARIANT=1 and primarily affecting rms_norm, the question is: should vLLM IR own batch-invariant dispatching as a first-class concern, or should it be handled outside the IR?

Background

The vLLM IR design introduces a unified dispatching and kernel registration mechanism to replace CustomOp. A key motivation for this replacement is that CustomOp dispatching was opaque and hard to reason about. A clear example is that @yewentao256 did not realize that forward_native (the torch-native implementation) was being used by default in batch-invariant mode (and I didn't know at all what was happening), meaning the batch-invariant triton kernel was not invoked, as it only existed inside forward_cuda and forward_hip. Multiple levels of initialization and dispatching made the actual execution path difficult to trace.

Under CustomOp, batch invariance for rms_norm intersects with three execution paths:

• forward_cuda / forward_hip — dispatching between the standard CUDA/HIP custom kernel and batch invariant Triton kernel (also AITER in forward_hip)
• forward_native — used by default during compilation; batch-invariant dispatching is absent here, so it is silently skipped. Instead we relied on Inductor to generate a batch-invariant Triton kernel.

Under vLLM IR, the forward_* methods are eliminated and replaced with a single dispatching mechanism. How to represent and route batch-invariant execution within that model is an open question.

Proposed Alternatives

There are basically two decisions embedded in this:

• Do we want to dispatch batch-invariant rms_norm to native torch ops when compiling the model, or just always use the batch-invariant Triton kernel?
• If yes (alternatives 1-4), how should this dispatch happen?

Alternatives 1-4 are various ways to dispatch to the native implementation, and alternative 5 just always dispatches to the batch-invariant Triton kernel.

Option 1: vLLM IR handles batch-invariant dispatching natively (original plan, PR #36816)

Batch invariance is a first-class property in the IR. Each implementation declares whether it is batch invariant, and IrOp.dispatch filters implementations accordingly when the batch_invariant flag is set to true.

Option 2: Control batch-invariant kernels via vLLM IR priority list (my new preference)

Batch invariance is not encoded in the IR itself. Instead, when VLLM_BATCH_INVARIANT=1, the KernelConfig initialization logic adjusts the priority list for affected ops to prefer batch-invariant implementations:

class CudaPlatform:
  ...

  @classmethod
  def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
    # Native used by default when compiling,
    # use vllm_c kernels where available when no codegen
    cc = vllm_config.compilation_config
    using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
    default = ["native"] if using_inductor else ["vllm_c", "native"]

    if using_inductor:
      rms_norm = ["native"]
    elif envs.VLLM_BATCH_INVARIANT:
      rms_norm = ["triton_batch_invariant", "native"]
    else:
      rms_norm = ["vllm_c", "native"]

    return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)

The IR dispatches according to whatever priority list is active and remains unaware of batch invariance as a concept.

Option 3: Separate rms_norm_batch_invariant IR op

Define a distinct IR op vllm.ir.ops.rms_norm_batch_invariant with its own semantics declaration and implementation registry. Layer code selects which op to invoke based on whether batch invariance is required. This makes the path separate and even visible in the FX graph, but that also means more duplication and manual effort.

Option 4: Separate dispatching mechanism outside vLLM IR

When VLLM_BATCH_INVARIANT=1, execution bypasses IR dispatching entirely and dispatches between the batch-invariant Triton kernel and native implementation directly. The IR is not involved in this path. This does not allow OOT extensibility or custom compile passes but it does at least keep the benefit of native performance, although the native implementation has to be duplicated as we’re not using the IR op (we could manually invoke ir.ops.rms_norm.impls[“native”]).

Option 5: Batch-invariant path always calls the Triton kernel (Wentao's preference)

When batch invariance is required, always dispatch directly to the batch-invariant Triton kernel, with no IR involvement and no priority list manipulation. This matches current behavior on A100 and earlier GPUs, where eager execution is already forced.

B200 perf results for option 5, comparing triton_batch_invariant and native implementations (this is before #40408 & #40413 which close roughly half of that gap). Llama was the easiest to demonstrate this on but this applies to Qwen and Deepseek even more as they have more norms per layer (there are some temporary confounding issues right now out of scope of this RFC):

# triton_batch_invariant
$ VLLM_BATCH_INVARIANT=1 vllm bench latency --model=RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --attention-backend=TRITON_ATTN --ir-op-priority.rms_norm=triton_batch_invariant --ir-op-priority.fused_add_rms_norm=triton_batch_invariant -cc.pass_config.fuse_norm_quant=False
Avg latency: 0.43808772719154754 seconds
10% percentile latency: 0.43741614669561385 seconds
25% percentile latency: 0.43829121452290565 seconds
50% percentile latency: 0.4384383929427713 seconds
75% percentile latency: 0.4386132051004097 seconds
90% percentile latency: 0.43903175396844746 seconds
99% percentile latency: 0.43941899626515807 seconds

# native
$ VLLM_BATCH_INVARIANT=1 vllm bench latency --model=RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --attention-backend=TRITON_ATTN --ir-op-priority.rms_norm=native --ir-op-priority.fused_add_rms_norm=native -cc.pass_config.fuse_norm_quant=False
Avg latency: 0.3949168597658475 seconds
10% percentile latency: 0.3945025376044214 seconds
25% percentile latency: 0.39466362388338894 seconds
50% percentile latency: 0.3952090779785067 seconds
75% percentile latency: 0.3955157147720456 seconds
90% percentile latency: 0.3957540170755237 seconds
99% percentile latency: 0.3962868914008141 seconds

Discussion

The core trade-off is between visible, explicit logic in layer/config code and a clean, unified IR dispatching model with superior performance.

Options 4 and 5 are the simplest from an IR perspective: no batch-invariance-related logic enters the core dispatch mechanism. However, both place the batch-invariant path outside the IR entirely, making it a second-class citizen. OOT platforms (that don’t support the Triton kernel) would need to implement their own separate routing logic and override the whole layer rather than using the standard register_impl interface.

Option 5 is the simplest, but it carries the classic downsides of a custom kernel: lower performance portability, no portability beyond Triton kernels, no automatic fusion, higher maintenance cost. Option 4 recovers torch.compile-kernel-gen benefits (including 5% E2E perf) but still lacks OOT extensibility and IR benefits (fusion passes & sequence parallelism/async TP, although those don’t currently work with batch invariance anyway).

Option 3 adds explicit IR visibility but nearly duplicates the rms_norm op declaration, creates maintenance burden, and requires conditional logic in the RMSNorm layer — complexity that grows if other ops ever need batch-invariant variants.

Option 1 (current plan) keeps the batch-invariant path fully inside the IR, giving it access to unified dispatching, OOT register_impl, and Inductor fusions. The cost is that the IR dispatch model must understand the concept of batch invariance, adding a layer of reasoning to the core.

Option 2 splits the difference: batch-invariant kernels are still registered through the IR and benefit from unified dispatch and OOT extensibility, but the IR itself has no knowledge of batch invariance — it just follows the active priority list. The batch-invariant override logic lives in CudaPlatform.get_default_ir_op_priority, which is explicit and easy to audit. The main downside relative to Option 1 is that there is no enforcement that a registered implementation actually satisfies batch invariance — correctness depends on the platform providing the right priority list.

Recommendation

Option 1 was the original plan and is the most principled long-term design. However, Option 2 is likely a better starting point: it does not preclude migrating to Option 1 in the future, keeps the IR dispatch model slightly simpler while the IR is still stabilizing, and avoids encoding a concept that currently only affects a single op (rms_norm). No other ops are currently known to require batch-invariant variants, so building first-class IR support for this now is premature. If batch invariance becomes relevant to additional ops, or Inductor fusions on the batch-invariant path prove impactful, migrating to Option 1 remains straightforward.

Feedback from RL post-training users and OOT platform maintainers would be especially valuable.

Feedback Period.

4/22 - 4/24, this is a high-priority blocker for vLLM IR enablement

CC List.

@yewentao256 @simon-mo @zhuohan123 @houseroad @tjtanaa @zou3519 @gmagogsfm @angelayi @PaulZhang12 @tlrmchlsmth