Upstream develop

.github/workflows/test-recipes.yml

@@ -7,6 +7,7 @@ on:
       - 'run-recipe.py'
       - 'run-recipe.sh'
       - 'launch-cluster.sh'
+      - 'mods/uma-fix/**'
       - 'recipes/**'
       - 'tests/**'
       - '.github/workflows/test-recipes.yml'
@@ -15,6 +16,7 @@ on:
       - 'run-recipe.py'
       - 'run-recipe.sh'
       - 'launch-cluster.sh'
+      - 'mods/uma-fix/**'
       - 'recipes/**'
       - 'tests/**'
       - '.github/workflows/test-recipes.yml'
@@ -52,6 +54,9 @@ jobs:
         ./tests/test_recipes.sh -v
         ./tests/test_launch_cluster_image_sync.sh
         ./tests/test_launch_cluster_vllm_pr.sh
+
+    - name: Run uma-fix compatibility test
+      run: python tests/test_uma_fix_mod.py -v
 
     - name: Verify all recipes with dry-run
       run: |

Dockerfile

@@ -60,6 +60,7 @@ RUN apt update && \
     python3-dev python3-pip git wget \
     libibverbs1 libibverbs-dev rdma-core \
     ccache devscripts debhelper fakeroot \
+    cuda-compat-13-3 \
     && rm -rf /var/lib/apt/lists/* \
     && pip install uv
 
@@ -90,6 +91,8 @@ ENV CCACHE_COMPRESS=1
 # Tell CMake to use ccache for compilation
 ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache
 ENV CMAKE_CUDA_COMPILER_LAUNCHER=ccache
+# Use CUDA 13.3 user-mode driver/JIT libraries on older host drivers.
+ENV LD_LIBRARY_PATH=/usr/local/cuda-13.3/compat:${LD_LIBRARY_PATH}
 
 # 2. Set Environment Variables
 ARG TORCH_CUDA_ARCH_LIST="12.1a"
@@ -718,6 +721,7 @@ ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
 ARG FLASHINFER_CUDA_ARCH_LIST="12.1a"
 ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST}
 ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
+ENV LD_LIBRARY_PATH=/usr/local/cuda-13.3/compat:${LD_LIBRARY_PATH}
 ENV TIKTOKEN_ENCODINGS_BASE=$VLLM_BASE_DIR/tiktoken_encodings
 ENV PATH=$VLLM_BASE_DIR:$PATH
 

mods/diffusiongemma/diffusiongemma-support.patch

@@ -3103,6 +3103,19 @@ index cf0c1d417..6a48b6282 100644
 
          # Dict to store valid cudagraph dispatching keys.
          self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = {
+@@ -215,7 +211,11 @@ class CudagraphDispatcher:
+             cudagraph_capture_sizes_for_decode = [
+                 x
+                 for x in self.compilation_config.cudagraph_capture_sizes
+-                if x <= max_num_tokens and x >= uniform_decode_query_len
++                if (
++                    x <= max_num_tokens
++                    and x >= uniform_decode_query_len
++                    and x % uniform_decode_query_len == 0
++                )
+             ]
+             for bs, num_active_loras in product(
+                 cudagraph_capture_sizes_for_decode, lora_cases
 diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
 index 08c814ab3..6258c8199 100644
 --- a/vllm/v1/engine/core.py

mods/exp-gb300-trtllm/README.md

@@ -0,0 +1,44 @@
+# Experimental GB300 TRTLLM startup mod
+
+This runtime mod accelerates the post-load TensorRT-LLM NVFP4 MoE weight-layout
+conversion when selected expert weights have been offloaded with vLLM's UVA
+backend on a Grace-Blackwell system.
+
+vLLM normally sees a UVA mapping as a CUDA tensor and therefore performs the
+TRTLLM expert permutation directly against host LPDDR. This mod detects the
+`FLASHINFER_TRTLLM` NVFP4 MoE backend and uses vLLM's internal
+`_vllm_is_uva_offloaded` parameter marker to clone every parameter that was
+actually UVA-offloaded for the current MoE layer into HBM. It runs the existing
+conversion unchanged, then lets vLLM copy the converted tensors back to pinned
+LPDDR before proceeding to the next layer.
+
+The internal marker is the source of truth after vLLM applies both
+`--cpu-offload-params` and `--cpu-offload-gb`. Consequently, the mod does not
+assume names such as `w13_weight` or `w2_weight`; it also handles broader
+selectors such as `experts`, future parameter names, and a final layer that is
+only partially offloaded because the byte budget was reached.
+
+The staging is bounded to one MoE layer at a time. It does not change runtime
+weight placement, KV-cache placement, model output, or the selected inference
+kernel. It does temporarily increase load-time HBM usage by the staged layer
+plus TRTLLM conversion intermediates.
+
+Apply it like the other spark-vllm mods:
+
+```bash
+./launch-cluster.sh --solo \
+  --apply-mod mods/exp-gb300-trtllm \
+  exec vllm serve nvidia/GLM-5.2-NVFP4 ...
+```
+
+The patch targets the source layout used by vLLM 0.25.1 and fails closed if its
+expected anchors are not found. It is idempotent. During model loading, each
+staged layer emits a line similar to:
+
+```text
+Staged 2 UVA-offloaded TRTLLM NVFP4 MoE parameters (4.50 GiB) in HBM for post-load layout conversion.
+```
+
+If loading runs out of HBM before KV-cache allocation, remove this mod; do not
+raise `--cpu-offload-gb` merely to compensate, because that does not increase
+the temporary HBM headroom available during conversion.

mods/exp-gb300-trtllm/patch_vllm.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+"""Patch vLLM to stage marked UVA-offloaded TRTLLM MoE weights in HBM."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import NoReturn
+
+
+PREFIX = "[exp-gb300-trtllm]"
+MARKER = "exp-gb300-trtllm: stage marked TRTLLM MoE UVA weights in HBM"
+TARGET_RELATIVE_PATH = Path("vllm/model_executor/model_loader/utils.py")
+
+
+def fail(message: str) -> NoReturn:
+    raise SystemExit(f"{PREFIX} {message}")
+
+
+def replace_once(text: str, old: str, new: str, description: str) -> str:
+    count = text.count(old)
+    if count != 1:
+        fail(
+            f"expected exactly one source anchor for {description}, found {count}; "
+            "refusing to patch an unknown vLLM source layout"
+        )
+    return text.replace(old, new, 1)
+
+
+def main() -> None:
+    if len(sys.argv) != 2:
+        fail("usage: patch_vllm.py SITE_PACKAGES")
+
+    site_packages = Path(sys.argv[1]).resolve()
+    target = site_packages / TARGET_RELATIVE_PATH
+    if not target.is_file():
+        fail(f"vLLM source file not found: {target}")
+
+    original = target.read_text()
+    if MARKER in original:
+        print(f"{PREFIX} Patch is already applied; skipping.")
+        return
+
+    text = original
+
+    text = replace_once(
+        text,
+        """            with device_loading_context(module, target_device):
+                quant_method.process_weights_after_loading(module)
+""",
+        f"""            # {MARKER}
+            # A CUDA UVA view reports itself as a CUDA tensor, so the generic
+            # loading context otherwise leaves these large, randomly indexed
+            # tensors in host LPDDR during TRTLLM layout conversion.
+            nvfp4_backend = getattr(quant_method, "nvfp4_backend", None)
+            stage_trtllm_moe_uva = (
+                getattr(nvfp4_backend, "name", None) == "FLASHINFER_TRTLLM"
+            )
+            with device_loading_context(
+                module,
+                target_device,
+                stage_trtllm_moe_uva=stage_trtllm_moe_uva,
+            ):
+                quant_method.process_weights_after_loading(module)
+""",
+        "quantized post-load device context",
+    )
+
+    text = replace_once(
+        text,
+        """@contextmanager
+def device_loading_context(module: torch.nn.Module, target_device: torch.device):
+""",
+        """@contextmanager
+def device_loading_context(
+    module: torch.nn.Module,
+    target_device: torch.device,
+    *,
+    stage_trtllm_moe_uva: bool = False,
+):
+""",
+        "device_loading_context signature",
+    )
+
+    text = replace_once(
+        text,
+        """    original_device_states: dict[str, torch.device] = {}
+    uva_offloaded_parameters: list[str] = []
+
+    # Store original device states and move parameters to GPU if they're on CPU
+""",
+        """    original_device_states: dict[str, torch.device] = {}
+    uva_offloaded_parameters: list[str] = []
+    staged_uva_backings: list[torch.Tensor] = []
+    staged_uva_bytes = 0
+    staged_uva_count = 0
+
+    # Keep the original UVA views alive until the converted HBM tensors have
+    # been copied back to new pinned CPU storage in the finally block.
+    # Store original device states and move parameters to GPU if they're on CPU.
+""",
+        "UVA staging state",
+    )
+
+    text = replace_once(
+        text,
+        """        if getattr(p, "_vllm_is_uva_offloaded", False):
+            uva_offloaded_parameters.append(name)
+        # Parameters already on target device are not touched
+
+    try:
+""",
+        """        if getattr(p, "_vllm_is_uva_offloaded", False):
+            uva_offloaded_parameters.append(name)
+            if stage_trtllm_moe_uva:
+                # A same-device .to() is a no-op for a UVA CUDA view. clone()
+                # allocates fresh CUDA storage, which is HBM on GB300, and copies
+                # exactly the parameters vLLM actually offloaded, including a
+                # partially offloaded final module, through the C2C link once.
+                staged_uva_backings.append(p.data)
+                p.data = p.data.clone(memory_format=torch.preserve_format)
+                delattr(p, "_vllm_is_uva_offloaded")
+                staged_uva_bytes += p.numel() * p.element_size()
+                staged_uva_count += 1
+        # Parameters already on target device are not touched.
+
+    if staged_uva_bytes:
+        logger.info(
+            "Staged %d UVA-offloaded TRTLLM NVFP4 MoE parameters (%.2f GiB) "
+            "in HBM for post-load layout conversion.",
+            staged_uva_count,
+            staged_uva_bytes / (1024**3),
+        )
+
+    try:
+""",
+        "marker-driven UVA-to-HBM staging loop",
+    )
+
+    # The existing finally block re-offloads a parameter whenever its original
+    # name was UVA-backed and the current Parameter no longer carries the UVA
+    # marker. Staging deliberately removes that marker, so both in-place and
+    # replacement-style quantization paths are restored without another change.
+
+    try:
+        compile(text, str(target), "exec")
+    except SyntaxError as exc:
+        fail(f"patched source failed Python syntax validation: {exc}")
+
+    target.write_text(text)
+
+    print(
+        f"{PREFIX} Applied marker-driven HBM staging for UVA-offloaded "
+        "TRTLLM NVFP4 MoE post-load conversion."
+    )
+
+
+if __name__ == "__main__":
+    main()

mods/exp-gb300-trtllm/run.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+set -euo pipefail
+
+PYTHON_ROOT="${PYTHON_ROOT:-${VLLM_SITE_PACKAGES:-/usr/local/lib/python3.12/dist-packages}}"
+MOD_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PREFIX="[exp-gb300-trtllm]"
+
+if ! command -v python3 >/dev/null 2>&1; then
+  echo "$PREFIX python3 is required to apply this mod." >&2
+  exit 1
+fi
+
+if [ ! -d "$PYTHON_ROOT/vllm" ]; then
+  echo "$PREFIX vLLM package not found at $PYTHON_ROOT/vllm" >&2
+  exit 1
+fi
+
+python3 "$MOD_DIR/patch_vllm.py" "$PYTHON_ROOT"
+
+# Do not leave bytecode compiled from the pre-patch source around. The patcher
+# syntax-checks the result, so clearing the generated cache is safe.
+find "$PYTHON_ROOT/vllm/model_executor/model_loader" \
+  -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true
+
+echo "$PREFIX Enabled GB300 HBM staging for TRTLLM NVFP4 MoE conversion."