@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b6cbba8bc8

[Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size (#48391)

Signed-off-by: oops-oom <73481342@qq.com> Signed-off-by: oops-oom <liubin8905@vip.qq.com> Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>

oops-oom committed Jul 28, 2026 at 22:24 UTC b6cbba8bc893c61e412a205533aafbee1ae6be31

6 files changed +189 -21

.buildkite/test_areas/misc.yaml

+7 -7

@@ -398,7 +398,7 @@ steps:

398

399 - label: Batch Invariance (A100)

400 key: batch-invariance-a100

401 - timeout_in_minutes: 40

401 + timeout_in_minutes: 60

402 device: a100

403 source_file_dependencies:

404 - vllm/v1/attention

@@ -408,11 +408,11 @@ steps:

408 - export VLLM_WORKER_MULTIPROC_METHOD=spawn

409 - pip install pytest-timeout pytest-forked

410 - pytest -v -s v1/determinism/test_batch_invariance.py

411 - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]

411 + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA

412

413 - label: Batch Invariance (H100)

414 key: batch-invariance-h100

415 - timeout_in_minutes: 40

415 + timeout_in_minutes: 60

416 device: h100

417 source_file_dependencies:

418 - vllm/v1/attention

@@ -423,8 +423,8 @@ steps:

423 - pip install pytest-timeout pytest-forked

424 - pytest -v -s v1/determinism/test_batch_invariance.py

425 - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py

426 - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]

427 - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]

426 + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA

427 + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN

428

429 - label: Batch Invariance (B200)

430 key: batch-invariance-b200

@@ -439,8 +439,8 @@ steps:

439 - pip install pytest-timeout pytest-forked

440 - pytest -v -s v1/determinism/test_batch_invariance.py

441 - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py

442 - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]

443 - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]

442 + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA

443 + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN

444 - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py

445 - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py

446 - pytest -v -s v1/determinism/test_matmul_batch_invariant.py

csrc/libtorch_stable/layernorm_kernels.cu

+10 -4

@@ -249,7 +249,9 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size]

249 int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;

250

251 // For large num_tokens, use smaller blocks to increase SM concurrency.

252 - const int max_block_size = (num_tokens < 256) ? 1024 : 256;

252 + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

253 + const int max_block_size =

254 + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);

255 dim3 grid(num_tokens);

256 const torch::stable::accelerator::DeviceGuard device_guard(

257 input.get_device_index());

@@ -325,8 +327,13 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]

327 /* This kernel is memory-latency bound in many scenarios.

328 When num_tokens is large, a smaller block size allows

329 for increased block occupancy on CUs and better latency

328 - hiding on global mem ops. */

329 - const int max_block_size = (num_tokens < 256) ? 1024 : 256;

330 + hiding on global mem ops. In batch-invariant mode the block size must

331 + not depend on num_tokens, otherwise the same token would use a different

332 + reduction width (and thus a different floating-point summation order)

333 + across batches; lock it to 1024 to keep results bit-exact. */

334 + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

335 + const int max_block_size =

336 + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);

337 dim3 block(std::min(hidden_size, max_block_size));

338 const torch::stable::accelerator::DeviceGuard device_guard(

339 input.get_device_index());

@@ -337,7 +344,6 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size]

344 auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());

345 bool offsets_are_multiple_of_vector_width =

346 hidden_size % vector_width == 0 && input_stride % vector_width == 0;

340 - bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

347 const bool has_weight = weight.has_value();

348 if (has_weight) {

349 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());

csrc/libtorch_stable/layernorm_quant_kernels.cu

+6 -3

@@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant(

215 int num_tokens = input.numel() / hidden_size;

216

217 // For large num_tokens, use smaller blocks to increase SM concurrency.

218 - const int max_block_size = (num_tokens < 256) ? 1024 : 256;

218 + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

219 + const int max_block_size =

220 + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);

221 dim3 grid(num_tokens);

222 const torch::stable::accelerator::DeviceGuard device_guard(

223 input.get_device_index());

@@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant(

281 When num_tokens is large, a smaller block size allows

282 for increased block occupancy on CUs and better latency

283 hiding on global mem ops. */

282 - const int max_block_size = (num_tokens < 256) ? 1024 : 256;

284 + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

285 + const int max_block_size =

286 + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);

