[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>

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")

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

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"

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

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