[Cleanup][SM70] Drop the forced compile-cache opt-out for the Flash-V100 graph

benchmarks/benchmark_sm70_decode.py

@@ -440,9 +440,6 @@ def _sm70_turbomind_policy() -> dict[str, Any]:
         "VLLM_SM70_AWQ_PRESERVE_DEFAULT_SPLITS_ONLY": os.environ.get(
             "VLLM_SM70_AWQ_PRESERVE_DEFAULT_SPLITS_ONLY"
         ),
-        "VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING": os.environ.get(
-            "VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING"
-        ),
         "awq_preserve_default_splits_effective": _env_bool(
             "VLLM_SM70_AWQ_PRESERVE_DEFAULT_SPLITS",
             True,
@@ -451,10 +448,6 @@ def _sm70_turbomind_policy() -> dict[str, Any]:
             "VLLM_SM70_AWQ_PRESERVE_DEFAULT_SPLITS_ONLY",
             False,
         ),
-        "allow_compile_cache_for_profiling_effective": _env_bool(
-            "VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING",
-            False,
-        ),
         "VLLM_SM70_AWQ_DENSE_TUNE_MAX_M": os.environ.get(
             "VLLM_SM70_AWQ_DENSE_TUNE_MAX_M"
         ),

tools/torch_patches/0001-aot-compile-serialize-triton-kernel-side-table.patch

@@ -0,0 +1,139 @@
+--- a/torch/_dynamo/aot_compile_types.py
++++ b/torch/_dynamo/aot_compile_types.py
+@@ -1,10 +1,78 @@
+ import abc
++import importlib
+ import pickle
+ from typing import Any
+
+ import torch
+
+
++def _serialize_triton_kernel(kernel: Any) -> tuple[str, str]:
++    """
++    Serialize a triton kernel by extracting its module path and function name.
++    Returns (module_path, function_name) tuple.
++
++    Triton JITFunction objects contain unpicklable _thread.RLock objects, so we
++    serialize the import path instead and reimport on load.
++
++    Raises:
++        RuntimeError: If the kernel cannot be serialized (missing attributes).
++    """
++    fn = getattr(kernel, "fn", None)
++    if fn is None:
++        raise RuntimeError(
++            f"Kernel {kernel} has no 'fn' attribute. "
++            f"Cannot serialize for precompilation."
++        )
++    module_path = getattr(fn, "__module__", None)
++    func_name = getattr(fn, "__name__", None)
++    if module_path is None or func_name is None:
++        raise RuntimeError(
++            f"Kernel fn missing __module__ or __name__: "
++            f"module={module_path}, name={func_name}. "
++            f"Cannot serialize for precompilation."
++        )
++    return (module_path, func_name)
++
++
++def _deserialize_triton_kernel(kernel_info: tuple[str, str]) -> Any:
++    """
++    Deserialize a triton kernel by reimporting from its module.
++    kernel_info is (module_path, function_name) tuple.
++    """
++    module_path, func_name = kernel_info
++    module = importlib.import_module(module_path)
++    kernel = getattr(module, func_name)
++    return kernel
++
++
++# Note: [Triton Kernel Side Table Serialization]
++#
++# When dynamo captures user-defined triton kernels, it creates FX graph nodes
++# (triton_kernel_wrapper_mutation/functional) with a `kernel_idx` parameter that
++# references the global `kernel_side_table` in triton_kernel_wrap.py. This side
++# table maps integer indices to actual triton kernel objects.
++#
++# For kernels that go through inductor's codegen path, this is fine - inductor
++# looks up the kernel from the side table at codegen time and embeds the kernel
++# source code directly into the generated wrapper. The compiled code doesn't
++# need the side table at runtime.
++#
++# However, not all triton kernels go through inductor codegen. When using
++# regional_inductor, only annotated regions are compiled by inductor. Triton
++# kernels outside these regions are executed via the FX interpreter, which
++# calls the higher-order op directly and needs the kernel to be in the side
++# table at runtime.
++#
++# When serializing/deserializing bundled AOT artifacts across process boundaries,
++# the kernel_side_table is empty in the new process, causing:
++#   AssertionError: Kernel index X not found in id_to_kernel
++#
++# To fix this, we capture the kernel_side_table state during serialization and
++# restore it during deserialization. Kernels are serialized by their import path
++# (module_path, function_name) since triton JITFunction objects contain
++# unpicklable RLock objects.
++
++
+ class SerializableCallable(abc.ABC):
+     @classmethod
+     @abc.abstractmethod
+@@ -46,8 +114,23 @@
+     def serialize_compile_artifacts(
+         cls, fn: "BundledAOTAutogradSerializableCallable"
+     ) -> bytes:
++        from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
++
++        # See Note: [Triton Kernel Side Table Serialization]
++        # Capture triton kernel side table state BEFORE serialization.
++        triton_kernels: dict[int, tuple[str, str]] = {
++            idx: _serialize_triton_kernel(kernel)
++            for idx, kernel in kernel_side_table.id_to_kernel.items()
++        }
++        triton_constant_args: dict[int, dict[str, Any]] = dict(
++            kernel_side_table.constant_args
++        )
++
+         with torch._functorch.config.patch("bundled_autograd_cache", True):
+-            result = pickle.dumps(fn.compiled_fn.serialize())
++            serialized_entry = fn.compiled_fn.serialize()
++            # Bundle the triton kernel side table with the serialized entry
++            bundle = (serialized_entry, triton_kernels, triton_constant_args)
++            result = pickle.dumps(bundle)
+             return result
+
+     @classmethod
+@@ -55,8 +138,31 @@
+         from torch._functorch._aot_autograd.aot_autograd_result import (
+             deserialize_bundled_cache_entry,
+         )
++        from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
++
++        bundle = pickle.loads(data)
++
++        # Handle both old format (just entry) and new format (entry, kernels, const_args)
++        if isinstance(bundle, tuple) and len(bundle) == 3:
++            entry, triton_kernels, triton_constant_args = bundle
++        else:
++            # Backwards compatibility with old serialized artifacts
++            entry = bundle
++            triton_kernels = {}
++            triton_constant_args = {}
++
++        # See Note: [Triton Kernel Side Table Serialization]
++        # Restore triton kernel side table BEFORE deserializing the compiled function.
++        # The compiled function may reference kernels by index if any triton kernels
++        # don't go through inductor codegen (e.g., triton kernels outside of
++        # regional_inductor compiled regions).
++        for idx, kernel_info in triton_kernels.items():
++            kernel = _deserialize_triton_kernel(kernel_info)
++            kernel_side_table.id_to_kernel[idx] = kernel
++            kernel_side_table.kernel_to_id[kernel] = idx
+
+-        entry = pickle.loads(data)
++        for idx, args in triton_constant_args.items():
++            kernel_side_table.constant_args[idx] = args
+
+         compiled_fn = deserialize_bundled_cache_entry(entry)
+         return cls(compiled_fn)

