From 3e267bae70d85785ede4ae724f87f89edd6905a9 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Wed, 16 Sep 2026 20:03:34 -0500 Subject: [PATCH] [ROCm] triton+triton_kernels 3.8 mxfp4 MoE support (gpt-oss + DeepSeek-V4) (#55934) Signed-off-by: Rohan Potdar Signed-off-by: Rohan138 Co-authored-by: Claude Opus 4.8 Co-authored-by: Matt <156021403+mawong-amd@users.noreply.github.com> --- docker/Dockerfile.rock | 8 +- .../moe/test_gpt_oss_triton_kernels.py | 3 +- .../model_executor/layers/fused_moe/config.py | 11 +- .../experts/aiter_mxfp4_w4a16_moe.py | 31 +-- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 65 +----- .../experts/gpt_oss_triton_kernels_moe.py | 214 ++++++++++++++---- .../layers/fused_moe/oracle/mxfp4.py | 37 +-- .../layers/quantization/utils/mxfp4_utils.py | 28 ++- vllm/utils/import_utils.py | 20 ++ 9 files changed, 270 insertions(+), 147 deletions(-) diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index 14ba523b07..32623ff1a9 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -5,7 +5,7 @@ ARG VLLM_BRANCH="main" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base -ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7 +ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557 ARG INSTALL_LMCACHE=false # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types @@ -322,8 +322,9 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && export CCACHE_BASEDIR="$PWD" \ && echo "=== ccache stats before ROCm native build ===" \ && (ccache --show-stats || true) \ @@ -669,8 +670,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Install the packaging helpers before modifying requirements/rocm.txt. RUN --mount=type=cache,target=/root/.cache/uv \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ vllm/cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rock ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ && cd vllm \ && uv pip install --system setuptools_scm regex diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 4afcba706c..86028a7feb 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -384,9 +384,8 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): make_routing_data = gptoss_moe.make_routing_data routing_data_from_sparse_topk = gptoss_moe.routing_data_from_sparse_topk - use_legacy_triton_kernels = gptoss_moe.use_legacy_triton_kernels - if use_legacy_triton_kernels: + if gptoss_moe.triton_kernels_version == "3.5.1": pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+") from triton_kernels.topk import topk as topk_fn diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index ccca24cc6f..9a2f8229b1 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -17,14 +17,19 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version from vllm.utils.math_utils import cdiv logger = init_logger(__name__) -if has_triton_kernels(): +_triton_kernels_version = get_triton_kernels_version() + +if _triton_kernels_version is not None: try: - from triton_kernels.matmul_ogs import PrecisionConfig + if _triton_kernels_version == "3.8": + from triton_kernels.matmul import PrecisionConfig + else: + from triton_kernels.matmul_ogs import PrecisionConfig except (ImportError, AttributeError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " 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 757d1d8295..f6afd6fbfa 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 @@ -12,9 +12,6 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, RoutingMethodType, ) -from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( - patch_gating_output, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -26,17 +23,6 @@ __all__ = [ ] -def _aiter_raw(t): - """Unwrap the `triton_kernels.tensor.Tensor` that `_swizzle_mxfp4` returns.""" - if t is None or isinstance(t, torch.Tensor): - return t - assert hasattr(t, "storage"), ( - f"expected a triton_kernels wrapped tensor with a .storage attribute, " - f"got {type(t)}" - ) - return t.storage.data - - def _aiter_w4a16_silu_via_a8w4( hidden_states: torch.Tensor, w1_data, @@ -168,11 +154,6 @@ def aiter_triton_kernel_w4a16_moe_forward( _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing - # See context in #50859. - # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged, - # AITER released, and AITER pin in vLLM increased from 0.1.19. - gating_output = patch_gating_output(gating_output, global_num_experts) - if 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( @@ -201,10 +182,14 @@ def aiter_triton_kernel_w4a16_moe_forward( assert quant_config.w1_precision is not None assert quant_config.w2_precision is not None - w1_data = _aiter_raw(w1) - w2_data = _aiter_raw(w2) - w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) - w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + w1_data = w1.storage.data + w2_data = w2.storage.data + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + weight_mx_scale, + ) + + w1_wscale = weight_mx_scale(quant_config.w1_precision).storage.data + w2_wscale = weight_mx_scale(quant_config.w2_precision).storage.data gammas = routing_data.gate_scal if routing_data else None 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 0691d9fed4..b327cbeb59 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 @@ -5,7 +5,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops -from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -25,60 +24,6 @@ __all__ = [ ] -# TODO(akaratza): Remove -# _get_padding_mask and patch_gating_output after -# https://github.com/ROCm/aiter/pull/4530 is released and consumed by vLLM. -def _get_padding_mask() -> torch.Tensor | None: - """Retrieves a boolean mask with non-padding (0) and padding (1) tokens. - - `slot_mapping < 0` comes from: - - slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1) - - in gpu_model_runner.py. - - Direct kernel callers do not have model-runner padding metadata, so no - mask is needed when there is no forward context. - """ - if not is_forward_context_available(): - return None - - forward_context = get_forward_context() - - # model runner v2. - if forward_context.is_padding is not None: - return forward_context.is_padding - - # model runner v1. - slot_mapping = forward_context.slot_mapping - - if isinstance(slot_mapping, list): - slot_mapping_dict = slot_mapping[0] - else: - slot_mapping_dict = slot_mapping - - if isinstance(slot_mapping_dict, dict): - slot_mapping_sample = next(iter(slot_mapping_dict.values()), None) - - if isinstance(slot_mapping_sample, torch.Tensor): - return slot_mapping_sample < 0 - else: - return None - - -def patch_gating_output( - gating_output: torch.Tensor, global_num_experts: int -) -> torch.Tensor: - if global_num_experts != 128: - is_padding = _get_padding_mask() - - if is_padding is not None: - assert is_padding.shape[0] == gating_output.shape[0] - gating_output = gating_output.masked_fill(is_padding[:, None], 0.0) - - return gating_output - - def aiter_triton_kernel_w4a8_moe_forward( hidden_states: torch.Tensor, w1, # Tensor or triton_kernels.Tensor @@ -112,11 +57,6 @@ def aiter_triton_kernel_w4a8_moe_forward( _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing - # See context in #50859. - # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged, - # AITER released, and AITER pin in vLLM increased from 0.1.19. - gating_output = patch_gating_output(gating_output, global_num_experts) - routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize ) @@ -191,6 +131,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( should_use_cdna4_mx_scale_swizzle, + weight_mx_scale, ) _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None @@ -210,7 +151,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( hidden_states, w1.storage.data, None, - quant_config.w1_precision.weight_scale.storage.data, + weight_mx_scale(quant_config.w1_precision).storage.data, quant_config.w1_precision.flex_ctx.lhs_data.scale, quant_config.w2_precision.flex_ctx.lhs_data.scale, quant_config.w1_bias, @@ -230,7 +171,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( intermediate_cache1, w2.storage.data, None, - quant_config.w2_precision.weight_scale.storage.data, + weight_mx_scale(quant_config.w2_precision).storage.data, quant_config.w2_precision.flex_ctx.lhs_data.scale, None, quant_config.w2_bias, diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index e2b4987bd0..27fdcf8bb2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -26,7 +26,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels from ..utils import swiglu_limit_func @@ -441,48 +441,59 @@ def _patch_legacy_routing_for_nonpow2_topk() -> None: _routing.sort_tokens = _sort_tokens_pow2 -# Two API generations of triton_kernels are supported: -# - v3.5.1 (the version bundled with vLLM): exposes `routing()` and -# `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix` -# constructor takes a `scratchpad` argument. -# - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix` -# based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only -# when the user has triton_kernels installed system-wide at v3.6.0+. -# -# `use_legacy_triton_kernels` selects between them at import time based on -# whether `SparseMatrix` is importable. -use_legacy_triton_kernels = False +triton_kernels_version = get_triton_kernels_version() -if has_triton_kernels(): +if triton_kernels_version is not None: try: import triton_kernels.swiglu - from triton_kernels.matmul_ogs import ( - FnSpecs, - FusedActivation, - GatherIndx, - RoutingData, - ScatterIndx, - matmul_ogs, - ) - from triton_kernels.tensor import ( - BIT, - Bitmatrix, - ) + if triton_kernels_version == "3.8": + from triton_kernels.matmul import FnSpecs, FusedActivation, matmul + from triton_kernels.matmul_details.opt_flags import ( + update_opt_flags_constraints, + ) + from triton_kernels.tensor import make_ragged_tensor_metadata + + from vllm.platforms.rocm import on_gfx950 + + if on_gfx950(): + # block_n=64/block_k=256 fit gfx950's 160KB LDS + the CDNA4 + # scale swizzle; block_m=32 avoids an autotuner block_m=128 + # padding cliff (~15x) at decode-sized ragged MoE shapes. + update_opt_flags_constraints( + {"block_m": 32, "block_n": 64, "block_k": 256} + ) + else: + from triton_kernels.matmul_ogs import ( + FnSpecs, + FusedActivation, + GatherIndx, + RoutingData, + ScatterIndx, + matmul_ogs, + ) + + if triton_kernels_version == "3.6": + from triton_kernels.tensor import ( + SparseMatrix, + make_ragged_tensor_metadata, + ) + + # BIT/Bitmatrix exist in v3.5.1/v3.6, gone in v3.8. try: from triton_kernels.tensor import ( - SparseMatrix, - make_ragged_tensor_metadata, + BIT, + Bitmatrix, ) except ImportError: - # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once - # the gpt-oss perf regression in v3.6.0+ is resolved upstream. - # Tracking: https://github.com/triton-lang/triton/issues/9969 - use_legacy_triton_kernels = True - if not use_legacy_triton_kernels: + BIT = None + Bitmatrix = None + + # Patch the routing kernels for non-pow2 top_k (DeepSeek-V4 top_k=6); + # v3.8 rebuilds routing from topk ids and needs no patch. + if triton_kernels_version == "3.6": _patch_make_bitmatrix_metadata() - else: - # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4). + elif triton_kernels_version == "3.5.1": _patch_legacy_routing_for_nonpow2_topk() except (AttributeError, ImportError) as e: logger.error( @@ -492,6 +503,21 @@ if has_triton_kernels(): ) +if triton_kernels_version == "3.8": + # 3.8's matmul dropped matmul_ogs's in-kernel scatter+topk-reduce, so the + # forward combines externally via index_add; this is the ragged routing + # bundle the (inlined) 3.8 paths pass around, in place of matmul_ogs's. + class RoutingData: # type: ignore[no-redef] + __slots__ = ("ragged", "gate_scal", "gather_tok", "n_expts_act", "n_valid") + + def __init__(self, ragged, gate_scal, gather_tok, n_expts_act, n_valid): + self.ragged = ragged + self.gate_scal = gate_scal + self.gather_tok = gather_tok + self.n_expts_act = n_expts_act + self.n_valid = n_valid + + @triton.jit def pack_bitmatrix( bitmatrix, @@ -562,7 +588,7 @@ def triton_kernel_moe_forward( # three separate kernels used by the generic path below. # Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines # equivalent logic via SparseMatrix in `make_routing_data`. - if use_legacy_triton_kernels and expert_map is None: + if triton_kernels_version == "3.5.1" and expert_map is None: from triton_kernels.routing import routing as fused_routing routing_data, gather_idx, scatter_idx = fused_routing( @@ -662,6 +688,52 @@ def triton_kernel_fused_experts( assert activation == MoEActivation.SWIGLUOAI, ( "Only SWIGLUOAI activation is supported" ) + + if triton_kernels_version == "3.8": + # 3.8 has no in-kernel scatter+topk-reduce, so the two matmuls are + # followed by an external index_add combine. + assert quant_config is not None + M, K = hidden_states.shape[-2:] + a_ragged = routing_data.ragged + gather_tok = routing_data.gather_tok + gammas = routing_data.gate_scal + act = FusedActivation( + FnSpecs( + "swiglu", + triton_kernels.swiglu.swiglu_fn, + ("alpha", "limit"), + reduction_n=2, + ), + (swiglu_alpha, swiglu_limit), + ) + inter = matmul( + hidden_states, + w1, + quant_config.w1_bias, + a_ragged, + None, + gather_tok, + None, + quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=act, + ) + y = matmul( + inter, + w2, + quant_config.w2_bias, + a_ragged, + None, + None, + None, + quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + ) + acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device) + acc.index_add_(0, gather_tok.to(torch.int64), y.to(torch.float32)) + out = output_tensor.view(M, K) + out.copy_(acc.to(out.dtype)) + return out assert quant_config is not None # type check, uint8 means mxfp4 @@ -704,7 +776,7 @@ def triton_kernel_fused_experts( ), (swiglu_alpha, swiglu_limit), ) - if not use_legacy_triton_kernels + if triton_kernels_version != "3.5.1" else FusedActivation( FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), (swiglu_alpha, swiglu_limit), @@ -744,6 +816,24 @@ def make_routing_data( topk_weights: torch.Tensor, num_local_experts: int, ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: + if triton_kernels_version == "3.8": + _, topk = topk_ids.shape + flat_e = topk_ids.reshape(-1) + valid = flat_e >= 0 + flat_e = flat_e.to(torch.int64).clamp(min=0) + flat_w = topk_weights.reshape(-1).to(torch.float32) * valid.to(torch.float32) + nrows = flat_e.shape[0] + order = torch.argsort(flat_e, stable=True) + col_sum = torch.zeros( + num_local_experts, dtype=torch.int32, device=flat_e.device + ) + col_sum.scatter_add_(0, flat_e, torch.ones_like(flat_e, dtype=torch.int32)) + ragged = make_ragged_tensor_metadata(col_sum, nrows) + gather_tok = (order // topk).to(torch.int32) + gammas = flat_w[order] + routing = RoutingData(ragged, gammas, gather_tok, topk, nrows) + return routing, gather_tok, gather_tok + topk_ids = topk_ids.to(torch.int16) topk_weights = topk_weights.to(torch.bfloat16) @@ -774,7 +864,7 @@ def make_routing_data( Bitmatrix( bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max ) - if not use_legacy_triton_kernels + if triton_kernels_version != "3.5.1" else Bitmatrix( bitmatrix, shape=bitmatrix_shape, @@ -786,7 +876,7 @@ def make_routing_data( # matmul_ogs expects invalid topk_weights to be -1s topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights) - if use_legacy_triton_kernels: + if triton_kernels_version == "3.5.1": from triton_kernels.routing import routing_from_bitmatrix return routing_from_bitmatrix( @@ -809,6 +899,10 @@ def routing_data_from_sparse_topk( call already produced; when the topk ids are used as-is (no expert-map remapping) that work is redundant. """ + if triton_kernels_version == "3.8": + # 3.8 routing is rebuilt from raw ids; no SparseMatrix reuse. + return make_routing_data(sparse_topk.indx, sparse_topk.vals, n_expts) + dispatch_indx = sparse_topk.mask_metadata.row_sorted_indx combine_indx = sparse_topk.mask_metadata.col_sorted_indx ragged_batch_metadata = make_ragged_tensor_metadata( @@ -1197,6 +1291,50 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): topk = topk_ids.size(1) + if triton_kernels_version == "3.8": + # Activation and topk-reduce kept unfused from the matmuls (3.8 has + # no in-kernel scatter+reduce); combine via external index_add. + assert self._lora_context is None, ( + "tk38 unfused MoE path does not support LoRA yet" + ) + act_out_dim = self.adjust_N_for_activation(w1.shape[2], activation) + M, K = hidden_states.shape[-2:] + a_ragged = routing_data.ragged + gather_tok = routing_data.gather_tok + gammas = routing_data.gate_scal + inter = matmul( + hidden_states, + w1, + quant_config.w1_bias, + a_ragged, + None, + gather_tok, + None, + quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + ) + act_out = torch.empty( + (routing_data.n_valid, act_out_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self.activation(activation, act_out, inter) + down = matmul( + act_out, + w2, + quant_config.w2_bias, + a_ragged, + None, + None, + None, + quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + ) + acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device) + acc.index_add_(0, gather_tok.to(torch.int64), down.to(torch.float32)) + output.view(M, K).copy_(acc.to(output.dtype)) + return + # type check, uint8 means mxfp4 assert hidden_states.dtype == torch.bfloat16 assert ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 912b713918..ce19ba39bf 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -45,7 +45,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d from vllm.platforms import current_platform -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version from vllm.utils.math_utils import round_up if TYPE_CHECKING: @@ -54,9 +54,14 @@ if TYPE_CHECKING: logger = init_logger(__name__) -if has_triton_kernels(): +triton_kernels_version = get_triton_kernels_version() + +if triton_kernels_version is not None: try: - from triton_kernels.matmul_ogs import PrecisionConfig + if triton_kernels_version == "3.8": + from triton_kernels.matmul import FlexCtx, PrecisionConfig + else: + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig except (ImportError, AttributeError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " @@ -65,6 +70,15 @@ if has_triton_kernels(): ) +def _mx_scale_kwargs(scale): + """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size, + 3.5.1/3.6 use weight_scale. + """ + if triton_kernels_version == "3.8": + return {"b_mx_scale": scale, "b_microblock_size": 32} + return {"weight_scale": scale} + + def _pack_deepgemm_mxfp4_scales( w13_weight: torch.Tensor, w2_weight: torch.Tensor, @@ -1182,7 +1196,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: # W4A8: MXFP4 weights + static FP8 activations (triton kernel) - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig from triton_kernels.numerics import InFlexData if w13_bias is not None: @@ -1217,11 +1230,11 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( # Create PrecisionConfig with both weight and activation info w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, + **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13), ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, + **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), ) @@ -1238,8 +1251,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend in TRITON_BACKENDS: - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig - if w13_bias is not None: w13_bias = w13_bias.to(torch.float32) if w2_bias is not None: @@ -1255,10 +1266,10 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( ) w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex) + **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex) ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex) + **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) ) # The original mxfp4 block scales have been swizzled into the @@ -1653,8 +1664,6 @@ def convert_weight_to_mxfp4_moe_kernel_format( elif mxfp4_backend in TRITON_BACKENDS or ( mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 ): - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig - if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: # AITER moe_gemm_a16w4 needs gate/up interleaved def interleave_gate_up(w: torch.Tensor) -> torch.Tensor: @@ -1698,10 +1707,10 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex) + **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex) ) w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex) + **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) ) # The original mxfp4 block scales have been swizzled into the diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 1bd6b3f931..a1d0501775 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -7,7 +7,7 @@ import torch from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_triton_kernels +from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer logger = init_logger(__name__) @@ -27,10 +27,23 @@ def should_use_cdna4_mx_scale_swizzle() -> bool: return on_gfx950() and get_tensor_model_parallel_world_size() <= 2 +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. + """ + if get_triton_kernels_version() == "3.8": + return precision_config.b_mx_scale + return precision_config.weight_scale + + def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): """Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel.""" assert has_triton_kernels() - import triton_kernels.matmul_ogs_details.opt_flags as opt_flags + if get_triton_kernels_version() == "3.8": + # 3.8: matmul_ogs_details -> matmul_details + import triton_kernels.matmul_details.opt_flags as opt_flags + else: + import triton_kernels.matmul_ogs_details.opt_flags as opt_flags from triton_kernels.numerics import InFlexData from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout @@ -51,6 +64,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): ) value_layout = StridedLayout scale_layout = StridedLayout + elif current_platform.is_rocm() and get_triton_kernels_version() == "3.8": + # StridedLayout(-2) is triton_kernels' default CDNA mxfp4 weight layout + # (make_default_matmul_mxfp4_w_layout); 3.8 wants layout instances. + from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout + + value_layout = StridedLayout(-2) + scale_layout = ( + CDNA4MXScaleLayout() + if should_use_cdna4_mx_scale_swizzle() + else StridedLayout() + ) elif current_platform.is_rocm(): value_layout = StridedLayout if should_use_cdna4_mx_scale_swizzle(): diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index d6a2aabbb4..4e90367a3e 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -527,6 +527,26 @@ def has_triton_kernels() -> bool: return is_available +@cache +def get_triton_kernels_version() -> str | None: + """The triton_kernels MoE-API generation ("3.5.1"/"3.6"/"3.8"), or None. + + Inferred by capability since the package exposes no usable version: 3.8 + replaced ``matmul_ogs`` with ``matmul``, and 3.5.1 predates ``SparseMatrix``. + """ + if not has_triton_kernels(): + return None + try: + import triton_kernels.matmul_ogs # noqa: F401 + except ImportError: + return "3.8" + try: + from triton_kernels.tensor import SparseMatrix # noqa: F401 + except ImportError: + return "3.5.1" + return "3.6" + + @cache def has_tilelang() -> bool: """Whether the optional `tilelang` package is available.