Your current environment

Reproducible on L4 (SM89) GPUs in CI. Environment details to be added once reproduced on an L4 node.

🐛 Describe the bug

When VLLM_BATCH_INVARIANT=1 is set with torch.compile active (enforce_eager=False), RMSNorm's batch invariant path is bypassed. This causes batch-size-dependent numerical differences that break speculative decoding determinism on L4 (SM89).

Root Cause

vLLM's batch invariance registers aten::IMPL overrides in enable_batch_invariant_mode() (e.g., aten::mean.dim → deterministic Triton kernel with fixed reduction order). However, under torch.compile with custom_ops=["none"] (the default when Inductor is active):

1. CustomOp.dispatch_forward() → RMSNorm.enabled() returns False (because default_on() returns False when custom_ops=["none"])
2. RMSNorm dispatches to forward_native instead of forward_cuda
3. forward_native calls ir.ops.rms_norm → x.pow(2).mean(dim=-1, keepdim=True) (pure PyTorch ops)
4. Inductor lowers these to its own IR and generates its own Triton reduction kernels, bypassing the aten::mean.dim override at the dispatch table
5. The Inductor-generated reduction kernel produces batch-size-dependent results on L4

In contrast, forward_cuda calls rms_norm_batch_invariant() directly, which is deterministic regardless of batch size.

Why this wasn't caught earlier

The e2e batch invariant tests (tests/v1/determinism/test_batch_invariance.py) set enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, so on L4 (SM89) they run in eager mode → custom_ops=["all"] → forward_cuda is used. The issue was only exposed by test_eagle_dp which uses enforce_eager=False unconditionally.

Minimal repro

import torch
import os

os.environ["VLLM_BATCH_INVARIANT"] = "1"

from vllm.model_executor.layers.batch_invariant import enable_batch_invariant_mode
from vllm.model_executor.layers.layernorm import RMSNorm

enable_batch_invariant_mode()

hidden_size = 4096
eps = 1e-5
norm = RMSNorm(hidden_size, eps=eps).cuda()

torch.manual_seed(42)
x_single = torch.randn(1, hidden_size, dtype=torch.bfloat16, device="cuda")
x_batch = torch.randn(4, hidden_size, dtype=torch.bfloat16, device="cuda")
x_batch[0] = x_single[0]

# forward_cuda path (batch invariant) — should match
out_cuda_single = norm.forward_cuda(x_single.clone())
out_cuda_batch = norm.forward_cuda(x_batch.clone())
print("forward_cuda match:", torch.equal(out_cuda_single[0], out_cuda_batch[0]))

# forward_native path (compiled by Inductor) — may NOT match on L4
out_native_single = norm.forward_native(x_single.clone())
out_native_batch = norm.forward_native(x_batch.clone())
print("forward_native match:", torch.equal(out_native_single[0], out_native_batch[0]))

On L4, forward_cuda should produce matching results while forward_native may not.

Current workaround

In PR [#38938], the EAGLE DP test forces +rms_norm via compilation_config={"custom_ops": ["none", "+rms_norm"]} to ensure forward_cuda is used.

Related

• [#31913](https://github.com/vllm-project/vllm/issues/31913) — flaky test_eagle_dp
• [pytorch/pytorch#170563](https://github.com/pytorch/pytorch/issues/170563)
• [#38566](https://github.com/vllm-project/vllm/pull/38566) — workaround PR
• [#32992](https://github.com/vllm-project/vllm/issues/32992) — torch.compile + batch invariance

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.