tools/torch_patches/README.md

@@ -0,0 +1,14 @@
+# torch backports for the pinned torch 2.10.0
+
+1Cat pins torch 2.10.0 (the last cu128 wheel with Volta, sm_70). Fixes that landed
+in later torch releases and that this stack needs are kept here as patches and
+applied to the installed torch with `apply.sh`. Each patch names its upstream
+source. `apply.sh` is idempotent and refuses files it does not recognise.
+
+| patch | upstream | why |
+|---|---|---|
+| 0001 aot_compile_types.py | pytorch/pytorch #173556 (main dffe73e2, in 2.11+) | AOT compile artifacts (`VLLM_USE_AOT_COMPILE`, compile cache on) reference Triton kernels through indices of a process-local side table; a fresh process has an empty table. Loading such an artifact fails with an empty assertion (recompile) or, seen with Qwen3.8-Flash-Next on an RTX 8000 stage, resolves a wrong kernel and dies with an illegal memory access. The fix serialises the table into the artifact and restores it on load; old artifacts without a table are still read the old way, so delete `torch_aot_compile/` under the cache root once after applying. |
+
+Usage after installing or reinstalling torch in a venv:
+
+    tools/torch_patches/apply.sh /path/to/venv/bin/python

tools/torch_patches/apply.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+# Apply the torch 2.10.0 backports in this directory to the torch of a given interpreter.
+# Idempotent: a file already at the patched hash is skipped, the pristine 2.10.0 file is
+# patched (original kept as .orig-2.10.0), anything else aborts loudly. Re-run after every
+# reinstall of torch (pip install -e . does not touch torch, a venv rebuild does).
+#   tools/torch_patches/apply.sh [python]     default: python from PATH
+set -euo pipefail
+PY=${1:-python}
+HERE=$(cd "$(dirname "$0")" && pwd)
+TORCH=$("$PY" -c 'import os, torch; print(os.path.dirname(torch.__file__))')
+VER=$("$PY" -c 'import torch; print(torch.__version__)')
+case "$VER" in 2.10.0*) ;; *) echo "torch $VER: these backports are for 2.10.0 only (2.11+ has them upstream)"; exit 2 ;; esac
+sha() { sha256sum "$1" | cut -c1-64; }
+# --- 0001: torch #173556, serialize the triton kernel side table into bundled AOT artifacts ---
+F="$TORCH/_dynamo/aot_compile_types.py"
+ORIG=93f2529e50a3fa31dd0a4dfcd51f1e907dbf74fe782e0d0819d78ef365fbe74e
+PATCHED=d3acc67b813f260f8989156ae5fbdcaf1c04c8eee236b567c8681dd377675f35
+case "$(sha "$F")" in
+  $PATCHED) echo "0001 aot_compile_types.py: already applied" ;;
+  $ORIG)    [ -e "$F.orig-2.10.0" ] || cp "$F" "$F.orig-2.10.0"; patch -p1 -d "$(dirname "$TORCH")" --forward --silent < "$HERE/0001-aot-compile-serialize-triton-kernel-side-table.patch"
+            [ "$(sha "$F")" = "$PATCHED" ] || { echo "0001: hash after patch unexpected"; exit 1; }; echo "0001 aot_compile_types.py: applied" ;;
+  *)        echo "0001: $F is neither pristine 2.10.0 nor patched; refusing"; exit 1 ;;
+esac

