@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 46f02e00f2

Fork 0

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 5, 14:45 UTC

[Bugfix] Fix AWQ models batch invariance issues (#38670)

Signed-off-by: yusuf <yusuf@deeplearningmachine.mynet> Signed-off-by: <> Co-authored-by: yusuf <yusuf@deeplearningmachine.mynet>

Yusuf Mohammad committed Apr 3, 2026 at 15:54 UTC 46f02e00f2d99944ecd24000bcc8be2a9504eafe

4 files changed +27 -10

tests/v1/determinism/test_batch_invariance.py

+5 -5

@@ -187,7 +187,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(

187 tensor_parallel_size=tp_size,

188 max_num_seqs=128,

189 max_model_len=8192,

190 - dtype="bfloat16", # not everything is supported

190 + dtype="auto", # not everything is supported

191 gpu_memory_utilization=0.9,

192 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,

193 attention_config={"backend": backend},

@@ -400,7 +400,7 @@ def test_simple_generation(backend):

400 tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),

401 gpu_memory_utilization=0.9,

402 max_model_len=2048,

403 - dtype="bfloat16",

403 + dtype="auto",

404 enable_prefix_caching=False,

405 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,

406 attention_config={"backend": backend},

@@ -466,7 +466,7 @@ def test_logprobs_without_batch_invariance_should_fail(

466 tensor_parallel_size=tp_size,

467 max_num_seqs=32,

468 max_model_len=8192,

469 - dtype="bfloat16",

469 + dtype="auto",

470 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,

471 attention_config={"backend": backend},

472 )

@@ -686,7 +686,7 @@ def test_decode_logprobs_match_prefill_logprobs(

686 tensor_parallel_size=tp_size,

687 max_num_seqs=32,

688 max_model_len=8192,

689 - dtype="bfloat16",

689 + dtype="auto",

690 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,

691 attention_config={"backend": backend},

692 )

@@ -931,7 +931,7 @@ def LLM_with_max_seqs(

931 max_num_seqs=max_num_seqs,

932 gpu_memory_utilization=gpu_memory_utilization,

933 max_model_len=max_model_len,

934 - dtype="bfloat16",

934 + dtype="auto",

935 tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),

936 enable_prefix_caching=False,

937 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,

vllm/model_executor/layers/batch_invariant.py

+12 -3

@@ -10,6 +10,7 @@ import vllm.envs as envs

10 from vllm.logger import init_logger

11 from vllm.platforms import current_platform

12 from vllm.triton_utils import tl, triton

13 +from vllm.utils.mem_utils import get_max_shared_memory_bytes

14 from vllm.utils.platform_utils import num_compute_units

15 from vllm.utils.torch_utils import is_torch_equal_or_newer

16 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -177,7 +178,7 @@ def matmul_persistent(

178 },

179 torch.float16: {

180 "BLOCK_SIZE_M": 128,

180 - "BLOCK_SIZE_N": 256,

181 + "BLOCK_SIZE_N": _fp16_block_size_n,

182 "BLOCK_SIZE_K": 64,

183 "GROUP_SIZE_M": 8,

184 "num_stages": 3,

@@ -700,7 +701,7 @@ def bmm_batch_invariant(a, b, *, out=None):

701 },

702 torch.float16: {

703 "BLOCK_SIZE_M": 128,

703 - "BLOCK_SIZE_N": 256,

704 + "BLOCK_SIZE_N": _fp16_block_size_n,

705 "BLOCK_SIZE_K": 64,

706 "num_stages": 3,

707 "num_warps": 8,

@@ -752,7 +753,8 @@ def addmm_batch_invariant(bias, a, b):

753

754

755 def _log_softmax_batch_invariant(input, dim, _half_to_float):

755 - assert not _half_to_float, "not implemented"

756 + if _half_to_float:

757 + return log_softmax(input.float(), dim=dim)

758 return log_softmax(input, dim=dim)

759

760

@@ -923,12 +925,15 @@ _original_fp16_reduction_precision = None

925 _original_bf16_reduction_precision = None

926 _original_cublas_workspace_cfg = None

927 _original_cublaslt_workspace_size = None

928 +_fp16_block_size_n = 256

929

930

931 def enable_batch_invariant_mode():

932 global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm

933 global _original_fp16_reduction_precision, _original_bf16_reduction_precision

934 global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size

935 + global _fp16_block_size_n

936 +

937 if _batch_invariant_MODE:

938 return

939

@@ -944,6 +949,10 @@ def enable_batch_invariant_mode():

949 _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")

950 _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")

951 _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")

952 +

953 + # Query the shared memory size and set block size

954 + # accordingly to avoid triton OutOfResources

955 + _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128

956 else:

957 # Only source of batch invariance for Hopper is split-k, can disable through

958 # cuBLAS workspace config

vllm/model_executor/layers/quantization/awq.py

+4 -2

@@ -8,6 +8,7 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE

8 from transformers import PretrainedConfig

10 from vllm import _custom_ops as ops

11 +from vllm import envs

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.fused_moe.layer import FusedMoE

14 from vllm.model_executor.layers.linear import (

@@ -273,8 +274,9 @@ class AWQLinearMethod(LinearMethodBase):

274

275 # num_tokens >= threshold

276 FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256

276 -

277 - if FP16_MATMUL_HEURISTIC_CONDITION:

277 + # Batch invariant mode requires torch.matmul path

278 + # for Triton override

279 + if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT:

280 out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)

281 out = torch.matmul(reshaped_x, out)

282 else:

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -10,6 +10,7 @@ from transformers import PretrainedConfig

11 import vllm.model_executor.layers.fused_moe # noqa

12 from vllm import _custom_ops as ops

13 +from vllm import envs

14 from vllm.logger import init_logger

15 from vllm.model_executor.kernels.linear import (

16 MPLinearLayerConfig,

@@ -233,6 +234,11 @@ class AWQMarlinConfig(QuantizationConfig):

234 def override_quantization_method(

235 cls, hf_quant_cfg, user_quant

236 ) -> "QuantizationMethods | None":

237 + # Skip override to marlin kernels, as they are not

238 + # batch invariant

239 + if envs.VLLM_BATCH_INVARIANT:

240 + return None

241 +

242 can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)

243 is_valid_user_quant = (

244 user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"