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 6, 14:30 UTC
[Bugfix] opcheck false mutation error in rms_norm_per_block_quant (#36688) (#36779)
Signed-off-by: Krish Gupta <krishom70@gmail.com>
Krish Gupta committed
Mar 17, 2026 at 02:41 UTC
c0f011918da543f1323833c8ee2bfcac99e0452a
2 files changed
+19
-9
csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+9
@@ -286,6 +286,15 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
286
"Outer scale stride must be 1 when scales are not transposed");
287
}
288
289
+ int64_t hidden_size = input.size(-1);
290
+ TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
291
+ "hidden_size must be a positive multiple of group_size");
292
+ int64_t num_tokens = input.numel() / hidden_size;
293
+ int64_t num_groups = hidden_size / group_size;
294
+ TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
295
+ "scales buffer too small: need ", num_tokens * num_groups,
296
+ " elements, got ", scales.numel());
297
+
298
rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
299
var_epsilon, scale_ub, residual,
300
is_scale_transposed);
tests/kernels/core/test_fused_quant_layernorm.py
+10
-9
@@ -280,21 +280,22 @@ def test_rms_norm(
280
assert torch.allclose(ref_residual, ops_residual)
281
282
output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
283
- scales = torch.empty(
284
- (x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
285
- )
286
-
283
if group_size is None:
284
+ scales = torch.empty(
285
+ (x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
286
+ )
287
opcheck(
288
torch.ops._C.rms_norm_dynamic_per_token_quant,
289
(output, x, layer.weight, scales, 1e-5, scale_ub, residual),
290
)
291
else:
293
- # TODO(luka/eliza) opcheck is broken?
294
- # Somehow the cloned args are getting mutated in-place,
295
- # which causes the opcheck to fail.
296
- # https://github.com/vllm-project/vllm/issues/36688
297
- return
292
+ assert hidden_size % group_size[1] == 0
293
+ num_groups = hidden_size // group_size[1]
294
+ scales = torch.empty(
295
+ (num_groups, num_tokens),
296
+ device=x.device,
297
+ dtype=torch.float32,
298
+ ).transpose(0, 1)
299
opcheck(
300
torch.ops._C.rms_norm_per_block_quant,
301
(