refactor(prims-ts): upstream attention precompile and sparse MLA caches

CLAUDE.md

@@ -738,6 +738,15 @@ Task-scheduled attention shares this cache through
 options, and embedded device-capacity constants. It hashes the attention source
 tree, including shared preparation, masks, and reductions. Occupancy queries
 are cached per device and cluster size by `_hardware.py`.
+PrimTS warm-up can use `_compile_cache.precompile_only()` to persist kernels
+without loading them, or `probe_cached()` to check their disk artifacts.
+`cache_compiled` bypasses the in-process memo in both modes; use it for new
+attention compiler entry points, including those returning kernel tuples.
+Sparse MLA caches validated geometry with an explicit device index. Native
+D512 single-token calls also reuse prepared pool/sink views after a full run;
+replanning clears those views, and dynamic query/index/output metadata is
+still checked with `validate=False`.
+
 
 ### Dispatch Macros
 

benchmarks/bench_prims_ts_sparse_mla_host.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 by FlashInfer team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Measure sparse MLA CPU launch and completed eager wall time.
+
+Run this same script with PYTHONPATH selecting the baseline and candidate
+checkouts. GPU work is included in completed_wall_us, but excluded from the
+post-loop synchronization in cpu_submit_us. Neither metric is pure kernel time.
+"""
+
+import argparse
+import json
+import statistics
+import time
+
+import torch
+
+from flashinfer.attention.prims_ts import mla_decode as mla
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=int, default=4)
+    parser.add_argument("--heads", type=int, default=16)
+    parser.add_argument("--topk", type=int, default=128)
+    parser.add_argument("--iterations", type=int, default=500)
+    args = parser.parse_args()
+    torch.manual_seed(0)
+    device = torch.device("cuda", torch.cuda.current_device())
+    query = torch.randn(
+        args.rows, 1, args.heads, 512, device=device, dtype=torch.bfloat16
+    )
+    pool = torch.randn(args.topk * 2, 1, 512, device=device, dtype=torch.bfloat16)
+    indices = torch.randint(
+        args.topk * 2, (args.rows, args.topk), device=device, dtype=torch.int32
+    )
+    lengths = torch.full((args.rows,), args.topk, device=device, dtype=torch.int32)
+    output = torch.empty_like(query)
+    wrapper = mla.BatchSparseMLADecodePagedTSWrapper()
+    wrapper.plan(device, args.rows, args.heads, max_topk=args.topk)
+
+    def run():
+        wrapper.run(query, pool, indices, lengths, out=output, validate=False)
+
+    for _ in range(20):
+        run()
+    submit, wall, plan = [], [], []
+    for _ in range(5):
+        torch.cuda.synchronize()
+        start = time.perf_counter()
+        for _ in range(args.iterations):
+            run()
+        submitted = time.perf_counter()
+        torch.cuda.synchronize()
+        done = time.perf_counter()
+        submit.append((submitted - start) * 1e6 / args.iterations)
+        wall.append((done - start) * 1e6 / args.iterations)
+        start = time.perf_counter()
+        for _ in range(args.iterations):
+            mla.get_prims_ts_sparse_mla_decode_workspace_size(
+                device, args.rows, args.heads, max_topk=args.topk
+            )
+        plan.append((time.perf_counter() - start) * 1e6 / args.iterations)
+    print(
+        json.dumps(
+            dict(
+                gpu=torch.cuda.get_device_name(),
+                module=mla.__file__,
+                rows=args.rows,
+                heads=args.heads,
+                topk=args.topk,
+                cpu_submit_us=statistics.median(submit),
+                completed_wall_us=statistics.median(wall),
+                workspace_query_us=statistics.median(plan),
+            ),
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

flashinfer/attention/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
 
 import torch
 
@@ -554,7 +553,7 @@ def compile_kernel():
         )
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_block_sparse(
     key: _BlockSparseCompileKey,
 ) -> Callable[..., object]:

flashinfer/attention/prims_ts/_compile_cache.py

@@ -1,3 +1,4 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 # Copyright (c) 2026 by FlashInfer team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,10 +16,13 @@
 """Persistent compilation shared by the PrimTS attention adapters."""
 
 from collections.abc import Callable
+import contextlib
+import contextvars
 import dataclasses
 import enum
 import functools
 import hashlib
+import os
 from pathlib import Path
 from typing import Any
 
@@ -64,6 +68,149 @@ def _compile_cache_key(spec: object, options: str, extra: tuple = ()) -> str | N
     return hashlib.sha256(encoded.encode()).hexdigest()
 
 
+# Warm-up precompile modes. Under ``precompile_only()`` the attention kernels
+# are compiled through ``JitSpecCuteDsl.compile_and_persist()`` (the
+# ``cute.compile`` runs outside the module-level file lock that
+# ``build_and_load()`` holds for a whole build, so sibling processes compile
+# different kernels concurrently), committed to disk, and not loaded: the
+# serving process loads them afterwards. Under ``probe_cached()`` nothing is
+# compiled; a kernel missing from the cache raises ``KernelNotCached``. Both
+# modes return a placeholder in place of the kernel.
+_PRECOMPILE_MODE: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "prims_ts_precompile_mode", default=None
+)
+_SKIP_BUSY: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "prims_ts_precompile_skip_busy", default=False
+)
+
+
+class KernelNotCached(Exception):
+    """Raised under ``probe_cached()`` for a kernel that is not on disk."""
+
+
+class KernelBusy(Exception):
+    """Raised under ``precompile_only(skip_busy=True)`` while another process compiles the kernel."""
+
+
+@contextlib.contextmanager
+def precompile_only(*, skip_busy: bool = False):
+    """Compile and persist kernels without loading them (see module note).
+
+    With ``skip_busy`` a kernel whose per-kernel lock another process holds
+    raises ``KernelBusy`` instead of waiting, so a worker can defer it and
+    compile something else first.
+    """
+    token = _PRECOMPILE_MODE.set("persist")
+    busy_token = _SKIP_BUSY.set(skip_busy)
+    try:
+        yield
+    finally:
+        _SKIP_BUSY.reset(busy_token)
+        _PRECOMPILE_MODE.reset(token)
+
+
+@contextlib.contextmanager
+def probe_cached():
+    """Report, without compiling, whether every kernel a plan needs is on disk."""
+    token = _PRECOMPILE_MODE.set("probe")
+    try:
+        yield
+    finally:
+        _PRECOMPILE_MODE.reset(token)
+
+
+def expected_cache_meta() -> dict:
+    """The on-disk module metadata this process would write (arch, DSL version, source hash)."""
+    return _make_spec(
+        "prims_ts", "metadata", lambda: None, _kernel_source_files()
+    ).expected_meta
+
+
+def _unloaded_kernel(*args, **kwargs):
+    raise RuntimeError(
+        "PrimTS kernel was compiled in precompile-only mode and never loaded; "
+        "load it through a normal plan() outside precompile_only()."
+    )
+
+
+def is_placeholder(compiled) -> bool:
+    """True for the stand-in returned under the precompile modes."""
+    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 ...jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_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:
+    """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
+    waits on the per-kernel lock (or raises ``KernelBusy`` under ``skip_busy``)
+    and then finds it on disk. The lock lives in FlashInfer's tmp directory,
+    not in the module directory: creating that directory before the first
+    export would let ``compile_and_persist()`` of a sibling read a missing
+    meta.json as a stale module and wipe the exports already there.
+    """
+    from filelock import FileLock, Timeout
+    from ...jit.core import get_tmpdir
+
+    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,
+    )
+    if _SKIP_BUSY.get():
+        try:
+            lock.acquire(timeout=0)
+        except Timeout:
+            raise KernelBusy(spec.name) from None
+    else:
+        lock.acquire()
+    try:
+        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()
+
+
+def _probe_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None:
+    spec = _make_spec(module_name, kernel_name, compile_fn, key_files)
+    if not spec.is_compiled:
+        raise KernelNotCached(spec.name)
+
+
 @functools.cache
 def _kernel_source_files() -> tuple[str, ...]:
     # Family adapters share masks, reductions, sparse/Sage preparation and
@@ -87,18 +234,48 @@ def build_and_load_attention_kernel(
         cute_dsl_cache_disabled,
     )
 
+    mode = _PRECOMPILE_MODE.get()
+    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(
+                "PrimTS warm-up precompile needs the on-disk CuTe DSL cache "
+                "(FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 is set)."
+            )
         return compile_fn()
 
     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."
+            )
         return compile_fn()
     import torch
 
     # A family target such as sm100f can cover physical GPUs with different
     # host-selected tuning defaults. The module metadata keys the target;
     # the specialization must also distinguish the physical capability.
     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(),
+        )
+        return _unloaded_kernel
+    if mode == "probe":
+        _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}",
         f"{kind}_sm{major}{minor}_{key}",

flashinfer/attention/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 .vc_attention import (
@@ -1970,7 +1970,7 @@ def _paged_context_compile_spec(
     )
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_context(
     compile_spec: _ContextCompileSpec,
 ):
@@ -2325,7 +2325,7 @@ def fake_compact(dtype, shape, assumed_align):
     return compiled, policy
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_paged_context(
     compile_spec: _PagedContextCompileSpec,
 ):

flashinfer/attention/prims_ts/decode.py

@@ -24,7 +24,7 @@
 
 import torch
 
-from ._compile_cache import build_and_load_attention_kernel
+from ._compile_cache import build_and_load_attention_kernel, cache_compiled
 
 from flashinfer.api_logging import flashinfer_experimental_api
 from flashinfer.utils import (
@@ -1564,7 +1564,7 @@ def _make_decode_compile_spec(
     )
 
 
-@functools.cache
+@cache_compiled
 def _get_compiled_decode(
     compile_spec: _DecodeCompileSpec,
 ):