@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f080a83511
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 13, 13:15 UTC

[RFC][ROCm][AITER] Keep all AITER kernels in `_aiter_ops` class like `_custom_ops` and `_ipex_ops` (#24490)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

vllmellm committed Nov 10, 2025 at 17:20 UTC f080a83511511a9c0a222451a752a1623aec095d
25 files changed +1194 -925
docs/design/moe_kernel_features.md
+1 -1
@@ -97,7 +97,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` sub-class, MoE kernels
97 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
98 | pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
99 | iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
100 -| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_moe_impl] |
100 +| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
101 | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
102 | naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
103
tests/kernels/moe/test_moe.py
+6 -5
@@ -6,6 +6,8 @@ Run `pytest tests/kernels/test_moe.py`.
6 """
7
8 import functools
9 +import importlib
10 +import sys
11 from collections.abc import Callable
12 from dataclasses import dataclass
13 from typing import Any
@@ -20,6 +22,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
22 import vllm.model_executor.layers.fused_moe # noqa
23 from tests.kernels.moe.utils import fused_moe
24 from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
25 +from vllm._aiter_ops import rocm_aiter_ops
26 from vllm.config import VllmConfig, set_current_vllm_config
27 from vllm.distributed.parallel_state import init_distributed_environment
28 from vllm.forward_context import set_forward_context
@@ -412,14 +415,12 @@ def test_mixtral_moe(
415 huggingface."""
416
417 # clear the cache before every test
415 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
416 - is_rocm_aiter_moe_enabled,
417 - )
418 + # Force reload aiter_ops to pick up the new environment variables.
419 + if "rocm_aiter_ops" in sys.modules:
420 + importlib.reload(rocm_aiter_ops)
421
419 - is_rocm_aiter_moe_enabled.cache_clear()
422 if use_rocm_aiter:
423 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
422 -
424 if dtype == torch.float32:
425 pytest.skip("AITER ROCm test skip for float32")
426
tests/model_executor/test_enabled_custom_ops.py
+14 -27
@@ -4,6 +4,7 @@
4 import pytest
5 import torch
6
7 +from vllm._aiter_ops import rocm_aiter_ops
8 from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
9 from vllm.model_executor.custom_op import CustomOp
10 from vllm.model_executor.layers.activation import (
@@ -15,9 +16,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
16 dispatch_topk_func,
17 vllm_topk_softmax,
18 )
18 -from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
19 - is_rocm_aiter_moe_enabled,
20 -)
19 from vllm.model_executor.layers.layernorm import (
20 RMSNorm,
21 dispatch_rocm_rmsnorm_func,
@@ -126,50 +124,39 @@ def test_enabled_ops_invalid(env: str):
124 RMSNorm(1024).enabled()
125
126
129 -@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
130 -def test_topk_dispatch(use_rocm_aiter: str, monkeypatch):
131 - monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
132 - topk_func = dispatch_topk_func()
133 - is_rocm_aiter_moe_enabled.cache_clear()
134 - if current_platform.is_rocm() and int(use_rocm_aiter):
135 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
136 - rocm_aiter_topk_softmax,
137 - )
127 +@pytest.mark.parametrize(
128 + "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
129 +)
130 +def test_topk_dispatch(use_rocm_aiter: bool):
131 + topk_func = dispatch_topk_func(use_rocm_aiter)
132
139 - assert topk_func == rocm_aiter_topk_softmax
133 + if current_platform.is_rocm() and use_rocm_aiter:
134 + assert topk_func == rocm_aiter_ops.topk_softmax
135 else:
136 assert topk_func == vllm_topk_softmax
137
138
139 @pytest.mark.parametrize("add_residual", [True, False])
140 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
146 -@pytest.mark.parametrize("use_rocm_aiter", ["0", "1"])
147 -@pytest.mark.parametrize("use_rocm_aiter_norm", ["0", "1"])
141 +@pytest.mark.parametrize("use_rocm_aiter", [True, False])
142 @pytest.mark.skipif(
143 not current_platform.is_rocm(), reason="AITER is a feature exclusive for ROCm"
144 )
145 def test_rms_norm_dispatch(
152 - add_residual: bool,
153 - dtype: torch.dtype,
154 - use_rocm_aiter: str,
155 - use_rocm_aiter_norm: str,
156 - monkeypatch,
146 + add_residual: bool, dtype: torch.dtype, use_rocm_aiter: bool
147 ):
158 - monkeypatch.setenv("VLLM_ROCM_USE_AITER", use_rocm_aiter)
159 - monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", use_rocm_aiter_norm)
160 - rms_norm_func = dispatch_rocm_rmsnorm_func(add_residual, dtype)
148 + rms_norm_func = dispatch_rocm_rmsnorm_func(add_residual, dtype, use_rocm_aiter)
149
150 should_use_rocm_aiter = (
151 current_platform.is_rocm()
164 - and int(use_rocm_aiter)
165 - and int(use_rocm_aiter_norm)
152 + and use_rocm_aiter
153 and dtype in RMS_NORM_SUPPORTED_DTYPES
154 )
155
156 if add_residual and should_use_rocm_aiter:
170 - assert rms_norm_func == torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add
157 + assert rms_norm_func == rocm_aiter_ops.rms_norm2d_with_add
158 elif should_use_rocm_aiter:
172 - assert rms_norm_func == torch.ops.vllm.rocm_aiter_rms_norm
159 + assert rms_norm_func == rocm_aiter_ops.rms_norm
160 elif add_residual:
161 assert rms_norm_func == fused_add_rms_norm
162 else:
vllm/_aiter_ops.py new
+941
@@ -0,0 +1,941 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 +from collections.abc import Callable
5 +
6 +import torch
7 +
8 +import vllm.envs as envs
9 +from vllm.platforms import current_platform
10 +from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
11 +
12 +
13 +def is_aiter_found() -> bool:
14 + from importlib.util import find_spec
15 +
16 + return find_spec("aiter") is not None
17 +
18 +
19 +# `find_spec` is not torch.compile compatible.
20 +# In cases where aiter availability might have
21 +# been checked in forward passes that are torch compiled.
22 +# we keep this global outside to not cause torch compile breaks.
23 +IS_AITER_FOUND = is_aiter_found()
24 +
25 +
26 +def if_aiter_supported(func: Callable) -> Callable:
27 + """Decorator that only executes the function if
28 + ROCm AITER package is supported on gfx9 archs.
29 + """
30 +
31 + @functools.wraps(func)
32 + def wrapper(*args, **kwargs):
33 + # checks the platform, device arch and aiter library existance.
34 +
35 + from vllm.platforms.rocm import on_gfx9
36 +
37 + if current_platform.is_rocm() and on_gfx9() and IS_AITER_FOUND:
38 + return func(*args, **kwargs)
39 + else:
40 + # Return None or do nothing if not supported
41 + return None
42 +
43 + return wrapper
44 +
45 +
46 +def _rocm_aiter_fused_moe_impl(
47 + hidden_states: torch.Tensor,
48 + w1: torch.Tensor,
49 + w2: torch.Tensor,
50 + topk_weight: torch.Tensor,
51 + topk_ids: torch.Tensor,
52 + expert_mask: torch.Tensor | None = None,
53 + activation_method: int = 0,
54 + quant_method: int = 0,
55 + doweight_stage1: bool = False,
56 + w1_scale: torch.Tensor | None = None,
57 + w2_scale: torch.Tensor | None = None,
58 + a1_scale: torch.Tensor | None = None,
59 + a2_scale: torch.Tensor | None = None,
60 +) -> torch.Tensor:
61 + from aiter import ActivationType, QuantType
62 + from aiter.fused_moe import fused_moe
63 +
64 + activation = ActivationType(activation_method)
65 + quant_type = QuantType(quant_method)
66 +
67 + return fused_moe(
68 + hidden_states,
69 + w1,
70 + w2,
71 + topk_weight,
72 + topk_ids,
73 + expert_mask,
74 + activation,
75 + quant_type,
76 + doweight_stage1,
77 + w1_scale,
78 + w2_scale,
79 + a1_scale,
80 + a2_scale,
81 + )
82 +
83 +
84 +def _rocm_aiter_fused_moe_fake(
85 + hidden_states: torch.Tensor,
86 + w1: torch.Tensor,
87 + w2: torch.Tensor,
88 + topk_weight: torch.Tensor,
89 + topk_ids: torch.Tensor,
90 + expert_mask: torch.Tensor | None = None,
91 + activation_method: int = 0,
92 + quant_method: int = 0,
93 + doweight_stage1: bool = False,
94 + w1_scale: torch.Tensor | None = None,
95 + w2_scale: torch.Tensor | None = None,
96 + a1_scale: torch.Tensor | None = None,
97 + a2_scale: torch.Tensor | None = None,
98 +) -> torch.Tensor:
99 + return torch.empty_like(hidden_states)
100 +
101 +
102 +def _rocm_aiter_asm_moe_tkw1_impl(
103 + hidden_states: torch.Tensor,
104 + w1: torch.Tensor,
105 + w2: torch.Tensor,
106 + topk_weights: torch.Tensor,
107 + topk_ids: torch.Tensor,
108 + fc1_scale: torch.Tensor | None = None,
109 + fc2_scale: torch.Tensor | None = None,
110 + fc1_smooth_scale: torch.Tensor | None = None,
111 + fc2_smooth_scale: torch.Tensor | None = None,
112 + a16: bool = False,
113 + per_tensor_quant_scale: torch.Tensor | None = None,
114 + expert_mask: torch.Tensor | None = None,
115 + activation_method: int = 0,
116 +) -> torch.Tensor:
117 + from aiter import ActivationType
118 + from aiter.fused_moe_bf16_asm import asm_moe_tkw1
119 +
120 + activation = ActivationType(activation_method)
121 +
122 + return asm_moe_tkw1(
123 + hidden_states,
124 + w1,
125 + w2,
126 + topk_weights,
127 + topk_ids,
128 + fc1_scale=fc1_scale,
129 + fc2_scale=fc2_scale,
130 + fc1_smooth_scale=fc1_smooth_scale,
131 + fc2_smooth_scale=fc2_smooth_scale,
132 + a16=a16,
133 + per_tensor_quant_scale=per_tensor_quant_scale,
134 + expert_mask=expert_mask,
135 + activation=activation,
136 + )
137 +
138 +
139 +def _rocm_aiter_asm_moe_tkw1_fake(
140 + hidden_states: torch.Tensor,
141 + w1: torch.Tensor,
142 + w2: torch.Tensor,
143 + topk_weights: torch.Tensor,
144 + topk_ids: torch.Tensor,
145 + fc1_scale: torch.Tensor | None = None,
146 + fc2_scale: torch.Tensor | None = None,
147 + fc1_smooth_scale: torch.Tensor | None = None,
148 + fc2_smooth_scale: torch.Tensor | None = None,
149 + a16: bool = False,
150 + per_tensor_quant_scale: torch.Tensor | None = None,
151 + expert_mask: torch.Tensor | None = None,
152 + activation_method: int = 0,
153 +) -> torch.Tensor:
154 + return torch.empty_like(hidden_states)
155 +
156 +
157 +def _rocm_aiter_topk_softmax_impl(
158 + topk_weights: torch.Tensor,
159 + topk_indices: torch.Tensor,
160 + token_expert_indices: torch.Tensor,
161 + gating_output: torch.Tensor,
162 + renormalize: bool,
163 +) -> None:
164 + from aiter import topk_softmax
165 +
166 + topk_softmax(
167 + topk_weights, topk_indices, token_expert_indices, gating_output, renormalize
168 + )
169 +
170 +
171 +def _rocm_aiter_topk_softmax_fake(
172 + topk_weights: torch.Tensor,
173 + topk_indices: torch.Tensor,
174 + token_expert_indices: torch.Tensor,
175 + gating_output: torch.Tensor,
176 + renormalize: bool,
177 +) -> None:
178 + pass
179 +
180 +
181 +def _rocm_aiter_biased_grouped_topk_impl(
182 + gating_output: torch.Tensor,
183 + correction_bias: torch.Tensor,
184 + topk_weights: torch.Tensor,
185 + topk_ids: torch.Tensor,
186 + num_expert_group: int,
187 + topk_group: int,
188 + need_renorm: bool,
189 + routed_scaling_factor: float = 1.0, # mul to topk_weights
190 +) -> None:
191 + from aiter import biased_grouped_topk
192 +
193 + biased_grouped_topk(
194 + gating_output,
195 + correction_bias,
196 + topk_weights,
197 + topk_ids,
198 + num_expert_group,
199 + topk_group,
200 + need_renorm,
201 + routed_scaling_factor,
202 + )
203 +
204 +
205 +def _rocm_aiter_biased_grouped_topk_fake(
206 + gating_output: torch.Tensor,
207 + correction_bias: torch.Tensor,
208 + topk_weights: torch.Tensor,
209 + topk_ids: torch.Tensor,
210 + num_expert_group: int,
211 + topk_group: int,
212 + need_renorm: bool,
213 + routed_scaling_factor: float = 1.0, # mul to topk_weights
214 +) -> None:
215 + pass
216 +
217 +
218 +def _rocm_aiter_grouped_topk_impl(
219 + gating_output: torch.Tensor,
220 + topk_weights: torch.Tensor,
221 + topk_ids: torch.Tensor,
222 + num_expert_group: int,
223 + topk_group: int,
224 + need_renorm: bool,
225 + scoring_func: str = "softmax",
226 + routed_scaling_factor: float = 1.0, # mul to topk_weights
227 +) -> None:
228 + is_softmax = scoring_func == "softmax"
229 + from aiter import grouped_topk
230 +
231 + grouped_topk(
232 + gating_output,
233 + topk_weights,
234 + topk_ids,
235 + num_expert_group,
236 + topk_group,
237 + need_renorm,
238 + is_softmax,
239 + routed_scaling_factor,
240 + )
241 +
242 +
243 +def _rocm_aiter_grouped_topk_fake(
244 + gating_output: torch.Tensor,
245 + topk_weights: torch.Tensor,
246 + topk_ids: torch.Tensor,
247 + num_expert_group: int,
248 + topk_group: int,
249 + need_renorm: bool,
250 + scoring_func: str = "softmax",
251 + routed_scaling_factor: float = 1.0, # mul to topk_weights
252 +) -> None:
253 + pass
254 +
255 +
256 +def _rocm_aiter_mla_decode_fwd_impl(
257 + q: torch.Tensor,
258 + kv_buffer: torch.Tensor,
259 + o: torch.Tensor,
260 + qo_indptr: torch.Tensor,
261 + max_seqlen_qo: int,
262 + kv_indptr: torch.Tensor | None = None,
263 + kv_indices: torch.Tensor | None = None,
264 + kv_last_page_lens: torch.Tensor | None = None,
265 + sm_scale: float = 1.0,
266 + logit_cap: float = 0.0,
267 +) -> None:
268 + from aiter.mla import mla_decode_fwd
269 +
270 + mla_decode_fwd(
271 + q,
272 + kv_buffer.view(-1, 1, 1, q.shape[-1]),
273 + o,
274 + qo_indptr,
275 + kv_indptr,
276 + kv_indices,
277 + kv_last_page_lens,
278 + max_seqlen_qo,
279 + sm_scale=sm_scale,
280 + logit_cap=logit_cap,
281 + )
282 +
283 +
284 +def _rocm_aiter_mla_decode_fwd_fake(
285 + q: torch.Tensor,
286 + kv_buffer: torch.Tensor,
287 + o: torch.Tensor,
288 + qo_indptr: torch.Tensor,
289 + max_seqlen_qo: int,
290 + kv_indptr: torch.Tensor | None = None,
291 + kv_indices: torch.Tensor | None = None,
292 + kv_last_page_lens: torch.Tensor | None = None,
293 + sm_scale: float = 1.0,
294 + logit_cap: float = 0.0,
295 +) -> None:
296 + pass
297 +
298 +
299 +def _rocm_aiter_gemm_w8a8_impl(
300 + A: torch.Tensor,
301 + B: torch.Tensor,
302 + As: torch.Tensor,
303 + Bs: torch.Tensor,
304 + bias: torch.Tensor | None = None,
305 + output_dtype: torch.dtype = torch.float16,
306 +) -> torch.Tensor:
307 + from aiter import gemm_a8w8_CK
308 +
309 + # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
310 + # a to be [M, K]
311 + # b to be [N, K]
312 + # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
313 + return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)
314 +
315 +
316 +def _rocm_aiter_gemm_w8a8_fake(
317 + A: torch.Tensor,
318 + B: torch.Tensor,
319 + As: torch.Tensor,
320 + Bs: torch.Tensor,
321 + bias: torch.Tensor | None = None,
322 + output_dtype: torch.dtype = torch.float16,
323 +) -> torch.Tensor:
324 + m = A.shape[0]
325 + n = B.shape[0]
326 + Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
327 + return Y
328 +
329 +
330 +def _rocm_aiter_gemm_w8a8_blockscale_impl(
331 + A: torch.Tensor,
332 + B: torch.Tensor,
333 + As: torch.Tensor,
334 + Bs: torch.Tensor,
335 + output_dtype: torch.dtype = torch.float16,
336 +) -> torch.Tensor:
337 + from aiter import gemm_a8w8_blockscale
338 +
339 + return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
340 +
341 +
342 +def _rocm_aiter_gemm_w8a8_blockscale_fake(
343 + A: torch.Tensor,
344 + B: torch.Tensor,
345 + As: torch.Tensor,
346 + Bs: torch.Tensor,
347 + output_dtype: torch.dtype = torch.float16,
348 +) -> torch.Tensor:
349 + m = A.shape[0]
350 + n = B.shape[0]
351 + Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
352 + return Y
353 +
354 +
355 +def _rocm_aiter_rms_norm_impl(
356 + x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
357 +) -> torch.Tensor:
358 + from aiter import rms_norm
359 +
360 + if x.dim() > 2:
361 + x_original_shape = x.shape
362 + x = x.reshape(-1, x_original_shape[-1])
363 + x = rms_norm(x, weight, variance_epsilon)
364 + return x.reshape(x_original_shape)
365 +
366 + return rms_norm(x, weight, variance_epsilon)
367 +
368 +
369 +def _rocm_aiter_rms_norm_fake(
370 + x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
371 +) -> torch.Tensor:
372 + return torch.empty_like(x)
373 +
374 +
375 +def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
376 + x: torch.Tensor,
377 + residual: torch.Tensor,
378 + weight: torch.Tensor,
379 + variance_epsilon: float,
380 +) -> tuple[torch.Tensor, torch.Tensor]:
381 + from aiter import rmsnorm2d_fwd_with_add
382 +
383 + residual_out = torch.empty_like(residual)
384 + output = torch.empty_like(x)
385 + rmsnorm2d_fwd_with_add(
386 + output, # output
387 + x, # input
388 + residual, # residual input
389 + residual_out, # residual output
390 + weight,
391 + variance_epsilon,
392 + )
393 + return output, residual_out
394 +
395 +
396 +def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
397 + x: torch.Tensor,
398 + residual: torch.Tensor,
399 + weight: torch.Tensor,
400 + variance_epsilon: float,
401 +) -> tuple[torch.Tensor, torch.Tensor]:
402 + return torch.empty_like(x), torch.empty_like(residual)
403 +
404 +
405 +# Global flag to ensure ops are registered only once
406 +_OPS_REGISTERED = False
407 +
408 +
409 +class rocm_aiter_ops:
410 + _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
411 + _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
412 + _RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
413 + _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
414 + _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
415 + _PG_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
416 + _MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
417 + _TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
418 + _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
419 + _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
420 + _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
421 + _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
422 +
423 + @classmethod
424 + @if_aiter_supported
425 + def is_enabled(cls) -> bool:
426 + """Verifies device specs and availability of aiter main env variable."""
427 + return cls._AITER_ENABLED
428 +
429 + @classmethod
430 + @if_aiter_supported
431 + def is_linear_enabled(cls) -> bool:
432 + """ "Verifies device specs and availability of env variable."""
433 + return cls._AITER_ENABLED and cls._LINEAR_ENABLED
434 +
435 + @classmethod
436 + @if_aiter_supported
437 + def is_linear_fp8_enaled(cls) -> bool:
438 + """ "Verifies device specs and availability of env variable."""
439 + return cls.is_linear_enabled() and current_platform.is_fp8_fnuz()
440 +
441 + @classmethod
442 + @if_aiter_supported
443 + def is_rmsnorm_enabled(cls) -> bool:
444 + """ "Verifies device specs and availability of env variable."""
445 + return cls._AITER_ENABLED and cls._RMSNORM_ENABLED
446 +
447 + @classmethod
448 + @if_aiter_supported
449 + def is_fused_moe_enabled(cls) -> bool:
450 + """ "Verifies device specs and availability of env variable."""
451 + return cls._AITER_ENABLED and cls._FMOE_ENABLED
452 +
453 + @classmethod
454 + @if_aiter_supported
455 + def is_fusion_moe_shared_experts_enabled(cls) -> bool:
456 + return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED
457 +
458 + @classmethod
459 + @if_aiter_supported
460 + def is_mla_enabled(cls) -> bool:
461 + """ "Verifies device specs and availability of env variable."""
462 + return cls._AITER_ENABLED and cls._MLA_ENABLED
463 +
464 + @classmethod
465 + @if_aiter_supported
466 + def is_mha_enabled(cls) -> bool:
467 + """ "Verifies device specs and availability of env variable."""
468 + return cls._AITER_ENABLED and cls._MHA_ENABLED
469 +
470 + @classmethod
471 + @if_aiter_supported
472 + def is_pa_attn_enabled(cls) -> bool:
473 + """ "Verifies device specs and availability of env variable."""
474 + return cls._AITER_ENABLED and cls._PG_ATTN_ENABLED
475 +
476 + @classmethod
477 + @if_aiter_supported
478 + def is_triton_unified_attn_enabled(cls) -> bool:
479 + """ "Verifies device specs and availability of env variable."""
480 + return cls._AITER_ENABLED and cls._TRITON_UNIFIED_ATTN_ENABLED
481 +
482 + @classmethod
483 + @if_aiter_supported
484 + def is_fp8bmm_enabled(cls) -> bool:
485 + return cls._AITER_ENABLED and cls._FP8BMM_ENABLED
486 +
487 + @classmethod
488 + @if_aiter_supported
489 + def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
490 + return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM
491 +
492 + @classmethod
493 + @if_aiter_supported
494 + def is_triton_rotary_embed_enabled(cls) -> bool:
495 + return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED
496 +
497 + @staticmethod
498 + @if_aiter_supported
499 + def register_ops_once() -> None:
500 + global _OPS_REGISTERED
501 + if not _OPS_REGISTERED:
502 + tags = (
503 + tuple()
504 + if is_torch_equal_or_newer("2.7.0")
505 + else (torch.Tag.needs_fixed_stride_order,)
506 + )
507 +
508 + # register all the custom ops here
509 + direct_register_custom_op(
510 + op_name="rocm_aiter_asm_moe_tkw1",
511 + op_func=_rocm_aiter_asm_moe_tkw1_impl,
512 + mutates_args=[],
513 + fake_impl=_rocm_aiter_asm_moe_tkw1_fake,
514 + dispatch_key=current_platform.dispatch_key,
515 + )
516 +
517 + direct_register_custom_op(
518 + op_name="rocm_aiter_fused_moe",
519 + op_func=_rocm_aiter_fused_moe_impl,
520 + mutates_args=[],
521 + fake_impl=_rocm_aiter_fused_moe_fake,
522 + dispatch_key=current_platform.dispatch_key,
523 + )
524 +
525 + direct_register_custom_op(
526 + op_name="rocm_aiter_topk_softmax",
527 + op_func=_rocm_aiter_topk_softmax_impl,
528 + mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
529 + fake_impl=_rocm_aiter_topk_softmax_fake,
530 + dispatch_key=current_platform.dispatch_key,
531 + )
532 +
533 + direct_register_custom_op(
534 + op_name="rocm_aiter_biased_grouped_topk",
535 + op_func=_rocm_aiter_biased_grouped_topk_impl,
536 + mutates_args=["topk_weights", "topk_ids"],
537 + fake_impl=_rocm_aiter_biased_grouped_topk_fake,
538 + dispatch_key=current_platform.dispatch_key,
539 + )
540 +
541 + direct_register_custom_op(
542 + op_name="rocm_aiter_grouped_topk",
543 + op_func=_rocm_aiter_grouped_topk_impl,
544 + mutates_args=["topk_weights", "topk_ids"],
545 + fake_impl=_rocm_aiter_grouped_topk_fake,
546 + dispatch_key=current_platform.dispatch_key,
547 + )
548 +
549 + direct_register_custom_op(
550 + op_name="rocm_aiter_mla_decode_fwd",
551 + op_func=_rocm_aiter_mla_decode_fwd_impl,
552 + mutates_args=["o"],
553 + fake_impl=_rocm_aiter_mla_decode_fwd_fake,
554 + tags=tags,
555 + )
556 +
557 + direct_register_custom_op(
558 + op_name="rocm_aiter_gemm_w8a8",
559 + op_func=_rocm_aiter_gemm_w8a8_impl,
560 + mutates_args=[],
561 + fake_impl=_rocm_aiter_gemm_w8a8_fake,
562 + dispatch_key=current_platform.dispatch_key,
563 + )
564 +
565 + direct_register_custom_op(
566 + op_name="rocm_aiter_gemm_w8a8_blockscale",
567 + op_func=_rocm_aiter_gemm_w8a8_blockscale_impl,
568 + mutates_args=[],
569 + fake_impl=_rocm_aiter_gemm_w8a8_blockscale_fake,
570 + dispatch_key=current_platform.dispatch_key,
571 + )
572 +
573 + direct_register_custom_op(
574 + op_name="rocm_aiter_rms_norm",
575 + op_func=_rocm_aiter_rms_norm_impl,
576 + mutates_args=[],
577 + fake_impl=_rocm_aiter_rms_norm_fake,
578 + dispatch_key=current_platform.dispatch_key,
579 + )
580 +
581 + direct_register_custom_op(
582 + op_name="rocm_aiter_rmsnorm2d_fwd_with_add",
583 + op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
584 + mutates_args=[],
585 + fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
586 + dispatch_key=current_platform.dispatch_key,
587 + )
588 +
589 + _OPS_REGISTERED = True
590 +
591 + @staticmethod
592 + def rms_norm2d_with_add(
593 + x: torch.Tensor,
594 + residual: torch.Tensor,
595 + weight: torch.Tensor,
596 + variance_epsilon: float,
597 + ) -> tuple[torch.Tensor, torch.Tensor]:
598 + return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add(
599 + x, residual, weight, variance_epsilon
600 + )
601 +
602 + @staticmethod
603 + def rms_norm(
604 + x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
605 + ) -> torch.Tensor:
606 + return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon)
607 +
608 + @staticmethod
609 + def gemm_w8a8(
610 + A: torch.Tensor,
611 + B: torch.Tensor,
612 + As: torch.Tensor,
613 + Bs: torch.Tensor,
614 + bias: torch.Tensor | None = None,
615 + output_dtype: torch.dtype = torch.float16,
616 + ) -> torch.Tensor:
617 + return torch.ops.vllm.rocm_aiter_gemm_w8a8(A, B, As, Bs, bias, output_dtype)
618 +
619 + @staticmethod
620 + def gemm_w8a8_blockscale(
621 + A: torch.Tensor,
622 + B: torch.Tensor,
623 + As: torch.Tensor,
624 + Bs: torch.Tensor,
625 + block_size: list[int],
626 + output_dtype: torch.dtype = torch.float16,
627 + ) -> torch.Tensor:
628 + return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale(
629 + A, B, As, Bs, output_dtype
630 + )
631 +
632 + @staticmethod
633 + def fused_moe(
634 + hidden_states: torch.Tensor,
635 + w1: torch.Tensor,
636 + w2: torch.Tensor,
637 + topk_weight: torch.Tensor,
638 + topk_ids: torch.Tensor,
639 + expert_mask: torch.Tensor | None = None,
640 + activation_method: int = 0,
641 + quant_method: int = 0,
642 + doweight_stage1: bool = False,
643 + w1_scale: torch.Tensor | None = None,
644 + w2_scale: torch.Tensor | None = None,
645 + a1_scale: torch.Tensor | None = None,
646 + a2_scale: torch.Tensor | None = None,
647 + ) -> torch.Tensor:
648 + return torch.ops.vllm.rocm_aiter_fused_moe(
649 + hidden_states,
650 + w1,
651 + w2,
652 + topk_weight,
653 + topk_ids,
654 + expert_mask,
655 + activation_method,
656 + quant_method,
657 + doweight_stage1,
658 + w1_scale,
659 + w2_scale,
660 + a1_scale,
661 + a2_scale,
662 + )
663 +
664 + @staticmethod
665 + def asm_moe_tkw1(
666 + hidden_states: torch.Tensor,
667 + w1: torch.Tensor,
668 + w2: torch.Tensor,
669 + topk_weights: torch.Tensor,
670 + topk_ids: torch.Tensor,
671 + fc1_scale: torch.Tensor | None = None,
672 + fc2_scale: torch.Tensor | None = None,
673 + fc1_smooth_scale: torch.Tensor | None = None,
674 + fc2_smooth_scale: torch.Tensor | None = None,
675 + a16: bool = False,
676 + per_tensor_quant_scale: torch.Tensor | None = None,
677 + expert_mask: torch.Tensor | None = None,
678 + activation_method: int = 0,
679 + ) -> torch.Tensor:
680 + return torch.ops.vllm.rocm_aiter_asm_moe_tkw1(
681 + hidden_states,
682 + w1,
683 + w2,
684 + topk_weights,
685 + topk_ids,
686 + fc1_scale,
687 + fc2_scale,
688 + fc1_smooth_scale,
689 + fc2_smooth_scale,
690 + a16,
691 + per_tensor_quant_scale,
692 + expert_mask,
693 + activation_method,
694 + )
695 +
696 + @staticmethod
697 + def topk_softmax(
698 + topk_weights: torch.Tensor,
699 + topk_indices: torch.Tensor,
700 + token_expert_indices: torch.Tensor,
701 + gating_output: torch.Tensor,
702 + renormalize: bool,
703 + ) -> tuple[torch.Tensor, ...]:
704 + torch.ops.vllm.rocm_aiter_topk_softmax(
705 + topk_weights, topk_indices, token_expert_indices, gating_output, renormalize
706 + )
707 + return topk_weights, topk_indices
708 +
709 + @staticmethod
710 + def biased_grouped_topk(
711 + gating_output: torch.Tensor,
712 + correction_bias: torch.Tensor,
713 + topk_weights: torch.Tensor,
714 + topk_ids: torch.Tensor,
715 + num_expert_group: int,
716 + topk_group: int,
717 + need_renorm: bool,
718 + routed_scaling_factor: float = 1.0,
719 + ) -> None:
720 + torch.ops.vllm.rocm_aiter_biased_grouped_topk(
721 + gating_output,
722 + correction_bias,
723 + topk_weights,
724 + topk_ids,
725 + num_expert_group,
726 + topk_group,
727 + need_renorm,
728 + routed_scaling_factor,
729 + )
730 +
731 + @staticmethod
732 + def grouped_topk(
733 + gating_output: torch.Tensor,
734 + topk_weights: torch.Tensor,
735 + topk_ids: torch.Tensor,
736 + num_expert_group: int,
737 + topk_group: int,
738 + need_renorm: bool,
739 + scoring_func: str = "softmax",
740 + routed_scaling_factor: float = 1.0,
741 + ) -> None:
742 + torch.ops.vllm.rocm_aiter_grouped_topk(
743 + gating_output,
744 + topk_weights,
745 + topk_ids,
746 + num_expert_group,
747 + topk_group,
748 + need_renorm,
749 + scoring_func,
750 + routed_scaling_factor,
751 + )
752 +
753 + @staticmethod
754 + def mla_decode_fwd(
755 + q: torch.Tensor,
756 + kv_buffer: torch.Tensor,
757 + o: torch.Tensor,
758 + sm_scale: float,
759 + qo_indptr: torch.Tensor,
760 + max_seqlen_qo: int,
761 + kv_indptr: torch.Tensor | None = None,
762 + kv_indices: torch.Tensor | None = None,
763 + kv_last_page_lens: torch.Tensor | None = None,
764 + logit_cap: float = 0.0,
765 + ):
766 + torch.ops.vllm.rocm_aiter_mla_decode_fwd(
767 + q,
768 + kv_buffer.view(-1, 1, 1, q.shape[-1]),
769 + o,
770 + qo_indptr,
771 + max_seqlen_qo,
772 + kv_indptr,
773 + kv_indices,
774 + kv_last_page_lens,
775 + sm_scale=sm_scale,
776 + logit_cap=logit_cap,
777 + )
778 +
779 + @staticmethod
780 + def triton_fp4_gemm_dynamic_qaunt(
781 + x: torch.Tensor,
782 + weight: torch.Tensor,
783 + weight_scale: torch.Tensor,
784 + out_dtype: torch.dtype | None = torch.bfloat16,
785 + x_scales: torch.Tensor | None = None,
786 + ) -> torch.Tensor:
787 + from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
788 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
789 +
790 + if x_scales is None:
791 + x_q, x_s = dynamic_mxfp4_quant(x)
792 + else:
793 + x_q = x
794 + x_s = x_scales
795 +
796 + y = torch.empty(
797 + x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
798 + )
799 +
800 + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
801 + return y
802 +
803 + @staticmethod
804 + def triton_rotary_embed(
805 + positions: torch.Tensor,
806 + query: torch.Tensor,
807 + key: torch.Tensor,
808 + cos_sin_cache: torch.Tensor,
809 + head_size: int,
810 + rotary_dim: int,
811 + is_neox_style: bool,
812 + ):
813 + from aiter.ops.triton.rope import rope_cached_thd_positions_2c_fwd_inplace
814 +
815 + num_tokens = positions.numel()
816 + cos, sin = cos_sin_cache.chunk(2, dim=-1)
817 + query_shape = query.shape
818 + key_shape = key.shape
819 + rotate_style = 0 if is_neox_style else 1
820 +
821 + query = query.view(num_tokens, -1, head_size)
822 + key = key.view(num_tokens, -1, head_size)
823 + query_ = query[..., :rotary_dim]
824 + key_ = key[..., :rotary_dim]
825 + positions = positions.view(*query.shape[:1])
826 + rope_cached_thd_positions_2c_fwd_inplace(
827 + positions,
828 + sin,
829 + cos,
830 + query_,
831 + key_,
832 + rotate_style,
833 + reuse_freqs_front_part=True,
834 + is_nope_first=False,
835 + )
836 + query = query.view(query_shape)
837 + key = key.view(key_shape)
838 +
839 + @staticmethod
840 + def triton_fp8_bmm(
841 + X: torch.Tensor,
842 + WQ: torch.Tensor,
843 + w_scale: torch.Tensor,
844 + group_size: int = 128,
845 + bias: torch.Tensor | None = None,
846 + dtype: torch.dtype | None = torch.bfloat16,
847 + splitK: int | None = None,
848 + YQ: torch.Tensor | None = None,
849 + transpose_bm: bool | None = False,
850 + config: dict | None = None,
851 + ) -> torch.Tensor:
852 + # ruff: noqa: E501 # isort: skip
853 + from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
854 + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm,
855 + )
856 +
857 + return aiter_triton_fp8_bmm(
858 + X,
859 + WQ,
860 + w_scale,
861 + group_size=group_size,
862 + bias=bias,
863 + dtype=dtype,
864 + splitK=splitK,
865 + YQ=YQ,
866 + transpose_bm=transpose_bm,
867 + config=config,
868 + )
869 +
870 + @staticmethod
871 + def triton_gemm_a8w8_blockscale(
872 + A: torch.Tensor,
873 + B: torch.Tensor,
874 + As: torch.Tensor,
875 + Bs: torch.Tensor,
876 + block_size: list[int],
877 + output_dtype: torch.dtype = torch.float16,
878 + ) -> torch.Tensor:
879 + from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale
880 +
881 + return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
882 +
883 + @staticmethod
884 + def per_1x128_fp8_quant(
885 + input_2d: torch.Tensor,
886 + ) -> tuple[torch.Tensor, ...]:
887 + """Only applies quantization method for fp8 data type only."""
888 + from aiter import QuantType, dtypes, get_hip_quant
889 +
890 + aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)
891 + return aiter_per1x128_quant(input_2d.contiguous(), quant_dtype=dtypes.fp8)
892 +
893 + @staticmethod
894 + def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:
895 + return (n, k) in [
896 + (1024, 8192),
897 + (2112, 7168),
898 + (3072, 1536),
899 + (32768, 8192),
900 + (4096, 7168),
901 + (4608, 7168),
902 + (512, 7168),
903 + (7168, 2048),
904 + (7168, 256),
905 + (8192, 1024),
906 + (8192, 32768),
907 + ]
908 +
909 + @staticmethod
910 + def shuffle_weight(
911 + self, tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
912 + ) -> torch.Tensor:
913 + from aiter.ops.shuffle import shuffle_weight
914 +
915 + return shuffle_weight(tensor, layout=layout)
916 +
917 + @staticmethod
918 + def shuffle_weights(
919 + *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
920 + ) -> tuple[torch.Tensor, ...]:
921 + """
922 + Applies shuffle_weight function from AITER to each
923 + input tensor and returns them.
924 +
925 + Rearranges (shuffles) the input tensor/s
926 + into a specified block layout for optimized computation.
927 +
928 + Args:
929 + *tensors: Variable number of torch.Tensor objects.
930 + layout: A pair of integers specifying the block sizes used to divide
931 + the tensors during shuffling. Default is (16, 16).
932 +
933 + Returns:
934 + A Tuple of shuffled tensors.
935 + """
936 + from aiter.ops.shuffle import shuffle_weight
937 +
938 + return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
939 +
940 +
941 +rocm_aiter_ops.register_ops_once()
vllm/attention/ops/rocm_aiter_mla.py deleted
-105
@@ -1,105 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -
5 -import torch
6 -
7 -from vllm.platforms import current_platform
8 -from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
9 -
10 -
11 -def get_aiter_mla_metadata(
12 - max_batch_size: int, block_size: int, max_block_per_batch: int, device: torch.device
13 -) -> tuple[torch.Tensor, ...]:
14 - paged_kv_indices = torch.zeros(
15 - max_batch_size * max_block_per_batch, dtype=torch.int32, device=device
16 - )
17 - paged_kv_indptr = torch.zeros(max_batch_size + 1, dtype=torch.int32, device=device)
18 - paged_kv_last_page_lens = torch.full(
19 - (max_batch_size,), block_size, dtype=torch.int32
20 - )
21 - qo_indptr = torch.zeros(max_batch_size + 1, dtype=torch.int, device=device)
22 - return paged_kv_indices, paged_kv_indptr, paged_kv_last_page_lens, qo_indptr
23 -
24 -
25 -def aiter_mla_decode_fwd(
26 - q: torch.Tensor,
27 - kv_buffer: torch.Tensor,
28 - o: torch.Tensor,
29 - sm_scale: float,
30 - qo_indptr: torch.Tensor,
31 - max_seqlen_qo: int,
32 - kv_indptr: torch.Tensor | None = None,
33 - kv_indices: torch.Tensor | None = None,
34 - kv_last_page_lens: torch.Tensor | None = None,
35 - logit_cap: float = 0.0,
36 -):
37 - torch.ops.vllm.rocm_aiter_mla_decode_fwd(
38 - q,
39 - kv_buffer.view(-1, 1, 1, q.shape[-1]),
40 - o,
41 - qo_indptr,
42 - max_seqlen_qo,
43 - kv_indptr,
44 - kv_indices,
45 - kv_last_page_lens,
46 - sm_scale=sm_scale,
47 - logit_cap=logit_cap,
48 - )
49 -
50 -
51 -def mla_decode_fwd_impl(
52 - q: torch.Tensor,
53 - kv_buffer: torch.Tensor,
54 - o: torch.Tensor,
55 - qo_indptr: torch.Tensor,
56 - max_seqlen_qo: int,
57 - kv_indptr: torch.Tensor | None = None,
58 - kv_indices: torch.Tensor | None = None,
59 - kv_last_page_lens: torch.Tensor | None = None,
60 - sm_scale: float = 1.0,
61 - logit_cap: float = 0.0,
62 -) -> None:
63 - from aiter.mla import mla_decode_fwd
64 -
65 - mla_decode_fwd(
66 - q,
67 - kv_buffer.view(-1, 1, 1, q.shape[-1]),
68 - o,
69 - qo_indptr,
70 - kv_indptr,
71 - kv_indices,
72 - kv_last_page_lens,
73 - max_seqlen_qo,
74 - sm_scale=sm_scale,
75 - logit_cap=logit_cap,
76 - )
77 -
78 -
79 -def mla_decode_fwd_fake(
80 - q: torch.Tensor,
81 - kv_buffer: torch.Tensor,
82 - o: torch.Tensor,
83 - qo_indptr: torch.Tensor,
84 - max_seqlen_qo: int,
85 - kv_indptr: torch.Tensor | None = None,
86 - kv_indices: torch.Tensor | None = None,
87 - kv_last_page_lens: torch.Tensor | None = None,
88 - sm_scale: float = 1.0,
89 - logit_cap: float = 0.0,
90 -) -> None:
91 - pass
92 -
93 -
94 -if current_platform.is_rocm():
95 - if is_torch_equal_or_newer("2.7.0"):
96 - tags = ()
97 - else:
98 - tags = ((torch.Tag.needs_fixed_stride_order,),)
99 - direct_register_custom_op(
100 - op_name="rocm_aiter_mla_decode_fwd",
101 - op_func=mla_decode_fwd_impl,
102 - mutates_args=["o"],
103 - fake_impl=mla_decode_fwd_fake,
104 - tags=tags,
105 - )
vllm/envs.py
+4 -4
@@ -109,7 +109,7 @@ if TYPE_CHECKING:
109 VLLM_ROCM_USE_AITER_MLA: bool = True
110 VLLM_ROCM_USE_AITER_MHA: bool = True
111 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
112 - VLLM_ROCM_USE_TRITON_ROPE: bool = False
112 + VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
113 VLLM_ROCM_USE_AITER_FP8BMM: bool = True
114 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False
115 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = True
@@ -926,8 +926,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
926 ),
927 # Whether to use aiter rope.
928 # By default is disabled.
929 - "VLLM_ROCM_USE_TRITON_ROPE": lambda: (
930 - os.getenv("VLLM_ROCM_USE_TRITON_ROPE", "False").lower() in ("true", "1")
929 + "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (
930 + os.getenv("VLLM_ROCM_USE_AITER_TRITON_ROPE", "False").lower() in ("true", "1")
931 ),
932 # Whether to use aiter triton fp8 bmm kernel
933 # By default is enabled.
@@ -1589,7 +1589,7 @@ def compute_hash() -> str:
1589 "VLLM_ROCM_USE_AITER_MLA",
1590 "VLLM_ROCM_USE_AITER_MHA",
1591 "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
1592 - "VLLM_ROCM_USE_TRITON_ROPE",
1592 + "VLLM_ROCM_USE_AITER_TRITON_ROPE",
1593 "VLLM_ROCM_USE_AITER_FP8BMM",
1594 "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION",
1595 "VLLM_ROCM_USE_AITER_TRITON_GEMM",
vllm/model_executor/layers/fused_moe/fused_moe.py
+7 -8
@@ -14,6 +14,7 @@ import torch.nn.functional as F
14 import vllm.envs as envs
15 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16 from vllm import _custom_ops as ops
17 +from vllm._aiter_ops import rocm_aiter_ops
18 from vllm.logger import init_logger
19 from vllm.model_executor.layers.batch_invariant import (
20 vllm_is_batch_invariant,
@@ -55,8 +56,6 @@ from vllm.triton_utils import tl, triton
56 from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
57 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
58
58 -from .rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled
59 -
59 logger = init_logger(__name__)
60
61
@@ -1089,11 +1088,11 @@ def vllm_topk_softmax(
1088 return topk_weights, topk_indices
1089
1090
1092 -def dispatch_topk_func() -> Callable[..., tuple[torch.Tensor, ...]]:
1093 - if is_rocm_aiter_moe_enabled():
1094 - from .rocm_aiter_fused_moe import rocm_aiter_topk_softmax
1095 -
1096 - return rocm_aiter_topk_softmax
1091 +def dispatch_topk_func(
1092 + use_rocm_aiter: bool = False,
1093 +) -> Callable[..., tuple[torch.Tensor, ...]]:
1094 + if use_rocm_aiter:
1095 + return rocm_aiter_ops.topk_softmax
1096 return vllm_topk_softmax
1097
1098
@@ -1121,7 +1120,7 @@ def fused_topk(
1120 M, topk, dtype=torch.int32, device=hidden_states.device
1121 )
1122
1124 - topk_func = dispatch_topk_func()
1123 + topk_func = dispatch_topk_func(use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled())
1124 topk_weights, topk_ids = topk_func(
1125 topk_weights, topk_ids, token_expert_indices, gating_output, renormalize
1126 )
vllm/model_executor/layers/fused_moe/layer.py
+43 -40
@@ -13,6 +13,7 @@ import torch.nn.functional as F
13 from torch.nn.parameter import UninitializedParameter
14
15 import vllm.envs as envs
16 +from vllm._aiter_ops import rocm_aiter_ops
17 from vllm.config import VllmConfig, get_current_vllm_config
18 from vllm.config.parallel import ExpertPlacementStrategy
19 from vllm.distributed import (
@@ -41,8 +42,6 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
42 )
43 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
44 init_aiter_topK_meta_data,
44 - is_rocm_aiter_fusion_shared_expert_enabled,
45 - is_rocm_aiter_moe_enabled,
45 )
46 from vllm.model_executor.layers.fused_moe.routing_simulator import RoutingSimulator
47 from vllm.model_executor.layers.quantization.base_config import (
@@ -92,13 +91,11 @@ else:
91 return topk_ids
92
93 eplb_map_to_physical_and_record = _eplb_map_to_physical_and_record
94 +from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
95 +from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
96 + rocm_aiter_grouped_topk,
97 +)
98
96 -if is_rocm_aiter_moe_enabled():
97 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
98 - rocm_aiter_grouped_topk as grouped_topk_aiter,
99 - )
100 -else:
101 - from vllm.model_executor.layers.fused_moe.fused_moe import grouped_topk
99 if current_platform.is_tpu():
100 from .moe_pallas import fused_moe as fused_moe_pallas
101 else:
@@ -463,7 +460,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
460
461 def __init__(self, moe: FusedMoEConfig):
462 super().__init__(moe)
466 - self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
463 +
464 + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
465 if self.rocm_aiter_moe_enabled:
466 from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
467
@@ -620,13 +618,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
618 # Padding the weight for better performance on ROCm
619 layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
620 layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
623 - # Lazy import to avoid importing triton.
624 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
625 - shuffle_weights,
626 - )
621
622 if self.rocm_aiter_moe_enabled:
629 - shuffled_w13, shuffled_w2 = shuffle_weights(
623 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
624 layer.w13_weight.data, layer.w2_weight.data
625 )
626
@@ -1002,6 +996,7 @@ def determine_expert_map(
996 global_num_experts: int,
997 expert_placement_strategy: ExpertPlacementStrategy = "linear",
998 num_fused_shared_experts: int = 0,
999 + return_expert_mask: bool = False,
1000 ) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
1001 """
1002 Calculates how many experts should be assigned to each rank for EP and
@@ -1064,7 +1059,7 @@ def determine_expert_map(
1059 )
1060
1061 expert_mask = None
1067 - if is_rocm_aiter_moe_enabled():
1062 + if return_expert_mask:
1063 expert_mask = torch.ones(
1064 (global_num_experts + num_fused_shared_experts + 1,), dtype=torch.int32
1065 )
@@ -1292,14 +1287,18 @@ class FusedMoE(CustomOp):
1287 self.logical_replica_count: torch.Tensor | None = None
1288
1289 # ROCm aiter shared experts fusion
1290 + self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
1291 + self.aiter_fmoe_shared_expert_enabled = (
1292 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
1293 + )
1294 +
1295 self.num_fused_shared_experts = (
1296 n_shared_experts
1297 - if n_shared_experts is not None
1298 - and is_rocm_aiter_fusion_shared_expert_enabled()
1297 + if n_shared_experts is not None and self.aiter_fmoe_shared_expert_enabled
1298 else 0
1299 )
1300 if (
1302 - not is_rocm_aiter_fusion_shared_expert_enabled()
1301 + not self.aiter_fmoe_shared_expert_enabled
1302 and self.num_fused_shared_experts != 0
1303 ):
1304 raise ValueError(
@@ -1346,6 +1345,7 @@ class FusedMoE(CustomOp):
1345 global_num_experts=self.global_num_experts,
1346 expert_placement_strategy=expert_placement_strategy,
1347 num_fused_shared_experts=self.num_fused_shared_experts,
1348 + return_expert_mask=self.rocm_aiter_fmoe_enabled,
1349 )
1350 self.local_num_experts = local_num_experts
1351 self.register_buffer("expert_map", expert_map)
@@ -1570,13 +1570,16 @@ class FusedMoE(CustomOp):
1570 ep_rank=self.ep_rank,
1571 global_num_experts=self.global_num_experts,
1572 num_fused_shared_experts=self.num_fused_shared_experts,
1573 + return_expert_mask=self.rocm_aiter_fmoe_enabled,
1574 )
1575 self.local_num_experts = local_num_experts
1576 self.register_buffer("expert_map", expert_map)
1577 self.register_buffer("expert_mask", expert_mask)
1577 - self._init_aiter_shared_experts_topK_buffer(
1578 - vllm_config=get_current_vllm_config(), dp_size=get_dp_group().world_size
1579 - )
1578 + if self.aiter_fmoe_shared_expert_enabled:
1579 + self._init_aiter_shared_experts_topK_buffer(
1580 + vllm_config=get_current_vllm_config(),
1581 + dp_size=get_dp_group().world_size,
1582 + )
1583
1584 def _load_per_tensor_weight_scale(
1585 self,
@@ -1753,20 +1756,19 @@ class FusedMoE(CustomOp):
1756 def _init_aiter_shared_experts_topK_buffer(
1757 self, vllm_config: VllmConfig, dp_size: int
1758 ):
1756 - if is_rocm_aiter_fusion_shared_expert_enabled():
1757 - if self.num_fused_shared_experts > 0:
1758 - init_aiter_topK_meta_data(
1759 - n_routed_experts=self.global_num_experts,
1760 - n_shared_experts=self.num_fused_shared_experts,
1761 - top_k=self.top_k,
1762 - tp_rank=self.ep_rank if self.use_ep else self.tp_rank,
1763 - tp_size=self.ep_size if self.use_ep else self.tp_size,
1764 - shared_experts_score=1.0,
1765 - max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens
1766 - * dp_size,
1767 - is_EP=self.use_ep,
1768 - )
1769 - self.local_num_experts += self.num_fused_shared_experts
1759 + if self.num_fused_shared_experts > 0:
1760 + init_aiter_topK_meta_data(
1761 + n_routed_experts=self.global_num_experts,
1762 + n_shared_experts=self.num_fused_shared_experts,
1763 + top_k=self.top_k,
1764 + tp_rank=self.ep_rank if self.use_ep else self.tp_rank,
1765 + tp_size=self.ep_size if self.use_ep else self.tp_size,
1766 + shared_experts_score=1.0,
1767 + max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens
1768 + * dp_size,
1769 + is_EP=self.use_ep,
1770 + )
1771 + self.local_num_experts += self.num_fused_shared_experts
1772
1773 @overload
1774 def weight_loader(
@@ -2208,15 +2210,16 @@ class FusedMoE(CustomOp):
2210 elif use_grouped_topk:
2211 assert topk_group is not None
2212 assert num_expert_group is not None
2211 - if is_rocm_aiter_moe_enabled():
2212 - if not is_rocm_aiter_fusion_shared_expert_enabled():
2213 + if rocm_aiter_ops.is_fused_moe_enabled():
2214 + if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
2215 assert num_fused_shared_experts == 0
2216 grouped_topk_impl = partial(
2215 - grouped_topk_aiter,
2217 + rocm_aiter_grouped_topk,
2218 num_fused_shared_experts=num_fused_shared_experts,
2219 )
2220 else:
2221 grouped_topk_impl = grouped_topk
2222 +
2223 topk_weights, topk_ids = grouped_topk_impl(
2224 hidden_states=hidden_states,
2225 gating_output=router_logits,
@@ -2448,7 +2451,7 @@ class FusedMoE(CustomOp):
2451 use_grouped_topk=self.use_grouped_topk,
2452 global_num_experts=self.global_num_experts,
2453 expert_map=self.expert_map
2451 - if not is_rocm_aiter_moe_enabled()
2454 + if not self.rocm_aiter_fmoe_enabled
2455 else self.expert_mask,
2456 topk_group=self.topk_group,
2457 num_expert_group=self.num_expert_group,
@@ -2612,7 +2615,7 @@ class FusedMoE(CustomOp):
2615 use_grouped_topk=self.use_grouped_topk,
2616 global_num_experts=self.global_num_experts,
2617 expert_map=self.expert_map
2615 - if not is_rocm_aiter_moe_enabled()
2618 + if not self.rocm_aiter_fmoe_enabled
2619 else self.expert_mask,
2620 topk_group=self.topk_group,
2621 num_expert_group=self.num_expert_group,
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+17 -312
@@ -1,17 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from enum import IntEnum
4 -from functools import cache, lru_cache
4 +from functools import lru_cache
5
6 import torch
7
8 -from vllm import envs
8 +from vllm._aiter_ops import rocm_aiter_ops
9 from vllm.model_executor.layers.fused_moe.config import (
10 FUSED_MOE_UNQUANTIZED_CONFIG,
11 FusedMoEQuantConfig,
12 )
13 -from vllm.platforms import current_platform
14 -from vllm.utils.torch_utils import direct_register_custom_op
13
14
15 class QuantMethod(IntEnum):
@@ -37,27 +35,6 @@ class ActivationMethod(IntEnum):
35 GELU = 1
36
37
40 -@cache
41 -def is_rocm_aiter_moe_enabled() -> bool:
42 - return (
43 - current_platform.is_rocm()
44 - and envs.VLLM_ROCM_USE_AITER_MOE
45 - and envs.VLLM_ROCM_USE_AITER
46 - )
47 -
48 -
49 -@cache
50 -def use_mxfp4_aiter_moe() -> bool:
51 - return current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER
52 -
53 -
54 -@cache
55 -def is_rocm_aiter_fusion_shared_expert_enabled() -> bool:
56 - return (
57 - envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS and is_rocm_aiter_moe_enabled()
58 - )
59 -
60 -
38 aiter_topK_meta_data = None
39
40
@@ -114,250 +91,6 @@ def init_aiter_topK_meta_data(
91 aiter_topK_meta_data = (total_topk_weights, total_topk_ids)
92
93
117 -def rocm_aiter_asm_moe_tkw1_impl(
118 - hidden_states: torch.Tensor,
119 - w1: torch.Tensor,
120 - w2: torch.Tensor,
121 - topk_weights: torch.Tensor,
122 - topk_ids: torch.Tensor,
123 - fc1_scale: torch.Tensor | None = None,
124 - fc2_scale: torch.Tensor | None = None,
125 - fc1_smooth_scale: torch.Tensor | None = None,
126 - fc2_smooth_scale: torch.Tensor | None = None,
127 - a16: bool = False,
128 - per_tensor_quant_scale: torch.Tensor | None = None,
129 - expert_mask: torch.Tensor | None = None,
130 - activation_method: int = ActivationMethod.SILU.value,
131 -) -> torch.Tensor:
132 - from aiter import ActivationType
133 - from aiter.fused_moe_bf16_asm import asm_moe_tkw1
134 -
135 - activation = ActivationType(activation_method)
136 -
137 - return asm_moe_tkw1(
138 - hidden_states,
139 - w1,
140 - w2,
141 - topk_weights,
142 - topk_ids,
143 - fc1_scale=fc1_scale,
144 - fc2_scale=fc2_scale,
145 - fc1_smooth_scale=fc1_smooth_scale,
146 - fc2_smooth_scale=fc2_smooth_scale,
147 - a16=a16,
148 - per_tensor_quant_scale=per_tensor_quant_scale,
149 - expert_mask=expert_mask,
150 - activation=activation,
151 - )
152 -
153 -
154 -def rocm_aiter_asm_moe_tkw1_fake(
155 - hidden_states: torch.Tensor,
156 - w1: torch.Tensor,
157 - w2: torch.Tensor,
158 - topk_weights: torch.Tensor,
159 - topk_ids: torch.Tensor,
160 - fc1_scale: torch.Tensor | None = None,
161 - fc2_scale: torch.Tensor | None = None,
162 - fc1_smooth_scale: torch.Tensor | None = None,
163 - fc2_smooth_scale: torch.Tensor | None = None,
164 - a16: bool = False,
165 - per_tensor_quant_scale: torch.Tensor | None = None,
166 - expert_mask: torch.Tensor | None = None,
167 - activation_method: int = ActivationMethod.SILU.value,
168 -) -> torch.Tensor:
169 - return torch.empty_like(hidden_states)
170 -
171 -
172 -def rocm_aiter_topk_softmax_impl(
173 - topk_weights: torch.Tensor,
174 - topk_indices: torch.Tensor,
175 - token_expert_indices: torch.Tensor,
176 - gating_output: torch.Tensor,
177 - renormalize: bool,
178 -) -> None:
179 - from aiter import topk_softmax
180 -
181 - topk_softmax(
182 - topk_weights, topk_indices, token_expert_indices, gating_output, renormalize
183 - )
184 -
185 -
186 -def rocm_aiter_topk_softmax_fake(
187 - topk_weights: torch.Tensor,
188 - topk_indices: torch.Tensor,
189 - token_expert_indices: torch.Tensor,
190 - gating_output: torch.Tensor,
191 - renormalize: bool,
192 -) -> None:
193 - pass
194 -
195 -
196 -def rocm_aiter_biased_grouped_topk_impl(
197 - gating_output: torch.Tensor,
198 - correction_bias: torch.Tensor,
199 - topk_weights: torch.Tensor,
200 - topk_ids: torch.Tensor,
201 - num_expert_group: int,
202 - topk_group: int,
203 - need_renorm: bool,
204 - routed_scaling_factor: float = 1.0, # mul to topk_weights
205 -) -> None:
206 - from aiter import biased_grouped_topk
207 -
208 - biased_grouped_topk(
209 - gating_output,
210 - correction_bias,
211 - topk_weights,
212 - topk_ids,
213 - num_expert_group,
214 - topk_group,
215 - need_renorm,
216 - routed_scaling_factor,
217 - )
218 -
219 -
220 -def rocm_aiter_biased_grouped_topk_fake(
221 - gating_output: torch.Tensor,
222 - correction_bias: torch.Tensor,
223 - topk_weights: torch.Tensor,
224 - topk_ids: torch.Tensor,
225 - num_expert_group: int,
226 - topk_group: int,
227 - need_renorm: bool,
228 - routed_scaling_factor: float = 1.0, # mul to topk_weights
229 -) -> None:
230 - pass
231 -
232 -
233 -def rocm_aiter_grouped_topk_impl(
234 - gating_output: torch.Tensor,
235 - topk_weights: torch.Tensor,
236 - topk_ids: torch.Tensor,
237 - num_expert_group: int,
238 - topk_group: int,
239 - need_renorm: bool,
240 - scoring_func: str = "softmax",
241 - routed_scaling_factor: float = 1.0, # mul to topk_weights
242 -) -> None:
243 - from aiter import grouped_topk
244 -
245 - grouped_topk(
246 - gating_output,
247 - topk_weights,
248 - topk_ids,
249 - num_expert_group,
250 - topk_group,
251 - need_renorm,
252 - scoring_func,
253 - routed_scaling_factor,
254 - )
255 -
256 -
257 -def rocm_aiter_grouped_topk_fake(
258 - gating_output: torch.Tensor,
259 - topk_weights: torch.Tensor,
260 - topk_ids: torch.Tensor,
261 - num_expert_group: int,
262 - topk_group: int,
263 - need_renorm: bool,
264 - scoring_func: str = "softmax",
265 - routed_scaling_factor: float = 1.0, # mul to topk_weights
266 -) -> None:
267 - pass
268 -
269 -
270 -def rocm_aiter_fused_moe_impl(
271 - hidden_states: torch.Tensor,
272 - w1: torch.Tensor,
273 - w2: torch.Tensor,
274 - topk_weight: torch.Tensor,
275 - topk_ids: torch.Tensor,
276 - expert_mask: torch.Tensor | None = None,
277 - activation_method: int = ActivationMethod.SILU.value,
278 - quant_method: int = QuantMethod.NO.value,
279 - doweight_stage1: bool = False,
280 - w1_scale: torch.Tensor | None = None,
281 - w2_scale: torch.Tensor | None = None,
282 - a1_scale: torch.Tensor | None = None,
283 - a2_scale: torch.Tensor | None = None,
284 -) -> torch.Tensor:
285 - from aiter import ActivationType, QuantType
286 - from aiter.fused_moe import fused_moe
287 -
288 - activation = ActivationType(activation_method)
289 - quant_type = QuantType(quant_method)
290 -
291 - return fused_moe(
292 - hidden_states,
293 - w1,
294 - w2,
295 - topk_weight,
296 - topk_ids,
297 - expert_mask,
298 - activation,
299 - quant_type,
300 - doweight_stage1,
301 - w1_scale,
302 - w2_scale,
303 - a1_scale,
304 - a2_scale,
305 - )
306 -
307 -
308 -def rocm_aiter_fused_moe_fake(
309 - hidden_states: torch.Tensor,
310 - w1: torch.Tensor,
311 - w2: torch.Tensor,
312 - topk_weight: torch.Tensor,
313 - topk_ids: torch.Tensor,
314 - expert_mask: torch.Tensor | None = None,
315 - activation_method: int = ActivationMethod.SILU.value,
316 - quant_method: int = QuantMethod.NO.value,
317 - doweight_stage1: bool = False,
318 - w1_scale: torch.Tensor | None = None,
319 - w2_scale: torch.Tensor | None = None,
320 - a1_scale: torch.Tensor | None = None,
321 - a2_scale: torch.Tensor | None = None,
322 -) -> torch.Tensor:
323 - return torch.empty_like(hidden_states)
324 -
325 -
326 -if current_platform.is_rocm():
327 - direct_register_custom_op(
328 - op_name="rocm_aiter_asm_moe_tkw1",
329 - op_func=rocm_aiter_asm_moe_tkw1_impl,
330 - fake_impl=rocm_aiter_asm_moe_tkw1_fake,
331 - )
332 -
333 - direct_register_custom_op(
334 - op_name="rocm_aiter_fused_moe",
335 - op_func=rocm_aiter_fused_moe_impl,
336 - fake_impl=rocm_aiter_fused_moe_fake,
337 - )
338 -
339 - direct_register_custom_op(
340 - op_name="rocm_aiter_topk_softmax",
341 - op_func=rocm_aiter_topk_softmax_impl,
342 - mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
343 - fake_impl=rocm_aiter_topk_softmax_fake,
344 - )
345 -
346 - direct_register_custom_op(
347 - op_name="rocm_aiter_biased_grouped_topk",
348 - op_func=rocm_aiter_biased_grouped_topk_impl,
349 - mutates_args=["topk_weights", "topk_ids"],
350 - fake_impl=rocm_aiter_biased_grouped_topk_fake,
351 - )
352 -
353 - direct_register_custom_op(
354 - op_name="rocm_aiter_grouped_topk",
355 - op_func=rocm_aiter_grouped_topk_impl,
356 - mutates_args=["topk_weights", "topk_ids"],
357 - fake_impl=rocm_aiter_grouped_topk_fake,
358 - )
359 -
360 -
94 def rocm_aiter_grouped_topk(
95 hidden_states: torch.Tensor,
96 gating_output: torch.Tensor,
@@ -372,7 +105,10 @@ def rocm_aiter_grouped_topk(
105 ) -> tuple[torch.Tensor, torch.Tensor]:
106 token = hidden_states.shape[0]
107 device = hidden_states.device
375 - if is_rocm_aiter_fusion_shared_expert_enabled() and num_fused_shared_experts > 0:
108 + if (
109 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
110 + and num_fused_shared_experts > 0
111 + ):
112 assert aiter_topK_meta_data is not None, (
113 "AITER topK meta data is not initialized. "
114 "Please ensure that init_aiter_topK_meta_data "
@@ -397,7 +133,7 @@ def rocm_aiter_grouped_topk(
133 topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device)
134
135 if e_score_correction_bias is not None:
400 - torch.ops.vllm.rocm_aiter_biased_grouped_topk(
136 + rocm_aiter_ops.biased_grouped_topk(
137 gating_output,
138 e_score_correction_bias.to(gating_output.dtype),
139 topk_weights,
@@ -409,7 +145,7 @@ def rocm_aiter_grouped_topk(
145 )
146 else:
147 assert scoring_func == "softmax" or scoring_func == "sigmoid"
412 - torch.ops.vllm.rocm_aiter_grouped_topk(
148 + rocm_aiter_ops.grouped_topk(
149 gating_output,
150 topk_weights,
151 topk_ids,
@@ -420,7 +156,10 @@ def rocm_aiter_grouped_topk(
156 routed_scaling_factor=routed_scaling_factor,
157 )
158
423 - if is_rocm_aiter_fusion_shared_expert_enabled() and num_fused_shared_experts > 0:
159 + if (
160 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
161 + and num_fused_shared_experts > 0
162 + ):
163 return total_topk_weights, total_topk_ids
164 return topk_weights, topk_ids
165
@@ -464,7 +203,7 @@ def rocm_aiter_fused_experts(
203 "Only support topk=1 when `apply_router_weight_on_input` is True"
204 )
205
467 - return torch.ops.vllm.rocm_aiter_asm_moe_tkw1(
206 + return rocm_aiter_ops.asm_moe_tkw1(
207 hidden_states,
208 w1,
209 w2,
@@ -482,7 +221,9 @@ def rocm_aiter_fused_experts(
221
222 else:
223 quant_method = QuantMethod.NO.value
485 -
224 + # quark moe for mxfp4 w_dtype
225 + if quant_config.use_mxfp4_w4a16:
226 + quant_method = QuantMethod.BLOCK_1X32.value
227 # w8a8 block-scaled
228 if quant_config.block_shape is not None and quant_config.use_fp8_w8a8:
229 assert not apply_router_weight_on_input, (
@@ -507,7 +248,7 @@ def rocm_aiter_fused_experts(
248 "Only support topk=1 when `apply_router_weight_on_input` is True"
249 )
250
510 - return torch.ops.vllm.rocm_aiter_fused_moe(
251 + return rocm_aiter_ops.fused_moe(
252 hidden_states,
253 w1,
254 w2,
@@ -522,39 +263,3 @@ def rocm_aiter_fused_experts(
263 a2_scale=quant_config.a2_scale,
264 doweight_stage1=apply_router_weight_on_input,
265 )
525 -
526 -
527 -def rocm_aiter_topk_softmax(
528 - topk_weights: torch.Tensor,
529 - topk_indices: torch.Tensor,
530 - token_expert_indices: torch.Tensor,
531 - gating_output: torch.Tensor,
532 - renormalize: bool,
533 -) -> tuple[torch.Tensor, ...]:
534 - torch.ops.vllm.rocm_aiter_topk_softmax(
535 - topk_weights, topk_indices, token_expert_indices, gating_output, renormalize
536 - )
537 - return topk_weights, topk_indices
538 -
539 -
540 -def shuffle_weights(
541 - *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
542 -) -> tuple[torch.Tensor, ...]:
543 - """
544 - Applies shuffle_weight function from AITER to each
545 - input tensor and returns them.
546 -
547 - Rearranges (shuffles) the input tensor/s
548 - into a specified block layout for optimized computation.
549 -
550 - Args:
551 - *tensors: Variable number of torch.Tensor objects.
552 - layout: A pair of integers specifying the block sizes used to divide
553 - the tensors during shuffling. Default is (16, 16).
554 -
555 - Returns:
556 - A Tuple of shuffled tensors.
557 - """
558 - from aiter.ops.shuffle import shuffle_weight
559 -
560 - return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
vllm/model_executor/layers/layernorm.py
+21 -69
@@ -6,18 +6,13 @@ import torch
6 import torch.nn as nn
7 import torch.nn.functional as F
8
9 -import vllm.envs as envs
9 +from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.model_executor.custom_op import CustomOp
11 from vllm.model_executor.layers.batch_invariant import (
12 rms_norm_batch_invariant,
13 vllm_is_batch_invariant,
14 )
15 from vllm.platforms import current_platform
16 -from vllm.utils.torch_utils import direct_register_custom_op
17 -
18 -
19 -def is_rocm_aiter_rmsnorm_enabled() -> bool:
20 - return envs.VLLM_ROCM_USE_AITER_RMSNORM and envs.VLLM_ROCM_USE_AITER
16
17
18 def rms_norm(
@@ -58,80 +53,34 @@ def fused_add_rms_norm(
53 return x, residual
54
55
61 -def rocm_aiter_rms_norm_impl(
62 - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
56 +def poly_norm(
57 + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, variance_epsilon: float
58 ) -> torch.Tensor:
64 - import aiter as rocm_aiter
65 -
66 - if x.dim() > 2:
67 - x_original_shape = x.shape
68 - x = x.reshape(-1, x_original_shape[-1])
69 - x = rocm_aiter.rms_norm(x, weight, variance_epsilon)
70 - return x.reshape(x_original_shape)
71 -
72 - return rocm_aiter.rms_norm(x, weight, variance_epsilon)
73 -
59 + from vllm import _custom_ops as ops
60
75 -def rocm_aiter_rmsnorm2d_fwd_with_add_impl(
76 - x: torch.Tensor,
77 - residual: torch.Tensor,
78 - weight: torch.Tensor,
79 - variance_epsilon: float,
80 -) -> tuple[torch.Tensor, torch.Tensor]:
81 - import aiter as rocm_aiter
82 -
83 - residual_out = torch.empty_like(residual)
84 - output = torch.empty_like(x)
85 - rocm_aiter.rmsnorm2d_fwd_with_add(
86 - output, # output
87 - x, # input
88 - residual, # residual input
89 - residual_out, # residual output
61 + out = torch.empty_like(x)
62 + ops.poly_norm(
63 + out,
64 + x,
65 weight,
66 + bias,
67 variance_epsilon,
68 )
93 - return output, residual_out
94 -
95 -
96 -def rocm_aiter_rms_norm_fake(
97 - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
98 -) -> torch.Tensor:
99 - return torch.empty_like(x)
100 -
101 -
102 -def rocm_aiter_rmsnorm2d_fwd_with_add_fake(
103 - x: torch.Tensor,
104 - residual: torch.Tensor,
105 - weight: torch.Tensor,
106 - variance_epsilon: float,
107 -) -> tuple[torch.Tensor, torch.Tensor]:
108 - return torch.empty_like(x), torch.empty_like(residual)
109 -
110 -
111 -if current_platform.is_rocm():
112 - direct_register_custom_op(
113 - op_name="rocm_aiter_rms_norm",
114 - op_func=rocm_aiter_rms_norm_impl,
115 - fake_impl=rocm_aiter_rms_norm_fake,
116 - )
117 -
118 - direct_register_custom_op(
119 - op_name="rocm_aiter_rmsnorm2d_fwd_with_add",
120 - op_func=rocm_aiter_rmsnorm2d_fwd_with_add_impl,
121 - fake_impl=rocm_aiter_rmsnorm2d_fwd_with_add_fake,
122 - )
69 + return out
70
71
125 -def dispatch_rocm_rmsnorm_func(with_fused_add: bool, dtype: torch.dtype):
126 - use_aiter = is_rocm_aiter_rmsnorm_enabled() and dtype in [
72 +def dispatch_rocm_rmsnorm_func(
73 + with_fused_add: bool, dtype: torch.dtype, use_aiter: bool = False
74 +):
75 + use_aiter = use_aiter and dtype in [
76 torch.float16,
77 torch.bfloat16,
78 ]
79
80 if use_aiter and with_fused_add:
132 - return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add
81 + return rocm_aiter_ops.rms_norm2d_with_add
82 if use_aiter:
134 - return torch.ops.vllm.rocm_aiter_rms_norm
83 + return rocm_aiter_ops.rms_norm
84
85 # fall back to CUDA implementation
86 if with_fused_add:
@@ -169,11 +118,14 @@ class RMSNorm(CustomOp):
118 self.weight = nn.Parameter(self.weight)
119
120 if current_platform.is_rocm():
121 + aiter_rmsnorm_enabled = rocm_aiter_ops.is_rmsnorm_enabled()
122 self.rocm_norm_func = dispatch_rocm_rmsnorm_func(
173 - with_fused_add=False, dtype=weight_dtype
123 + with_fused_add=False,
124 + dtype=weight_dtype,
125 + use_aiter=aiter_rmsnorm_enabled,
126 )
127 self.rocm_norm_func_with_add = dispatch_rocm_rmsnorm_func(
176 - with_fused_add=True, dtype=weight_dtype
128 + with_fused_add=True, dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled
129 )
130
131 @staticmethod
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+3 -9
@@ -12,6 +12,7 @@ from compressed_tensors.quantization import ActivationOrdering, QuantizationStra
12 import vllm.envs as envs
13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 from vllm import _custom_ops as ops
15 +from vllm._aiter_ops import rocm_aiter_ops
16 from vllm.distributed import get_tensor_model_parallel_world_size
17 from vllm.logger import init_logger
18 from vllm.model_executor.layers.fused_moe import (
@@ -582,11 +583,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
583 # Disable marlin for rocm
584 if current_platform.is_rocm():
585 self.use_marlin = False
585 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
586 - is_rocm_aiter_moe_enabled,
587 - )
586
589 - self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
587 + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
588
589 # cutlass path
590 self.is_fp8_w8a8_sm100 = quant_config._is_fp8_w8a8_sm100(
@@ -829,12 +827,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
827
828 # Property to determine if AITER is used
829 if self.rocm_aiter_moe_enabled:
832 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501
833 - shuffle_weights,
834 - )
835 -
830 # reshaping weights is required for aiter moe kernel.
837 - shuffled_w13, shuffled_w2 = shuffle_weights(
831 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
832 layer.w13_weight.data, layer.w2_weight.data
833 )
834
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+2 -2
@@ -7,12 +7,12 @@ import torch
7 from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
8 from torch.nn import Parameter
9
10 +from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
12 CompressedTensorsScheme,
13 )
14 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
15 W8A8BlockFp8LinearOp,
15 - check_aiter_fp8_linear_support,
16 create_fp8_input_scale,
17 create_fp8_scale_parameter,
18 create_fp8_weight_parameter,
@@ -61,7 +61,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
61 )
62
63 self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
64 - self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
64 + self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enaled()
65
66 if self.weight_block_size is not None:
67 assert not self.is_static_input_scheme
vllm/model_executor/layers/quantization/fp8.py
+6 -10
@@ -12,6 +12,7 @@ from torch.nn.parameter import Parameter
12 import vllm.envs as envs
13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 from vllm import _custom_ops as ops
15 +from vllm._aiter_ops import rocm_aiter_ops
16 from vllm.distributed import get_tensor_model_parallel_world_size
17 from vllm.logger import init_logger
18 from vllm.model_executor.layers.batch_invariant import (
@@ -56,7 +57,6 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
57 )
58 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
59 W8A8BlockFp8LinearOp,
59 - check_aiter_fp8_linear_support,
60 create_fp8_input_scale,
61 create_fp8_scale_parameter,
62 create_fp8_weight_parameter,
@@ -369,7 +369,7 @@ class Fp8LinearMethod(LinearMethodBase):
369 if vllm_is_batch_invariant():
370 self.use_marlin = False
371
372 - self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
372 + self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enaled()
373 self.use_deep_gemm = is_deep_gemm_supported()
374
375 self.weight_block_size = self.quant_config.weight_block_size
@@ -869,12 +869,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
869
870 def process_weights_after_loading(self, layer: Module) -> None:
871 # Lazy import to avoid importing triton too early.
872 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
873 - is_rocm_aiter_moe_enabled,
874 - shuffle_weights,
875 - )
872
877 - self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
873 + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
874
875 # TODO (rob): refactor block quant into separate class.
876 if self.block_quant:
@@ -916,7 +912,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
912 )
913 if self.rocm_aiter_moe_enabled:
914 # reshaping weights is required for aiter moe kernel.
919 - shuffled_w13, shuffled_w2 = shuffle_weights(
915 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
916 layer.w13_weight.data, layer.w2_weight.data
917 )
918
@@ -962,7 +958,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
958 layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
959 if self.rocm_aiter_moe_enabled:
960 # reshaping weights is required for aiter moe kernel.
965 - shuffled_w13, shuffled_w2 = shuffle_weights(
961 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
962 layer.w13_weight, layer.w2_weight
963 )
964
@@ -1042,7 +1038,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1038 start += shard_size
1039
1040 if self.rocm_aiter_moe_enabled:
1045 - shuffled_w13, shuffled_w2 = shuffle_weights(
1041 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
1042 layer.w13_weight, layer.w2_weight
1043 )
1044
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
+3 -45
@@ -4,54 +4,14 @@
4
5 import torch
6
7 -import vllm.envs as envs
7 from vllm import _custom_ops as ops
8 +from vllm._aiter_ops import rocm_aiter_ops
9 from vllm.platforms import current_platform
10 -from vllm.utils.torch_utils import direct_register_custom_op
10
11 from .cutlass import CutlassScaledMMLinearKernel
12 from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
13
14
16 -def rocm_aiter_gemm_w8a8_impl(
17 - A: torch.Tensor,
18 - B: torch.Tensor,
19 - As: torch.Tensor,
20 - Bs: torch.Tensor,
21 - bias: torch.Tensor | None = None,
22 - output_dtype: torch.dtype = torch.float16,
23 -) -> torch.Tensor:
24 - from aiter import gemm_a8w8_CK
25 -
26 - # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
27 - # a to be [M, K]
28 - # b to be [N, K]
29 - # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
30 - return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)
31 -
32 -
33 -def rocm_aiter_gemm_w8a8_fake(
34 - A: torch.Tensor,
35 - B: torch.Tensor,
36 - As: torch.Tensor,
37 - Bs: torch.Tensor,
38 - bias: torch.Tensor | None = None,
39 - output_dtype: torch.dtype = torch.float16,
40 -) -> torch.Tensor:
41 - m = A.shape[0]
42 - n = B.shape[0]
43 - Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
44 - return Y
45 -
46 -
47 -if current_platform.is_rocm():
48 - direct_register_custom_op(
49 - op_name="rocm_aiter_gemm_w8a8",
50 - op_func=rocm_aiter_gemm_w8a8_impl,
51 - fake_impl=rocm_aiter_gemm_w8a8_fake,
52 - )
53 -
54 -
15 class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
16 @classmethod
17 def get_min_capability(cls) -> int:
@@ -75,7 +35,7 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
35 + "installed on ROCm.",
36 )
37 # Check if rocm_aiter_gemm_w8a8_scaled_mm is enabled
78 - if not (envs.VLLM_ROCM_USE_AITER_LINEAR and envs.VLLM_ROCM_USE_AITER):
38 + if not (rocm_aiter_ops.is_linear_enabled()):
39 return (
40 False,
41 "AiterScaledMMLinearKernel is disabled. "
@@ -157,6 +117,4 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
117 # a to be [M, K]
118 # b to be [N, K]
119 # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format
160 - return torch.ops.vllm.rocm_aiter_gemm_w8a8(
161 - x_q, w_q.t(), x_s, w_s, bias, out_dtype
162 - )
120 + return rocm_aiter_ops.gemm_w8a8(x_q, w_q.t(), x_s, w_s, bias, out_dtype)
vllm/model_executor/layers/quantization/quark/quark_moe.py
+16 -31
@@ -8,6 +8,7 @@ import torch
8
9 import vllm.envs as envs
10 from vllm import _custom_ops as ops
11 +from vllm._aiter_ops import rocm_aiter_ops
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.fused_moe import (
14 FusedMoE,
@@ -21,10 +22,6 @@ from vllm.model_executor.layers.fused_moe.config import (
22 ocp_mx_moe_quant_config,
23 )
24 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
24 -from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
25 - is_rocm_aiter_moe_enabled,
26 - use_mxfp4_aiter_moe,
27 -)
25 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
26 prepare_moe_fp8_layer_for_marlin,
27 )
@@ -122,7 +119,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
119 if current_platform.is_rocm():
120 self.use_marlin = False
121
125 - self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
122 + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
123
124 def create_weights(
125 self,
@@ -309,12 +306,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
306 )
307 # Property to determine if AITER is used
308 if self.rocm_aiter_moe_enabled:
312 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501
313 - shuffle_weights,
314 - )
315 -
309 # reshaping weights is required for aiter moe kernel.
317 - shuffled_w13, shuffled_w2 = shuffle_weights(
310 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
311 layer.w13_weight.data, layer.w2_weight.data
312 )
313
@@ -470,13 +463,15 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
463 "not implemented. Please open an issue."
464 )
465
466 + self.use_rocm_aiter_moe = rocm_aiter_ops.is_fused_moe_enabled()
467 +
468 self.emulate = not current_platform.supports_mx() or not (
474 - use_mxfp4_aiter_moe() and self.ocp_mx_scheme == "w_mxfp4_a_mxfp4"
469 + self.use_rocm_aiter_moe and self.ocp_mx_scheme == "w_mxfp4_a_mxfp4"
470 )
471 if self.emulate:
472 logger.warning_once(
473 f"The current mode (supports_mx={current_platform.supports_mx()}, "
479 - f"use_mxfp4_aiter_moe={use_mxfp4_aiter_moe()}, "
474 + f"use_mxfp4_aiter_moe={self.use_rocm_aiter_moe}, "
475 f"ocp_mx_scheme={self.ocp_mx_scheme}) "
476 "does not support native MXFP4/MXFP6 "
477 "computation. Simulated weight dequantization and activation "
@@ -656,28 +651,18 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
651 )
652
653 if not self.emulate:
659 - from aiter import ActivationType, QuantType
660 - from aiter.fused_moe import fused_moe
661 -
662 - aiter_acts = {
663 - ActivationType.No.name.lower(): ActivationType.No,
664 - ActivationType.Silu.name.lower(): ActivationType.Silu,
665 - ActivationType.Gelu.name.lower(): ActivationType.Gelu,
666 - }
667 - assert activation in aiter_acts, (
668 - f"Aiter CK fp4 MoE doesn't support activation {activation}"
669 - )
670 - out = fused_moe(
654 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
655 + rocm_aiter_fused_experts,
656 + )
657 +
658 + out = rocm_aiter_fused_experts(
659 x,
660 layer.w13_weight,
661 layer.w2_weight,
674 - topk_weights,
675 - topk_ids,
676 - quant_type=QuantType.per_1x32,
677 - w1_scale=layer.w13_weight_scale,
678 - w2_scale=layer.w2_weight_scale,
679 - activation=aiter_acts[activation],
680 - doweight_stage1=False,
662 + topk_weights=topk_weights,
663 + topk_ids=topk_ids,
664 + activation=activation,
665 + quant_config=self.moe_quant_config,
666 )
667 else:
668 from vllm.model_executor.layers.fused_moe import fused_experts
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+7
@@ -31,6 +31,13 @@ from .quark_scheme import QuarkScheme
31 logger = init_logger(__name__)
32
33
34 +# TODO: move registration of custom op to aiter_ops.py
35 +# `from vllm._aiter_ops import rocm_aiter_ops`
36 +# use `rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()`
37 +# for envs checks which does not require @cache anymore.
38 +# triton kernel is torch compile compatible.
39 +# does not require direct registeration.
40 +# use `rocm_aiter_ops.triton_fp4_gemm_dynamic_qaunt`.
41 @cache
42 def is_rocm_aiter_fp4_asm_gemm_enabled() -> bool:
43 return (
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+35 -89
@@ -12,6 +12,7 @@ import torch
12
13 import vllm.envs as envs
14 from vllm import _custom_ops as ops
15 +from vllm._aiter_ops import rocm_aiter_ops
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
18 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -68,78 +69,6 @@ def cutlass_scaled_mm(
69 )
70
71
71 -def rocm_aiter_gemm_w8a8_blockscale_impl(
72 - input_2d: torch.Tensor,
73 - weight: torch.Tensor,
74 - input_scale: torch.Tensor,
75 - weight_scale: torch.Tensor,
76 - group_size: int,
77 - output_dtype: torch.dtype = torch.float16,
78 -) -> torch.Tensor:
79 - def is_aiter_triton_kernel_tuned(n, k):
80 - return (n, k) in [
81 - (1024, 8192),
82 - (2112, 7168),
83 - (3072, 1536),
84 - (32768, 8192),
85 - (4096, 7168),
86 - (4608, 7168),
87 - (512, 7168),
88 - (7168, 2048),
89 - (7168, 256),
90 - (8192, 1024),
91 - (8192, 32768),
92 - ]
93 -
94 - n, k = weight.shape
95 - if input_scale is not None:
96 - q_input = input_2d
97 - elif not current_platform.is_fp8_fnuz() and is_aiter_triton_kernel_tuned(n, k):
98 - from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale
99 -
100 - # MI350 case uses triton kernel
101 - q_input, input_scale = per_token_group_quant_fp8(
102 - input_2d,
103 - group_size,
104 - column_major_scales=False,
105 - use_ue8m0=False,
106 - )
107 - else:
108 - # MI300 uses tuned AITER ASM/C++ kernel
109 - import aiter as rocm_aiter
110 - from aiter import gemm_a8w8_blockscale, get_hip_quant
111 -
112 - aiter_per1x128_quant = get_hip_quant(rocm_aiter.QuantType.per_1x128)
113 - q_input, input_scale = aiter_per1x128_quant(
114 - input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8
115 - )
116 -
117 - return gemm_a8w8_blockscale(
118 - q_input, weight, input_scale, weight_scale, dtype=output_dtype
119 - )
120 -
121 -
122 -def rocm_aiter_gemm_w8a8_blockscale_fake(
123 - input_2d: torch.Tensor,
124 - weight: torch.Tensor,
125 - input_scale: torch.Tensor,
126 - weight_scale: torch.Tensor,
127 - group_size: int,
128 - output_dtype: torch.dtype = torch.float16,
129 -) -> torch.Tensor:
130 - m = input_2d.shape[0]
131 - n = weight.shape[0]
132 - return torch.empty(m, n, dtype=output_dtype, device=input_2d.device)
133 -
134 -
135 -if current_platform.is_rocm():
136 - direct_register_custom_op(
137 - op_name="rocm_aiter_gemm_w8a8_blockscale",
138 - op_func=rocm_aiter_gemm_w8a8_blockscale_impl,
139 - fake_impl=rocm_aiter_gemm_w8a8_blockscale_fake,
140 - )
141 -
142 -
72 # TODO we should be able to change the type of block_size to GroupShape
73 # after we resolve GroupShape compilation issue
74 # https://github.com/vllm-project/vllm/issues/25270
@@ -385,14 +314,40 @@ class W8A8BlockFp8LinearOp:
314 input_scale: torch.Tensor | None = None,
315 ) -> torch.Tensor:
316 assert self.act_quant_group_shape == GroupShape(1, 128)
388 - return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale(
389 - input_2d,
390 - weight,
391 - input_scale,
392 - weight_scale,
393 - self.act_quant_group_shape.col,
394 - input_2d.dtype,
395 - )
317 +
318 + n, k = weight.shape
319 + if input_scale is not None:
320 + q_input = input_2d
321 +
322 + # MI350 case uses triton kernel
323 + if (
324 + not current_platform.is_fp8_fnuz()
325 + and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
326 + ):
327 + q_input, input_scale = per_token_group_quant_fp8(
328 + input_2d,
329 + self.act_quant_group_shape.col,
330 + column_major_scales=False,
331 + use_ue8m0=False,
332 + )
333 + return rocm_aiter_ops.triton_gemm_a8w8_blockscale(
334 + q_input,
335 + weight,
336 + input_scale,
337 + weight_scale,
338 + input_2d.dtype,
339 + )
340 +
341 + # MI300 uses tuned AITER ASM/C++ kernel
342 + else:
343 + q_input, input_scale = rocm_aiter_ops.per_1x128_fp8_quant(input_2d)
344 + return rocm_aiter_ops.gemm_w8a8_blockscale(
345 + q_input,
346 + weight,
347 + input_scale,
348 + weight_scale,
349 + input_2d.dtype,
350 + )
351
352 def _run_triton(
353 self,
@@ -971,15 +926,6 @@ def requant_weight_ue8m0_inplace(
926 s_old.copy_(s_requant)
927
928
974 -def check_aiter_fp8_linear_support() -> bool:
975 - """AITER is only supported on ROCm for MI3XX"""
976 - return (
977 - current_platform.is_rocm()
978 - and envs.VLLM_ROCM_USE_AITER
979 - and envs.VLLM_ROCM_USE_AITER_LINEAR
980 - )
981 -
982 -
929 def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
930 """Pad the weight tensor. This is an optimization on ROCm platform, which
931 can benefit from tensors located far enough from one another in memory"""
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+1 -1
@@ -472,7 +472,7 @@ class Fp8LinearOp:
472 # Example:
473 # When the number of token is 1, per-token scale is [[1]]
474 # When per-tensor scale is [1] or ().
475 - per_tensor_weights = (weight_scale.numel() == 1) and weight_scale.dim() < 2
475 + per_tensor_weights = weight_scale.numel() == 1
476 per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2
477
478 # TODO(luka) do this dispatch during init (after ScaledMM refactor)
vllm/model_executor/layers/rotary_embedding/base.py
+5 -8
@@ -4,13 +4,10 @@
4
5 import torch
6
7 +from vllm._aiter_ops import rocm_aiter_ops
8 from vllm.model_executor.custom_op import CustomOp
9
10 from .common import apply_rotary_emb_torch
10 -from .rocm_aiter_rope_ops import (
11 - is_rocm_triton_rotary_embedding_enabled,
12 - rocm_aiter_rotary_emb,
13 -)
11
12
13 @CustomOp.register("rotary_embedding")
@@ -48,8 +45,8 @@ class RotaryEmbeddingBase(CustomOp):
45 cache = cache.to(dtype)
46 self.cos_sin_cache: torch.Tensor
47 self.register_buffer("cos_sin_cache", cache, persistent=False)
51 - self.is_rocm_triton_rotary_embedding_enabled = (
52 - is_rocm_triton_rotary_embedding_enabled()
48 + self.is_rocm_triton_rotary_embed_enabled = (
49 + rocm_aiter_ops.is_triton_rotary_embed_enabled()
50 )
51
52 def _compute_inv_freq(self, base: float) -> torch.Tensor:
@@ -169,9 +166,9 @@ class RotaryEmbedding(RotaryEmbeddingBase):
166 query: torch.Tensor,
167 key: torch.Tensor | None = None,
168 ) -> tuple[torch.Tensor, torch.Tensor | None]:
172 - if self.is_rocm_triton_rotary_embedding_enabled:
169 + if self.is_rocm_triton_rotary_embed_enabled:
170 self._match_cos_sin_cache_dtype(query)
174 - rocm_aiter_rotary_emb(
171 + rocm_aiter_ops.triton_rotary_embed(
172 positions,
173 query,
174 key,
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+9
@@ -146,6 +146,15 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
146 key = key_rot
147 return query, key
148
149 + def forward_hip(
150 + self,
151 + positions: torch.Tensor,
152 + query: torch.Tensor,
153 + key: torch.Tensor | None = None,
154 + offsets: torch.Tensor | None = None,
155 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
156 + return self.forward_native(positions, query, key, offsets)
157 +
158 def forward_cuda(
159 self,
160 positions: torch.Tensor,
vllm/model_executor/layers/rotary_embedding/rocm_aiter_rope_ops.py deleted
-94
@@ -1,94 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch
5 -
6 -import vllm.envs as envs
7 -from vllm.platforms import current_platform
8 -from vllm.utils.torch_utils import direct_register_custom_op
9 -
10 -
11 -def is_rocm_triton_rotary_embedding_enabled() -> bool:
12 - return (
13 - current_platform.is_rocm()
14 - and envs.VLLM_ROCM_USE_AITER
15 - and envs.VLLM_ROCM_USE_TRITON_ROPE
16 - )
17 -
18 -
19 -def rocm_aiter_rotary_emb_with_key_forward_triton_impl(
20 - positions: torch.Tensor,
21 - sin: torch.Tensor,
22 - cos: torch.Tensor,
23 - query: torch.Tensor,
24 - key: torch.Tensor,
25 - rotate_style: int = 0,
26 - is_nope_first: bool = False,
27 -) -> None:
28 - import aiter.ops.triton.rope as ops
29 -
30 - ops.rope_cached_thd_positions_2c_fwd_inplace(
31 - query,
32 - key,
33 - cos,
34 - sin,
35 - positions,
36 - rotate_style,
37 - reuse_freqs_front_part=True,
38 - nope_first=is_nope_first,
39 - )
40 -
41 -
42 -def rocm_aiter_rotary_emb_with_key_forward_triton_fake(
43 - positions: torch.Tensor,
44 - sin: torch.Tensor,
45 - cos: torch.Tensor,
46 - query: torch.Tensor,
47 - key: torch.Tensor,
48 - rotate_style: int = 0,
49 - is_nope_first: bool = False,
50 -) -> None:
51 - pass
52 -
53 -
54 -if is_rocm_triton_rotary_embedding_enabled():
55 - direct_register_custom_op(
56 - op_name="rocm_aiter_rotary_emb_with_key_forward_triton",
57 - op_func=rocm_aiter_rotary_emb_with_key_forward_triton_impl,
58 - mutates_args=["key", "query"],
59 - fake_impl=rocm_aiter_rotary_emb_with_key_forward_triton_fake,
60 - dispatch_key=current_platform.dispatch_key,
61 - )
62 -
63 -
64 -def rocm_aiter_rotary_emb(
65 - positions: torch.Tensor,
66 - query: torch.Tensor,
67 - key: torch.Tensor,
68 - cos_sin_cache: torch.Tensor,
69 - head_size: int,
70 - rotary_dim: int,
71 - is_neox_style: bool,
72 -):
73 - num_tokens = positions.numel()
74 - cos, sin = cos_sin_cache.chunk(2, dim=-1)
75 - query_shape = query.shape
76 - key_shape = key.shape
77 - rotate_style = 0 if is_neox_style else 1
78 -
79 - query = query.view(num_tokens, -1, head_size)
80 - key = key.view(num_tokens, -1, head_size)
81 - query_ = query[..., :rotary_dim]
82 - key_ = key[..., :rotary_dim]
83 - positions = positions.view(*query.shape[:1])
84 - torch.ops.vllm.rocm_aiter_rotary_emb_with_key_forward_triton(
85 - positions,
86 - sin,
87 - cos,
88 - query_,
89 - key_,
90 - rotate_style,
91 - False,
92 - )
93 - query = query.view(query_shape)
94 - key = key.view(key_shape)
vllm/model_executor/models/deepseek_v2.py
+12 -15
@@ -33,6 +33,7 @@ import torch
33 from torch import nn
34 from transformers import DeepseekV2Config, DeepseekV3Config
35
36 +from vllm._aiter_ops import rocm_aiter_ops
37 from vllm.attention import Attention
38 from vllm.attention.backends.abstract import AttentionBackend
39 from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
@@ -50,10 +51,6 @@ from vllm.logger import init_logger
51 from vllm.model_executor.layers.activation import SiluAndMul
52 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
53 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
53 -from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
54 - is_rocm_aiter_fusion_shared_expert_enabled,
55 - is_rocm_aiter_moe_enabled,
56 -)
54 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
55 from vllm.model_executor.layers.linear import (
56 ColumnParallelLinear,
@@ -294,10 +291,8 @@ class DeepseekV2MoE(nn.Module):
291 self.physical_expert_start + self.n_local_physical_experts
292 )
293
297 - if (
298 - config.n_shared_experts is None
299 - or is_rocm_aiter_fusion_shared_expert_enabled()
300 - ):
294 + self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
295 + if config.n_shared_experts is None or self.is_rocm_aiter_moe_enabled:
296 self.shared_experts = None
297 else:
298 intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -330,14 +325,14 @@ class DeepseekV2MoE(nn.Module):
325 # we do scaling outside, set factor to 1.0 to avoid double mul
326 # aiter applies routed_scaling_factor internally
327 routed_scaling_factor=1.0
333 - if not is_rocm_aiter_moe_enabled()
328 + if not self.is_rocm_aiter_moe_enabled
329 else self.routed_scaling_factor,
330 e_score_correction_bias=self.gate.e_score_correction_bias,
331 enable_eplb=self.enable_eplb,
332 num_redundant_experts=self.n_redundant_experts,
333 is_sequence_parallel=self.is_sequence_parallel,
334 n_shared_experts=config.n_shared_experts
340 - if is_rocm_aiter_fusion_shared_expert_enabled()
335 + if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
336 else None,
337 )
338
@@ -371,7 +366,7 @@ class DeepseekV2MoE(nn.Module):
366 # Fix FP16 overflow
367 # See DeepseekV2DecoderLayer for more details.
368 if hidden_states.dtype != torch.float16:
374 - if not is_rocm_aiter_moe_enabled():
369 + if not self.is_rocm_aiter_moe_enabled:
370 final_hidden_states *= self.routed_scaling_factor
371 elif self.shared_experts is not None:
372 assert shared_output is not None
@@ -1428,6 +1423,9 @@ class DeepseekV2ForCausalLM(
1423 )
1424
1425 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1426 + rocm_aiter_moe_shared_expert_enabled = (
1427 + rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
1428 + )
1429 stacked_params_mapping = [
1430 # (param_name, shard_name, shard_id)
1431 ("gate_up_proj", "gate_proj", 0),
@@ -1456,7 +1454,7 @@ class DeepseekV2ForCausalLM(
1454 num_experts=self.config.n_routed_experts
1455 + (
1456 self.config.n_shared_experts
1459 - if is_rocm_aiter_fusion_shared_expert_enabled()
1457 + if rocm_aiter_moe_shared_expert_enabled
1458 else 0
1459 ),
1460 num_redundant_experts=self.num_redundant_experts,
@@ -1472,9 +1470,8 @@ class DeepseekV2ForCausalLM(
1470 if spec_layer is not None:
1471 continue # skip spec decode layers for main model
1472
1475 - is_fuse_shared_experts_layer = (
1476 - is_rocm_aiter_fusion_shared_expert_enabled()
1477 - and ("mlp.shared_experts" in name)
1473 + is_fuse_shared_experts_layer = rocm_aiter_moe_shared_expert_enabled and (
1474 + "mlp.shared_experts" in name
1475 )
1476
1477 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/platforms/rocm.py
+18 -9
@@ -142,6 +142,8 @@ def use_rocm_custom_paged_attention(
142 alibi_slopes: torch.Tensor | None = None,
143 sinks: torch.Tensor | None = None,
144 ) -> bool:
145 + from vllm._aiter_ops import rocm_aiter_ops
146 +
147 GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
148 ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
149 ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
@@ -157,7 +159,7 @@ def use_rocm_custom_paged_attention(
159 and (gqa_ratio >= 1 and gqa_ratio <= 16)
160 and max_seq_len <= 128 * 1024
161 and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
160 - and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN and envs.VLLM_ROCM_USE_AITER)
162 + and not (rocm_aiter_ops.is_pa_attn_enabled())
163 and sinks is None
164 )
165
@@ -202,12 +204,15 @@ class RocmPlatform(Platform):
204 ]
205
206 @classmethod
205 - def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":
207 + def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> _Backend:
208 from importlib.util import find_spec
209
210 + from vllm._aiter_ops import rocm_aiter_ops
211 from vllm.attention.backends.registry import _Backend
212
210 - if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
213 + if rocm_aiter_ops.is_mha_enabled():
214 + # Note: AITER FA is only supported for Qwen-VL models.
215 + # TODO: Add support for other VL models in their model class.
216 return _Backend.ROCM_AITER_FA
217
218 if on_gfx9() and find_spec("flash_attn") is not None:
@@ -228,19 +233,23 @@ class RocmPlatform(Platform):
233 has_sink,
234 use_sparse,
235 ) -> str:
236 + from vllm._aiter_ops import rocm_aiter_ops
237 from vllm.attention.backends.registry import _Backend
238
239 if use_sparse:
240 raise NotImplementedError("Sparse Attention is not supported on ROCm.")
235 - if use_mla:
236 - from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
237 - is_aiter_mla_enabled,
241 +
242 + if not use_v1:
243 + raise RuntimeError(
244 + "V0 attention backends have been removed. Set VLLM_USE_V1=1 "
245 + "to select a supported backend."
246 )
247
248 + if use_mla:
249 if selected_backend is None:
250 selected_backend = (
251 _Backend.ROCM_AITER_MLA
243 - if is_aiter_mla_enabled() or block_size == 1
252 + if rocm_aiter_ops.is_mla_enabled() or block_size == 1
253 else _Backend.TRITON_MLA
254 )
255
@@ -265,12 +274,12 @@ class RocmPlatform(Platform):
274 logger.info("Using FlexAttention backend.")
275 return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
276 if (
268 - envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()
277 + rocm_aiter_ops.is_mha_enabled()
278 ) or selected_backend == _Backend.ROCM_AITER_FA:
279 logger.info("Using Aiter Flash Attention backend.")
280 return "vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
281 if (
273 - envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
282 + rocm_aiter_ops.is_triton_unified_attn_enabled()
283 ) or selected_backend == _Backend.ROCM_AITER_UNIFIED_ATTN:
284 logger.info("Using Aiter Unified Attention backend.")
285 return (
vllm/v1/attention/backends/mla/common.py
+21 -34
@@ -198,6 +198,7 @@ from tqdm import tqdm
198
199 import vllm.envs as envs
200 from vllm import _custom_ops as ops
201 +from vllm._aiter_ops import rocm_aiter_ops
202 from vllm.attention.backends.abstract import (
203 AttentionBackend,
204 AttentionLayer,
@@ -270,28 +271,15 @@ except ImportError:
271 flashinfer_available = False
272
273
273 -def is_rocm_aiter_fp8bmm_enabled() -> bool:
274 - return (
275 - current_platform.is_rocm()
276 - and envs.VLLM_ROCM_USE_AITER_FP8BMM
277 - and envs.VLLM_ROCM_USE_AITER
278 - )
279 -
280 -
281 -if is_rocm_aiter_fp8bmm_enabled():
282 - from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( # noqa: E501
283 - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm, # noqa: E501
284 - )
285 -
286 - def dynamic_per_batched_tensor_quant(
287 - x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn
288 - ):
289 - DTYPE_MAX = torch.finfo(dtype).max
290 - min_val, max_val = x.aminmax()
291 - amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10)
292 - scale = DTYPE_MAX / amax
293 - x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX)
294 - return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
274 +def dynamic_per_batched_tensor_quant(
275 + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn
276 +):
277 + DTYPE_MAX = torch.finfo(dtype).max
278 + min_val, max_val = x.aminmax()
279 + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10)
280 + scale = DTYPE_MAX / amax
281 + x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX)
282 + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
283
284
285 logger = init_logger(__name__)
@@ -1109,6 +1097,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1097 self.kv_b_proj = kv_b_proj
1098 self.indexer = indexer
1099 self.q_pad_num_heads = q_pad_num_heads
1100 + self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
1101
1102 def process_weights_after_loading(self, act_dtype: torch.dtype):
1103 def get_layer_weight(layer):
@@ -1158,7 +1147,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1147 [self.qk_nope_head_dim, self.v_head_dim], dim=-1
1148 )
1149
1161 - if is_rocm_aiter_fp8bmm_enabled():
1150 + if self.is_aiter_triton_fp8_bmm_enabled:
1151 W_K = W_UK.transpose(0, 1) # 16 512 128
1152 W_V = W_UV.permute(1, 2, 0) # 16 128 512
1153 self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
@@ -1187,7 +1176,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1176 dtype=torch.bfloat16,
1177 device=self.W_K.device,
1178 )
1190 - aiter_triton_fp8_bmm(
1179 + rocm_aiter_ops.triton_fp8_bmm(
1180 x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
1181 )
1182
@@ -1196,7 +1185,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1185 dtype=torch.bfloat16,
1186 device=self.W_V.device,
1187 )
1199 - aiter_triton_fp8_bmm(
1188 + rocm_aiter_ops.triton_fp8_bmm(
1189 x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
1190 )
1191 else:
@@ -1208,10 +1197,9 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1197 def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1198 # Convert from (B, N, L) to (N, B, L)
1199 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1211 -
1212 - if is_rocm_aiter_fp8bmm_enabled():
1200 + if self.is_aiter_triton_fp8_bmm_enabled:
1201 # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1214 - x = aiter_triton_fp8_bmm(
1202 + x = rocm_aiter_ops.triton_fp8_bmm(
1203 x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
1204 )
1205 # Convert from (B, N, V) to (B, N * V)
@@ -1571,7 +1559,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1559 [self.qk_nope_head_dim, self.v_head_dim], dim=-1
1560 )
1561
1574 - if is_rocm_aiter_fp8bmm_enabled():
1562 + if self.is_aiter_triton_fp8_bmm_enabled:
1563 W_K = W_UK.transpose(0, 1) # 16 512 128
1564 W_V = W_UV.permute(1, 2, 0) # 16 128 512
1565 self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
@@ -1600,7 +1588,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1588 dtype=torch.bfloat16,
1589 device=self.W_K.device,
1590 )
1603 - aiter_triton_fp8_bmm(
1591 + rocm_aiter_ops.triton_fp8_bmm(
1592 x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
1593 )
1594
@@ -1609,7 +1597,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1597 dtype=torch.bfloat16,
1598 device=self.W_V.device,
1599 )
1612 - aiter_triton_fp8_bmm(
1600 + rocm_aiter_ops.triton_fp8_bmm(
1601 x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
1602 )
1603 else:
@@ -1958,7 +1946,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1946 # Convert from (B, N, P) to (N, B, P)
1947 decode_q_nope = decode_q_nope.transpose(0, 1)
1948
1961 - # Pads the head_dim if necessary (for the underlying kernel)
1949 if self.q_pad_num_heads is not None:
1950 B, N, L = decode_q_pe.shape
1951 decode_pe_padded = decode_q_pe.new_empty((B, self.q_pad_num_heads, L))
@@ -1966,9 +1953,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1953 decode_pe_padded.copy_(decode_q_pe)
1954 decode_q_pe = decode_pe_padded
1955
1969 - if is_rocm_aiter_fp8bmm_enabled():
1956 + if self.is_aiter_triton_fp8_bmm_enabled:
1957 # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
1971 - decode_ql_nope = aiter_triton_fp8_bmm(
1958 + decode_ql_nope = rocm_aiter_ops.triton_fp8_bmm(
1959 decode_q_nope,
1960 self.W_K,
1961 self.W_K_scale,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+2 -7
@@ -6,9 +6,8 @@ from typing import ClassVar
6
7 import torch
8
9 -import vllm.envs as envs
9 +from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.attention.backends.abstract import AttentionLayer
11 -from vllm.attention.ops.rocm_aiter_mla import aiter_mla_decode_fwd
11 from vllm.config import VllmConfig
12 from vllm.utils.math_utils import cdiv
13 from vllm.v1.attention.backends.mla.common import (
@@ -22,10 +21,6 @@ from vllm.v1.attention.backends.utils import AttentionCGSupport
21 from vllm.v1.kv_cache_interface import AttentionSpec
22
23
25 -def is_aiter_mla_enabled() -> bool:
26 - return envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MLA
27 -
28 -
24 class AiterMLABackend(MLACommonBackend):
25 @staticmethod
26 def get_name() -> str:
@@ -284,7 +279,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
279 # max_seqlen_qo must be 1 except for MTP
280 # TODO: Find the best value for MTP
281 max_seqlen_qo = 1
287 - aiter_mla_decode_fwd(
282 + rocm_aiter_ops.mla_decode_fwd(
283 q,
284 kv_buffer,
285 o,