[None][refactor] Move PrimTS attention cache updates into upstream source

3rdparty/vendor_sources.lock.yaml

@@ -14,15 +14,15 @@ vendors:
     digest: sha256-tree-v1:12eb5c011c3fef1831aeba02691df8e2f5bff323f9da9c58c0f31e80883a9f06
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261107
-    commit: 911d6a9d62308b6c8a5249e8dc4aabc7518020f1
+    branch: prims-ts-upstream-refactor--backport
+    commit: 52be15bbc078ad86710fa2a92367c80bcead3a76
     source: flashinfer/attention/prims_ts
     destination: tensorrt_llm/_torch/attention/backends/prims_ts
     include:
     - '**/*.py'
     patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
-    patch_digest: sha256:f35ef4d096314a21c710b849b0a1f731cf1edbf50d32cdee46b635f45fe6b47a
-    digest: sha256-tree-v1:bf74c7664256a1c7064163ca417a6c589f8ab8f47bb478444bb1b3c402c856ee
+    patch_digest: sha256:eaf642dd5d4b5370a055a94f637925aecba61d01329510dde88d55ef3a192be1
+    digest: sha256-tree-v1:c507cdb90c90db4a13f4f042b6bb961752ee5ad44681e8bbcacda8ec0b8c7d70
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py

@@ -15,9 +15,8 @@
 """Compilation and caching for PrimTS block-sparse attention adapters."""
 
 from collections.abc import Callable
-import functools
 
-from .._compile_cache import build_and_load_attention_kernel
+from .._compile_cache import build_and_load_attention_kernel, cache_compiled
 from .._scales import fake_scale
 
 import torch
@@ -573,7 +572,7 @@ def compile_kernel():
         )
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_block_sparse(
     key: _BlockSparseCompileKey,
 ) -> Callable[..., object]:

tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py

@@ -121,17 +121,9 @@ def probe_cached():
 
 def expected_cache_meta() -> dict:
     """The on-disk module metadata this process would write (arch, DSL version, source hash)."""
-    from flashinfer.jit.cute_dsl_core import (
-        _get_compile_arch,
-        _get_cute_dsl_version,
-        _hash_source_files,
-    )
-
-    return {
-        "arch": _get_compile_arch(),
-        "cute_dsl_version": _get_cute_dsl_version(),
-        "source_sha256": _hash_source_files(tuple(_kernel_source_files())),
-    }
+    return _make_spec(
+        "prims_ts", "metadata", lambda: None, _kernel_source_files()
+    ).expected_meta
 
 
 def _unloaded_kernel(*args, **kwargs):
@@ -146,13 +138,37 @@ def is_placeholder(compiled) -> bool:
     return compiled is _unloaded_kernel
 
 
+def cache_compiled(compile_fn):
+    """Memoize loaded kernels, bypassing the memo for probe/persist plans.
+
+    A plan may return a kernel tuple or a callable. Bypassing the whole memo
+    keeps placeholders out of either shape and probes disk even after a normal
+    plan has loaded a kernel into this process.
+    """
+    cached = functools.cache(compile_fn)
+
+    @functools.wraps(compile_fn)
+    def wrapper(*args, **kwargs):
+        if _PRECOMPILE_MODE.get() is not None:
+            return compile_fn(*args, **kwargs)
+        return cached(*args, **kwargs)
+
+    wrapper.cache_clear = cached.cache_clear
+    wrapper.cache_info = cached.cache_info
+    return wrapper
+
+
 def _make_spec(module_name: str, kernel_name: str, compile_fn, key_files):
     from flashinfer.jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_files
 
