wehub/vllm-project--vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
7ce4c8e27e7e5d9be714f2f0899a2b62a63ee155
vllm-project--vllm/vllm/_aiter_ops.py
T
wehub-resource-sync 7ce4c8e27e
pre-commit / pre-run-check (push) Has been cancelled
Details
pre-commit / pre-commit (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 12:55:37 +08:00

3033 lines
96 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import functools
from collections.abc import Callable
import torch
from torch._ops import OpOverload
import vllm.envs as envs
from vllm.platforms import current_platform
from vllm.utils.import_utils import PlaceholderModule
from vllm.utils.torch_utils import direct_register_custom_op
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
rocm_aiter_sparse_attn_indexer,
rocm_aiter_sparse_attn_indexer_fake,
)
try:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
# fp8_dtype is not cached.
# on ROCm the fp8_dtype always calls is_fp8_fnuz
# which is a host op, so we cache it once here.
FP8_DTYPE = current_platform.fp8_dtype()
_HIPB_MM_INITIALIZED_DEVICES: set[int] = set()
def _ensure_hipb_mm_extension_initialized() -> None:
import aiter
device = torch.accelerator.current_device_index()
if device not in _HIPB_MM_INITIALIZED_DEVICES:
aiter.hipb_create_extension()
_HIPB_MM_INITIALIZED_DEVICES.add(device)
def is_aiter_found() -> bool:
from importlib.util import find_spec
return find_spec("aiter") is not None
# `find_spec` is not torch.compile compatible.
# In cases where aiter availability might have
# been checked in forward passes that are torch compiled.
# we keep this global outside to not cause torch compile breaks.
IS_AITER_FOUND = is_aiter_found()
def is_aiter_found_and_supported() -> bool:
"""Check if AITER library is available and platform supports it.
Checks: platform (ROCm), device arch (gfx9), and library existence.
Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled().
This function determines if aiter CAN be used, not if it SHOULD be used.
Separation of concerns:
- This function: Can aiter work on this system? (platform + library availability)
- rocm_aiter_ops.is_enabled(): Should aiter be used by default? (adds env var check)
- Backend selection: Can explicitly request aiter regardless of env var
This allows explicit backend selection via attention_config to work even when
VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.
"""
if current_platform.is_rocm() and IS_AITER_FOUND:
from vllm.platforms.rocm import on_mi3xx
return on_mi3xx()
return False
@functools.cache
def _load_gemm_tuned_configs(
q_dtype_w: torch.dtype, csv_path: str
) -> set[tuple[int, int, int]]:
try:
df = pd.read_csv(csv_path).drop_duplicates()
df = df[df["q_dtype_w"] == str(q_dtype_w)]
return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int)))
except Exception:
return set()
def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:
configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)
l_m = (
[1, 2, 4]
+ list(range(8, 513, 8))
+ [1024, 1536]
+ [2**i for i in range(11, 19)]
)
return any((N, K, M) in configs for M in l_m)
def if_aiter_supported(func: Callable) -> Callable:
"""Decorator that only executes the function if
ROCm AITER package is supported and enabled on gfx9 archs.
"""
@functools.wraps(func)
def wrapper(*args, **kwargs):
if is_aiter_found_and_supported():
return func(*args, **kwargs)
return None
return wrapper
def _rocm_aiter_fused_moe_impl(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
expert_mask: torch.Tensor | None = None,
activation_method: int = 0,
quant_method: int = 0,
doweight_stage1: bool = False,
w1_scale: torch.Tensor | None = None,
w2_scale: torch.Tensor | None = None,
a1_scale: torch.Tensor | None = None,
a2_scale: torch.Tensor | None = None,
num_local_tokens: torch.Tensor | None = None,
output_dtype: torch.dtype | None = None,
hidden_pad: int = 0,
intermediate_pad: int = 0,
gate_mode: str = "",
bias1: torch.Tensor | None = None,
bias2: torch.Tensor | None = None,
moe_sorting_dispatch_policy: int = 0,
swiglu_limit: float = 0.0,
) -> torch.Tensor:
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe
activation = ActivationType(activation_method)
quant_type = QuantType(quant_method)
extra_kwargs: dict = {}
if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode():
extra_kwargs["gate_mode"] = gate_mode
return fused_moe(
hidden_states,
w1,
w2,
topk_weight,
topk_ids,
expert_mask,
activation,
quant_type,
doweight_stage1,
w1_scale,
w2_scale,
a1_scale,
a2_scale,
num_local_tokens=num_local_tokens,
dtype=output_dtype,
hidden_pad=hidden_pad,
intermediate_pad=intermediate_pad,
bias1=bias1,
bias2=bias2,
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
swiglu_limit=swiglu_limit,
**extra_kwargs,
)
def _rocm_aiter_fused_moe_fake(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
expert_mask: torch.Tensor | None = None,
activation_method: int = 0,
quant_method: int = 0,
doweight_stage1: bool = False,
w1_scale: torch.Tensor | None = None,
w2_scale: torch.Tensor | None = None,
a1_scale: torch.Tensor | None = None,
a2_scale: torch.Tensor | None = None,
num_local_tokens: torch.Tensor | None = None,
output_dtype: torch.dtype | None = None,
hidden_pad: int = 0,
intermediate_pad: int = 0,
gate_mode: str = "",
bias1: torch.Tensor | None = None,
bias2: torch.Tensor | None = None,
moe_sorting_dispatch_policy: int = 0,
swiglu_limit: float = 0.0,
) -> torch.Tensor:
if output_dtype is not None:
return torch.empty_like(hidden_states, dtype=output_dtype)
return torch.empty_like(hidden_states)
def _rocm_aiter_asm_moe_tkw1_impl(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
fc1_scale: torch.Tensor | None = None,
fc2_scale: torch.Tensor | None = None,
fc1_smooth_scale: torch.Tensor | None = None,
fc2_smooth_scale: torch.Tensor | None = None,
a16: bool = False,
per_tensor_quant_scale: torch.Tensor | None = None,
expert_mask: torch.Tensor | None = None,
activation_method: int = 0,
) -> torch.Tensor:
from aiter import ActivationType
from aiter.fused_moe_bf16_asm import asm_moe_tkw1
activation = ActivationType(activation_method)
return asm_moe_tkw1(
hidden_states,
w1,
w2,
topk_weights,
topk_ids,
fc1_scale=fc1_scale,
fc2_scale=fc2_scale,
fc1_smooth_scale=fc1_smooth_scale,
fc2_smooth_scale=fc2_smooth_scale,
a16=a16,
per_tensor_quant_scale=per_tensor_quant_scale,
expert_mask=expert_mask,
activation=activation,
)
def _rocm_aiter_asm_moe_tkw1_fake(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
fc1_scale: torch.Tensor | None = None,
fc2_scale: torch.Tensor | None = None,
fc1_smooth_scale: torch.Tensor | None = None,
fc2_smooth_scale: torch.Tensor | None = None,
a16: bool = False,
per_tensor_quant_scale: torch.Tensor | None = None,
expert_mask: torch.Tensor | None = None,
activation_method: int = 0,
) -> torch.Tensor:
return torch.empty_like(hidden_states)
def _rocm_aiter_topk_softmax_impl(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool,
num_shared_experts: int = 0,
shared_expert_scoring_func: str = "",
) -> None:
from aiter import topk_softmax
topk_softmax(
topk_weights,
topk_indices,
token_expert_indices,
gating_output,
renormalize,
num_shared_experts,
shared_expert_scoring_func,
)
def _rocm_aiter_topk_softmax_fake(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool,
num_shared_experts: int = 0,
shared_expert_scoring_func: str = "",
) -> None:
pass
def _rocm_aiter_topk_sigmoid_impl(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
gating_output: torch.Tensor,
) -> None:
from aiter import topk_sigmoid
topk_sigmoid(topk_weights, topk_indices, gating_output)
def _rocm_aiter_topk_sigmoid_fake(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
gating_output: torch.Tensor,
) -> None:
pass
def _rocm_aiter_biased_grouped_topk_impl(
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
routed_scaling_factor: float = 1.0, # mul to topk_weights
) -> None:
from aiter import biased_grouped_topk
biased_grouped_topk(
gating_output,
correction_bias,
topk_weights,
topk_ids,
num_expert_group,
topk_group,
need_renorm,
routed_scaling_factor,
)
def _rocm_aiter_biased_grouped_topk_fake(
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
routed_scaling_factor: float = 1.0, # mul to topk_weights
) -> None:
pass
def _rocm_aiter_grouped_topk_impl(
gating_output: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
scoring_func: str = "softmax",
routed_scaling_factor: float = 1.0, # mul to topk_weights
) -> None:
is_softmax = scoring_func == "softmax"
from aiter import grouped_topk
grouped_topk(
gating_output,
topk_weights,
topk_ids,
num_expert_group,
topk_group,
need_renorm,
is_softmax,
routed_scaling_factor,
)
def _rocm_aiter_grouped_topk_fake(
gating_output: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
scoring_func: str = "softmax",
routed_scaling_factor: float = 1.0, # mul to topk_weights
) -> None:
pass
def _rocm_aiter_fused_topk_impl(
x: torch.Tensor,
router_logits: torch.Tensor,
top_k: int,
gate_up: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.fused_moe import fused_topk
# fused_topk returns (topk_weights, topk_indices)
return fused_topk(x, router_logits, top_k, gate_up)
def _rocm_aiter_fused_topk_fake(
x: torch.Tensor,
router_logits: torch.Tensor,
top_k: int,
gate_up: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
num_tokens = x.shape[0]
topk_weights = torch.empty(
(num_tokens, top_k), dtype=torch.float32, device=x.device
)
topk_indices = torch.empty((num_tokens, top_k), dtype=torch.int32, device=x.device)
return topk_weights, topk_indices
# Cache whether aiter supports FP8 MLA parameters
_AITER_MLA_SUPPORTS_FP8: bool | None = None
_AITER_HAS_FUSED_QK_RMSNORM: bool | None = None
def check_aiter_fused_qk_rmsnorm() -> bool:
"""Check if aiter provides fused_qk_rmsnorm.
Supports both the new private name ``_fused_qk_rmsnorm``
(AITER >= PR #2958) and the old public name ``fused_qk_rmsnorm``
(AITER >= PR #2442).
TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
"""
global _AITER_HAS_FUSED_QK_RMSNORM
if _AITER_HAS_FUSED_QK_RMSNORM is None:
try:
from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
_fused_qk_rmsnorm,
)
_AITER_HAS_FUSED_QK_RMSNORM = True
except (ImportError, ModuleNotFoundError, AttributeError):
try:
from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
fused_qk_rmsnorm,
)
_AITER_HAS_FUSED_QK_RMSNORM = True
except (ImportError, ModuleNotFoundError, AttributeError):
_AITER_HAS_FUSED_QK_RMSNORM = False
return _AITER_HAS_FUSED_QK_RMSNORM
def _check_aiter_mla_fp8_support() -> bool:
"""Check if aiter.mla.mla_decode_fwd supports q_scale and kv_scale parameters."""
global _AITER_MLA_SUPPORTS_FP8
if _AITER_MLA_SUPPORTS_FP8 is None:
try:
import inspect
from aiter.mla import mla_decode_fwd
sig = inspect.signature(mla_decode_fwd)
_AITER_MLA_SUPPORTS_FP8 = (
"q_scale" in sig.parameters and "kv_scale" in sig.parameters
)
except (
ImportError,
ModuleNotFoundError,
AttributeError,
ValueError,
TypeError,
):
# ImportError/ModuleNotFoundError: aiter.mla module not available
# AttributeError: mla_decode_fwd doesn't exist
# ValueError: mla_decode_fwd has no signature (e.g., built-in)
# TypeError: mla_decode_fwd is not a callable
_AITER_MLA_SUPPORTS_FP8 = False
return _AITER_MLA_SUPPORTS_FP8
def _rocm_aiter_mla_decode_fwd_impl(
q: torch.Tensor,
kv_buffer: torch.Tensor,
o: torch.Tensor,
qo_indptr: torch.Tensor,
max_seqlen_qo: int,
kv_indptr: torch.Tensor | None = None,
kv_indices: torch.Tensor | None = None,
kv_last_page_lens: torch.Tensor | None = None,
sm_scale: float = 1.0,
logit_cap: float = 0.0,
q_scale: torch.Tensor | None = None,
kv_scale: torch.Tensor | None = None,
work_meta_data: torch.Tensor | None = None,
work_indptr: torch.Tensor | None = None,
work_info_set: torch.Tensor | None = None,
reduce_indptr: torch.Tensor | None = None,
reduce_final_map: torch.Tensor | None = None,
reduce_partial_map: torch.Tensor | None = None,
) -> None:
from aiter.mla import mla_decode_fwd
kwargs: dict[str, float | torch.Tensor | None] = {
"sm_scale": sm_scale,
"logit_cap": logit_cap,
}
# Only pass q_scale and kv_scale if the aiter library supports them
if _check_aiter_mla_fp8_support():
kwargs["q_scale"] = q_scale
kwargs["kv_scale"] = kv_scale
if work_meta_data is not None:
assert work_indptr is not None, (
"work_indptr must be provided with work_meta_data"
)
assert work_info_set is not None, (
"work_info_set must be provided with work_meta_data"
)
assert reduce_indptr is not None, (
"reduce_indptr must be provided with work_meta_data"
)
assert reduce_final_map is not None, (
"reduce_final_map must be provided with work_meta_data"
)
assert reduce_partial_map is not None, (
"reduce_partial_map must be provided with work_meta_data"
)
kwargs["work_meta_data"] = work_meta_data
kwargs["work_indptr"] = work_indptr
kwargs["work_info_set"] = work_info_set
kwargs["reduce_indptr"] = reduce_indptr
kwargs["reduce_final_map"] = reduce_final_map
kwargs["reduce_partial_map"] = reduce_partial_map
mla_decode_fwd(
q,
kv_buffer.view(-1, 1, 1, q.shape[-1]),
o,
qo_indptr,
kv_indptr,
kv_indices,
kv_last_page_lens,
max_seqlen_qo,
**kwargs,
)
def _rocm_aiter_mla_decode_fwd_fake(
q: torch.Tensor,
kv_buffer: torch.Tensor,
o: torch.Tensor,
qo_indptr: torch.Tensor,
max_seqlen_qo: int,
kv_indptr: torch.Tensor | None = None,
kv_indices: torch.Tensor | None = None,
kv_last_page_lens: torch.Tensor | None = None,
sm_scale: float = 1.0,
logit_cap: float = 0.0,
q_scale: torch.Tensor | None = None,
kv_scale: torch.Tensor | None = None,
work_meta_data: torch.Tensor | None = None,
work_indptr: torch.Tensor | None = None,
work_info_set: torch.Tensor | None = None,
reduce_indptr: torch.Tensor | None = None,
reduce_final_map: torch.Tensor | None = None,
reduce_partial_map: torch.Tensor | None = None,
) -> None:
pass
def _rocm_aiter_w8a8_gemm_impl(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
from aiter import gemm_a8w8_CK
# gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
# a to be [M, K]
# b to be [N, K]
# CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)
def _rocm_aiter_w8a8_gemm_fake(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
m = A.shape[0]
n = B.shape[0]
Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
return Y
def _rocm_aiter_preshuffled_per_token_w8a8_gemm_impl(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
from aiter import gemm_a8w8_bpreshuffle
output = gemm_a8w8_bpreshuffle(A, B, As, Bs, None, output_dtype)
if bias is not None:
output.add_(bias)
return output
def _rocm_aiter_preshuffled_per_token_w8a8_gemm_fake(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
m = A.shape[0]
n = B.shape[0]
return torch.empty(m, n, dtype=output_dtype, device=A.device)
def _rocm_aiter_hipb_mm_fp8_impl(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
from aiter import hipb_mm
_ensure_hipb_mm_extension_initialized()
return hipb_mm(
A,
B,
solution_index=-1,
bias=bias,
out_dtype=output_dtype,
scaleA=As,
scaleB=Bs,
scaleOut=None,
bpreshuffle=True,
)
def _rocm_aiter_hipb_mm_fp8_fake(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
m = A.shape[0]
n = B.shape[1]
return torch.empty(m, n, dtype=output_dtype, device=A.device)
def _rocm_aiter_triton_gemm_a8w8_blockscale_impl(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale
return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
def _rocm_aiter_triton_gemm_a8w8_blockscale_fake(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
m = A.shape[0]
n = B.shape[0]
Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
return Y
def _rocm_aiter_gemm_a8w8_blockscale_impl(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
from aiter import gemm_a8w8_blockscale
return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
def _rocm_aiter_gemm_a8w8_blockscale_fake(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
m = A.shape[0]
n = B.shape[0]
Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
return Y
def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
quant_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
import aiter as rocm_aiter
assert quant_dtype in [torch.int8, FP8_DTYPE]
y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
residual_out = torch.empty_like(x)
rocm_aiter.rmsnorm2d_fwd_with_add_dynamicquant(
out,
x,
residual,
residual_out,
y_scale,
weight,
epsilon,
use_model_sensitive_rmsnorm=0,
)
return out, residual_out, y_scale
def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_fake(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
quant_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
residual_out = torch.empty_like(x)
return out, residual_out, y_scale
def _rocm_aiter_rmsnorm_fused_dynamic_quant_impl(
x: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
quant_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
import aiter as rocm_aiter
assert quant_dtype in [torch.int8, FP8_DTYPE]
y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
rocm_aiter.rmsnorm2d_fwd_with_dynamicquant(
out, x, y_scale, weight, epsilon, use_model_sensitive_rmsnorm=0
)
return out, y_scale
def _rocm_aiter_rmsnorm_fused_dynamic_quant_fake(
x: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
quant_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
return out, y_scale
def _rocm_aiter_fused_allreduce_rmsnorm_impl(
input_: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor]:
aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
assert aiter_ar is not None, "aiter allreduce must be initialized"
ca = aiter_ar.aiter_ca
total_bytes = input_.numel() * input_.element_size()
hidden_dim = input_.shape[-1]
token_num = input_.shape[0]
if input_.dtype in (torch.bfloat16, torch.float16):
pack_size = 16 // input_.element_size()
hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
else:
hidden_ok = False
token_ok = token_num <= 80
world_size = ca.world_size
full_nvlink = ca.fully_connected
if world_size == 2:
size_ok = True
elif full_nvlink and world_size <= 4:
size_ok = total_bytes < 256 * 1024
elif full_nvlink and world_size <= 8:
size_ok = total_bytes < 128 * 1024
else:
size_ok = False
use_1stage = hidden_ok and token_ok and size_ok
result = ca.custom_fused_ar_rms(
input_,
residual,
weight,
epsilon,
use_1stage=use_1stage,
)
assert result is not None
return result[0], result[1]
def _rocm_aiter_fused_allreduce_rmsnorm_fake(
input_: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.empty_like(input_), torch.empty_like(residual)
def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
input_: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Fused AllReduce + RMSNorm + per-group FP8 quant.
Mirrors the eligibility logic of ``_rocm_aiter_fused_allreduce_rmsnorm_impl``
for the 1-stage vs 2-stage AITER kernel dispatch (both variants run inside
AITER, the only choice we make here is the launcher to call into).
"""
aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
assert aiter_ar is not None, "aiter allreduce must be initialized"
ca = aiter_ar.aiter_ca
total_bytes = input_.numel() * input_.element_size()
hidden_dim = input_.shape[-1]
token_num = input_.shape[0]
if input_.dtype in (torch.bfloat16, torch.float16):
pack_size = 16 // input_.element_size()
hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
else:
hidden_ok = False
token_ok = token_num <= 80
world_size = ca.world_size
full_nvlink = ca.fully_connected
if world_size == 2:
size_ok = True
elif full_nvlink and world_size <= 4:
size_ok = total_bytes < 256 * 1024
elif full_nvlink and world_size <= 8:
size_ok = total_bytes < 128 * 1024
else:
size_ok = False
use_1stage = hidden_ok and token_ok and size_ok
result = ca.fused_ar_rms_per_group_quant(
input_,
residual,
w=weight,
eps=epsilon,
group_size=group_size,
registered=torch.cuda.is_current_stream_capturing(),
use_1stage=use_1stage,
)
assert result is not None
return result[0], result[1], result[2]
def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_fake(
input_: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
hidden_dim = input_.shape[-1]
num_groups = hidden_dim // group_size
quant_out = torch.empty(input_.shape, dtype=FP8_DTYPE, device=input_.device)
residual_out = torch.empty_like(residual)
scale_out = torch.empty(
input_.shape[:-1] + (num_groups,),
dtype=torch.float32,
device=input_.device,
)
return quant_out, residual_out, scale_out
def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
input_: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Fused AllReduce + add-RMSNorm + per-group FP8 quant + bf16 normed act.
Wraps the same AITER launcher as ``_rocm_aiter_fused_allreduce_rmsnorm_
quant_per_group_impl`` with ``emit_bf16=True``, emitting the pre-quant
bf16/fp16 normed activation for a parallel consumer (DeepSeek V3.2 sparse
indexer ``wk_weights_proj``).
"""
aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
assert aiter_ar is not None, "aiter allreduce must be initialized"
ca = aiter_ar.aiter_ca
total_bytes = input_.numel() * input_.element_size()
hidden_dim = input_.shape[-1]
token_num = input_.shape[0]
if input_.dtype in (torch.bfloat16, torch.float16):
pack_size = 16 // input_.element_size()
hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
else:
hidden_ok = False
token_ok = token_num <= 80
world_size = ca.world_size
full_nvlink = ca.fully_connected
if world_size == 2:
size_ok = True
elif full_nvlink and world_size <= 4:
size_ok = total_bytes < 256 * 1024
elif full_nvlink and world_size <= 8:
size_ok = total_bytes < 128 * 1024
else:
size_ok = False
use_1stage = hidden_ok and token_ok and size_ok
result = ca.fused_ar_rms_per_group_quant(
input_,
residual,
w=weight,
eps=epsilon,
group_size=group_size,
registered=torch.cuda.is_current_stream_capturing(),
use_1stage=use_1stage,
emit_bf16=True,
)
assert result is not None
assert len(result) == 4, "emit_bf16=True must return four tensors from aiter"
return result[0], result[1], result[2], result[3]
def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake(
input_: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
hidden_dim = input_.shape[-1]
num_groups = hidden_dim // group_size
quant_out = torch.empty(input_.shape, dtype=FP8_DTYPE, device=input_.device)
residual_out = torch.empty_like(residual)
scale_out = torch.empty(
input_.shape[:-1] + (num_groups,),
dtype=torch.float32,
device=input_.device,
)
bf16_norm_out = torch.empty_like(input_)
return quant_out, residual_out, scale_out, bf16_norm_out
def _rocm_aiter_per_tensor_quant_impl(
out: torch.Tensor,
x: torch.Tensor,
scale: torch.Tensor,
is_dynamic: bool,
) -> None:
from aiter.ops.quant import dynamic_per_tensor_quant, static_per_tensor_quant
if is_dynamic:
dynamic_per_tensor_quant(out, x, scale)
else:
static_per_tensor_quant(out, x, scale)
def _rocm_aiter_per_tensor_quant_fake(
out: torch.Tensor,
x: torch.Tensor,
scale: torch.Tensor,
is_dynamic: bool,
) -> None:
pass
def _rocm_aiter_per_token_quant_impl(
x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.ops.quant import dynamic_per_token_scaled_quant
assert quant_dtype in [torch.int8, FP8_DTYPE]
out_shape = x.shape
out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
if scale is None:
scale = torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device)
dynamic_per_token_scaled_quant(
out,
x,
scale,
scale_ub=None,
shuffle_scale=False,
num_rows=None,
num_rows_factor=1,
)
return out, scale
def _rocm_aiter_per_token_quant_fake(
x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
out_shape = x.shape
return (
torch.empty(x.shape, dtype=quant_dtype, device=x.device),
torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device),
)
def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
(x_quant, x_quant_scales), _, _, res = fused_rms_fp8_group_quant(
x,
weight,
variance_epsilon,
None,
None,
None,
group_size=group_size,
dtype_quant=FP8_DTYPE,
res1=residual,
)
return (
x_quant,
res,
x_quant_scales,
)
def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
M, N = x.shape
scale_shape = (M, (N + group_size - 1) // group_size)
return (
torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
torch.empty_like(residual, device=residual.device),
torch.empty(scale_shape, dtype=torch.float32, device=x.device),
)
def _rocm_aiter_rmsnorm_fp8_group_quant_impl(
x: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
(x_quant, x_quant_scales), _, _, res = fused_rms_fp8_group_quant(
x,
weight,
variance_epsilon,
None,
None,
None,
group_size=group_size,
dtype_quant=FP8_DTYPE,
res1=None,
)
return (x_quant, x_quant_scales)
def _rocm_aiter_rmsnorm_fp8_group_quant_fake(
x: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
M, N = x.shape
scale_shape = (M, (N + group_size - 1) // group_size)
return (
torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
torch.empty(scale_shape, dtype=torch.float32, device=x.device),
)
def _rocm_aiter_fused_rms_gated_fp8_group_quant_impl(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
z: torch.Tensor,
eps: float,
norm_before_gate: bool,
activation: str,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Fused gated-RMSNorm + FP8 group quantization via aiter Triton kernel."""
from aiter.ops.triton.quant import fused_rms_gated_fp8_group_quant
return fused_rms_gated_fp8_group_quant(
x,
weight,
bias,
z,
eps,
norm_before_gate=norm_before_gate,
activation=activation,
out_dtype=FP8_DTYPE,
group_size=group_size,
)
def _rocm_aiter_fused_rms_gated_fp8_group_quant_fake(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor | None,
z: torch.Tensor,
eps: float,
norm_before_gate: bool,
activation: str,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
M, N = x.shape
scale_shape = (M, (N + group_size - 1) // group_size)
return (
torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
torch.empty(scale_shape, dtype=torch.float32, device=x.device),
)
def _rocm_aiter_group_fp8_quant_impl(
x: torch.Tensor,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
assert x.shape[-1] % group_size == 0, "Input shape must be divisible by group size"
from aiter import QuantType, get_hip_quant
aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)
return aiter_per1x128_quant(x.contiguous(), quant_dtype=FP8_DTYPE)
def _rocm_aiter_group_fp8_quant_fake(
x: torch.Tensor,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
M, N = x.shape
x_fp8 = torch.empty((M, N), dtype=FP8_DTYPE, device=x.device)
out_bs = torch.empty(
(
M,
(N + group_size - 1) // group_size,
),
dtype=torch.float32,
device=x.device,
)
return x_fp8, out_bs
def _rocm_aiter_act_mul_and_fp8_group_quant_impl(
x: torch.Tensor,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.ops.triton.activation import act_mul_and_fp8_group_quant
return act_mul_and_fp8_group_quant(
x,
activation="silu",
group_size=group_size,
dtype_quant=FP8_DTYPE,
)
def _rocm_aiter_act_mul_and_fp8_group_quant_fake(
x: torch.Tensor,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
M, N = x.shape
assert N % 2 == 0
N_half = N // 2
x_fp8 = torch.empty((M, N_half), dtype=FP8_DTYPE, device=x.device)
out_bs = torch.empty(
(
M,
(N_half + group_size - 1) // group_size,
),
dtype=torch.float32,
device=x.device,
)
return x_fp8, out_bs
def _rocm_aiter_triton_add_rmsnorm_pad_impl(
x: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float,
residual: torch.Tensor,
x_pad_to_multiple: int,
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.ops.triton.fused_add_rmsnorm_pad import fused_add_rmsnorm_pad
return fused_add_rmsnorm_pad(
x,
weight,
variance_epsilon,
residual,
x_pad_to_multiple=x_pad_to_multiple,
)
def _rocm_aiter_triton_add_rmsnorm_pad_fake(
x: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float,
residual: torch.Tensor,
x_pad_to_multiple: int,
) -> tuple[torch.Tensor, torch.Tensor]:
M, N = x.shape
if x_pad_to_multiple > 0:
N_out = (N + x_pad_to_multiple - 1) // x_pad_to_multiple * x_pad_to_multiple
else:
N_out = N
out = torch.empty((M, N_out), dtype=x.dtype, device=x.device)
residual_out = torch.empty_like(residual)
return out, residual_out
def _fused_mla_dual_rms_norm_impl(
x1: torch.Tensor,
x1_weight: torch.Tensor,
x2: torch.Tensor,
x2_weight: torch.Tensor,
x1_epsilon: float,
x2_epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor]:
try:
import aiter.ops.fused_qk_norm_rope_cache_quant as aiter_ops
except (ImportError, ModuleNotFoundError, AttributeError) as exc:
raise ImportError(
"fused_qk_rmsnorm requires AITer >= PR #2442. "
"Please upgrade aiter or disable the "
"fuse_mla_dual_rms_norm pass."
) from exc
if hasattr(aiter_ops, "_fused_qk_rmsnorm"):
return aiter_ops._fused_qk_rmsnorm(
q_out=None,
q=x1,
q_weight=x1_weight,
q_eps=x1_epsilon,
k_out=None,
k=x2,
k_weight=x2_weight,
k_eps=x2_epsilon,
)
# TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
# AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
if hasattr(aiter_ops, "fused_qk_rmsnorm"):
return aiter_ops.fused_qk_rmsnorm(
q=x1,
q_weight=x1_weight,
q_eps=x1_epsilon,
k=x2,
k_weight=x2_weight,
k_eps=x2_epsilon,
)
raise ImportError(
"fused_qk_rmsnorm requires AITer >= PR #2442. "
"Please upgrade aiter or disable the "
"fuse_mla_dual_rms_norm pass."
)
def _fused_mla_dual_rms_norm_fake(
x1: torch.Tensor,
x1_weight: torch.Tensor,
x2: torch.Tensor,
x2_weight: torch.Tensor,
x1_epsilon: float,
x2_epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor]:
return (torch.empty_like(x1), torch.empty_like(x2))
def _fused_mla_dual_rms_norm_per_token_quant_impl(
q: torch.Tensor,
q_weight: torch.Tensor,
kv: torch.Tensor,
kv_weight: torch.Tensor,
q_epsilon: float,
kv_epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Fused MLA q/kv RMSNorm (+ FP8 per-token quant on q) via AITER.
Backs the ``fused_mla_dual_rms_norm_per_token_quant`` custom op used by the
MLA FP8 attention fusion when the q latent is quantized *per token* (a single
``(M, 1)`` scale). Only the *q* latent is FP8 quantized (it feeds the
FP8 ``q_b_proj`` GEMM); the *kv* latent is RMS-normed and consumed by attention as bf16.
"""
from aiter.ops.fused_qk_rmsnorm_group_quant import (
fused_qk_rmsnorm_per_token_quant,
)
mq, nq = q.shape
q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
q_scale = torch.empty((mq, 1), dtype=torch.float32, device=q.device)
kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
# q -> RMSNorm + FP8 per-token quant (q slot); kv -> RMSNorm only (k slot).
# `split` views are accepted directly (unit inner stride); the kernel
# handles strided inputs, matching the aiter op-test usage.
fused_qk_rmsnorm_per_token_quant(
q_out_quantized=q_out,
q_out_scale=q_scale,
q=q,
q_weight=q_weight,
q_epsilon=q_epsilon,
k_out=kv_normed,
k=kv,
k_weight=kv_weight,
k_epsilon=kv_epsilon,
gemma_norm=False,
)
return q_out, q_scale, kv_normed
def _fused_mla_dual_rms_norm_per_token_quant_fake(
q: torch.Tensor,
q_weight: torch.Tensor,
kv: torch.Tensor,
kv_weight: torch.Tensor,
q_epsilon: float,
kv_epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
mq, nq = q.shape
q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
q_scale = torch.empty((mq, 1), dtype=torch.float32, device=q.device)
kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
return q_out, q_scale, kv_normed
def _rocm_aiter_gemm_a8wfp4_impl(
x: torch.Tensor,
w: torch.Tensor,
x_scales: torch.Tensor,
w_scales: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
from aiter.ops.triton.gemm_a8wfp4 import gemm_a8wfp4
M, N = x.shape[0], w.shape[0]
y = torch.empty(M, N, dtype=out_dtype, device=x.device)
gemm_a8wfp4(
x=x,
w=w,
y=y,
x_scales=x_scales,
w_scales=w_scales,
dtype=out_dtype,
config=None,
)
return y
def _rocm_aiter_gemm_a8wfp4_fake(
x: torch.Tensor,
w: torch.Tensor,
x_scales: torch.Tensor,
w_scales: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
return torch.empty(x.shape[0], w.shape[0], dtype=out_dtype, device=x.device)
def _triton_rotary_embedding_impl(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox: bool,
offsets: torch.Tensor | None = None,
) -> None:
# Modifies query and key in-place
from aiter.ops.triton.rope.rope import (
rope_cached_thd_positions_offsets_2c_fwd_inplace,
)
num_tokens = positions.numel()
cos, sin = cos_sin_cache.chunk(2, dim=-1)
query_shape = query.shape
key_shape = key.shape
rotate_style = 0 if is_neox else 1
rotary_dim = head_size
query = query.view(num_tokens, -1, head_size)
key = key.view(num_tokens, -1, head_size)
query_ = query[..., :rotary_dim]
key_ = key[..., :rotary_dim]
positions = positions.view(*query.shape[:1])
rope_cached_thd_positions_offsets_2c_fwd_inplace(
query_,
key_,
cos,
sin,
positions,
offsets,
rotate_style,
reuse_freqs_front_part=True,
nope_first=False,
)
query = query.view(query_shape)
key = key.view(key_shape)
def _triton_rotary_embedding_fake(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox_style: bool,
offsets: torch.Tensor | None = None,
) -> None:
return
# Global flag to ensure ops are registered only once
_OPS_REGISTERED = False
class rocm_aiter_ops:
"""ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.
This class centralizes the import and registration of AITER ops,
and provides a unified interface for checking if AITER is enabled.
Operations are only available on supported gfx9
architectures when aiter is installed.
The class uses environment variables to control which features are enabled,
allowing fine-grained control over which AITER optimizations are used.
Environment Variables:
VLLM_ROCM_USE_AITER: Main toggle for all AITER operations.
VLLM_ROCM_USE_AITER_LINEAR: Controls GEMM and quantization ops.
VLLM_ROCM_USE_AITER_RMSNORM: Controls RMSNorm operations.
VLLM_ROCM_USE_AITER_MOE: Controls MoE (Mixture of Experts) ops.
VLLM_ROCM_USE_AITER_MLA: Controls MLA (Multi-head Latent Attention) ops.
VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen.
VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention.
VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.
Note:
The environment variables are assigned when the module is imported,
so you can't change the environment variables after the module is imported.
This is done out of performance consideration. Accessing environment variables
is expensive as described in issue https://github.com/vllm-project/vllm/issues/17067
so we don't want to do it repeatedly, especially in the hot path (the forward pass).
You can call the refresh_env_variables() function to reload the env variables
after monkey patching the env variables in the unit test.
Check Functions:
All check functions (is_*_enabled) are decorated with @if_aiter_supported,
which verifies: (1) platform is ROCm, (2) device arch is gfx9, and
(3) aiter library is installed. The check function then also verifies
the corresponding environment variable is enabled.
i.e. ___
is_enabled() == current_platform.is_rocm() and | checked by
current_platform.is_on_gfx9() and | @if_aiter_supported
IS_AITER_FOUND and _______________|
cls._AITER_ENABLED -----> Check by the logic in `is_enabled()`
Example:
from vllm._aiter_ops import rocm_aiter_ops
# Check if aiter is enabled before using operations
if rocm_aiter_ops.is_enabled():
result = rocm_aiter_ops.per_token_quant(x, FP8_DTYPE)
Operations:
- GEMM operations: gemm_a8w8, gemm_a8w8_blockscale
- Fused MoE: fused_moe, asm_moe_tkw1
- Routing: topk_softmax, biased_grouped_topk, grouped_topk
- MLA decode: mla_decode_fwd
- Quantization: per_tensor_quant, per_token_quant, group_fp8_quant
- Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
"""
_MOE_DISPATCH_POLICY: int | None = None
@classmethod
@if_aiter_supported
def get_moe_dispatch_policy(cls) -> int:
"""Cached MoE sorting dispatch policy."""
if cls._MOE_DISPATCH_POLICY is None:
import vllm.envs as envs
cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY
return cls._MOE_DISPATCH_POLICY
# Check if the env variable is set
_AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
_CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR
_LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
_FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
_MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
_MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
_SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT
_TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
# TODO: Consolidate under _LINEAR_ENABLED
_FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
_FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
_LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
# TODO: Consolidate under _LINEAR_ENABLED
_FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
# TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
_TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
# TODO: Consolidate under _LINEAR_ENABLED
_TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
# Lazily probed: whether aiter.topk_softmax supports the
# num_shared_experts / shared_expert_scoring_func args (7-arg form).
_TOPK_SOFTMAX_FUSED_SIGMOID: bool | None = None
@classmethod
def refresh_env_variables(cls):
"""
Since the environment variables are assigned when the module is imported,
This is a helper function to reload all the env variables from
the environment variables.
for example, after monkey patching the env variables in the unit test,
you can call this function to reload the env variables.
"""
cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
cls._CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR
cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
cls._SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT
cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM
cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
@staticmethod
def get_aiter_activation_type(activation_str: str):
"""
Given an activation type as a string, returns the corresponding aiter ActivationType enum.
Supported activation types: "no", "none", "silu", "gelu", "swiglu".
Returns None if the mapping fails.
Args:
activation_str (str): Activation type as string.
Returns:
Aiter ActivationType enum value, or None if not found.
"""
# Import only locally, since aiter may not always be available.
try:
from aiter import ActivationType
except ImportError:
return None
if not isinstance(activation_str, str):
return None
name = activation_str.strip().lower()
mapping = {
"none": ActivationType.No,
"no": ActivationType.No,
"silu": ActivationType.Silu,
"gelu": ActivationType.Gelu,
"swiglu": ActivationType.Swiglu,
}
return mapping.get(name)
@staticmethod
def get_aiter_quant_type(quant_type_str: str):
"""
Given a quantization type as a string, returns the corresponding aiter QuantType enum.
Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".
Returns None if the mapping fails.
Args:
quant_type_str (str): Quantization type as string.
Returns:
Aiter QuantType enum value, or None if not found.
"""
try:
from aiter import QuantType
except ImportError:
return None
if not isinstance(quant_type_str, str):
return None
name = quant_type_str.strip().lower()
mapping = {
"no": QuantType.No,
"per_tensor": QuantType.per_Tensor,
"per_token": QuantType.per_Token,
"per_1x32": QuantType.per_1x32,
"per_1x128": QuantType.per_1x128,
"per_128x128": QuantType.per_128x128,
}
return mapping.get(name)
@classmethod
@if_aiter_supported
def is_enabled(cls) -> bool:
return cls._AITER_ENABLED
@classmethod
@if_aiter_supported
def is_linear_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._LINEAR_ENABLED
@classmethod
@if_aiter_supported
def is_linear_fp8_enabled(cls) -> bool:
return cls.is_linear_enabled()
@classmethod
@if_aiter_supported
def is_fused_moe_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._FMOE_ENABLED
@classmethod
@if_aiter_supported
def is_fusion_moe_shared_experts_enabled(cls) -> bool:
return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED
@classmethod
@if_aiter_supported
def topk_softmax_supports_fused_sigmoid(cls) -> bool:
"""Check if topk_softmax supports fused shared expert activation."""
if cls._TOPK_SOFTMAX_FUSED_SIGMOID is None:
try:
import inspect
from aiter import topk_softmax
params = inspect.signature(topk_softmax).parameters
if "num_shared_experts" in params:
cls._TOPK_SOFTMAX_FUSED_SIGMOID = True
else:
# @compile_ops wrapper loses the original signature.
# Fall back to the torch custom op schema.
import torch
schema = getattr(
getattr(torch.ops.aiter, "topk_softmax", None), "default", None
)
schema_str = str(getattr(schema, "_schema", ""))
cls._TOPK_SOFTMAX_FUSED_SIGMOID = "num_shared_experts" in schema_str
except (ImportError, ValueError):
cls._TOPK_SOFTMAX_FUSED_SIGMOID = False
return cls._TOPK_SOFTMAX_FUSED_SIGMOID
@classmethod
@if_aiter_supported
def fuse_sigmoid_in_kernel(cls, aiter_topK_meta_data: object) -> bool:
"""Whether fused shared-expert sigmoid in the topk kernel is usable.
Combines the cached static capability checks (FSE enabled, fused-moe
enabled, topk_softmax supports fused sigmoid) with the runtime
readiness check (topK meta-data buffer initialized).
``aiter_topK_meta_data`` is accepted as a parameter rather than
imported internally so callers cannot hit initialization-order
issues where the module-level global has not been set yet.
"""
return (
cls.is_fusion_moe_shared_experts_enabled()
and cls.topk_softmax_supports_fused_sigmoid()
and aiter_topK_meta_data is not None
)
@classmethod
@if_aiter_supported
def is_mla_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._MLA_ENABLED
@classmethod
@if_aiter_supported
def is_mha_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._MHA_ENABLED
@classmethod
@if_aiter_supported
def is_custom_all_reduce_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._CUSTOM_ALL_REDUCE_ENABLED
@classmethod
@if_aiter_supported
def is_shuffle_kv_cache_enabled(cls) -> bool:
return cls._SHUFFLE_KV_CACHE_ENABLED
@classmethod
@if_aiter_supported
def is_triton_unified_attn_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._TRITON_UNIFIED_ATTN_ENABLED
@classmethod
@if_aiter_supported
def is_fp8bmm_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._FP8BMM_ENABLED
@classmethod
@if_aiter_supported
def is_fp4bmm_enabled(cls) -> bool:
from vllm.platforms.rocm import on_gfx950
return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950()
@classmethod
@if_aiter_supported
def is_linear_hipbmm_enabled(cls) -> bool:
from vllm.platforms.rocm import on_mi3xx
return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED
@classmethod
@if_aiter_supported
def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:
from vllm.platforms.rocm import on_gfx950
return cls._AITER_ENABLED and cls._FP4_GEMM_DYNAMIC_QUANT_ASM and on_gfx950()
@classmethod
@if_aiter_supported
def is_triton_rotary_embed_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED
@classmethod
@if_aiter_supported
def is_triton_gemm_enabled(cls) -> bool:
return cls._AITER_ENABLED and cls._TRITON_UNQUANT_GEMM
@classmethod
@if_aiter_supported
def is_tgemm_enabled(cls) -> bool:
from vllm.platforms.rocm import on_gfx950
return cls.is_linear_enabled() and on_gfx950()
@classmethod
def get_aiter_allreduce(cls):
"""Return the TP device communicator's AITER custom-allreduce if it has
one, return None otherwise
"""
from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
AiterCustomAllreduce,
)
from vllm.distributed.parallel_state import get_tp_group
device_comm = get_tp_group().device_communicator
aiter_ar_comm = getattr(device_comm, "aiter_ar_comm", None)
return (
aiter_ar_comm if isinstance(aiter_ar_comm, AiterCustomAllreduce) else None
)
@classmethod
@if_aiter_supported
def are_gdn_triton_kernels_available(cls) -> bool:
"""Check if AITER Triton kernels for GDN attention are importable.
These are optional Triton kernels (conv1d fast-path, gated delta net)
used by GatedDeltaNetAttention's decode fast-path. They may be absent
in older aiter builds.
"""
if not cls._AITER_ENABLED:
return False
try:
import aiter.ops.triton.causal_conv1d_update_single_token # noqa: F401
import aiter.ops.triton.gated_delta_net # noqa: F401
from aiter.ops.triton.quant import ( # noqa: F401
fused_rms_gated_fp8_group_quant,
)
return True
except (ImportError, ModuleNotFoundError):
return False
@classmethod
@if_aiter_supported
@functools.cache
def fused_moe_supports_gate_mode(cls) -> bool:
"""Probe whether the installed aiter.fused_moe accepts `gate_mode`.
Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14).
Builds with older AITER must omit this argument.
"""
import inspect
from aiter.fused_moe import fused_moe
return "gate_mode" in inspect.signature(fused_moe).parameters
@staticmethod
@if_aiter_supported
def register_ops_once() -> None:
global _OPS_REGISTERED
if not _OPS_REGISTERED:
# register all the custom ops here
direct_register_custom_op(
op_name="rocm_aiter_asm_moe_tkw1",
op_func=_rocm_aiter_asm_moe_tkw1_impl,
mutates_args=[],
fake_impl=_rocm_aiter_asm_moe_tkw1_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_moe",
op_func=_rocm_aiter_fused_moe_impl,
mutates_args=[],
fake_impl=_rocm_aiter_fused_moe_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_topk_softmax",
op_func=_rocm_aiter_topk_softmax_impl,
mutates_args=["topk_weights", "topk_indices", "token_expert_indices"],
fake_impl=_rocm_aiter_topk_softmax_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_topk_sigmoid",
op_func=_rocm_aiter_topk_sigmoid_impl,
mutates_args=["topk_weights", "topk_indices"],
fake_impl=_rocm_aiter_topk_sigmoid_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_biased_grouped_topk",
op_func=_rocm_aiter_biased_grouped_topk_impl,
mutates_args=["topk_weights", "topk_ids"],
fake_impl=_rocm_aiter_biased_grouped_topk_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_grouped_topk",
op_func=_rocm_aiter_grouped_topk_impl,
mutates_args=["topk_weights", "topk_ids"],
fake_impl=_rocm_aiter_grouped_topk_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_topk",
op_func=_rocm_aiter_fused_topk_impl,
mutates_args=[],
fake_impl=_rocm_aiter_fused_topk_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_mla_decode_fwd",
op_func=_rocm_aiter_mla_decode_fwd_impl,
mutates_args=["o"],
fake_impl=_rocm_aiter_mla_decode_fwd_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_w8a8_gemm",
op_func=_rocm_aiter_w8a8_gemm_impl,
fake_impl=_rocm_aiter_w8a8_gemm_fake,
)
direct_register_custom_op(
op_name="_rocm_aiter_preshuffled_per_token_w8a8_gemm",
op_func=_rocm_aiter_preshuffled_per_token_w8a8_gemm_impl,
fake_impl=_rocm_aiter_preshuffled_per_token_w8a8_gemm_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_hipb_mm_fp8",
op_func=_rocm_aiter_hipb_mm_fp8_impl,
fake_impl=_rocm_aiter_hipb_mm_fp8_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_triton_gemm_a8w8_blockscale",
op_func=_rocm_aiter_triton_gemm_a8w8_blockscale_impl,
fake_impl=_rocm_aiter_triton_gemm_a8w8_blockscale_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_gemm_a8w8_blockscale",
op_func=_rocm_aiter_gemm_a8w8_blockscale_impl,
fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_rmsnorm_fused_dynamic_quant",
op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl,
fake_impl=_rocm_aiter_rmsnorm_fused_dynamic_quant_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_rmsnorm_fused_add_dynamic_quant",
op_func=_rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl,
fake_impl=_rocm_aiter_rmsnorm_fused_add_dynamic_quant_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_rmsnorm_fp8_group_quant",
op_func=_rocm_aiter_rmsnorm_fp8_group_quant_impl,
fake_impl=_rocm_aiter_rmsnorm_fp8_group_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_rms_gated_fp8_group_quant",
op_func=_rocm_aiter_fused_rms_gated_fp8_group_quant_impl,
fake_impl=_rocm_aiter_fused_rms_gated_fp8_group_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_rmsnorm_with_add_fp8_group_quant",
op_func=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl,
fake_impl=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_act_mul_and_fp8_group_quant",
op_func=_rocm_aiter_act_mul_and_fp8_group_quant_impl,
fake_impl=_rocm_aiter_act_mul_and_fp8_group_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_triton_add_rmsnorm_pad",
op_func=_rocm_aiter_triton_add_rmsnorm_pad_impl,
fake_impl=_rocm_aiter_triton_add_rmsnorm_pad_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_group_fp8_quant",
op_func=_rocm_aiter_group_fp8_quant_impl,
fake_impl=_rocm_aiter_group_fp8_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_per_tensor_quant",
op_func=_rocm_aiter_per_tensor_quant_impl,
mutates_args=["out", "scale"],
fake_impl=_rocm_aiter_per_tensor_quant_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_per_token_quant",
op_func=_rocm_aiter_per_token_quant_impl,
fake_impl=_rocm_aiter_per_token_quant_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_sparse_attn_indexer",
op_func=rocm_aiter_sparse_attn_indexer,
mutates_args=["topk_indices_buffer"],
fake_impl=rocm_aiter_sparse_attn_indexer_fake,
dispatch_key=current_platform.dispatch_key,
)
direct_register_custom_op(
op_name="rocm_aiter_gemm_a8wfp4",
op_func=_rocm_aiter_gemm_a8wfp4_impl,
mutates_args=[],
fake_impl=_rocm_aiter_gemm_a8wfp4_fake,
dispatch_key=current_platform.dispatch_key,
)
# Register rocm aiter rotary embedding custom op
direct_register_custom_op(
op_name="rocm_aiter_triton_rotary_embedding",
op_func=_triton_rotary_embedding_impl,
mutates_args=["query", "key"], # These tensors are modified in-place
fake_impl=_triton_rotary_embedding_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_allreduce_rmsnorm",
op_func=_rocm_aiter_fused_allreduce_rmsnorm_impl,
fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_allreduce_rmsnorm_quant_per_group",
op_func=(_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl),
fake_impl=(_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_fake),
)
direct_register_custom_op(
op_name="rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm", # noqa: E501
op_func=_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl, # noqa: E501
fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake, # noqa: E501
)
direct_register_custom_op(
op_name="fused_mla_dual_rms_norm",
op_func=_fused_mla_dual_rms_norm_impl,
mutates_args=[],
fake_impl=_fused_mla_dual_rms_norm_fake,
)
direct_register_custom_op(
op_name="fused_mla_dual_rms_norm_per_token_quant",
op_func=_fused_mla_dual_rms_norm_per_token_quant_impl,
mutates_args=[],
fake_impl=_fused_mla_dual_rms_norm_per_token_quant_fake,
)
_OPS_REGISTERED = True
@staticmethod
def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant.default
@staticmethod
def get_rmsnorm_fused_dynamic_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant.default
@staticmethod
def get_rmsnorm_group_fused_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
@staticmethod
def get_fused_rms_gated_fp8_group_quant_op() -> OpOverload:
"""Return the fused gated-RMSNorm + FP8 group quant custom op."""
return torch.ops.vllm.rocm_aiter_fused_rms_gated_fp8_group_quant.default
@staticmethod
def get_rmsnorm_group_add_fused_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant.default
@staticmethod
def get_per_token_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_per_token_quant.default
@staticmethod
def get_group_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_group_fp8_quant.default
@staticmethod
def get_act_mul_fused_fp8_group_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant.default
@staticmethod
def get_triton_add_rmsnorm_pad_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_triton_add_rmsnorm_pad.default
@staticmethod
def get_triton_rotary_embedding_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default
@staticmethod
def get_fused_allreduce_rmsnorm_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm.default
@staticmethod
def get_fused_allreduce_rmsnorm_quant_per_group_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group.default
@staticmethod
def get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() -> OpOverload: # noqa: E501
return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm.default # noqa: E501
@staticmethod
def get_fused_mla_dual_rms_norm_op() -> OpOverload:
return torch.ops.vllm.fused_mla_dual_rms_norm.default
@staticmethod
def get_fused_mla_dual_rms_norm_per_token_quant_op() -> OpOverload:
return torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default
@staticmethod
def w8a8_gemm(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_w8a8_gemm(A, B, As, Bs, bias, output_dtype)
@staticmethod
def preshuffled_per_token_w8a8_gemm(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
return torch.ops.vllm._rocm_aiter_preshuffled_per_token_w8a8_gemm(
A, B, As, Bs, bias, output_dtype
)
@staticmethod
def hipb_mm_fp8(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_hipb_mm_fp8(A, B, As, Bs, bias, output_dtype)
@staticmethod
def triton_gemm_a8w8_blockscale(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_triton_gemm_a8w8_blockscale(
A, B, As, Bs, output_dtype
)
@staticmethod
def gemm_a8w8_blockscale(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_gemm_a8w8_blockscale(
A, B, As, Bs, output_dtype
)
@staticmethod
def fused_moe(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weight: torch.Tensor,
topk_ids: torch.Tensor,
expert_mask: torch.Tensor | None = None,
activation_method: int = 0,
quant_method: int = 0,
doweight_stage1: bool = False,
w1_scale: torch.Tensor | None = None,
w2_scale: torch.Tensor | None = None,
a1_scale: torch.Tensor | None = None,
a2_scale: torch.Tensor | None = None,
num_local_tokens: torch.Tensor | None = None,
output_dtype: torch.dtype | None = None,
hidden_pad: int = 0,
intermediate_pad: int = 0,
gate_mode: str = "",
bias1: torch.Tensor | None = None,
bias2: torch.Tensor | None = None,
moe_sorting_dispatch_policy: int = 0,
swiglu_limit: float = 0.0,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_fused_moe(
hidden_states,
w1,
w2,
topk_weight,
topk_ids,
expert_mask,
activation_method,
quant_method,
doweight_stage1,
w1_scale,
w2_scale,
a1_scale,
a2_scale,
num_local_tokens,
output_dtype,
hidden_pad,
intermediate_pad,
gate_mode,
bias1,
bias2,
moe_sorting_dispatch_policy,
swiglu_limit,
)
@staticmethod
def asm_moe_tkw1(
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
fc1_scale: torch.Tensor | None = None,
fc2_scale: torch.Tensor | None = None,
fc1_smooth_scale: torch.Tensor | None = None,
fc2_smooth_scale: torch.Tensor | None = None,
a16: bool = False,
per_tensor_quant_scale: torch.Tensor | None = None,
expert_mask: torch.Tensor | None = None,
activation_method: int = 0,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_asm_moe_tkw1(
hidden_states,
w1,
w2,
topk_weights,
topk_ids,
fc1_scale,
fc2_scale,
fc1_smooth_scale,
fc2_smooth_scale,
a16,
per_tensor_quant_scale,
expert_mask,
activation_method,
)
@staticmethod
def topk_softmax(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool,
num_shared_experts: int = 0,
shared_expert_scoring_func: str = "",
) -> tuple[torch.Tensor, ...]:
torch.ops.vllm.rocm_aiter_topk_softmax(
topk_weights,
topk_indices,
token_expert_indices,
gating_output,
renormalize,
num_shared_experts,
shared_expert_scoring_func,
)
return topk_weights, topk_indices
@staticmethod
def topk_sigmoid(
topk_weights: torch.Tensor,
topk_indices: torch.Tensor,
token_expert_indices: torch.Tensor,
gating_output: torch.Tensor,
renormalize: bool,
) -> tuple[torch.Tensor, ...]:
torch.ops.vllm.rocm_aiter_topk_sigmoid(
topk_weights, topk_indices, gating_output
)
return topk_weights, topk_indices
@staticmethod
def biased_grouped_topk(
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
routed_scaling_factor: float = 1.0,
) -> None:
if correction_bias.dtype != gating_output.dtype:
correction_bias = correction_bias.to(gating_output.dtype)
torch.ops.vllm.rocm_aiter_biased_grouped_topk(
gating_output,
correction_bias,
topk_weights,
topk_ids,
num_expert_group,
topk_group,
need_renorm,
routed_scaling_factor,
)
@staticmethod
def grouped_topk(
gating_output: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
num_expert_group: int,
topk_group: int,
need_renorm: bool,
scoring_func: str = "softmax",
routed_scaling_factor: float = 1.0,
) -> None:
torch.ops.vllm.rocm_aiter_grouped_topk(
gating_output,
topk_weights,
topk_ids,
num_expert_group,
topk_group,
need_renorm,
scoring_func,
routed_scaling_factor,
)
@staticmethod
def fused_topk(
x: torch.Tensor,
router_logits: torch.Tensor,
top_k: int,
gate_up: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
@staticmethod
def mla_decode_fwd(
q: torch.Tensor,
kv_buffer: torch.Tensor,
o: torch.Tensor,
sm_scale: float,
qo_indptr: torch.Tensor,
max_seqlen_qo: int,
kv_indptr: torch.Tensor | None = None,
kv_indices: torch.Tensor | None = None,
kv_last_page_lens: torch.Tensor | None = None,
logit_cap: float = 0.0,
q_scale: torch.Tensor | None = None,
kv_scale: torch.Tensor | None = None,
work_meta_data: torch.Tensor | None = None,
work_indptr: torch.Tensor | None = None,
work_info_set: torch.Tensor | None = None,
reduce_indptr: torch.Tensor | None = None,
reduce_final_map: torch.Tensor | None = None,
reduce_partial_map: torch.Tensor | None = None,
):
torch.ops.vllm.rocm_aiter_mla_decode_fwd(
q,
kv_buffer.view(-1, 1, 1, q.shape[-1]),
o,
qo_indptr,
max_seqlen_qo,
kv_indptr,
kv_indices,
kv_last_page_lens,
sm_scale=sm_scale,
logit_cap=logit_cap,
q_scale=q_scale,
kv_scale=kv_scale,
work_meta_data=work_meta_data,
work_indptr=work_indptr,
work_info_set=work_info_set,
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
)
@staticmethod
def per_tensor_quant(
x: torch.Tensor,
quant_dtype: torch.dtype,
scale: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
out = torch.empty_like(x, dtype=quant_dtype)
is_dynamic = scale is None
if is_dynamic:
scale = torch.empty(1, dtype=torch.float32, device=x.device)
torch.ops.vllm.rocm_aiter_per_tensor_quant(out, x, scale, is_dynamic)
return out, scale
@staticmethod
def per_token_quant(
x: torch.Tensor,
quant_dtype: torch.dtype,
scale: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops.vllm.rocm_aiter_per_token_quant(x, quant_dtype, scale)
@staticmethod
def gemm_a8wfp4(
x: torch.Tensor,
w: torch.Tensor,
x_scales: torch.Tensor,
w_scales: torch.Tensor,
out_dtype: torch.dtype,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_gemm_a8wfp4(
x, w, x_scales, w_scales, out_dtype
)
@staticmethod
def triton_fp4_gemm_dynamic_quant(
x: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
out_dtype: torch.dtype | None = torch.bfloat16,
x_scales: torch.Tensor | None = None,
) -> torch.Tensor:
from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4
from aiter.ops.triton.quant import dynamic_mxfp4_quant
if x_scales is None:
x_q, x_s = dynamic_mxfp4_quant(x)
else:
x_q = x
x_s = x_scales
y = torch.empty(
x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
)
gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
return y
@staticmethod
def triton_rope_and_cache(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
layer_slot_mapping: torch.Tensor,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
flash_layout: bool,
apply_scale: bool,
):
from aiter.ops.triton.fused_kv_cache import fused_qk_rope_reshape_and_cache
cos, sin = cos_sin_cache.chunk(2, dim=-1)
fused_qk_rope_reshape_and_cache(
query,
key,
value,
key_cache,
value_cache,
layer_slot_mapping,
positions,
cos,
sin,
k_scale,
v_scale,
is_neox,
flash_layout=flash_layout,
apply_scale=apply_scale,
q_out=query,
k_out=key,
output_zeros=False,
)
@staticmethod
def batched_gemm_a16wfp4(
X: torch.Tensor,
W: torch.Tensor,
w_scale: torch.Tensor,
Y: torch.Tensor,
transpose_bm: bool | None = False,
prequant: bool | None = False,
y_scale: torch.Tensor | None = None,
) -> torch.Tensor:
# ruff: noqa: E501 # isort: skip
from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4
return batched_gemm_a16wfp4(
X,
W,
w_scale,
y=Y,
transpose_bm=transpose_bm,
prequant=prequant,
y_scale=y_scale,
)
@staticmethod
def triton_fp8_bmm(
X: torch.Tensor,
WQ: torch.Tensor,
w_scale: torch.Tensor,
group_size: int = 128,
bias: torch.Tensor | None = None,
dtype: torch.dtype | None = torch.bfloat16,
splitK: int | None = None,
YQ: torch.Tensor | None = None,
transpose_bm: bool | None = False,
config: dict | None = None,
) -> torch.Tensor:
# ruff: noqa: E501 # isort: skip
from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm,
)
return aiter_triton_fp8_bmm(
X,
WQ,
w_scale,
group_size=group_size,
bias=bias,
dtype=dtype,
splitK=splitK,
YQ=YQ,
transpose_bm=transpose_bm,
config=config,
)
@staticmethod
def group_fp8_quant(
input_2d: torch.Tensor,
group_size: int = 128,
) -> tuple[torch.Tensor, torch.Tensor]:
assert group_size == 128, "Group size must be 128"
return torch.ops.vllm.rocm_aiter_group_fp8_quant(input_2d, group_size)
@staticmethod
def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:
return (n, k) in [
(1024, 8192),
(2112, 7168),
(3072, 1536),
(32768, 8192),
(4096, 7168),
(4608, 7168),
(512, 7168),
(7168, 2048),
(7168, 256),
(8192, 1024),
(8192, 32768),
]
@staticmethod
def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool:
return (n, k) in [
(8192, 4096),
(1280, 8192),
(16384, 53248),
(106496, 16384),
(57344, 8192),
(8192, 2048),
(2560, 8192),
(10240, 8192),
(16384, 16384),
(8192, 28672),
(28672, 8192),
(18432, 16384),
(8192, 1024),
(7168, 8192),
(5120, 8192),
(8192, 8192),
(8192, 7168),
(14336, 8192),
(8192, 14336),
(8192, 3584),
]
@staticmethod
def is_shuffled_per_token_w8a8_gemm_tuned(
N: int, K: int, q_dtype_w: torch.dtype
) -> bool:
import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
csv_path = (
aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE
)
return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
@staticmethod
def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool:
import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE
return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
@staticmethod
def shuffle_weight(
tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
) -> torch.Tensor:
from aiter.ops.shuffle import shuffle_weight
return shuffle_weight(tensor, layout=layout)
@staticmethod
def shuffle_weight_a16w4(
tensor: "torch.Tensor",
nLane: int,
gate_up: bool,
) -> "torch.Tensor":
"""
Shuffles the weight tensor into (A16W4) layout for AITER kernels.
Args:
tensor: The input weight tensor to be shuffled.
layout: The block layout to use, defaults to (16, 4).
Returns:
torch.Tensor: The shuffled tensor.
"""
from aiter.ops.shuffle import shuffle_weight_a16w4
return shuffle_weight_a16w4(tensor, nLane, gate_up)
@staticmethod
def shuffle_scale_a16w4(
tensor: "torch.Tensor",
num_experts: int,
gate_up: bool,
) -> "torch.Tensor":
"""
Shuffles the scale tensor into (A16W4) layout for AITER kernels.
Args:
tensor: The input scale tensor to be shuffled.
num_experts: Number of experts, needed for reshaping logic.
gate_up: Whether the scale is for w13 (True) or w2 (False).
Returns:
torch.Tensor: The shuffled scale tensor.
"""
from aiter.ops.shuffle import shuffle_scale_a16w4
return shuffle_scale_a16w4(tensor, num_experts, gate_up)
@staticmethod
def shuffle_weights(
*tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
) -> tuple[torch.Tensor, ...]:
"""
Applies shuffle_weight function from AITER to each
input tensor and returns them.
Rearranges (shuffles) the input tensor/s
into a specified block layout for optimized computation.
Args:
*tensors: Variable number of torch.Tensor objects.
layout: A pair of integers specifying the block sizes used to divide
the tensors during shuffling. Default is (16, 16).
Returns:
A Tuple of shuffled tensors.
"""
from aiter.ops.shuffle import shuffle_weight
return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)
@staticmethod
def shuffle_mxfp8_moe_weights(
w13: torch.Tensor,
w2: torch.Tensor,
w13_scale: torch.Tensor,
w2_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Preshuffle MXFP8 MoE weights + E8M0 scales into AITER's FlyDSL layout:
gate/up-interleaved weights, interleaved scale for w13 (gate/up), plain
scale for w2 (the interleaved variant is gate/up-only and misaligns w2).
"""
from aiter.ops.shuffle import shuffle_scale, shuffle_weight
num_experts = w13.shape[0]
w13 = shuffle_weight(w13, is_guinterleave=True, gate_up=True)
w2 = shuffle_weight(w2, is_guinterleave=True, gate_up=False)
w13_scale = shuffle_scale(
w13_scale.reshape(-1, w13_scale.shape[-1]),
num_experts,
is_guinterleave=True,
gate_up=True,
)
w2_scale = shuffle_scale(w2_scale.reshape(-1, w2_scale.shape[-1]))
return w13, w2, w13_scale, w2_scale
@staticmethod
def flash_attn_varlen_func(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int,
max_seqlen_k: int,
min_seqlen_q: int | None = None,
dropout_p: float = 0.0,
softmax_scale: float | None = None,
causal: bool = False,
window_size: tuple[int, int] | None = None,
alibi_slopes: torch.Tensor | None = None,
return_lse: bool = False,
out: torch.Tensor | None = None,
sink_ptr: torch.Tensor | None = None,
):
"""
Flash attention with variable length sequences.
This function is NOT wrapped with @is_aiter_supported decorator
to allow explicit backend selection via attention_config to work
even when VLLM_ROCM_USE_AITER=0.
Note: This performs lazy import of aiter.flash_attn_varlen_func
"""
from aiter import flash_attn_varlen_func
return flash_attn_varlen_func(
q=q,
k=k,
v=v,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
min_seqlen_q=min_seqlen_q,
dropout_p=dropout_p,
softmax_scale=softmax_scale,
causal=causal,
window_size=window_size,
alibi_slopes=alibi_slopes,
return_lse=return_lse,
out=out,
sink_ptr=sink_ptr,
)
@staticmethod
def pa_fwd_asm(
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
block_tables: torch.Tensor,
context_lens: torch.Tensor,
block_tables_stride0: int,
K_QScale: torch.Tensor,
V_QScale: torch.Tensor,
out_: torch.Tensor,
):
"""
Paged attention forward pass using assembly kernel.
This function is NOT wrapped with @is_aiter_supported decorator
to allow explicit backend selection via attention_config to work
even when VLLM_ROCM_USE_AITER=0.
Note: This performs lazy import of aiter.pa_fwd_asm
"""
from aiter import pa_fwd_asm
return pa_fwd_asm(
Q=Q,
K=K,
V=V,
block_tables=block_tables,
context_lens=context_lens,
block_tables_stride0=block_tables_stride0,
K_QScale=K_QScale,
V_QScale=V_QScale,
out_=out_,
)
@staticmethod
def paged_attention_common(
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
tmp_out: torch.Tensor,
max_logits: torch.Tensor,
exp_sums: torch.Tensor,
max_seq_len: int,
block_tables: torch.Tensor,
context_lens: torch.Tensor,
block_tables_stride0: int,
scale: float,
K_QScale_hip: torch.Tensor,
V_QScale_hip: torch.Tensor,
K_QScale_asm: torch.Tensor,
V_QScale_asm: torch.Tensor,
out_: torch.Tensor,
kv_cache_dtype: str,
):
"""
Paged attention common function.
This function is NOT wrapped with @is_aiter_supported decorator
to allow explicit backend selection via attention_config to work
even when VLLM_ROCM_USE_AITER=0.
Note: This performs lazy import of aiter.paged_attention_common
"""
from aiter import paged_attention_common
return paged_attention_common(
Q=Q,
K=K,
V=V,
tmp_out=tmp_out,
max_logits=max_logits,
exp_sums=exp_sums,
max_seq_len=max_seq_len,
block_tables=block_tables,
context_lens=context_lens,
block_tables_stride0=block_tables_stride0,
scale=scale,
K_QScale_hip=K_QScale_hip,
V_QScale_hip=V_QScale_hip,
K_QScale_asm=K_QScale_asm,
V_QScale_asm=V_QScale_asm,
out_=out_,
kv_cache_dtype=kv_cache_dtype,
)
@staticmethod
def mhc_pre(
residual: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Forward pass for mHC pre block.
Args:
residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
hc_scale: shape (3,), dtype torch.float32
hc_base: shape (hc_mult3,), dtype torch.float32
rms_eps: RMS normalization epsilon
hc_pre_eps: pre-mix epsilon
hc_sinkhorn_eps: sinkhorn epsilon
hc_post_mult_value: post-mix multiplier value
sinkhorn_repeat: number of sinkhorn iterations
n_splits: split-k factor;
Returns:
post_mix: shape (..., hc_mult), dtype torch.float32
comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
layer_input: shape (..., hidden_size), dtype torch.bfloat16
"""
from aiter.ops.mhc import mhc_pre
# Validate shapes
assert residual.dtype == torch.bfloat16
assert fn.dtype == torch.float32
assert hc_scale.dtype == torch.float32
assert hc_base.dtype == torch.float32
hc_mult = residual.shape[-2]
hidden_size = residual.shape[-1]
hc_mult2 = hc_mult * hc_mult
hc_mult3 = hc_mult * 2 + hc_mult2
hc_hidden_size = hc_mult * hidden_size
assert fn.shape[0] == hc_mult3
assert fn.shape[1] == hc_hidden_size
assert hc_scale.shape == (3,)
assert hc_base.shape == (hc_mult3,)
outer_shape = residual.shape[:-2]
residual_flat = residual.view(-1, hc_mult, hidden_size)
num_tokens = residual_flat.shape[0]
if num_tokens == 0:
return (
torch.empty(
num_tokens,
hc_mult,
1,
dtype=torch.float32,
device=residual_flat.device,
),
torch.empty(
num_tokens,
hc_mult,
hc_mult,
dtype=torch.float32,
device=residual_flat.device,
),
torch.empty(
num_tokens,
hidden_size,
dtype=torch.bfloat16,
device=residual_flat.device,
),
)
# AITER's Python wrapper allocates intermediate/output tensors without
# explicit device arguments, so run it under the residual tensor's device.
with torch.device(residual_flat.device):
post_mix, comb_mix, layer_input = mhc_pre(
residual_flat,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
)
return (
post_mix.view(*outer_shape, hc_mult, 1),
comb_mix.view(*outer_shape, hc_mult, hc_mult),
layer_input.view(*outer_shape, hidden_size),
)
@staticmethod
def hc_head(
hs_flat: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
out: torch.Tensor,
hidden_size: int,
rms_eps: float,
hc_eps: float,
hc_mult: int,
) -> None:
"""Run hc_head through AITER mhc_pre and write the result to out."""
assert hs_flat.dtype == torch.bfloat16
assert fn.dtype == torch.float32
assert hc_scale.dtype == torch.float32
assert hc_base.dtype == torch.float32
assert hs_flat.shape[-2:] == (hc_mult, hidden_size)
assert fn.shape == (hc_mult, hc_mult * hidden_size)
assert hc_scale.shape == (1,)
assert hc_base.shape == (hc_mult,)
num_tokens = hs_flat.shape[0]
if num_tokens == 0:
return
hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
full_fn = torch.zeros(
hc_mult3,
hc_mult * hidden_size,
dtype=fn.dtype,
device=fn.device,
)
full_fn[:hc_mult] = fn
full_base = torch.zeros(hc_mult3, dtype=hc_base.dtype, device=hc_base.device)
full_base[:hc_mult] = hc_base
full_scale = torch.zeros(3, dtype=hc_scale.dtype, device=hc_scale.device)
full_scale[0] = hc_scale[0]
_, _, layer_input = rocm_aiter_ops.mhc_pre(
hs_flat,
full_fn,
full_scale,
full_base,
rms_eps,
hc_eps,
0.0,
1.0,
0,
)
out.copy_(layer_input)
@staticmethod
def mhc_post(
x: torch.Tensor,
residual: torch.Tensor,
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor:
from aiter.ops.mhc import mhc_post
hc_mult = residual.shape[-2]
hidden_size = residual.shape[-1]
residual_flat = residual.view(-1, hc_mult, hidden_size)
num_tokens = residual_flat.shape[0]
out = torch.empty_like(residual_flat)
mhc_post(
out,
x.view(num_tokens, hidden_size),
residual_flat,
post_layer_mix.view(num_tokens, hc_mult, 1),
comb_res_mix.view(num_tokens, hc_mult, hc_mult),
)
return out.view_as(residual)
rocm_aiter_ops.register_ops_once()
Reference in New Issue View Git Blame Copy Permalink