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 14, 05:30 UTC
[ROCm] Small functional changes for gptoss (#25201)
Signed-off-by: jpvillam <jpvillam@amd.com> Co-authored-by: jpvillam <jpvillam@amd.com>
Juan Villamizar committed
Sep 23, 2025 at 18:39 UTC
bde2a1a8a442cd594eec35f5e356bddb5982de45
3 files changed
+26
-6
vllm/model_executor/layers/quantization/mxfp4.py
+6
-3
@@ -212,12 +212,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
212
intermediate_size_per_partition_after_pad = round_up(
213
intermediate_size_per_partition, 256)
214
hidden_size = round_up(hidden_size, 256)
215
- elif current_platform.is_rocm() or (
216
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
217
- or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16):
215
+ elif (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
216
+ or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16):
217
intermediate_size_per_partition_after_pad = round_up(
218
intermediate_size_per_partition, 128)
219
hidden_size = round_up(hidden_size, 128)
220
+ elif current_platform.is_rocm():
221
+ intermediate_size_per_partition_after_pad = round_up(
222
+ intermediate_size_per_partition, 256)
223
+ hidden_size = round_up(hidden_size, 256)
224
else:
225
intermediate_size_per_partition_after_pad = round_up(
226
intermediate_size_per_partition, 64)
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+14
-3
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Callable, Optional
3
+from typing import Any, Callable, Optional
4
5
import torch
6
@@ -21,6 +21,10 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
21
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
22
from triton_kernels.tensor_details import layout
23
from triton_kernels.tensor_details.layout import StridedLayout
24
+
25
+ value_layout_opts: dict[str, Any] = {}
26
+ scale_layout_opts: dict[str, Any] = {}
27
+
28
if (current_platform.is_cuda()
29
and current_platform.is_device_capability(90)
30
and not is_torch_equal_or_newer("2.8.1")):
@@ -28,8 +32,15 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
32
"Mxfp4 on hopper is running on torch < 2.8.1, "
33
"this cause swizling to be disabled, which may "
34
"cause performance degradation. Please upgrade to torch nightly")
31
- value_layout, value_layout_opts = StridedLayout, dict()
32
- scale_layout, scale_layout_opts = StridedLayout, dict()
35
+ value_layout = StridedLayout
36
+ scale_layout = StridedLayout
37
+ elif current_platform.is_rocm():
38
+ from triton_kernels.tensor_details.layout import (GFX950MXScaleLayout,
39
+ StridedLayout)
40
+
41
+ from vllm.platforms.rocm import on_gfx950
42
+ value_layout = StridedLayout
43
+ scale_layout = GFX950MXScaleLayout if on_gfx950() else StridedLayout
44
else:
45
value_layout, value_layout_opts = \
46
layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
vllm/platforms/rocm.py
+6
@@ -118,6 +118,12 @@ def on_gfx9() -> bool:
118
return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
119
120
121
+@cache
122
+def on_gfx950() -> bool:
123
+ GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
124
+ return any(arch in GPU_ARCH for arch in ["gfx950"])
125
+
126
+
127
@cache
128
def use_rocm_custom_paged_attention(
129
qtype: torch.dtype,