287 dim3 block(std::min(hidden_size, max_block_size));

288 const torch::stable::accelerator::DeviceGuard device_guard(

289 input.get_device_index());

@@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant(

300 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());

301 bool ptrs_are_aligned =

302 inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;

299 - bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

303 if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&

304 !batch_invariant_launch) {

305 LAUNCH_FUSED_ADD_RMS_NORM(8);

csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

+4 -1

@@ -2,6 +2,7 @@

2 #include "../../torch_utils.h"

4 #include "../../dispatch_utils.h"

5 +#include "../../../core/batch_invariant.hpp"

6 #include "layernorm_utils.cuh"

7 #include "quant_conversions.cuh"

@@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch(

232 auto num_tokens = input.numel() / hidden_size;

233

234 dim3 grid(num_tokens);

234 - const int max_block_size = (num_tokens <= 256) ? 512 : 256;

235 + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();

236 + const int max_block_size =

237 + batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256);

238 dim3 block(std::min(hidden_size, max_block_size));

239 const torch::stable::accelerator::DeviceGuard device_guard(

240 input.get_device_index());

tests/v1/determinism/test_batch_invariance.py

+18

@@ -27,8 +27,10 @@ from vllm.platforms import current_platform

27 "backend",

28 BACKENDS,

29 )

30 +@pytest.mark.parametrize("rms_norm_impl", ["default", "vllm_c"])

31 def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(

32 backend,

33 + rms_norm_impl,

34 ):

35 """

36 Ensures that the same request (the 'needle' prompt) yields identical output

@@ -60,6 +62,16 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(

62 random.seed(seed)

64 attention_config = {"backend": backend}

65 + # Force the C++ RMSNorm implementation so we actually exercise the

66 + # num_tokens-dependent block-size branches.

67 + kernel_config = None

68 + if rms_norm_impl == "vllm_c":

69 + kernel_config = {

70 + "ir_op_priority": {

71 + "rms_norm": ["vllm_c"],

72 + "fused_add_rms_norm": ["vllm_c"],

73 + }

74 + }

75 # Allow overrides from environment (useful for CI tuning)

76 # "facebook/opt-125m" is too small, doesn't reliably test determinism

77 model = TEST_MODEL

@@ -96,6 +108,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(

108 gpu_memory_utilization=gpu_mem_util,

109 max_model_len=max_model_len,

110 attention_config=attention_config,

111 + kernel_config=kernel_config,

112 )

113

114 # Baseline generation for the needle prompt alone.

@@ -923,11 +936,15 @@ def LLM_with_max_seqs(

936 gpu_memory_utilization: float,

937 max_model_len: int,

938 attention_config: dict | None = None,

939 + kernel_config: dict | None = None,

940 ) -> LLM:

941 """

942 Helper to construct an LLM with a specific max_num_seqs (batch-size limit)

943 using the high-level v1 LLM API, while constraining memory usage.

944 """

945 + extra_kwargs: dict = {}

946 + if kernel_config is not None:

947 + extra_kwargs["kernel_config"] = kernel_config

948 return LLM(

949 model=model,

950 max_num_seqs=max_num_seqs,

@@ -939,4 +956,5 @@ def LLM_with_max_seqs(

956 attention_config=attention_config,

957 # Enable for MOE models

958 # enable_expert_parallel=True,

959 + **extra_kwargs,

960 )

tests/v1/determinism/test_rms_norm_batch_invariant.py

+144 -6

@@ -28,16 +28,18 @@ def _rms_norm_reference(

30 @skip_if_not_cuda

31 -@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])

31 +@pytest.mark.parametrize("batch_size", [1, 4, 64, 300])

32 @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])

33 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

34 @pytest.mark.parametrize("eps", [1e-6, 1e-5])

35 +@pytest.mark.parametrize("seed", list(range(4)))

36 def test_rms_norm_batch_invariant_vs_reference(

37 default_vllm_config,

38 batch_size: int,

39 hidden_size: int,

40 dtype: torch.dtype,

41 eps: float,

42 + seed: int,

43 ):

44 """

