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 13, 15:45 UTC
[amd][gptoss] Perf gain because of block alignment (#28024)
Signed-off-by: Smit Kadvani <smit.kadvani@gmail.com> Co-authored-by: Smit Shaileshbhai Kadvani <kadvani@meta.com>
smit kadvani committed
Nov 6, 2025 at 21:27 UTC
11fd69dd54060a59c6f62a6d217e1ecc47d74a68
2 files changed
+13
-2
vllm/model_executor/layers/quantization/mxfp4.py
+4
-2
@@ -43,6 +43,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
43
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
44
_can_support_mxfp4,
45
_swizzle_mxfp4,
46
+ get_padding_alignment,
47
)
48
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
49
from vllm.model_executor.utils import set_weight_attrs
@@ -282,10 +283,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
283
)
284
hidden_size = round_up(hidden_size, 128)
285
elif current_platform.is_rocm():
286
+ pad_align = get_padding_alignment()
287
intermediate_size_per_partition_after_pad = round_up(
286
- intermediate_size_per_partition, 256
288
+ intermediate_size_per_partition, pad_align
289
)
288
- hidden_size = round_up(hidden_size, 256)
290
+ hidden_size = round_up(hidden_size, pad_align)
291
else:
292
intermediate_size_per_partition_after_pad = round_up(
293
intermediate_size_per_partition, 64
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+9
@@ -7,6 +7,7 @@ import torch
7
8
from vllm.logger import init_logger
9
from vllm.platforms import current_platform
10
+from vllm.triton_utils import triton
11
from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
12
13
logger = init_logger(__name__)
@@ -99,6 +100,14 @@ def _can_support_mxfp4(
100
)
101
102
103
+def get_padding_alignment():
104
+ return (
105
+ 256
106
+ if triton.runtime.driver.active.get_current_target().arch in ("gfx950",)
107
+ else 128
108
+ )
109
+
110
+
111
def _dequant_mxfp4(
112
x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
113
) -> torch.Tensor: