[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>
Rohan Potdar committed
Sep 16, 2026 at 20:03 UTC
3e267bae70d85785ede4ae724f87f89edd6905a9
9 files changed
+270
-147
docker/Dockerfile.rock
+5
-3
@@ -5,7 +5,7 @@ ARG VLLM_BRANCH="main"
5
ARG COMMON_WORKDIR=/app
6
ARG BASE_IMAGE=rocm/vllm-dev:base
7
ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
8
-ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7
8
+ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557
9
ARG INSTALL_LMCACHE=false
10
# NIC backend for MoRI RDMA support.
11
# 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"
322
323
RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
324
--mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
325
- grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
325
+ { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
326
cmake/external_projects/triton_kernels.cmake \
327
+ || 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"; } \
328
&& export CCACHE_BASEDIR="$PWD" \
329
&& echo "=== ccache stats before ROCm native build ===" \
330
&& (ccache --show-stats || true) \
@@ -669,8 +670,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
670
671
# Install the packaging helpers before modifying requirements/rocm.txt.
672
RUN --mount=type=cache,target=/root/.cache/uv \
672
- grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
673
+ { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \
674
vllm/cmake/external_projects/triton_kernels.cmake \
675
+ || 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"; } \
676
&& cd vllm \
677
&& uv pip install --system setuptools_scm regex
678
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+1
-2
@@ -384,9 +384,8 @@ def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk):
384
385
make_routing_data = gptoss_moe.make_routing_data
386
routing_data_from_sparse_topk = gptoss_moe.routing_data_from_sparse_topk
387
- use_legacy_triton_kernels = gptoss_moe.use_legacy_triton_kernels
387
389
- if use_legacy_triton_kernels:
388
+ if gptoss_moe.triton_kernels_version == "3.5.1":
389
pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+")
390
391
from triton_kernels.topk import topk as topk_fn
vllm/model_executor/layers/fused_moe/config.py
+8
-3
@@ -17,14 +17,19 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
17
)
18
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
19
from vllm.platforms import current_platform
20
-from vllm.utils.import_utils import has_triton_kernels
20
+from vllm.utils.import_utils import get_triton_kernels_version
21
from vllm.utils.math_utils import cdiv
22
23
logger = init_logger(__name__)
24
25
-if has_triton_kernels():
25
+_triton_kernels_version = get_triton_kernels_version()
26
+
27
+if _triton_kernels_version is not None:
28
try:
27
- from triton_kernels.matmul_ogs import PrecisionConfig
29
+ if _triton_kernels_version == "3.8":
30
+ from triton_kernels.matmul import PrecisionConfig
31
+ else:
32
+ from triton_kernels.matmul_ogs import PrecisionConfig
33
except (ImportError, AttributeError) as e:
34
logger.error(
35
"Failed to import Triton kernels. Please make sure your triton "
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py
+8
-23
@@ -12,9 +12,6 @@ from vllm.model_executor.layers.fused_moe.config import (
12
FusedMoEQuantConfig,
13
RoutingMethodType,
14
)
15
-from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
16
- patch_gating_output,
17
-)
15
from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
QuantKey,
17
kMxfp4Static,
@@ -26,17 +23,6 @@ __all__ = [
23
]
24
25
29
-def _aiter_raw(t):
30
- """Unwrap the `triton_kernels.tensor.Tensor` that `_swizzle_mxfp4` returns."""
31
- if t is None or isinstance(t, torch.Tensor):
32
- return t
33
- assert hasattr(t, "storage"), (
34
- f"expected a triton_kernels wrapped tensor with a .storage attribute, "
35
- f"got {type(t)}"
36
- )
37
- return t.storage.data
38
-
39
-
26
def _aiter_w4a16_silu_via_a8w4(
27
hidden_states: torch.Tensor,
28
w1_data,
@@ -168,11 +154,6 @@ def aiter_triton_kernel_w4a16_moe_forward(
154
_routing_mod.is_tdm_avail = lambda: False
155
aiter_routing = _routing_mod.routing
156
171
- # See context in #50859.
172
- # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged,
173
- # AITER released, and AITER pin in vLLM increased from 0.1.19.
174
- gating_output = patch_gating_output(gating_output, global_num_experts)
175
-
157
if score_mode is not None:
158
use_grouped_topk = num_expert_group is not None and num_expert_group > 1
159
routing_data, gather_idx, scatter_idx = aiter_routing(
@@ -201,10 +182,14 @@ def aiter_triton_kernel_w4a16_moe_forward(
182
assert quant_config.w1_precision is not None
183
assert quant_config.w2_precision is not None
184
204
- w1_data = _aiter_raw(w1)
205
- w2_data = _aiter_raw(w2)
206
- w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale)
207
- w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale)
185
+ w1_data = w1.storage.data
186
+ w2_data = w2.storage.data
187
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
188
+ weight_mx_scale,
189
+ )
190
+
191
+ w1_wscale = weight_mx_scale(quant_config.w1_precision).storage.data
192
+ w2_wscale = weight_mx_scale(quant_config.w2_precision).storage.data
193
194
gammas = routing_data.gate_scal if routing_data else None
195
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+3
-62
@@ -5,7 +5,6 @@ import torch
5
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
from vllm._aiter_ops import rocm_aiter_ops
8
-from vllm.forward_context import get_forward_context, is_forward_context_available
8
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
from vllm.model_executor.layers.fused_moe.config import (
10
FusedMoEConfig,
@@ -25,60 +24,6 @@ __all__ = [
24
]
25
26
28
-# TODO(akaratza): Remove
29
-# _get_padding_mask and patch_gating_output after
30
-# https://github.com/ROCm/aiter/pull/4530 is released and consumed by vLLM.
31
-def _get_padding_mask() -> torch.Tensor | None:
32
- """Retrieves a boolean mask with non-padding (0) and padding (1) tokens.
33
-
34
- `slot_mapping < 0` comes from:
35
-
36
- slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)
37
-
38
- in gpu_model_runner.py.
39
-
40
- Direct kernel callers do not have model-runner padding metadata, so no
41
- mask is needed when there is no forward context.
42
- """
43
- if not is_forward_context_available():
44
- return None
45
-
46
- forward_context = get_forward_context()
47
-
48
- # model runner v2.
49
- if forward_context.is_padding is not None:
50
- return forward_context.is_padding
51
-
52
- # model runner v1.
53
- slot_mapping = forward_context.slot_mapping
54
-
55
- if isinstance(slot_mapping, list):
56
- slot_mapping_dict = slot_mapping[0]
57
- else:
58
- slot_mapping_dict = slot_mapping
59
-
60
- if isinstance(slot_mapping_dict, dict):
61
- slot_mapping_sample = next(iter(slot_mapping_dict.values()), None)
62
-
63
- if isinstance(slot_mapping_sample, torch.Tensor):
64
- return slot_mapping_sample < 0
65
- else:
66
- return None
67
-
68
-
69
-def patch_gating_output(
70
- gating_output: torch.Tensor, global_num_experts: int
71
-) -> torch.Tensor:
72
- if global_num_experts != 128:
73
- is_padding = _get_padding_mask()
74
-
75
- if is_padding is not None:
76
- assert is_padding.shape[0] == gating_output.shape[0]
77
- gating_output = gating_output.masked_fill(is_padding[:, None], 0.0)
78
-
79
- return gating_output
80
-
81
-
27
def aiter_triton_kernel_w4a8_moe_forward(
28
hidden_states: torch.Tensor,
29
w1, # Tensor or triton_kernels.Tensor
@@ -112,11 +57,6 @@ def aiter_triton_kernel_w4a8_moe_forward(
57
_routing_mod.is_tdm_avail = lambda: False
58
aiter_routing = _routing_mod.routing
59
115
- # See context in #50859.
116
- # TODO: Remove once https://github.com/ROCm/aiter/pull/4530 is merged,
117
- # AITER released, and AITER pin in vLLM increased from 0.1.19.
118
- gating_output = patch_gating_output(gating_output, global_num_experts)
119
-
60
routing_data, gather_idx, scatter_idx = aiter_routing(
61
gating_output, topk, sm_first=not renormalize
62
)
@@ -191,6 +131,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
131
132
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
133
should_use_cdna4_mx_scale_swizzle,
134
+ weight_mx_scale,
135
)
136
137
_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(
151
hidden_states,
152
w1.storage.data,
153
None,
213
- quant_config.w1_precision.weight_scale.storage.data,
154
+ weight_mx_scale(quant_config.w1_precision).storage.data,
155
quant_config.w1_precision.flex_ctx.lhs_data.scale,
156
quant_config.w2_precision.flex_ctx.lhs_data.scale,
157
quant_config.w1_bias,
@@ -230,7 +171,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
171
intermediate_cache1,
172
w2.storage.data,
173
None,
233
- quant_config.w2_precision.weight_scale.storage.data,
174
+ weight_mx_scale(quant_config.w2_precision).storage.data,
175
quant_config.w2_precision.flex_ctx.lhs_data.scale,
176
None,
177
quant_config.w2_bias,
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+176
-38
@@ -26,7 +26,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
26
)
27
from vllm.platforms import current_platform
28
from vllm.triton_utils import tl, triton
29
-from vllm.utils.import_utils import has_triton_kernels
29
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
30
31
from ..utils import swiglu_limit_func
32
@@ -441,48 +441,59 @@ def _patch_legacy_routing_for_nonpow2_topk() -> None:
441
_routing.sort_tokens = _sort_tokens_pow2
442
443
444
-# Two API generations of triton_kernels are supported:
445
-# - v3.5.1 (the version bundled with vLLM): exposes `routing()` and
446
-# `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix`
447
-# constructor takes a `scratchpad` argument.
448
-# - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix`
449
-# based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only
450
-# when the user has triton_kernels installed system-wide at v3.6.0+.
451
-#
452
-# `use_legacy_triton_kernels` selects between them at import time based on
453
-# whether `SparseMatrix` is importable.
454
-use_legacy_triton_kernels = False
444
+triton_kernels_version = get_triton_kernels_version()
445
456
-if has_triton_kernels():
446
+if triton_kernels_version is not None:
447
try:
448
import triton_kernels.swiglu
459
- from triton_kernels.matmul_ogs import (
460
- FnSpecs,
461
- FusedActivation,
462
- GatherIndx,
463
- RoutingData,
464
- ScatterIndx,
465
- matmul_ogs,
466
- )
467
- from triton_kernels.tensor import (
468
- BIT,
469
- Bitmatrix,
470
- )
449
450
+ if triton_kernels_version == "3.8":
451
+ from triton_kernels.matmul import FnSpecs, FusedActivation, matmul
452
+ from triton_kernels.matmul_details.opt_flags import (
453
+ update_opt_flags_constraints,
454
+ )
455
+ from triton_kernels.tensor import make_ragged_tensor_metadata
456
+
457
+ from vllm.platforms.rocm import on_gfx950
458
+
459
+ if on_gfx950():
460
+ # block_n=64/block_k=256 fit gfx950's 160KB LDS + the CDNA4
461
+ # scale swizzle; block_m=32 avoids an autotuner block_m=128
462
+ # padding cliff (~15x) at decode-sized ragged MoE shapes.
463
+ update_opt_flags_constraints(
464
+ {"block_m": 32, "block_n": 64, "block_k": 256}
465
+ )
466
+ else:
467
+ from triton_kernels.matmul_ogs import (
468
+ FnSpecs,
469
+ FusedActivation,
470
+ GatherIndx,
471
+ RoutingData,
472
+ ScatterIndx,
473
+ matmul_ogs,
474
+ )
475
+
476
+ if triton_kernels_version == "3.6":
477
+ from triton_kernels.tensor import (
478
+ SparseMatrix,
479
+ make_ragged_tensor_metadata,
480
+ )
481
+
482
+ # BIT/Bitmatrix exist in v3.5.1/v3.6, gone in v3.8.
483
try:
484
from triton_kernels.tensor import (
474
- SparseMatrix,
475
- make_ragged_tensor_metadata,
485
+ BIT,
486
+ Bitmatrix,
487
)
488
except ImportError:
478
- # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once
479
- # the gpt-oss perf regression in v3.6.0+ is resolved upstream.
480
- # Tracking: https://github.com/triton-lang/triton/issues/9969
481
- use_legacy_triton_kernels = True
482
- if not use_legacy_triton_kernels:
489
+ BIT = None
490
+ Bitmatrix = None
491
+
492
+ # Patch the routing kernels for non-pow2 top_k (DeepSeek-V4 top_k=6);
493
+ # v3.8 rebuilds routing from topk ids and needs no patch.
494
+ if triton_kernels_version == "3.6":
495
_patch_make_bitmatrix_metadata()
484
- else:
485
- # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4).
496
+ elif triton_kernels_version == "3.5.1":
497
_patch_legacy_routing_for_nonpow2_topk()
498
except (AttributeError, ImportError) as e:
499
logger.error(
@@ -492,6 +503,21 @@ if has_triton_kernels():
503
)
504
505
506
+if triton_kernels_version == "3.8":
507
+ # 3.8's matmul dropped matmul_ogs's in-kernel scatter+topk-reduce, so the
508
+ # forward combines externally via index_add; this is the ragged routing
509
+ # bundle the (inlined) 3.8 paths pass around, in place of matmul_ogs's.
510
+ class RoutingData: # type: ignore[no-redef]
511
+ __slots__ = ("ragged", "gate_scal", "gather_tok", "n_expts_act", "n_valid")
512
+
513
+ def __init__(self, ragged, gate_scal, gather_tok, n_expts_act, n_valid):
514
+ self.ragged = ragged
515
+ self.gate_scal = gate_scal
516
+ self.gather_tok = gather_tok
517
+ self.n_expts_act = n_expts_act
518
+ self.n_valid = n_valid
519
+
520
+
521
@triton.jit
522
def pack_bitmatrix(
523
bitmatrix,
@@ -562,7 +588,7 @@ def triton_kernel_moe_forward(
588
# three separate kernels used by the generic path below.
589
# Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines
590
# equivalent logic via SparseMatrix in `make_routing_data`.
565
- if use_legacy_triton_kernels and expert_map is None:
591
+ if triton_kernels_version == "3.5.1" and expert_map is None:
592
from triton_kernels.routing import routing as fused_routing
593
594
routing_data, gather_idx, scatter_idx = fused_routing(
@@ -662,6 +688,52 @@ def triton_kernel_fused_experts(
688
assert activation == MoEActivation.SWIGLUOAI, (
689
"Only SWIGLUOAI activation is supported"
690
)
691
+
692
+ if triton_kernels_version == "3.8":
693
+ # 3.8 has no in-kernel scatter+topk-reduce, so the two matmuls are
694
+ # followed by an external index_add combine.
695
+ assert quant_config is not None
696
+ M, K = hidden_states.shape[-2:]
697
+ a_ragged = routing_data.ragged
698
+ gather_tok = routing_data.gather_tok
699
+ gammas = routing_data.gate_scal
700
+ act = FusedActivation(
701
+ FnSpecs(
702
+ "swiglu",
703
+ triton_kernels.swiglu.swiglu_fn,
704
+ ("alpha", "limit"),
705
+ reduction_n=2,
706
+ ),
707
+ (swiglu_alpha, swiglu_limit),
708
+ )
709
+ inter = matmul(
710
+ hidden_states,
711
+ w1,
712
+ quant_config.w1_bias,
713
+ a_ragged,
714
+ None,
715
+ gather_tok,
716
+ None,
717
+ quant_config.w1_precision,
718
+ gammas=gammas if apply_router_weight_on_input else None,
719
+ fused_activation=act,
720
+ )
721
+ y = matmul(
722
+ inter,
723
+ w2,
724
+ quant_config.w2_bias,
725
+ a_ragged,
726
+ None,
727
+ None,
728
+ None,
729
+ quant_config.w2_precision,
730
+ gammas=None if apply_router_weight_on_input else gammas,
731
+ )
732
+ acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device)
733
+ acc.index_add_(0, gather_tok.to(torch.int64), y.to(torch.float32))
734
+ out = output_tensor.view(M, K)
735
+ out.copy_(acc.to(out.dtype))
736
+ return out
737
assert quant_config is not None
738
739
# type check, uint8 means mxfp4
@@ -704,7 +776,7 @@ def triton_kernel_fused_experts(
776
),
777
(swiglu_alpha, swiglu_limit),
778
)
707
- if not use_legacy_triton_kernels
779
+ if triton_kernels_version != "3.5.1"
780
else FusedActivation(
781
FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
782
(swiglu_alpha, swiglu_limit),
@@ -744,6 +816,24 @@ def make_routing_data(
816
topk_weights: torch.Tensor,
817
num_local_experts: int,
818
) -> tuple["RoutingData", torch.Tensor, torch.Tensor]:
819
+ if triton_kernels_version == "3.8":
820
+ _, topk = topk_ids.shape
821
+ flat_e = topk_ids.reshape(-1)
822
+ valid = flat_e >= 0
823
+ flat_e = flat_e.to(torch.int64).clamp(min=0)
824
+ flat_w = topk_weights.reshape(-1).to(torch.float32) * valid.to(torch.float32)
825
+ nrows = flat_e.shape[0]
826
+ order = torch.argsort(flat_e, stable=True)
827
+ col_sum = torch.zeros(
828
+ num_local_experts, dtype=torch.int32, device=flat_e.device
829
+ )
830
+ col_sum.scatter_add_(0, flat_e, torch.ones_like(flat_e, dtype=torch.int32))
831
+ ragged = make_ragged_tensor_metadata(col_sum, nrows)
832
+ gather_tok = (order // topk).to(torch.int32)
833
+ gammas = flat_w[order]
834
+ routing = RoutingData(ragged, gammas, gather_tok, topk, nrows)
835
+ return routing, gather_tok, gather_tok
836
+
837
topk_ids = topk_ids.to(torch.int16)
838
topk_weights = topk_weights.to(torch.bfloat16)
839
@@ -774,7 +864,7 @@ def make_routing_data(
864
Bitmatrix(
865
bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
866
)
777
- if not use_legacy_triton_kernels
867
+ if triton_kernels_version != "3.5.1"
868
else Bitmatrix(
869
bitmatrix,
870
shape=bitmatrix_shape,
@@ -786,7 +876,7 @@ def make_routing_data(
876
# matmul_ogs expects invalid topk_weights to be -1s
877
topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
878
789
- if use_legacy_triton_kernels:
879
+ if triton_kernels_version == "3.5.1":
880
from triton_kernels.routing import routing_from_bitmatrix
881
882
return routing_from_bitmatrix(
@@ -809,6 +899,10 @@ def routing_data_from_sparse_topk(
899
call already produced; when the topk ids are used as-is (no expert-map
900
remapping) that work is redundant.
901
"""
902
+ if triton_kernels_version == "3.8":
903
+ # 3.8 routing is rebuilt from raw ids; no SparseMatrix reuse.
904
+ return make_routing_data(sparse_topk.indx, sparse_topk.vals, n_expts)
905
+
906
dispatch_indx = sparse_topk.mask_metadata.row_sorted_indx
907
combine_indx = sparse_topk.mask_metadata.col_sorted_indx
908
ragged_batch_metadata = make_ragged_tensor_metadata(
@@ -1197,6 +1291,50 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1291
1292
topk = topk_ids.size(1)
1293
1294
+ if triton_kernels_version == "3.8":
1295
+ # Activation and topk-reduce kept unfused from the matmuls (3.8 has
1296
+ # no in-kernel scatter+reduce); combine via external index_add.
1297
+ assert self._lora_context is None, (
1298
+ "tk38 unfused MoE path does not support LoRA yet"
1299
+ )
1300
+ act_out_dim = self.adjust_N_for_activation(w1.shape[2], activation)
1301
+ M, K = hidden_states.shape[-2:]
1302
+ a_ragged = routing_data.ragged
1303
+ gather_tok = routing_data.gather_tok
1304
+ gammas = routing_data.gate_scal
1305
+ inter = matmul(
1306
+ hidden_states,
1307
+ w1,
1308
+ quant_config.w1_bias,
1309
+ a_ragged,
1310
+ None,
1311
+ gather_tok,
1312
+ None,
1313
+ quant_config.w1_precision,
1314
+ gammas=gammas if apply_router_weight_on_input else None,
1315
+ )
1316
+ act_out = torch.empty(
1317
+ (routing_data.n_valid, act_out_dim),
1318
+ dtype=hidden_states.dtype,
1319
+ device=hidden_states.device,
1320
+ )
1321
+ self.activation(activation, act_out, inter)
1322
+ down = matmul(
1323
+ act_out,
1324
+ w2,
1325
+ quant_config.w2_bias,
1326
+ a_ragged,
1327
+ None,
1328
+ None,
1329
+ None,
1330
+ quant_config.w2_precision,
1331
+ gammas=None if apply_router_weight_on_input else gammas,
1332
+ )
1333
+ acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device)
1334
+ acc.index_add_(0, gather_tok.to(torch.int64), down.to(torch.float32))
1335
+ output.view(M, K).copy_(acc.to(output.dtype))
1336
+ return
1337
+
1338
# type check, uint8 means mxfp4
1339
assert hidden_states.dtype == torch.bfloat16
1340
assert (
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+23
-14
@@ -45,7 +45,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
)
46
from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d
47
from vllm.platforms import current_platform
48
-from vllm.utils.import_utils import has_triton_kernels
48
+from vllm.utils.import_utils import get_triton_kernels_version
49
from vllm.utils.math_utils import round_up
50
51
if TYPE_CHECKING:
@@ -54,9 +54,14 @@ if TYPE_CHECKING:
54
55
logger = init_logger(__name__)
56
57
-if has_triton_kernels():
57
+triton_kernels_version = get_triton_kernels_version()
58
+
59
+if triton_kernels_version is not None:
60
try:
59
- from triton_kernels.matmul_ogs import PrecisionConfig
61
+ if triton_kernels_version == "3.8":
62
+ from triton_kernels.matmul import FlexCtx, PrecisionConfig
63
+ else:
64
+ from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
65
except (ImportError, AttributeError) as e:
66
logger.error(
67
"Failed to import Triton kernels. Please make sure your triton "
@@ -65,6 +70,15 @@ if has_triton_kernels():
70
)
71
72
73
+def _mx_scale_kwargs(scale):
74
+ """PrecisionConfig weight-scale kwargs: 3.8 uses b_mx_scale/b_microblock_size,
75
+ 3.5.1/3.6 use weight_scale.
76
+ """
77
+ if triton_kernels_version == "3.8":
78
+ return {"b_mx_scale": scale, "b_microblock_size": 32}
79
+ return {"weight_scale": scale}
80
+
81
+
82
def _pack_deepgemm_mxfp4_scales(
83
w13_weight: torch.Tensor,
84
w2_weight: torch.Tensor,
@@ -1182,7 +1196,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1196
1197
elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
1198
# W4A8: MXFP4 weights + static FP8 activations (triton kernel)
1185
- from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1199
from triton_kernels.numerics import InFlexData
1200
1201
if w13_bias is not None:
@@ -1217,11 +1230,11 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1230
1231
# Create PrecisionConfig with both weight and activation info
1232
w13_precision_config = PrecisionConfig(
1220
- weight_scale=w13_scale,
1233
+ **_mx_scale_kwargs(w13_scale),
1234
flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
1235
)
1236
w2_precision_config = PrecisionConfig(
1224
- weight_scale=w2_scale,
1237
+ **_mx_scale_kwargs(w2_scale),
1238
flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
1239
)
1240
@@ -1238,8 +1251,6 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1251
)
1252
1253
elif mxfp4_backend in TRITON_BACKENDS:
1241
- from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1242
-
1254
if w13_bias is not None:
1255
w13_bias = w13_bias.to(torch.float32)
1256
if w2_bias is not None:
@@ -1255,10 +1266,10 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1266
)
1267
1268
w13_precision_config = PrecisionConfig(
1258
- weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
1269
+ **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex)
1270
)
1271
w2_precision_config = PrecisionConfig(
1261
- weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
1272
+ **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
1273
)
1274
1275
# The original mxfp4 block scales have been swizzled into the
@@ -1653,8 +1664,6 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1664
elif mxfp4_backend in TRITON_BACKENDS or (
1665
mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
1666
):
1656
- from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1657
-
1667
if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16:
1668
# AITER moe_gemm_a16w4 needs gate/up interleaved
1669
def interleave_gate_up(w: torch.Tensor) -> torch.Tensor:
@@ -1698,10 +1707,10 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1707
)
1708
1709
w13_precision_config = PrecisionConfig(
1701
- weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
1710
+ **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex)
1711
)
1712
w2_precision_config = PrecisionConfig(
1704
- weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
1713
+ **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex)
1714
)
1715
1716
# The original mxfp4 block scales have been swizzled into the
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+26
-2
@@ -7,7 +7,7 @@ import torch
7
from vllm.logger import init_logger
8
from vllm.platforms import current_platform
9
from vllm.triton_utils import tl, triton
10
-from vllm.utils.import_utils import has_triton_kernels
10
+from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels
11
from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
12
13
logger = init_logger(__name__)
@@ -27,10 +27,23 @@ def should_use_cdna4_mx_scale_swizzle() -> bool:
27
return on_gfx950() and get_tensor_model_parallel_world_size() <= 2
28
29
30
+def weight_mx_scale(precision_config):
31
+ """mxfp4 weight scale from a triton_kernels PrecisionConfig: 3.8 exposes it
32
+ as b_mx_scale, 3.5.1/3.6 as weight_scale.
33
+ """
34
+ if get_triton_kernels_version() == "3.8":
35
+ return precision_config.b_mx_scale
36
+ return precision_config.weight_scale
37
+
38
+
39
def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
40
"""Weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel."""
41
assert has_triton_kernels()
33
- import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
42
+ if get_triton_kernels_version() == "3.8":
43
+ # 3.8: matmul_ogs_details -> matmul_details
44
+ import triton_kernels.matmul_details.opt_flags as opt_flags
45
+ else:
46
+ import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
47
from triton_kernels.numerics import InFlexData
48
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
49
from triton_kernels.tensor_details import layout
@@ -51,6 +64,17 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
64
)
65
value_layout = StridedLayout
66
scale_layout = StridedLayout
67
+ elif current_platform.is_rocm() and get_triton_kernels_version() == "3.8":
68
+ # StridedLayout(-2) is triton_kernels' default CDNA mxfp4 weight layout
69
+ # (make_default_matmul_mxfp4_w_layout); 3.8 wants layout instances.
70
+ from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout
71
+
72
+ value_layout = StridedLayout(-2)
73
+ scale_layout = (
74
+ CDNA4MXScaleLayout()
75
+ if should_use_cdna4_mx_scale_swizzle()
76
+ else StridedLayout()
77
+ )
78
elif current_platform.is_rocm():
79
value_layout = StridedLayout
80
if should_use_cdna4_mx_scale_swizzle():
vllm/utils/import_utils.py
+20
@@ -527,6 +527,26 @@ def has_triton_kernels() -> bool:
527
return is_available
528
529
530
+@cache
531
+def get_triton_kernels_version() -> str | None:
532
+ """The triton_kernels MoE-API generation ("3.5.1"/"3.6"/"3.8"), or None.
533
+
534
+ Inferred by capability since the package exposes no usable version: 3.8
535
+ replaced ``matmul_ogs`` with ``matmul``, and 3.5.1 predates ``SparseMatrix``.
536
+ """
537
+ if not has_triton_kernels():
538
+ return None
539
+ try:
540
+ import triton_kernels.matmul_ogs # noqa: F401
541
+ except ImportError:
542
+ return "3.8"
543
+ try:
544
+ from triton_kernels.tensor import SparseMatrix # noqa: F401
545
+ except ImportError:
546
+ return "3.5.1"
547
+ return "3.6"
548
+
549
+
550
@cache
551
def has_tilelang() -> bool:
552
"""Whether the optional `tilelang` package is available.