From e440b75bb6d5d9f23a5403fab22fa8886c91b0a8 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 25 Sep 2026 14:55:52 -0500 Subject: [PATCH] [ROCm] Bump torch 2.13, triton 3.8, torchaudio, torchvision (#50605) Signed-off-by: Rohan Potdar Signed-off-by: Rohan138 Signed-off-by: jpvillam Co-authored-by: Claude Fable 5 Co-authored-by: jpvillam --- cmake/external_projects/triton_kernels.cmake | 4 +- docker/Dockerfile.rocm | 2 +- docker/Dockerfile.rocm_base | 6 +-- .../kernels/linear/mxfp8/rocm_native.py | 8 ++-- .../experts/gpt_oss_triton_kernels_moe.py | 46 +++++++++++++++++-- 5 files changed, 52 insertions(+), 14 deletions(-) diff --git a/cmake/external_projects/triton_kernels.cmake b/cmake/external_projects/triton_kernels.cmake index 8312c1f7ca..1906e6da5c 100644 --- a/cmake/external_projects/triton_kernels.cmake +++ b/cmake/external_projects/triton_kernels.cmake @@ -12,8 +12,8 @@ if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) else() if (VLLM_TARGET_DEVICE STREQUAL "rocm") set(TRITON_GIT "https://github.com/ROCm/triton.git") - # Pinned from release/internal/3.6.x - set(TRITON_KERNELS_TAG "0f380657dbf3ee86eb57558ff71df24f03b5d4e7") + # Pinned from release/internal/3.8.x + set(TRITON_KERNELS_TAG "669b31acc1dd1b3fd93286afd5db67f65d9f7557") else() set(TRITON_GIT "https://github.com/triton-lang/triton.git") set(TRITON_KERNELS_TAG "v3.5.1") diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 5170a71081..ecc2600a71 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -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 diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index eb4bb5cc97..750d4ce4d9 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,9 +1,9 @@ ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="f0b55c0" # release/internal/3.7.x as of 08/18 +ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17 ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="6bbd260" # release/2.12 as of 08/02 +ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.27.1" +ARG PYTORCH_VISION_BRANCH="v0.28.0" ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" ARG PYTORCH_AUDIO_BRANCH="v2.11.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 1f2e8fdf48..faf1e180fa 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -120,12 +120,10 @@ def _mxfp8_dot_scaled_linear( return out -# Triton 3.8 enables TRITON_HIP_USE_ASYNC_COPY by default on gfx950; its extra LDS -# buffer leaves no room for num_stages=3 at BLOCK_K=256. -# TODO(rasmith)(Rohan138): Remove the 3.8 check once -# https://github.com/vllm-project/vllm/pull/50605 merges. +# gfx950's async-copy LDS (default in the triton 3.8 this branch pins) leaves no +# room for num_stages=3 at BLOCK_K=256; use 2 there. _BK256_STAGES = 3 -if triton.__version__.startswith("3.8") and current_platform.is_rocm(): +if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx950 if on_gfx950(): 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 27fdcf8bb2..5bf5065858 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 @@ -1294,9 +1294,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): 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 @@ -1313,6 +1310,23 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): quant_config.w1_precision, gammas=gammas if apply_router_weight_on_input else None, ) + + lora_context = self._lora_context + if lora_context is not None: + s_ids = e_ids = n_pad = t_map = None + s_ids, e_ids, n_pad, t_map = self.apply_w13_lora( + lora_context, + y=inter, + x=hidden_states, + topk_ids=global_topk_ids, + topk_weights=topk_weights, + expert_map=expert_map, + w1=w1, + w2=w2, + num_tokens=M, + top_k_num=topk, + ) + act_out = torch.empty( (routing_data.n_valid, act_out_dim), dtype=hidden_states.dtype, @@ -1330,6 +1344,32 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): quant_config.w2_precision, gammas=None if apply_router_weight_on_input else gammas, ) + + if lora_context is not None: + # w2 LoRA works in (num_tokens, topk, K) token-topk order. Scatter + # the ragged rows there via the same expert-sort permutation + # (flat idx = token*topk + slot = argsort(flat_e)), apply, then + # gather the LoRA-updated rows back to ragged order for the combine. + flat_e = topk_ids.reshape(-1).to(torch.int64).clamp(min=0) + order = torch.argsort(flat_e, stable=True) + down_tt = down.new_zeros((M * topk, K)) + down_tt[order] = down + self.apply_w2_lora( + lora_context, + y=down_tt.view(M, topk, K), + x=act_out, + topk_weights=topk_weights, + sorted_token_ids_lora=s_ids, + expert_ids_lora=e_ids, + num_tokens_post_padded_lora=n_pad, + token_lora_mapping=t_map, + num_tokens=M, + w1=w1, + w2=w2, + top_k_num=topk, + ) + down = down_tt[order] + 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))