[ROCm] Bump torch 2.13, triton 3.8, torchaudio, torchvision (#50605)
Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Signed-off-by: jpvillam <juan.villamizar@amd.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: jpvillam <juan.villamizar@amd.com>
Rohan Potdar committed
Sep 25, 2026 at 14:55 UTC
e440b75bb6d5d9f23a5403fab22fa8886c91b0a8
5 files changed
+52
-14
cmake/external_projects/triton_kernels.cmake
+2
-2
@@ -12,8 +12,8 @@ if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
12
else()
13
if (VLLM_TARGET_DEVICE STREQUAL "rocm")
14
set(TRITON_GIT "https://github.com/ROCm/triton.git")
15
- # Pinned from release/internal/3.6.x
16
- set(TRITON_KERNELS_TAG "0f380657dbf3ee86eb57558ff71df24f03b5d4e7")
15
+ # Pinned from release/internal/3.8.x
16
+ set(TRITON_KERNELS_TAG "669b31acc1dd1b3fd93286afd5db67f65d9f7557")
17
else()
18
set(TRITON_GIT "https://github.com/triton-lang/triton.git")
19
set(TRITON_KERNELS_TAG "v3.5.1")
docker/Dockerfile.rocm
+1
-1
@@ -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
docker/Dockerfile.rocm_base
+3
-3
@@ -1,9 +1,9 @@
1
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
2
-ARG TRITON_BRANCH="f0b55c0" # release/internal/3.7.x as of 08/18
2
+ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17
3
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
4
-ARG PYTORCH_BRANCH="6bbd260" # release/2.12 as of 08/02
4
+ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17
5
ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
6
-ARG PYTORCH_VISION_BRANCH="v0.27.1"
6
+ARG PYTORCH_VISION_BRANCH="v0.28.0"
7
ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
8
ARG PYTORCH_AUDIO_BRANCH="v2.11.0"
9
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+3
-5
@@ -120,12 +120,10 @@ def _mxfp8_dot_scaled_linear(
120
return out
121
122
123
-# Triton 3.8 enables TRITON_HIP_USE_ASYNC_COPY by default on gfx950; its extra LDS
124
-# buffer leaves no room for num_stages=3 at BLOCK_K=256.
125
-# TODO(rasmith)(Rohan138): Remove the 3.8 check once
126
-# https://github.com/vllm-project/vllm/pull/50605 merges.
123
+# gfx950's async-copy LDS (default in the triton 3.8 this branch pins) leaves no
124
+# room for num_stages=3 at BLOCK_K=256; use 2 there.
125
_BK256_STAGES = 3
128
-if triton.__version__.startswith("3.8") and current_platform.is_rocm():
126
+if current_platform.is_rocm():
127
from vllm.platforms.rocm import on_gfx950
128
129
if on_gfx950():
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+43
-3
@@ -1294,9 +1294,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
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
- )
1297
act_out_dim = self.adjust_N_for_activation(w1.shape[2], activation)
1298
M, K = hidden_states.shape[-2:]
1299
a_ragged = routing_data.ragged
@@ -1313,6 +1310,23 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1310
quant_config.w1_precision,
1311
gammas=gammas if apply_router_weight_on_input else None,
1312
)
1313
+
1314
+ lora_context = self._lora_context
1315
+ if lora_context is not None:
1316
+ s_ids = e_ids = n_pad = t_map = None
1317
+ s_ids, e_ids, n_pad, t_map = self.apply_w13_lora(
1318
+ lora_context,
1319
+ y=inter,
1320
+ x=hidden_states,
1321
+ topk_ids=global_topk_ids,
1322
+ topk_weights=topk_weights,
1323
+ expert_map=expert_map,
1324
+ w1=w1,
1325
+ w2=w2,
1326
+ num_tokens=M,
1327
+ top_k_num=topk,
1328
+ )
1329
+
1330
act_out = torch.empty(
1331
(routing_data.n_valid, act_out_dim),
1332
dtype=hidden_states.dtype,
@@ -1330,6 +1344,32 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1344
quant_config.w2_precision,
1345
gammas=None if apply_router_weight_on_input else gammas,
1346
)
1347
+
1348
+ if lora_context is not None:
1349
+ # w2 LoRA works in (num_tokens, topk, K) token-topk order. Scatter
1350
+ # the ragged rows there via the same expert-sort permutation
1351
+ # (flat idx = token*topk + slot = argsort(flat_e)), apply, then
1352
+ # gather the LoRA-updated rows back to ragged order for the combine.
1353
+ flat_e = topk_ids.reshape(-1).to(torch.int64).clamp(min=0)
1354
+ order = torch.argsort(flat_e, stable=True)
1355
+ down_tt = down.new_zeros((M * topk, K))
1356
+ down_tt[order] = down
1357
+ self.apply_w2_lora(
1358
+ lora_context,
1359
+ y=down_tt.view(M, topk, K),
1360
+ x=act_out,
1361
+ topk_weights=topk_weights,
1362
+ sorted_token_ids_lora=s_ids,
1363
+ expert_ids_lora=e_ids,
1364
+ num_tokens_post_padded_lora=n_pad,
1365
+ token_lora_mapping=t_map,
1366
+ num_tokens=M,
1367
+ w1=w1,
1368
+ w2=w2,
1369
+ top_k_num=topk,
1370
+ )
1371
+ down = down_tt[order]
1372
+
1373
acc = torch.zeros((M, K), dtype=torch.float32, device=hidden_states.device)
1374
acc.index_add_(0, gather_tok.to(torch.int64), down.to(torch.float32))
1375
output.view(M, K).copy_(acc.to(output.dtype))