45 Compare batch-invariant Triton RMS norm against a PyTorch reference.

@@ -48,7 +50,7 @@ def test_rms_norm_batch_invariant_vs_reference(

50 device = torch.device(DEVICE_TYPE)

52 # Create test input and weight

51 - torch.manual_seed(42)

53 + torch.manual_seed(seed)

54 input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)

55 weight = torch.randn(hidden_size, dtype=dtype, device=device)

@@ -71,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_reference(

73 atol=atol,

74 msg=f"RMS norm mismatch for batch_size={batch_size}, "

75 f"hidden_size={hidden_size}, "

74 - f"dtype={dtype}, eps={eps}",

76 + f"dtype={dtype}, eps={eps}, seed={seed}",

77 )

@@ -79,17 +81,21 @@ def test_rms_norm_batch_invariant_vs_reference(

81 @pytest.mark.parametrize("hidden_size", [512, 4096])

82 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

83 @pytest.mark.parametrize("eps", [1e-6])

84 +@pytest.mark.parametrize("n_extra", [3, 299])

85 +@pytest.mark.parametrize("seed", list(range(16)))

86 def test_fused_add_rms_norm_batch_invariant_residual_path(

87 hidden_size: int,

88 dtype: torch.dtype,

89 eps: float,

90 + n_extra: int,

91 + seed: int,

92 ):

93 """

94 Test the batch-invariant fused residual-add + RMSNorm helper directly.

95 """

96 device = torch.device(DEVICE_TYPE)

92 - torch.manual_seed(42)

98 + torch.manual_seed(seed)

99 x_single = torch.randn(1, hidden_size, dtype=dtype, device=device)

100 residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device)

101 weight = torch.randn(hidden_size, dtype=dtype, device=device)

@@ -97,14 +103,14 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(

103 x_batch = torch.cat(

104 [

105 x_single,

100 - torch.randn(3, hidden_size, dtype=dtype, device=device),

106 + torch.randn(n_extra, hidden_size, dtype=dtype, device=device),

107 ],

108 dim=0,

109 )

110 residual_batch = torch.cat(

111 [

112 residual_single,

107 - torch.randn(3, hidden_size, dtype=dtype, device=device),

113 + torch.randn(n_extra, hidden_size, dtype=dtype, device=device),

114 ],

115 dim=0,

116 )

@@ -168,6 +174,138 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(

174 )

175

176

177 +FP8_DTYPE = current_platform.fp8_dtype()

178 +

179 +# The large launch (num_tokens=300 >= 256) drops an un-pinned kernel to block

180 +# 256, while the small launch (255 rows) stays under the threshold and keeps the

181 +# larger block (1024, or 512 for per-block quant). Under the pin the two launches

182 +# use the same block, so the shared first 255 rows must match bit-for-bit; 255 is

183 +# the most rows a single small launch can hold (< 256, and <= 256 for per-block).

184 +_LARGE_TOKENS = 300

185 +_SMALL_TOKENS = 255

186 +

187 +

188 +def _assert_rows_bit_identical(small, large, msg):

189 + if small.dtype == FP8_DTYPE:

190 + assert torch.equal(small.view(torch.uint8), large.view(torch.uint8)), msg

191 + else:

192 + torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg)

193 +

194 +

195 +@skip_if_not_cuda

196 +@pytest.mark.parametrize("hidden_size", [512, 4096])

197 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

198 +@pytest.mark.parametrize("seed", list(range(4)))

199 +def test_rms_norm_batch_invariant_nonresidual_kernel(

200 + hidden_size: int, dtype: torch.dtype, seed: int

201 +):

202 + """C++ ``rms_norm`` (no residual) must be batch invariant across the block

203 + threshold. Reached in compiled mode with ``ir_op_priority.rms_norm=["vllm_c"]``

204 + (default priority is ``native``/inductor codegen when compiling).

205 + """

206 + import vllm._custom_ops as ops

207 +

208 + device = torch.device(DEVICE_TYPE)

209 + torch.manual_seed(seed)

210 + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)

211 + weight = torch.randn(hidden_size, dtype=dtype, device=device)

212 +

213 + def rms_norm(x):

214 + out = torch.empty_like(x)

