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, 04:30 UTC
[ROCm] Fix fused_moe_fake signature mismatch and other AITER bugs (#36100)
Signed-off-by: Li <chuali@amd.com>
Chuan (Richard) Li committed
Mar 23, 2026 at 00:48 UTC
e99fb98867c222b73922f0120ee7a09e442bc8ea
4 files changed
+16
-26
vllm/_aiter_ops.py
+5
-1
@@ -137,6 +137,10 @@ def _rocm_aiter_fused_moe_fake(
137
a2_scale: torch.Tensor | None = None,
138
num_local_tokens: torch.Tensor | None = None,
139
output_dtype: torch.dtype | None = None,
140
+ hidden_pad: int = 0,
141
+ intermediate_pad: int = 0,
142
+ bias1: torch.Tensor | None = None,
143
+ bias2: torch.Tensor | None = None,
144
) -> torch.Tensor:
145
if output_dtype is not None:
146
return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -1700,7 +1704,7 @@ class rocm_aiter_ops:
1704
)
1705
1706
@staticmethod
1703
- def triton_fp4_gemm_dynamic_qaunt(
1707
+ def triton_fp4_gemm_dynamic_quant(
1708
x: torch.Tensor,
1709
weight: torch.Tensor,
1710
weight_scale: torch.Tensor,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
-1
@@ -765,7 +765,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
765
if self.emulate:
766
logger.warning_once(
767
f"The current mode (supports_mx={current_platform.supports_mx()}, "
768
- f"use_mxfp4_aiter_moe={self.use_rocm_aiter_moe}, "
768
+ f"use_rocm_aiter_moe={self.use_rocm_aiter_moe}, "
769
f"ocp_mx_scheme={self.ocp_mx_scheme}) "
770
"does not support native MXFP4/MXFP6 "
771
"computation. Simulated weight dequantization and activation "
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+5
-20
@@ -3,13 +3,12 @@
3
4
from collections.abc import Callable
5
from fractions import Fraction
6
-from functools import cache, partial
6
+from functools import partial
7
from typing import Any
8
9
import torch
10
import torch.nn.functional as F
11
12
-from vllm import envs
12
from vllm._aiter_ops import rocm_aiter_ops
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
@@ -37,22 +36,6 @@ from .quark_scheme import QuarkScheme
36
logger = init_logger(__name__)
37
38
40
-# TODO: move registration of custom op to aiter_ops.py
41
-# `from vllm._aiter_ops import rocm_aiter_ops`
42
-# use `rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()`
43
-# for envs checks which does not require @cache anymore.
44
-# triton kernel is torch compile compatible.
45
-# does not require direct registration.
46
-# use `rocm_aiter_ops.triton_fp4_gemm_dynamic_qaunt`.
47
-@cache
48
-def is_rocm_aiter_fp4_asm_gemm_enabled() -> bool:
49
- return (
50
- current_platform.is_rocm()
51
- and envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
52
- and envs.VLLM_ROCM_USE_AITER
53
- )
54
-
55
-
39
try:
40
from aiter.ops.shuffle import shuffle_weight
41
from aiter.ops.triton.gemm_afp4wfp4 import (
@@ -63,7 +46,7 @@ try:
46
47
from vllm.utils.torch_utils import direct_register_custom_op
48
66
- if is_rocm_aiter_fp4_asm_gemm_enabled():
49
+ if rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled():
50
from aiter import gemm_a4w4, per_1x32_f4_quant_hip
51
52
def gemm_with_dynamic_quant(
@@ -233,7 +216,9 @@ class QuarkOCP_MX(QuarkScheme):
216
self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
217
)
218
236
- self.rocm_use_aiter_fp4_asm_gemm = is_rocm_aiter_fp4_asm_gemm_enabled()
219
+ self.rocm_use_aiter_fp4_asm_gemm = (
220
+ rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()
221
+ )
222
223
if not self.emulate and (dynamic_mxfp4_quant is None or gemm_afp4wfp4 is None):
224
# Currently need these kernels if not emulating
vllm/v1/attention/backends/rocm_aiter_fa.py
+5
-4
@@ -157,13 +157,13 @@ if current_platform.is_rocm():
157
total_tokens: int,
158
):
159
assert kv_cache_layout in ["NHD", "SHUFFLE"], (
160
- "kv_cache_layout only support NHD, SHUFFLE"
160
+ "kv_cache_layout only supports NHD, SHUFFLE"
161
)
162
head_dim = key.shape[2]
163
x = 16 // key_cache.element_size()
164
# assert dequant is True, "Currently, we only support "\
165
# "gather cache with dequant"
166
- # For k cache layout: [num_blocks, num_heads, page_size, head_dim]
166
+ # For k cache layout: [num_blocks, page_size, num_heads, head_dim]
167
assert head_dim == key_cache.shape[3], (
168
"We assume your kv cache layout is [num_blocks, "
169
"page_size, num_heads, head_dim], but got otherwise"
@@ -832,7 +832,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
832
833
if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
834
raise NotImplementedError(
835
- "Encoder self-attention is not implemented for FlashAttentionImpl"
835
+ "Encoder self-attention is not implemented for AiterFlashAttentionImpl"
836
)
837
838
def extend_for_sliding_window(
@@ -1047,7 +1047,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1047
1048
if output_scale is not None or output_block_scale is not None:
1049
raise NotImplementedError(
1050
- "fused output quantization is not yet supported for FlashAttentionImpl"
1050
+ "fused output quantization is not yet supported "
1051
+ "for AiterFlashAttentionImpl"
1052
)
1053
1054
if attn_metadata is None: