From e711723352670e811b0e904300d771471e28930d Mon Sep 17 00:00:00 2001 From: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:51:57 -0500 Subject: [PATCH] [ROCm] gfx1250 optimizations and model enablement (#58800) Signed-off-by: jpvillam Signed-off-by: Jaden Mathias Signed-off-by: Juan Villamizar Signed-off-by: jpvillam Signed-off-by: Omar Muhammad Signed-off-by: Daniel Signed-off-by: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> Co-authored-by: Jaden Mathias Co-authored-by: jpvillam Co-authored-by: Omar Muhammad Co-authored-by: Daniel Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Daniel Chan --- docker/Dockerfile.rocm_base_gfx1250 | 38 +- docker/Dockerfile.rocm_gfx1250 | 8 +- tests/kernels/moe/test_ocp_mx_moe.py | 15 +- vllm/_aiter_ops.py | 63 +++ vllm/config/kernel.py | 3 + .../kernels/linear/scaled_mm/aiter.py | 39 +- .../experts/aiter_mxfp4_w4a16_moe.py | 10 +- .../fused_moe/experts/aiter_mxfp4_w4a4_moe.py | 373 ++++++++++++++++++ .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 10 +- vllm/model_executor/layers/fused_moe/layer.py | 1 + .../layers/fused_moe/modular_kernel.py | 14 + .../layers/fused_moe/oracle/mxfp4.py | 81 +++- .../layers/fused_moe/routed_experts.py | 2 + .../layers/quantization/quark/quark_moe.py | 4 +- .../layers/quantization/utils/mxfp4_utils.py | 28 ++ vllm/model_executor/layers/utils.py | 13 +- vllm/models/deepseek_v4/amd/model.py | 19 +- vllm/models/deepseek_v4/amd/rocm.py | 26 +- .../backends/rocm_aiter_unified_attn.py | 62 ++- vllm/v1/sample/ops/topk_topp_sampler.py | 48 ++- vllm/v1/sample/sampler.py | 61 +++ 21 files changed, 845 insertions(+), 73 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py diff --git a/docker/Dockerfile.rocm_base_gfx1250 b/docker/Dockerfile.rocm_base_gfx1250 index 3fad6800ab..ae41096a6e 100644 --- a/docker/Dockerfile.rocm_base_gfx1250 +++ b/docker/Dockerfile.rocm_base_gfx1250 @@ -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 \ diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index 0609c73c04..634e522e0d 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -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"] diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 0a8f5c29d0..968781a188 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -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 diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 41a76bfc60..ee146f2a57 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -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: diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 5c4b877fd7..a1b141237f 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -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) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 03e9be51cf..23a5886ab7 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index fc9f72ad10..4706fb50b7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -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, diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py new file mode 100644 index 0000000000..a40d8184fb --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a4_moe.py @@ -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, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index bc1a4af6c0..f7d2be94aa 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -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" diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index a51a1e9d0b..7be89aa4c8 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -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: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 99928afb13..ef3319bf12 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -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, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3abc0cbede..ae79298eea 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -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( diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 89a401eba3..2d7ab1e6aa 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -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() diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 849b0c9638..439b48a3e6 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -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}" diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 45562b12b1..074ee9c95f 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -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 ) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index b89a6443ad..ef4594e113 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -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) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index e7894f319f..f933cffcda 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -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): diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3dad09dd08..5353b13ebc 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -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, diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 9f8e9e9ab1..4a9f67a637 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -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, diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index daa94ebc3a..d38497b2ff 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -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( diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index d19f65f6d8..a9cb2f1a23 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -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.