[ROCm] triton+triton_kernels 3.8 mxfp4 MoE support (gpt-oss + DeepSeek-V4) (#55934)

Signed-off-by: Rohan Potdar <rohan.potdar@amd.com>
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Matt <156021403+mawong-amd@users.noreply.github.com>

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
 

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

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 "

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
 

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,

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 (

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

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():

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.