[Docker] Fix CUTLASS DSL cu13 install order in Dockerfile (#45204)

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
(cherry picked from commit 40e065e86a)

docker/Dockerfile

@@ -218,6 +218,10 @@ COPY requirements/common.txt requirements/common.txt
 COPY requirements/cuda.txt requirements/cuda.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
+# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
+# share paths with different content. uv can extract them in either order,
+# leaving base files that break CUDA 13 CuTe DSL JIT.
+# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
         sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
@@ -234,6 +238,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     else \
         uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+    fi \
+    && if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
+        CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
+        if [ -n "$CUTLASS_DSL_VERSION" ]; then \
+            uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \
+                "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
+        fi; \
     fi
 
 # Track PyTorch lib versions used during build and match in downstream instances.
@@ -745,6 +756,10 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
 ARG PYTORCH_CUDA_INDEX_BASE_URL
 COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
+# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
+# share paths with different content. uv can extract them in either order,
+# leaving base files that break CUDA 13 CuTe DSL JIT.
+# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
         sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
@@ -752,6 +767,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     fi && \
     uv pip install --system -r /tmp/requirements-cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
+    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
+        CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
+        if [ -n "$CUTLASS_DSL_VERSION" ]; then \
+            uv pip install --system --force-reinstall --no-deps \
+                "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
+        fi; \
+    fi && \
     rm /tmp/requirements-cuda.txt /tmp/common.txt
 
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
@@ -842,6 +864,19 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
     uv pip install --system ep_kernels/dist/*.whl --verbose \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
+# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
+# share paths with different content. Force -libs-cu13 last after runtime
+# dependency installs so uv cannot leave base files behind.
+# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
+RUN --mount=type=cache,target=/opt/uv/cache \
+    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
+        CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
+        if [ -n "$CUTLASS_DSL_VERSION" ]; then \
+            uv pip install --system --force-reinstall --no-deps \
+                "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
+        fi; \
+    fi
+
 # Download FlashInfer precompiled cubins AFTER all pip installs are done.
 # This must run after the vLLM wheel and EP kernels installs above, because
 # those can reinstall/touch flashinfer packages. Downloading cubins earlier

requirements/rocm.txt

@@ -23,6 +23,7 @@ timm>=1.0.17
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
 tilelang==0.1.10
-
+# Required apache-tvm-ffi matching tilelang version
+apache-tvm-ffi==0.1.10 
 # Required for faster safetensors model loading
 fastsafetensors >= 0.3.2

requirements/test/rocm.txt

@@ -44,6 +44,7 @@ anyio==4.13.0
     #   watchfiles
 apache-tvm-ffi==0.1.10
     # via
+    #   -c requirements/rocm.txt
     #   tilelang
     #   xgrammar
 arctic-inference==0.1.1

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

@@ -340,6 +340,16 @@ def rocm_aiter_fused_experts(
             moe_config.intermediate_size_per_partition
             - moe_config.intermediate_size_per_partition_unpadded
         )
+        # Round hidden_pad/intermediate_pad to match AITER's CK/FlyDSL MoE
+        # dispatch (currently pinned to v0.1.13.post1):
+        # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1073
+        # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099
+        # TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes
+        # for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401
+        hidden_pad = hidden_pad // 128 * 128
+        intermediate_pad = (
+            intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
+        )
 
         return rocm_aiter_ops.fused_moe(
             hidden_states,
@@ -357,8 +367,8 @@ def rocm_aiter_fused_experts(
             doweight_stage1=apply_router_weight_on_input,
             num_local_tokens=num_local_tokens,
             output_dtype=output_dtype,
-            hidden_pad=hidden_pad // 128 * 128,
-            intermediate_pad=intermediate_pad // 64 * 64 * 2,
+            hidden_pad=hidden_pad,
+            intermediate_pad=intermediate_pad,
             bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
             bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
             moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -242,7 +242,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 vllm_config.model_config.max_model_len,
                 vllm_config.scheduler_config.max_num_batched_tokens,
             )
-            self._init_fp8_prefill_ps_buffers(max_num_reqs, max_prefill_qlen, device)
+            self._init_fp8_prefill_ps_buffers(
+                max_num_reqs,
+                max_prefill_qlen,
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                device,
+            )
 
         if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
             self.paged_kv_indptr = torch.zeros(
@@ -257,21 +262,29 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         self,
         max_num_reqs: int,
         max_prefill_qlen: int,
+        max_num_batched_tokens: int,
         device: torch.device,
     ) -> None:
         """Pre-allocate persistent buffers for FP8 MLA prefill PS metadata.
 
         Uses ``get_ps_metadata_info_v1`` with max values so the buffers are
         large enough for any batch.  ``get_ps_metadata_v1`` fills them
-        per-batch in ``build()``.
+        per-batch in ``build()``.  The FP8 prefill forward path also uses the
+        global workspace manager for per-call scratch, so reserve its maximum
+        shape here before the workspace manager is locked after warmup.
 
         Args:
             max_num_reqs: Maximum number of concurrent requests.
             max_prefill_qlen: Maximum Q-length for a single request in one
                 prefill batch.  Should be ``min(max_model_len,
-                max_num_batched_tokens)`` — the chunked-prefill scheduler
-                never emits more than ``max_num_batched_tokens`` new tokens
-                per batch.
+                max_num_batched_tokens)`` — a single request never exceeds
+                ``max_model_len`` tokens, nor the per-batch token budget.
+            max_num_batched_tokens: Maximum number of tokens scheduled in one
+                batch.  The ``final_lse`` scratch is sized by ``total_q`` (the
+                summed Q-length over all prefill requests in the batch), which
+                is bounded by this budget rather than by a single request's
+                ``max_prefill_qlen`` — concurrent requests can sum to more than
+                ``max_model_len`` when ``max_model_len < max_num_batched_tokens``.
             device: Target device for the buffers.
         """
         from aiter import get_ps_metadata_info_v1
@@ -279,6 +292,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         # After kv_b_proj decompression, K has num_heads heads (same as Q).
         # So gqa_ratio=1 and num_head_k=num_heads for the PS kernel.
         num_head_k = self.num_heads
+        v_head_dim = self.mla_dims.v_head_dim
         # gqa_ratio = 1
         # qlen_granularity = _FP8_PREFILL_TILE_Q // max(gqa_ratio, 1)
         qlen_granularity = _FP8_PREFILL_TILE_Q
@@ -318,6 +332,21 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             device=device,
         )
 
+        from vllm.v1.worker.workspace import current_workspace_manager
+
+        max_num_partial_tiles = reduce_partial_map_size
+        current_workspace_manager().get_simultaneous(
+            (
+                (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k, v_head_dim),
+                torch.float32,
+            ),
+            (
+                (max_num_partial_tiles * _FP8_PREFILL_TILE_Q, num_head_k),
+                torch.float32,
+            ),
+            ((max_num_batched_tokens, num_head_k), torch.float32),
+        )
+
         logger.info(
             "FP8 MLA prefill PS buffers allocated "
             "(max_batch=%d, max_qlen=%d, num_head_k=%d)",