iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm] gfx1250 optimizations and model enablement (#58800)

Browse Source 
Signed-off-by: jpvillam <juan.villamizar@amd.com>
Signed-off-by: Jaden Mathias <jaden.mathias@amd.com>
Signed-off-by: Juan Villamizar <jpvillam@amd.com>
Signed-off-by: jpvillam <jpvillam@amd.com>
Signed-off-by: Omar Muhammad <omar.muhammad@amd.com>
Signed-off-by: Daniel <danichan@amd.com>
Signed-off-by: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com>
Co-authored-by: Jaden Mathias <jaden.mathias@amd.com>
Co-authored-by: jpvillam <jpvillam@amd.com>
Co-authored-by: Omar Muhammad <omar.muhammad@amd.com>
Co-authored-by: Daniel <danichan@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Daniel Chan <Daniel.Chan@amd.com>
This commit is contained in:
Juan Villamizar
2026-10-15 14:51:57 -07:00
committed by GitHub
co-authored by Jaden Mathias jpvillam Omar Muhammad Daniel Claude Opus 5 Daniel Chan
parent 9b99ccb56b
commit e711723352
21 changed files with 845 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/Dockerfile.rocm_base_gfx1250
+14 -24
View File
@@ -1,16 +1,16 @@
ARG BASE_IMAGE=ubuntu:24.04
ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/
ARG ROCM_SDK_VERSION=7.14.0a20260623
ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623
ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623
ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623
ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623
ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623
ARG ROCM_WHEEL_INDEX=https://rocm.nightlies.amd.com/whl-multi-arch/
ARG ROCM_SDK_VERSION=10.1.0a20260822
ARG TORCH_VERSION=2.11.0+rocm10.1.0a20260822
ARG TORCHVISION_VERSION=0.26.0+rocm10.1.0a20260822
ARG TORCHAUDIO_VERSION=2.11.0+rocm10.1.0a20260822
ARG TRITON_VERSION=3.8.0+git4cff872c.rocm10.1.0a20260822
ARG APEX_VERSION=1.11.0+rocm10.1.0a20260822
ARG FA_BRANCH="jpvillam/gfx1250_wip"
ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git"
ARG AITER_BRANCH="main"
ARG AITER_BRANCH="f85b24b"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -119,10 +119,10 @@ ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi
# which keeps the set consistent and avoids pip backtracking.
RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \
--extra-index-url https://pypi.org/simple \
"torch[device-all]==${TORCH_VERSION}" \
"torch[device-gfx1250]==${TORCH_VERSION}" \
"torchvision==${TORCHVISION_VERSION}" \
"torchaudio==${TORCHAUDIO_VERSION}" \
"rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \
"rocm[libraries,devel,device-gfx1250]==${ROCM_SDK_VERSION}" && \
rocm-sdk init
# Torch runtime deps that may not be published on the ROCm wheel index;
@@ -200,15 +200,6 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
###
### AMD SMI Build
###
FROM base AS build_amdsmi
RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \
&& pip wheel . --wheel-dir=dist
RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install
###
### MORI Build TODO(Build needs fixing)
###
@@ -253,7 +244,10 @@ FROM base AS build_aiter
ARG AITER_BRANCH
ARG AITER_REPO
ARG USE_SCCACHE
RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
RUN git clone ${AITER_REPO} \
&& cd aiter \
&& git checkout --detach ${AITER_BRANCH}
RUN cd aiter \
&& git submodule update --init --recursive \
&& pip install -r requirements.txt
@@ -278,8 +272,6 @@ FROM base AS debs_wheel_release
RUN mkdir /app/debs
RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
cp /install/*.whl /app/debs
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
cp /install/*.whl /app/debs
@@ -288,8 +280,6 @@ FROM base AS debs
RUN mkdir /app/debs
RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
cp /install/*.whl /app/debs
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
cp /install/*.whl /app/debs
RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
docker/Dockerfile.rocm_gfx1250
+5 -3
View File
@@ -15,6 +15,7 @@ ARG NIC_BACKEND=all
# Users can specify a custom version compatible with their host drivers.
# The default version has been tested with ionic-dkms=26.03.3.001
# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC.
# TODO (jmathias) : Check if we can revert this Pinned AINIC version
ARG AINIC_VERSION=1.117.5-a-77
ARG UBUNTU_CODENAME=jammy
@@ -115,8 +116,8 @@ WORKDIR ${COMMON_WORKDIR}
FROM base AS fetch_vllm_0
ONBUILD COPY ./ vllm/
FROM base AS fetch_vllm_1
ARG VLLM_REPO="https://github.com/ROCm/vllm.git"
ARG VLLM_BRANCH="455_wip"
ARG VLLM_REPO="https://github.com/vllm-project/vllm.git"
ARG VLLM_BRANCH="main"
ENV VLLM_REPO=${VLLM_REPO}
ENV VLLM_BRANCH=${VLLM_BRANCH}
ONBUILD RUN git clone ${VLLM_REPO} \
@@ -726,13 +727,14 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
# Download bench scripts for gfx1250
RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh
RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_benchserve.sh
### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better
RUN pip3 uninstall -y triton && \
git clone https://github.com/triton-lang/triton.git && \
cd triton && \
git checkout c517f38c && \
git checkout aedd1f6 && \
TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install .
CMD ["/bin/bash"]
tests/kernels/moe/test_ocp_mx_moe.py
+14 -1
View File
@@ -42,13 +42,15 @@ ROCM_AVAILABLE = current_platform.is_rocm()
ROCM_TRITON_KERNELS_AVAILABLE = False
ROCM_AITER_AVAILABLE = is_aiter_found()
ROCM_GFX950 = False
ROCM_GFX1250 = False
if ROCM_AVAILABLE:
from vllm.platforms.rocm import on_gfx950
from vllm.platforms.rocm import on_gfx950, on_gfx1250
from vllm.utils.import_utils import has_triton_kernels
ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
ROCM_GFX950 = on_gfx950()
ROCM_GFX1250 = on_gfx1250()
if ROCM_AITER_AVAILABLE:
from aiter.ops.triton.quant import dynamic_mxfp4_quant
@@ -1568,6 +1570,15 @@ ROCM_BACKEND_CONFIGS = {
"requires_aiter": True,
"requires_gfx950": True,
},
"AITER_TRITON_MXFP4_MXFP4": {
"activation": "SILU",
"act_type": "mxfp4",
"rtol": 1.0,
"percent": 0.8,
"requires_aiter": True,
"requires_gfx950": False,
"requires_gfx1250": True,
},
}
@@ -1607,6 +1618,8 @@ def test_rocm_mxfp4_moe_oracle(
pytest.skip(f"Backend {backend_name} requires AITER")
if config["requires_gfx950"] and not ROCM_GFX950:
pytest.skip(f"Backend {backend_name} requires GFX950")
if config.get("requires_gfx1250") and not ROCM_GFX1250:
pytest.skip(f"Backend {backend_name} requires GFX1250")
import vllm.distributed.parallel_state as ps
from vllm.config import VllmConfig, set_current_vllm_config
vllm/_aiter_ops.py
+63
View File
@@ -1010,6 +1010,39 @@ def _rocm_aiter_triton_gemm_a8w8_blockscale_fake(
return Y
def _rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_impl(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale_preshuffle
n, k = B.shape
return gemm_a8w8_blockscale_preshuffle(
A,
B.reshape(n // 16, k * 16),
As,
Bs,
output_dtype,
None, # y
None, # config
False, # skip_reduce
False, # is_x_scale_transposed (misspelled in aiter)
)
def _rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_fake(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
return torch.empty(A.shape[0], B.shape[0], dtype=output_dtype, device=A.device)
def _rocm_aiter_gemm_a8w8_blockscale_impl(
A: torch.Tensor,
B: torch.Tensor,
@@ -2984,6 +3017,12 @@ class rocm_aiter_ops:
fake_impl=_rocm_aiter_triton_gemm_a8w8_blockscale_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle",
op_func=_rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_impl,
fake_impl=_rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_gemm_a8w8_blockscale",
op_func=_rocm_aiter_gemm_a8w8_blockscale_impl,
@@ -3298,6 +3337,19 @@ class rocm_aiter_ops:
A, B, As, Bs, output_dtype
)
@staticmethod
def triton_gemm_a8w8_blockscale_preshuffle(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
return torch.ops.vllm.rocm_aiter_triton_gemm_a8w8_blockscale_preshuffle(
A, B, As, Bs, output_dtype
)
@staticmethod
def gemm_a8w8_blockscale(
A: torch.Tensor,
@@ -4074,6 +4126,17 @@ class rocm_aiter_ops:
except (AssertionError, ImportError):
return False
@staticmethod
@if_aiter_supported
def is_triton_blockscale_bpreshuffle_gemm_tuned(n: int, k: int) -> bool:
from aiter.ops.triton.utils.gemm_config_utils import (
get_gemm_config,
)
# Fake M value for now
_, is_tuned = get_gemm_config("GEMM-A8W8_BLOCKSCALE_PRESHUFFLED", 1, n, k)
return is_tuned
@staticmethod
@functools.cache
def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool:
vllm/config/kernel.py
+3
View File
@@ -139,6 +139,7 @@ MoEBackend = Literal[
"aiter",
"aiter_mega_moe",
"aiter_triton_mxfp4_bf16",
"aiter_triton_mxfp4_mxfp4",
"flydsl",
"hpc",
"emulation",
@@ -329,6 +330,8 @@ class KernelConfig:
DeepSeek-V4 only)
- "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16
(moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250)
- "aiter_triton_mxfp4_mxfp4": Use the AITER Triton/Gluon MXFP4 W4A4
(moe_gemm_a4w4) MoE kernel (ROCm gfx1250)
- "flydsl": Use AMD FlyDSL kernels (ROCm only)
- "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only)
- "hpc": Use HPC kernels (FP8 and Hopper only)
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+38 -1
View File
@@ -9,6 +9,7 @@ from vllm._aiter_ops import (
rocm_aiter_ops,
)
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
_upcast_e8m0_to_fp32,
)
@@ -399,6 +400,26 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250
)
self.use_bpreshuffle = (
_on_gfx1250
and n % 16 == 0
and k % 128 == 0
and bool(rocm_aiter_ops.is_triton_blockscale_bpreshuffle_gemm_tuned(n, k))
)
if self.use_bpreshuffle:
# The tuned preshuffle configs expect a column-major group scale.
act_scale_descriptor = config.activation_quant_key.scale
self.quant_fp8 = QuantFP8(
static=act_scale_descriptor.static,
group_shape=act_scale_descriptor.group_shape,
num_token_padding=self.get_output_padding(),
use_ue8m0=False,
column_major_scales=True,
)
logger.info_once(
"Using aiter A8W8 blockscale B-preshuffle GEMM for N=%d K=%d", n, k
)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
super().process_weights_after_loading(layer)
@@ -410,6 +431,18 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
if ws is not None and ws.dtype == torch.float8_e8m0fnu:
replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous())
if getattr(layer, "is_bmm", False):
self.use_bpreshuffle = False
if self.use_bpreshuffle:
w = FP8BlockParams.from_layer(layer).weight
replace_parameter(
layer,
FP8BlockParams.WEIGHT,
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
layer.aiter_bpreshuffled = True
@classmethod
def is_supported(cls, compute_capability=None):
if (
@@ -464,7 +497,11 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
Bs = Bs.to(torch.float32)
out_dtype = self.config.out_dtype
if self.use_triton:
if self.use_bpreshuffle:
gemm_a8w8_blockscale_op = (
rocm_aiter_ops.triton_gemm_a8w8_blockscale_preshuffle
)
elif self.use_triton:
gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
else:
gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
+7 -3
View File
@@ -53,10 +53,10 @@ def _aiter_w4a16_silu_via_a8w4(
from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
should_use_cdna4_mx_scale_swizzle,
get_aiter_mx_scale_swizzle,
)
swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
swz = get_aiter_mx_scale_swizzle()
quant_dtype = torch.float8_e4m3fn
g1_gammas = gammas if apply_router_weight_on_input else None
@@ -181,7 +181,11 @@ def aiter_triton_kernel_w4a16_moe_forward(
gating_output,
topk,
score_mode=score_mode,
bias=e_score_correction_bias,
bias=(
e_score_correction_bias.float()
if e_score_correction_bias is not None
else None
),
renorm=renormalize,
routed_scaling_factor=scale,
use_grouped_topk=use_grouped_topk,
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py
+373
View File
@@ -0,0 +1,373 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm._aiter_ops import rocm_aiter_ops
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
FusedMoEQuantConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
kMxfp4Dynamic,
kMxfp4Static,
)
__all__ = [
"AiterW4A4ExpertsMonolithic",
"aiter_triton_kernel_w4a4_moe_forward",
]
def _aiter_raw(t):
if t is None or isinstance(t, torch.Tensor):
return t
return t.storage.data if hasattr(t, "storage") else t
def aiter_triton_kernel_w4a4_moe_forward(
hidden_states: torch.Tensor,
w1,
w2,
gating_output: torch.Tensor,
topk: int,
renormalize: bool,
activation: MoEActivation = MoEActivation.SWIGLUOAI,
quant_config: FusedMoEQuantConfig | None = None,
apply_router_weight_on_input: bool = False,
global_num_experts: int = -1,
expert_map: torch.Tensor | None = None,
unpadded_N_w1=None,
unpadded_K_w1=None,
unpadded_N_w2=None,
unpadded_K_w2=None,
num_expert_group: int | None = None,
topk_group: int | None = None,
e_score_correction_bias: torch.Tensor | None = None,
routed_scaling_factor: float | None = None,
score_mode: str | None = None,
input_ids: torch.Tensor | None = None,
hash_indices_table: torch.Tensor | None = None,
):
assert quant_config is not None and rocm_aiter_ops.is_enabled()
from aiter.ops.triton.moe.moe_op_gemm_a4w4 import moe_gemm_a4w4, mxfp4_quant
try:
from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
except ImportError:
from aiter.ops.triton.moe_routing import routing as _routing_mod
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
get_aiter_mx_scale_swizzle,
)
from vllm.platforms.rocm import on_gfx1250
if on_gfx1250():
_routing_mod.is_tdm_avail = lambda: False
aiter_routing = _routing_mod.routing
gating_output = torch.nan_to_num(gating_output, nan=0.0, posinf=0.0, neginf=0.0)
if hash_indices_table is not None:
assert input_ids is not None, "hash routing requires input_ids"
n_tokens, n_expts_tot = gating_output.shape
tokens_per_expt = max(1, n_tokens * topk // n_expts_tot)
block_m = max(16, min(1 << (tokens_per_expt - 1).bit_length(), 128))
routing_data, gather_idx, scatter_idx = _routing_mod.routing_from_hash(
gating_output,
hash_indices_table,
input_ids.to(hash_indices_table.dtype),
topk,
block_m,
score_mode=score_mode or "sqrtsoftplus",
renorm=renormalize,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
)
elif score_mode is not None:
use_grouped_topk = num_expert_group is not None and num_expert_group > 1
routing_data, gather_idx, scatter_idx = aiter_routing(
gating_output,
topk,
score_mode=score_mode,
bias=(
e_score_correction_bias.float()
if e_score_correction_bias is not None
else None
),
renorm=renormalize,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
use_grouped_topk=use_grouped_topk,
num_expert_group=num_expert_group,
topk_group=topk_group,
)
else:
routing_data, gather_idx, scatter_idx = aiter_routing(
gating_output, topk, sm_first=not renormalize
)
if on_gfx1250():
gather_src = gather_idx.to(torch.long) // topk
hidden_states = hidden_states[gather_src]
gather_idx = None
w1_data = _aiter_raw(w1)
w2_data = _aiter_raw(w2)
w1_scale = _aiter_raw(quant_config.w1_scale)
w2_scale = _aiter_raw(quant_config.w2_scale)
gammas = routing_data.gate_scal if routing_data else None
swiglu_limit = (
quant_config.gemm1_clamp_limit
if quant_config.gemm1_clamp_limit is not None
else 7.0
)
# SWIGLUOAI needs the alpha/residual swiglu fused into the GEMM1 epilogue,
# which reads gate/up interleaved along N (see the oracle's weight prep).
fused_swiglu = activation in (
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
)
swiglu_kwargs: dict[str, object] = {}
if fused_swiglu:
assert quant_config.gemm1_beta in (None, 1.0), (
"aiter's fused swiglu hardcodes the residual to (up + 1); "
f"gemm1_beta={quant_config.gemm1_beta} cannot be expressed"
)
swiglu_kwargs = {
"alpha": (
quant_config.gemm1_alpha
if quant_config.gemm1_alpha is not None
else 1.0
),
"limit": swiglu_limit,
"swiglu_add_residual": quant_config.gemm1_beta == 1.0,
}
swizzle_mx_scale = get_aiter_mx_scale_swizzle()
x_q, x_scale = mxfp4_quant(hidden_states.to(torch.bfloat16))
# GEMM1: gate+up projection.
raw_intermediate = moe_gemm_a4w4(
x_q,
w1_data,
x_scale,
w1_scale,
bias=quant_config.w1_bias,
routing_data=routing_data,
gather_indx=gather_idx,
gammas=gammas if apply_router_weight_on_input else None,
swizzle_mx_scale=swizzle_mx_scale,
apply_swiglu=fused_swiglu,
**swiglu_kwargs,
)
if fused_swiglu:
# The swiglu epilogue already halved N.
intermediate = (
raw_intermediate
if unpadded_N_w1 is None
else raw_intermediate[:, : unpadded_N_w1 // 2]
)
else:
if unpadded_N_w1 is not None:
raw_intermediate = raw_intermediate[:, :unpadded_N_w1]
# SiLU(gate) * up on the concatenated [gate | up] halves
from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
half_n = raw_intermediate.shape[-1] // 2
intermediate = torch.empty(
raw_intermediate.shape[0],
half_n,
dtype=raw_intermediate.dtype,
device=raw_intermediate.device,
)
fused_clamp_act_mul(
raw_intermediate,
out=intermediate,
swiglu_limit=swiglu_limit,
activation="silu",
dtype_quant=None,
)
# GEMM2: down projection with scatter-reduce
mid_q, mid_scale = mxfp4_quant(intermediate.to(torch.bfloat16))
out = moe_gemm_a4w4(
mid_q,
w2_data,
mid_scale,
w2_scale,
bias=quant_config.w2_bias,
routing_data=routing_data,
scatter_indx=scatter_idx,
gammas=None if apply_router_weight_on_input else gammas,
swizzle_mx_scale=swizzle_mx_scale,
apply_swiglu=False,
)
return out
class AiterW4A4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
"""Monolithic MXFP4 W4A4 expert using AITER Triton/Gluon kernels (gfx1250).
Uses moe_gemm_a4w4 (auto-selects gluon backend on gfx1250) with dynamic
MXFP4 activation quantization via mxfp4_quant.
"""
def __init__(
self,
moe_config: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
):
super().__init__(moe_config, quant_config)
self.topk = moe_config.experts_per_token
self.renormalize = moe_config.routing_method in (
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
RoutingMethodType.DeepseekV4,
RoutingMethodType.DeepSeekV3,
RoutingMethodType.MiniMax2,
)
@staticmethod
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
@staticmethod
def _supports_hash_routing() -> bool:
return True
@staticmethod
def _supports_current_device() -> bool:
if not rocm_aiter_ops.is_enabled():
return False
from vllm.platforms.rocm import on_gfx1250
return on_gfx1250()
@staticmethod
def _supports_no_act_and_mul() -> bool:
return False
@staticmethod
def _supports_quant_scheme(
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic)
@staticmethod
def _supports_activation(activation: MoEActivation) -> bool:
return activation in (
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
MoEActivation.SILU,
)
@staticmethod
def _supports_parallel_config(
moe_parallel_config: FusedMoEParallelConfig,
) -> bool:
return (
not moe_parallel_config.use_all2all_kernels
and not moe_parallel_config.enable_eplb
and moe_parallel_config.dp_size <= 1
)
@staticmethod
def _supports_routing_method(
routing_method: RoutingMethodType,
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
return routing_method in [
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
RoutingMethodType.DeepseekV4,
RoutingMethodType.DeepSeekV3,
RoutingMethodType.MiniMax2,
]
@staticmethod
def _supports_router_logits_dtype(
router_logits_dtype: torch.dtype | None,
routing_method: RoutingMethodType,
) -> bool:
return True
@property
def expects_unquantized_inputs(self) -> bool:
return True
def apply(
self,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
router_logits: torch.Tensor,
activation: MoEActivation,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
apply_router_weight_on_input: bool,
num_expert_group: int | None = None,
e_score_correction_bias: torch.Tensor | None = None,
routed_scaling_factor: float | None = None,
topk_group: int | None = None,
routing_replay_out: torch.Tensor | None = None,
input_ids: torch.Tensor | None = None,
hash_indices_table: torch.Tensor | None = None,
) -> torch.Tensor:
assert self.moe_config.intermediate_size_per_partition_unpadded is not None
assert self.moe_config.hidden_dim_unpadded is not None
routing_method = self.moe_config.routing_method
if routing_method == RoutingMethodType.DeepseekV4:
score_mode = "sqrtsoftplus"
elif routing_method in (
RoutingMethodType.DeepSeekV3,
RoutingMethodType.MiniMax2,
):
# Sigmoid scores + routing bias; grouped for DeepSeek-R1/V3,
# ungrouped for MiniMax-M3.
score_mode = "sigmoid"
else:
score_mode = None
return aiter_triton_kernel_w4a4_moe_forward(
hidden_states=hidden_states,
w1=w1,
w2=w2,
gating_output=router_logits,
topk=self.topk,
renormalize=self.renormalize,
activation=activation,
global_num_experts=global_num_experts,
expert_map=expert_map,
quant_config=self.quant_config,
apply_router_weight_on_input=apply_router_weight_on_input,
unpadded_N_w1=(
self.moe_config.intermediate_size_per_partition_unpadded * 2
),
unpadded_K_w1=self.moe_config.hidden_dim_unpadded,
unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
unpadded_K_w2=(self.moe_config.intermediate_size_per_partition_unpadded),
num_expert_group=num_expert_group,
topk_group=topk_group,
e_score_correction_bias=e_score_correction_bias,
routed_scaling_factor=routed_scaling_factor,
score_mode=score_mode,
input_ids=input_ids,
hash_indices_table=hash_indices_table,
)
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+2 -8
View File
@@ -61,12 +61,6 @@ def aiter_triton_kernel_w4a8_moe_forward(
gating_output, topk, sm_first=not renormalize
)
# gfx1250: aiter's in-kernel gather is numerically broken
if on_gfx1250():
gather_src = gather_idx.to(torch.long) // topk
hidden_states = hidden_states[gather_src]
gather_idx = None
return triton_kernel_fused_mxfp4_w4a8_experts(
None,
hidden_states,
@@ -130,11 +124,11 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
from aiter.ops.triton.quant_moe import downcast_to_static_fp8
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
should_use_cdna4_mx_scale_swizzle,
get_aiter_mx_scale_swizzle,
weight_mx_scale,
)
_swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
_swizzle_mx_scale = get_aiter_mx_scale_swizzle()
assert quant_config.w1_precision is not None, (
"w1_precision in quant config can't be None"
vllm/model_executor/layers/fused_moe/layer.py
+1
View File
@@ -436,6 +436,7 @@ def FusedMoEFactory(
# TODO get from router? needs to be truncated?
e_score_correction_bias=e_score_correction_bias,
apply_router_weight_on_input=apply_router_weight_on_input,
hash_indices_table=hash_indices_table,
**routed_experts_args if routed_experts_args is not None else {},
)
if hash_indices_table is not None:
vllm/model_executor/layers/fused_moe/modular_kernel.py
+14
View File
@@ -668,6 +668,15 @@ class FusedMoEExperts(ABC):
"""
return True
@staticmethod
def _supports_hash_routing() -> bool:
"""Whether a monolithic kernel implements hash routing itself.
Kernels that return True receive `input_ids` and `hash_indices_table`
in `apply()` for layers with a hash-routing table.
"""
return False
@staticmethod
def _supports_router_logits_dtype(
router_logits_dtype: torch.dtype | None,
@@ -1623,6 +1632,11 @@ class FusedMoEKernelMonolithicImpl:
f"kernel {type(self.fused_experts).__name__}."
)
routing_replay_out = routing_sink.buffer[: len(a1q)]
hash_kwargs = (
{"input_ids": input_ids, "hash_indices_table": hash_indices_table}
if hash_indices_table is not None
else {}
)
fused_out = self.fused_experts.apply(
hidden_states=a1q,
w1=w1,
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+78 -3
View File
@@ -152,7 +152,9 @@ class Mxfp4MoeBackend(Enum):
# W4A16: aiter Triton moe_gemm_a16w4 kernel (gfx942/gfx950/gfx1250)
AITER_TRITON_MXFP4_BF16 = "AITER_TRITON_MXFP4_BF16"
AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel
AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel
AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel (gfx950)
# W4A4: aiter Triton/Gluon moe_gemm_a4w4 kernel (gfx1250)
AITER_TRITON_MXFP4_MXFP4 = "AITER_TRITON_MXFP4_MXFP4"
# Triton
TRITON = "TRITON"
TRITON_UNFUSED = "TRITON_UNFUSED"
@@ -310,6 +312,13 @@ def backend_to_kernel_cls(
return [AiterMxfp4Experts]
elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4:
from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a4_moe import (
AiterW4A4ExpertsMonolithic,
)
return [AiterW4A4ExpertsMonolithic]
elif backend == Mxfp4MoeBackend.XPU:
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4
@@ -370,9 +379,11 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16,
Mxfp4MoeBackend.AITER_MXFP4_FP8,
Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4,
],
"aiter_triton_mxfp4_bf16": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16],
"aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],
"aiter_triton_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4],
"aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4],
"xpu": [Mxfp4MoeBackend.XPU],
"cpu": [Mxfp4MoeBackend.CPU],
@@ -464,6 +475,7 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
return kFp8StaticTensorSym
if backend in (
Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4,
Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4,
):
return kMxfp4Dynamic
@@ -896,6 +908,18 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
return hidden_size, intermediate_size
def _interleave_gate_up_rows(t: torch.Tensor) -> torch.Tensor:
"""Repack [E, 2n, X] from [all gates | all ups] to [g0, u0, g1, u1, ...]."""
num_experts, two_n, inner = t.shape
assert two_n % 2 == 0, f"expected an even gate/up row count, got {two_n}"
return (
t.view(num_experts, 2, two_n // 2, inner)
.transpose(1, 2)
.contiguous()
.view(num_experts, two_n, inner)
)
def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
mxfp4_backend: Mxfp4MoeBackend,
layer: torch.nn.Module,
@@ -1192,9 +1216,56 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
w2_bias,
)
elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
from vllm._aiter_ops import rocm_aiter_ops
elif mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4:
if w13_bias is not None:
w13_bias = w13_bias.data.to(torch.float32)
if w2_bias is not None:
w2_bias = w2_bias.data.to(torch.float32)
# Triton moe_gemm_a4w4: uint8 column-major weights.
w13_data = w13_weight.data.view(torch.uint8)
w2_data = w2_weight.data.view(torch.uint8)
w13_scale_data = w13_weight_scale.data.view(torch.uint8)
# SILU applies the activation outside the GEMM and keeps the
# checkpoint's packed halves; the SWIGLUOAI epilogue needs them
# interleaved.
if getattr(layer, "activation", None) in (
MoEActivation.SWIGLUOAI,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
):
w13_data = _interleave_gate_up_rows(w13_data)
w13_scale_data = _interleave_gate_up_rows(w13_scale_data)
# Transpose to [E, K, N] — the view has stride(-2)==1 (column-major),
# which is what moe_gemm_a4w4 asserts.
w13_data = w13_data.transpose(1, 2)
w2_data = w2_data.transpose(1, 2)
# Scales are [E, N, K_scale]; both kernels index them as
# [E, K_scale, N] with K innermost, so transpose(1, 2)
# Keep them uint8: Triton moe_gemm_a4w4 takes e8m0 scales
w13_scale = w13_scale_data.transpose(1, 2)
w2_scale = w2_weight_scale.data.view(torch.uint8).transpose(1, 2)
from aiter.ops.triton.utils.shuffle import shuffle_scale_moe
w13_scale = shuffle_scale_moe(
w13_scale,
arch="gfx1250",
preshuffle_factor=32,
scale_kwidth=4,
)
w2_scale = shuffle_scale_moe(
w2_scale,
arch="gfx1250",
preshuffle_factor=32,
scale_kwidth=4,
)
return (w13_data, w2_data, w13_scale, w2_scale, w13_bias, w2_bias)
elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:
if w13_bias is not None:
w13_bias = w13_bias.data.to(torch.float32)
if w2_bias is not None:
@@ -1203,6 +1274,8 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
# e8m0_shuffle on weight scales (GFX950 swizzle layout)
from aiter.utility.fp4_utils import e8m0_shuffle
from vllm._aiter_ops import rocm_aiter_ops
s0, s1, _ = w13_weight_scale.shape
w13_weight_scale.data = e8m0_shuffle(w13_weight_scale.view(s0 * s1, -1)).view(
s0, s1, -1
@@ -1935,6 +2008,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
)
elif mxfp4_backend in (
Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4,
Mxfp4MoeBackend.AITER_MXFP4_FP8,
):
return convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
@@ -2088,6 +2162,7 @@ def make_mxfp4_moe_quant_config(
)
elif mxfp4_backend in (
Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
Mxfp4MoeBackend.AITER_TRITON_MXFP4_MXFP4,
Mxfp4MoeBackend.CUTLASS_MXFP4_MXFP4,
):
return ocp_mx_moe_quant_config(
vllm/model_executor/layers/fused_moe/routed_experts.py
+2
View File
@@ -110,6 +110,7 @@ class RoutedExperts(PluggableLayer):
swiglu_beta: float | None = None,
e_score_correction_bias: torch.Tensor | None = None,
apply_router_weight_on_input: bool = False,
hash_indices_table: torch.Tensor | None = None,
):
super().__init__()
self.layer_name = layer_name
@@ -155,6 +156,7 @@ class RoutedExperts(PluggableLayer):
# Set by bind_routed_experts_capturer for monolithic kernels.
self.routing_sink: RoutedExpertsSink | None = None
self.apply_router_weight_on_input = apply_router_weight_on_input
self.hash_indices_table = hash_indices_table
# End random parameters
self._loaded_expert_biases: set[str] = set()
self._fused_shared_expert_quantizer = self._get_fused_shared_expert_quantizer()
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3 -1
View File
@@ -1647,8 +1647,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
# If no native backend available, use emulation.
if self.mxfp4_backend is Mxfp4MoeBackend.NONE:
self.mxfp4_backend = Mxfp4MoeBackend.EMULATION
self.experts_cls = None
self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]
if self.experts_cls is None:
self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]
logger.info_once(
f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}"
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+28
View File
@@ -13,6 +13,14 @@ from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_
logger = init_logger(__name__)
def should_use_gfx1250_mx_scale_swizzle() -> bool:
if not current_platform.is_rocm():
return False
from vllm.platforms.rocm import on_gfx1250
return on_gfx1250()
def should_use_cdna4_mx_scale_swizzle() -> bool:
"""Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950.
@@ -27,6 +35,15 @@ def should_use_cdna4_mx_scale_swizzle() -> bool:
return on_gfx950() and get_tensor_model_parallel_world_size() <= 2
def get_aiter_mx_scale_swizzle() -> str | None:
"""AITER `swizzle_mx_scale` arg matching the load-time scale swizzle."""
if should_use_gfx1250_mx_scale_swizzle():
return "GFX1250_SCALE"
if should_use_cdna4_mx_scale_swizzle():
return "CDNA4_SCALE"
return None
def weight_mx_scale(precision_config):
"""mxfp4 weight scale from a triton_kernels PrecisionConfig: 3.8 exposes it
as b_mx_scale, 3.5.1/3.6 as weight_scale.
@@ -136,6 +153,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
# transpose the tensor so that the quantization axis is on dim1
quant_tensor = quant_tensor.transpose(-2, -1)
scale = scale.transpose(-2, -1)
if should_use_gfx1250_mx_scale_swizzle():
from aiter.ops.triton.utils.shuffle import shuffle_scale_moe
assert (
scale.dim() == 3 and scale.shape[-1] % 32 == 0 and scale.shape[-2] % 4 == 0
), f"GFX1250 scale swizzle needs (E,K_SCALE%4,N%32); got {tuple(scale.shape)}"
scale = shuffle_scale_moe(
scale, arch="gfx1250", preshuffle_factor=32, scale_kwidth=4
)
quant_tensor = convert_layout(
wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts
)
vllm/model_executor/layers/utils.py
+7 -6
View File
@@ -225,6 +225,8 @@ direct_register_custom_op(
def use_aiter_triton_gemm(n, m, k, dtype):
from vllm.platforms.rocm import on_gfx1250
if (
not rocm_aiter_ops.is_triton_gemm_enabled()
# MI300's - fp8nuz=True
@@ -233,6 +235,9 @@ def use_aiter_triton_gemm(n, m, k, dtype):
):
return False
if on_gfx1250():
return True
# use hipblaslt for the larger GEMMs
if n > 2048 and m > 512:
return False
@@ -333,7 +338,7 @@ def rocm_unquantized_gemm_impl(
return linear_batch_invariant(x, weight, bias)
from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx11, on_gfx950, on_gfx1250
from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx11, on_gfx950
n = x.numel() // x.size(-1)
m = weight.shape[0]
@@ -366,11 +371,7 @@ def rocm_unquantized_gemm_impl(
x_view = x.reshape(-1, x.size(-1)).contiguous()
return ops.wvSplitKrc(x_view, weight, cu_count, bias)
# gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires
# K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a
# partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b
# hidden), so skip aiter there and fall back to the torch GEMM path below.
if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0):
if use_aiter_triton_gemm(n, m, k, x.dtype):
from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
return gemm_a16w16(x, weight, bias)
vllm/models/deepseek_v4/amd/model.py
+13 -6
View File
@@ -80,7 +80,7 @@ from vllm.models.deepseek_v4.amd.mega_moe import (
)
from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
from vllm.platforms import current_platform
from vllm.platforms.rocm import on_gfx950
from vllm.platforms.rocm import on_gfx950, on_gfx1250
from vllm.sequence import IntermediateTensors
from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset
@@ -142,6 +142,11 @@ class DeepseekV4MLP(nn.Module):
# B-preshuffle the gate_up_proj weight in place (single weight).
if not self._gateup:
return
# aiter's gemm_a8w8_blockscale_bpreshuffle currently disabled on gfx1250.
# Leaving _gateup_scale as None -> "not preshuffled"
# forward() falls back to the standard gate_up_proj linear
if on_gfx1250():
return
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
_upcast_e8m0_to_fp32,
get_fp8_block_weight_scale,
@@ -157,11 +162,13 @@ class DeepseekV4MLP(nn.Module):
return
if ws.dtype == torch.float8_e8m0fnu:
ws = _upcast_e8m0_to_fp32(ws).contiguous()
replace_parameter(
self.gate_up_proj,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
# Skip when the linear kernel already B-preshuffled this weight.
if not getattr(self.gate_up_proj, "aiter_bpreshuffled", False):
replace_parameter(
self.gate_up_proj,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
self._gateup_scale = ws
def forward(self, x):
vllm/models/deepseek_v4/amd/rocm.py
+18 -8
View File
@@ -23,7 +23,7 @@ from vllm.models.deepseek_v4.sparse_mla import (
DeepseekV4SparseMLAMetadataBuilder,
)
from vllm.platforms import current_platform
from vllm.platforms.rocm import _ON_GFX950
from vllm.platforms.rocm import _ON_GFX950, on_gfx1250
from vllm.triton_utils import tl, triton
from vllm.utils.multi_stream_utils import execute_in_parallel
from vllm.v1.attention.backend import (
@@ -912,6 +912,11 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
if not rocm_aiter_ops.is_enabled():
return
# aiter's gemm_a8w8_blockscale_bpreshuffle currently disabled on gfx1250.
# Leaving the block scales as None -> "not preshuffled"
# _fused_wqa_wkv_gemm and _o_proj both fall back to the standard linear path
if on_gfx1250():
return
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
_upcast_e8m0_to_fp32,
get_fp8_block_weight_scale,
@@ -930,12 +935,14 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
return None
if ws.dtype == torch.float8_e8m0fnu:
ws = _upcast_e8m0_to_fp32(ws).contiguous()
# Shuffle the weight in place (single weight, no unshuffled copy).
replace_parameter(
linear,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
# Shuffle the weight in place (single weight, no unshuffled copy),
# unless the linear kernel already B-preshuffled it at load time.
if not getattr(linear, "aiter_bpreshuffled", False):
replace_parameter(
linear,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
return ws
self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
@@ -1300,9 +1307,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
fuse_inv_rope = self._wo_a_fp8_weight is None
rotated = 0
if num_decodes > 0:
fuse_decode_inv_rope = fuse_inv_rope and not on_gfx1250()
rotated = self._forward_decode(
q=q[:num_decode_tokens],
positions=positions[:num_decode_tokens] if fuse_inv_rope else None,
positions=(
positions[:num_decode_tokens] if fuse_decode_inv_rope else None
),
kv_cache=self_kv_cache,
swa_metadata=swa_metadata,
attn_metadata=rocm_metadata,
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+59 -3
View File
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
kFp8StaticTensorSym,
)
from vllm.platforms.rocm import on_gfx1250
from vllm.utils.torch_utils import is_quantized_kv_cache
from vllm.v1.attention.backend import (
AttentionLayer,
@@ -28,7 +29,12 @@ from vllm.v1.attention.backends.rocm_attn import (
RocmAttentionMetadata,
RocmAttentionMetadataBuilder,
)
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheLayout,
KVQuantMode,
get_kv_quant_mode,
)
logger = init_logger(__name__)
@@ -186,12 +192,34 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
# RDNA3 has no low precision MMA, preserve query precision
self.supports_quant_query_input = not on_gfx11()
# Layout flag: must mirror get_kv_cache_shape, which is attn_type
# agnostic because the allocation is shared across layers.
self._shuffle_kv_layout = (
rocm_aiter_ops.is_shuffle_kv_cache_enabled() and on_gfx1250()
)
# Kernel flag: the shuffled reader/writer is causal decoder only.
self._use_shuffle_kv = (
self._shuffle_kv_layout and self.attn_type != AttentionType.ENCODER_DECODER
)
def _split_kv_cache(
self, kv_cache: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
if self._shuffle_kv_layout:
# (2, B, N, H, hs) -> ((B, N, H, hs), (B, N, H, hs))
return kv_cache.unbind(0)
# (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
return kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
def _shuffle_kv_views(
self, key_cache: torch.Tensor, value_cache: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
nb, bs, nkv, hd = key_cache.shape
x = 16 // key_cache.element_size()
kc = key_cache.reshape(nb, nkv, hd // x, bs, x)
vc = value_cache.reshape(nb, nkv, bs // x, hd, x)
return kc, vc
def forward(
self,
layer: torch.nn.Module,
@@ -264,6 +292,9 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
key_cache = key_cache.view(self.fp8_dtype)
value_cache = value_cache.view(self.fp8_dtype)
if self._use_shuffle_kv:
key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache)
cu_seqlens_q = attn_metadata.query_start_loc
seqused_k = attn_metadata.seq_lens
max_seqlen_q = attn_metadata.max_query_len
@@ -291,6 +322,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
v_descale=layer._v_scale,
sinks=self.sinks,
output_scale=output_scale,
shuffled_kv_cache=self._use_shuffle_kv,
)
else:
# The aiter kernel is causal-only. Non-causal cross-attention
@@ -339,6 +371,27 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
return
key_cache, value_cache = self._split_kv_cache(kv_cache)
if self._use_shuffle_kv:
from vllm.v1.attention.backends.rocm_aiter_fa import (
reshape_and_cache_shuffle_triton,
)
if get_kv_quant_mode(self.kv_cache_dtype) == KVQuantMode.FP8_PER_TENSOR:
key_cache = key_cache.view(self.fp8_dtype)
value_cache = value_cache.view(self.fp8_dtype)
_dummy = torch.empty(1, dtype=torch.float32, device=key.device)
reshape_and_cache_shuffle_triton(
key,
value,
key_cache,
value_cache,
slot_mapping,
self.kv_cache_dtype,
_dummy,
_dummy,
)
return
# Reshape the input keys and values and store them in the cache.
ops.reshape_and_cache_flash(
key,
@@ -395,7 +448,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
k_scale=layer._k_scale_cpu,
v_scale=layer._v_scale_cpu,
kv_cache_dtype=self.kv_cache_dtype,
use_shuffle_layout=False,
use_shuffle_layout=self._use_shuffle_kv,
)
def do_qk_norm_mrope_kvcache_update(
@@ -456,13 +509,16 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
# we use direct Q, K, V tensors without caching
return
key_cache, value_cache = self._split_kv_cache(kv_cache)
flash_layout = True
is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
if is_fp8_kv_cache:
key_cache = key_cache.view(self.fp8_dtype)
value_cache = value_cache.view(self.fp8_dtype)
if self._use_shuffle_kv:
key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache)
flash_layout = not self._use_shuffle_kv
rocm_aiter_ops.triton_rope_and_cache(
query,
key,
vllm/v1/sample/ops/topk_topp_sampler.py
+42 -6
View File
@@ -34,11 +34,34 @@ def register_top_k_top_p_warmups() -> None:
_topp_split_mask.register_warmup()
def _skip_aiter_sampler_on_gfx1250() -> bool:
# Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU.
from vllm.platforms.rocm import on_gfx1250
def _aiter_temp_gumbel_sample(
logits: torch.Tensor, generators: dict[int, torch.Generator]
) -> torch.Tensor:
"""Fused temperature Gumbel-max sampling via aiter.
return on_gfx1250()
Logits are already temperature-scaled by the parent Sampler, so temps=1.
Returns int64 token ids of shape (num_tokens,).
"""
from aiter import mixed_sample_outer_exponential
n, vocab = logits.shape
out = torch.empty(n, dtype=torch.int32, device=logits.device)
if not generators:
# Common case: shared (1, vocab) exponential noise (cheap RNG).
exp = (
torch.empty((1, vocab), dtype=torch.float32, device=logits.device)
.exponential_()
.expand(n, vocab)
)
else:
exp = torch.empty((n, vocab), dtype=torch.float32, device=logits.device)
if len(generators) != n:
exp.exponential_()
for i, g in generators.items():
exp[i].exponential_(generator=g)
temps = torch.ones(n, dtype=torch.float32, device=logits.device)
mixed_sample_outer_exponential(out, logits, exp, temps, eps=1e-10)
return out.to(torch.int64)
def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None:
@@ -246,8 +269,12 @@ class TopKTopPSampler(nn.Module):
elif (
logprobs_mode not in PROCESSED_LOGPROBS_MODES
and rocm_aiter_ops.is_enabled()
and not _skip_aiter_sampler_on_gfx1250() # TODO (JPVILLAM): Enable
):
# Lazy ROCm-only import; keeps arch detection out of import time
# on CUDA/CPU.
from vllm.platforms.rocm import on_gfx1250
self.on_gfx1250 = on_gfx1250()
self.aiter_ops = None
self._aiter_ops_import_failed = False
logger.info_once(
@@ -375,7 +402,16 @@ class TopKTopPSampler(nn.Module):
k: torch.Tensor | None,
p: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Optimized ROCm/aiter path (same structure as forward_cuda)."""
"""Optimized ROCm/aiter path."""
if self.on_gfx1250:
if (
k is None
and p is None
and not self.use_fp64_gumbel
and self.logprobs_mode not in PROCESSED_LOGPROBS_MODES
):
return _aiter_temp_gumbel_sample(logits, generators), None
return self.forward_native(logits, generators, k, p)
if (k is None and p is None) or generators:
if generators:
logger.warning_once(
vllm/v1/sample/sampler.py
+61
View File
@@ -5,6 +5,7 @@
import torch
import torch.nn as nn
from vllm._aiter_ops import rocm_aiter_ops
from vllm.config.model import LogprobsMode
from vllm.platforms import current_platform
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
@@ -19,6 +20,61 @@ from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
_SAMPLING_EPS = 1e-5
def mixed_sample_eligible(sm: SamplingMetadata, predict_bonus_token: bool) -> bool:
"""Whether the fused aiter sample-from-raw-logits kernel can be used."""
if not rocm_aiter_ops.is_enabled():
return False
from vllm.platforms.rocm import on_gfx1250
if not on_gfx1250():
return False
if predict_bonus_token:
return False
if not sm.all_random or sm.temperature is None:
return False
if sm.top_k is not None or sm.top_p is not None:
return False
if sm.max_num_logprobs is not None or sm.logprob_token_ids:
return False
if (
not sm.no_penalties
or sm.allowed_token_ids_mask is not None
or sm.bad_words_token_ids
):
return False
lp = sm.logitsprocs
if lp.argmax_invariant or lp.non_argmax_invariant:
return False
h = sm.thinking_budget_state_holder
return h is None or not h.has_tracked_requests()
def _aiter_raw_fused_sample(logits: torch.Tensor, sm: SamplingMetadata) -> torch.Tensor:
"""Sample token ids from raw (unscaled) logits via a fused aiter kernel.
Returns int32 ids of shape (num_tokens, 1).
"""
from aiter import mixed_sample_outer_exponential
n, vocab = logits.shape
out = torch.empty(n, dtype=torch.int32, device=logits.device)
gens = sm.generators
if not gens:
exp = (
torch.empty((1, vocab), dtype=torch.float32, device=logits.device)
.exponential_()
.expand(n, vocab)
)
else:
exp = torch.empty((n, vocab), dtype=torch.float32, device=logits.device)
exp.exponential_()
for i, g in gens.items():
exp[i].exponential_(generator=g)
mixed_sample_outer_exponential(out, logits, exp, sm.temperature, eps=1e-10)
return out.unsqueeze(-1)
class Sampler(nn.Module):
"""A layer that samples the next tokens from the model's outputs
with the following steps in order:
@@ -77,6 +133,11 @@ class Sampler(nn.Module):
predict_bonus_token: bool = False,
logprobs_mode_override: LogprobsMode | None = None,
) -> SamplerOutput:
if mixed_sample_eligible(sampling_metadata, predict_bonus_token):
return SamplerOutput(
sampled_token_ids=_aiter_raw_fused_sample(logits, sampling_metadata),
logprobs_tensors=None,
)
logprobs_mode = logprobs_mode_override or self.logprobs_mode
# NOTE(woosuk): Use the original logits (before any penalties or
# temperature scaling) for the top-k logprobs.