-    return JitSpecCuteDsl(module_name, kernel_name, compile_fn, _hash_source_files(tuple(key_files)))
+    return JitSpecCuteDsl(
+        module_name, kernel_name, compile_fn, _hash_source_files(tuple(key_files))
+    )
 
 
-def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None:
+def _precompile_kernel(
+    module_name: str, kernel_name: str, compile_fn, key_files
+) -> None:
     """Compile one kernel into the on-disk cache unless it is already there.
 
     One compile per kernel per workspace: a sibling that wants the same kernel
@@ -168,7 +184,10 @@ def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files
     spec = _make_spec(module_name, kernel_name, compile_fn, key_files)
     if spec.is_compiled:
         return
-    lock = FileLock(get_tmpdir() / f"{spec.module_dir_name}.{spec.kernel_name}.lock", thread_local=False)
+    lock = FileLock(
+        get_tmpdir() / f"{spec.module_dir_name}.{spec.kernel_name}.lock",
+        thread_local=False,
+    )
     if _SKIP_BUSY.get():
         try:
             lock.acquire(timeout=0)
@@ -180,6 +199,8 @@ def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files
         if spec.is_compiled:
             return
         spec.compile_and_persist()
+        if not spec.is_compiled:
+            raise RuntimeError(f"PrimTS kernel was not persisted: {spec.name}")
     finally:
         lock.release()
 
@@ -214,8 +235,10 @@ def build_and_load_attention_kernel(
     )
 
     mode = _PRECOMPILE_MODE.get()
-    if mode is not None and os.environ.get("FLASHINFER_DISABLE_JIT"):
-        raise RuntimeError("PrimTS warm-up precompile cannot compile with FLASHINFER_DISABLE_JIT set.")
+    if mode == "persist" and os.environ.get("FLASHINFER_DISABLE_JIT"):
+        raise RuntimeError(
+            "PrimTS warm-up precompile cannot compile with FLASHINFER_DISABLE_JIT set."
+        )
     if cute_dsl_cache_disabled():
         if mode is not None:
             raise RuntimeError(
@@ -227,7 +250,9 @@ def build_and_load_attention_kernel(
     key = _compile_cache_key(spec, options, extra)
     if key is None:
         if mode is not None:
-            raise RuntimeError("PrimTS warm-up precompile cannot key this kernel for the on-disk cache.")
+            raise RuntimeError(
+                "PrimTS warm-up precompile cannot key this kernel for the on-disk cache."
+            )
         return compile_fn()
     import torch
 
@@ -237,11 +262,19 @@ def build_and_load_attention_kernel(
     major, minor = torch.cuda.get_device_capability()
     if mode == "persist":
         _precompile_kernel(
-            f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files()
+            f"prims_ts_{family}",
+            f"{kind}_sm{major}{minor}_{key}",
+            compile_fn,
+            _kernel_source_files(),
         )
         return _unloaded_kernel
     if mode == "probe":
-        _probe_kernel(f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files())
+        _probe_kernel(
+            f"prims_ts_{family}",
+            f"{kind}_sm{major}{minor}_{key}",
+            compile_fn,
+            _kernel_source_files(),
+        )
         return _unloaded_kernel
     return build_and_load_cute_dsl_kernel(
         f"prims_ts_{family}",

tensorrt_llm/_torch/attention/backends/prims_ts/context.py

@@ -43,7 +43,7 @@
 
 import torch
 
-from ._compile_cache import build_and_load_attention_kernel
+from ._compile_cache import build_and_load_attention_kernel, cache_compiled
 from ._hardware import get_max_active_clusters
 from ._scales import (
     GlobalScale,
@@ -1995,7 +1995,7 @@ def _paged_context_compile_spec(
     )
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_context(
     compile_spec: _ContextCompileSpec,
 ):
@@ -2389,7 +2389,7 @@ def fake_compact(dtype, shape, assumed_align):
     return compiled, policy
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_paged_context(
     compile_spec: _PagedContextCompileSpec,
 ):

tensorrt_llm/_torch/attention/backends/prims_ts/decode.py

@@ -25,7 +25,7 @@
 
 import torch
 
-from ._compile_cache import build_and_load_attention_kernel
+from ._compile_cache import build_and_load_attention_kernel, cache_compiled
 from ._scales import (
     GlobalScale,
     fake_scale,
@@ -1628,7 +1628,7 @@ def _make_decode_compile_spec(
     )
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_decode(
     compile_spec: _DecodeCompileSpec,
 ):

tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py

@@ -1,4 +1,3 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 # Copyright (c) 2026 by FlashInfer team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -118,6 +117,7 @@
 from ...sage.index import load_sage_q_scale, sage_section_word, sage_sequence_words
 from ...sage.layout import _FLOAT4_WORDS, _SageScaleLayout
 from ...sage.modes import SageKTransport, SageScoreScaling
+from .helpers_paged import _load_block_table_row_bounds, _load_runtime_seq_len_kv
 from .helpers import (
     add_f16x2_to_f32x2,
     _resolve_kv_tile_idx_context,
@@ -1479,7 +1479,7 @@ def compute_coords(
                 else:
                     kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds(
                         Int32(self.block_table_row_stride),
-                        self.cfg,
+                        self.cfg.num_tokens_per_page,
                         seqlen_k,
                         batch_coord,
                     )
@@ -1568,7 +1568,7 @@ def compute_page_coords(self, stage_info: StageInfo) -> tuple[Int32, Int32, Int3
             )
             kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds(
                 Int32(self.block_table_row_stride),
-                self.cfg,
+                self.cfg.num_tokens_per_page,
                 cached_seqlen_kv,
                 batch_coord,
             )
@@ -7815,13 +7815,3 @@ def tma_store(
                 prims.cp_async_bulk_commit_group()
                 if cutlass.const_expr(self.cfg.gmem_o_store_wait_after_write):
                     prims.cp_async_bulk_wait_group(0, read=True)
-
-
-# helpers_paged imports FmhaConfig from this module, so it is imported here,
-# after FmhaConfig is defined. The methods above call these helpers while being
-# traced inside the kernel's runtime loops, where the CuTe DSL rejects a first
-# import of a module; importing at module scope keeps those calls import-free.
-from .helpers_paged import (
-    _load_block_table_row_bounds,
-    _load_runtime_seq_len_kv,
-)

tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py

@@ -23,8 +23,6 @@
 import cutlass.cute as cute
 from cutlass import Int32
 
-from .fmha_resources import FmhaConfig
-
 
 @cute.jit
 def _load_runtime_seq_len_kv(
@@ -41,19 +39,21 @@ def _load_runtime_seq_len_kv(
 
 
 @cute.jit
-def _runtime_last_valid_page_idx(cfg: FmhaConfig, seq_len_kv: Int32) -> Int32:
-    num_pages = cute.ceil_div(seq_len_kv, cfg.num_tokens_per_page)
+def _runtime_last_valid_page_idx(
+    num_tokens_per_page: cutlass.Constexpr[int], seq_len_kv: Int32
+) -> Int32:
+    num_pages = cute.ceil_div(seq_len_kv, num_tokens_per_page)
     return cute.math.max(num_pages - Int32(1), Int32(0))
 
 
 @cute.jit
 def _load_block_table_row_bounds(
     block_table_row_stride: Int32,
-    cfg: FmhaConfig,
+    num_tokens_per_page: cutlass.Constexpr[int],
     seq_len_kv: Int32,
     batch_coord: Int32,
 ) -> tuple[Int32, Int32]:
     """Return one fixed-table row base and runtime-valid inclusive page bound."""
 
     row_begin = batch_coord * block_table_row_stride
-    return row_begin, _runtime_last_valid_page_idx(cfg, seq_len_kv)
+    return row_begin, _runtime_last_valid_page_idx(num_tokens_per_page, seq_len_kv)