215 + ops.rms_norm(out, x, weight, 1e-6)

216 + return out

217 +

218 + large = rms_norm(rows.clone())

219 + small = rms_norm(rows[:_SMALL_TOKENS].clone())

220 + _assert_rows_bit_identical(

221 + small,

222 + large[:_SMALL_TOKENS],

223 + "rms_norm output depends on num_tokens (block size)",

224 + )

225 +

226 +

227 +@skip_if_not_cuda

228 +@pytest.mark.parametrize("hidden_size", [512, 4096])

229 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

230 +@pytest.mark.parametrize("seed", list(range(4)))

231 +@pytest.mark.parametrize("add_residual", [False, True])

232 +def test_rms_norm_static_fp8_quant_batch_invariant(

233 + hidden_size: int, dtype: torch.dtype, seed: int, add_residual: bool

234 +):

235 + """C++ static per-tensor fp8-quant RMSNorm must be batch invariant across

236 + the block threshold. Covers ``rms_norm_static_fp8_quant`` and, with

237 + ``add_residual``, ``fused_add_rms_norm_static_fp8_quant`` (the compiled fp8

238 + path where ``RMSNormQuantFusionPass`` rewrites norm + quant into them).

239 + """

240 + device = torch.device(DEVICE_TYPE)

241 + torch.manual_seed(seed)

242 + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)

243 + residual = (

244 + torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)

245 + if add_residual

246 + else None

247 + )

248 + weight = torch.randn(hidden_size, dtype=dtype, device=device)

249 + quant_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

250 +

251 + def quant(x, res):

252 + out = torch.empty_like(x, dtype=FP8_DTYPE)

253 + if add_residual:

254 + torch.ops._C.fused_add_rms_norm_static_fp8_quant(

255 + out, x, res, weight, quant_scale, 1e-6

256 + )

257 + else:

258 + torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, quant_scale, 1e-6)

259 + return out

260 +

261 + large = quant(rows.clone(), residual.clone() if residual is not None else None)

262 + small = quant(

263 + rows[:_SMALL_TOKENS].clone(),

264 + residual[:_SMALL_TOKENS].clone() if residual is not None else None,

265 + )

266 + _assert_rows_bit_identical(

267 + small,

268 + large[:_SMALL_TOKENS],

269 + "static-fp8-quant RMSNorm output depends on num_tokens (block size)",

270 + )

271 +

272 +

273 +@skip_if_not_cuda

274 +@pytest.mark.parametrize("hidden_size", [512, 4096])

275 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

276 +@pytest.mark.parametrize("seed", list(range(4)))

277 +def test_rms_norm_per_block_quant_batch_invariant(

278 + hidden_size: int, dtype: torch.dtype, seed: int

279 +):

280 + """C++ ``rms_norm_per_block_quant`` must be batch invariant across the

281 + block threshold (compiled fp8 block-quant path; block pinned to 512)."""

282 + import vllm._custom_ops as ops

283 +

284 + device = torch.device(DEVICE_TYPE)

285 + torch.manual_seed(seed)

286 + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)

287 + weight = torch.randn(hidden_size, dtype=dtype, device=device)

288 + group_size = [1, 128]

289 +

290 + def per_block_quant(x):

291 + return ops.rms_norm_per_block_quant(x, weight, 1e-6, FP8_DTYPE, group_size)

292 +

293 + out_large, scale_large = per_block_quant(rows.clone())

294 + out_small, scale_small = per_block_quant(rows[:_SMALL_TOKENS].clone())

295 + _assert_rows_bit_identical(

296 + out_small,

297 + out_large[:_SMALL_TOKENS],

298 + "rms_norm_per_block_quant output depends on num_tokens (block size)",

299 + )

300 + torch.testing.assert_close(

301 + scale_small,

302 + scale_large[:_SMALL_TOKENS],

303 + rtol=0.0,

304 + atol=0.0,

305 + msg="rms_norm_per_block_quant scales depend on num_tokens (block size)",

306 + )

307 +

308 +

309 @skip_if_not_cuda

310 @pytest.mark.parametrize("batch_size", [1, 16, 128])

311 @pytest.mark.parametrize("seq_len", [1, 32, 512])