A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 15:15 UTC
[Bugfix][ROCm][GPT-OSS] Use old triton_kernels implementation on ROCm if the new API is not available (#34153)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Feb 9, 2026 at 17:38 UTC
c60f8e3b49eced1a17ba0e11da3f8c107b309df9
1 file changed
+52
-9
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+52
-9
@@ -19,11 +19,14 @@ from vllm.model_executor.layers.fused_moe.utils import _resize_cache
19
from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
QuantKey,
21
)
22
+from vllm.platforms import current_platform
23
from vllm.triton_utils import tl, triton
24
from vllm.utils.import_utils import has_triton_kernels
25
26
logger = init_logger(__name__)
27
28
+use_legacy_triton_kernels = False
29
+
30
if has_triton_kernels():
31
try:
32
import triton_kernels.swiglu
@@ -38,10 +41,20 @@ if has_triton_kernels():
41
from triton_kernels.tensor import (
42
BIT,
43
Bitmatrix,
41
- SparseMatrix,
42
- make_ragged_tensor_metadata,
44
)
45
from triton_kernels.topk import topk
46
+
47
+ try:
48
+ from triton_kernels.tensor import (
49
+ SparseMatrix,
50
+ make_ragged_tensor_metadata,
51
+ )
52
+ except ImportError:
53
+ if current_platform.is_rocm():
54
+ logger.warning_once("Using legacy triton_kernels on ROCm")
55
+ use_legacy_triton_kernels = True
56
+ else:
57
+ raise
58
except (AttributeError, ImportError) as e:
59
logger.error(
60
"Failed to import Triton kernels. Please make sure your triton "
@@ -101,6 +114,12 @@ def legacy_routing_from_bitmatrix(
114
Replacement for the removed triton_kernels.routing.routing_from_bitmatrix.
115
Creates routing data from a bitmatrix representation.
116
"""
117
+ if use_legacy_triton_kernels:
118
+ from triton_kernels.routing import routing_from_bitmatrix
119
+
120
+ return routing_from_bitmatrix(
121
+ bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act
122
+ )
123
sparse_logits = SparseMatrix(indx=expt_indx, vals=expt_scal, mask=bitmatrix)
124
dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
125
combine_indx = sparse_logits.mask_metadata.col_sorted_indx
@@ -130,6 +149,10 @@ def legacy_routing(
149
Replacement for the removed triton_kernels.routing.routing function.
150
Computes routing data from gating logits.
151
"""
152
+ if use_legacy_triton_kernels:
153
+ from triton_kernels.routing import routing
154
+
155
+ return routing(logits, n_expts_act, sm_first=sm_first)
156
if sm_first:
157
logits = torch.softmax(logits, dim=-1)
158
sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first)
@@ -231,11 +254,22 @@ def triton_kernel_fused_experts(
254
)
255
output_tensor = _resize_cache(output_tensor, (batch_dim, M, K))
256
234
- act = FusedActivation(
235
- FnSpecs(
236
- "swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit"), reduction_n=2
237
- ),
238
- (swiglu_alpha, swiglu_limit),
257
+ act = (
258
+ FusedActivation(
259
+ FnSpecs(
260
+ "swiglu",
261
+ triton_kernels.swiglu.swiglu_fn,
262
+ ("alpha", "limit"),
263
+ reduction_n=2,
264
+ ),
265
+ (swiglu_alpha, swiglu_limit),
266
+ )
267
+ if not use_legacy_triton_kernels
268
+ else FusedActivation(
269
+ FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
270
+ (swiglu_alpha, swiglu_limit),
271
+ 2,
272
+ )
273
)
274
gammas = routing_data.gate_scal if routing_data else None
275
@@ -296,8 +330,17 @@ def make_routing_data(
330
331
bitmatrix_shape = [n_rows, bm_cols * 32]
332
bitmatrix_shape_max = [n_rows, None]
299
- bitmatrix = Bitmatrix(
300
- bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
333
+ bitmatrix = (
334
+ Bitmatrix(
335
+ bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
336
+ )
337
+ if not use_legacy_triton_kernels
338
+ else Bitmatrix(
339
+ bitmatrix,
340
+ shape=bitmatrix_shape,
341
+ shape_max=bitmatrix_shape_max,
342
+ scratchpad=None,
343
+ )
344
)
345
346
# matmul_ogs expects invalid topk_weights to be -1s