vllm/config/vllm.py

@@ -2192,28 +2192,6 @@ def __post_init__(self):
                             "configuration: regular torch.compile reproduced "
                             "deterministic greedy token drift."
                         )
-                if envs.VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING:
-                    logger.warning_once(
-                        "VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING=1: "
-                        "leaving VLLM_DISABLE_COMPILE_CACHE unset for "
-                        "diagnostic profiling. This reuses compile artifacts "
-                        "and is not a quality-parity baseline."
-                    )
-                elif "VLLM_DISABLE_COMPILE_CACHE" not in os.environ:
-                    os.environ["VLLM_DISABLE_COMPILE_CACHE"] = "1"
-                    logger.info_once(
-                        "Auto-setting VLLM_DISABLE_COMPILE_CACHE=1 for SM70 "
-                        "Flash-V100 0.0.3 compile graph quality parity; "
-                        "decode throughput is preserved, but AOT artifact "
-                        "reload stays disabled until its token drift is fixed."
-                    )
-                elif os.environ.get("VLLM_DISABLE_COMPILE_CACHE") == "0":
-                    logger.warning_once(
-                        "VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH=1 with "
-                        "explicit VLLM_DISABLE_COMPILE_CACHE=0 is a "
-                        "diagnostic-only configuration: cached AOT artifact "
-                        "reload reproduced deterministic greedy token drift."
-                    )
                 self.compilation_config.inductor_compile_config["combo_kernels"] = True
                 self.compilation_config.inductor_compile_config[
                     "benchmark_combo_kernel"

vllm/envs.py

@@ -615,7 +615,6 @@
     VLLM_SM70_USE_BREAKABLE_CUDAGRAPH: bool = False
     VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH: bool = False
     VLLM_SM70_QWEN38_HYBRID_PLE: bool = False
-    VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING: bool = False
     VLLM_SM70_SYNC_BEFORE_COMPILE_GRAPH_FORWARD: bool = False
     VLLM_SM70_FLASH_V100_0DOT3_ELIMINATE_NOOPS: bool = False
     VLLM_SM70_FLASH_V100_0DOT3_BENCHMARK_COMBO_KERNEL: bool = False
@@ -856,12 +855,7 @@ def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None:
 
 
 def disable_compile_cache() -> bool:
-    sm70_compile_graph = os.getenv(
-        "VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH",
-        "0",
-    ).strip().lower() in ("1", "true", "yes", "on")
-    default_value = "1" if sm70_compile_graph else "0"
-    return bool(int(os.getenv("VLLM_DISABLE_COMPILE_CACHE", default_value)))
+    return bool(int(os.getenv("VLLM_DISABLE_COMPILE_CACHE", "0")))
 
 
 def use_aot_compile() -> bool:
@@ -3698,10 +3692,6 @@ def _resolve_rust_frontend_path() -> str | None:
     # Diagnostic-only profiling knob. The SM70 compile-graph quality profile
     # disables AOT cache reload by default due known token drift, but long
     # profiler runs need an explicit way to reuse compile artifacts.
-    "VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING": lambda: bool(
-        os.getenv("VLLM_SM70_ALLOW_COMPILE_CACHE_FOR_PROFILING", "0").strip().lower()
-        in ("1", "true", "yes", "on")
-    ),
     "VLLM_SM70_SYNC_BEFORE_COMPILE_GRAPH_FORWARD": lambda: bool(
         os.getenv(
             "VLLM_SM70_SYNC_BEFORE_COMPILE_GRAPH_FORWARD",