[yuxianq]

📌 Description

Moves general PrimTS functionality out of TensorRT-LLM's persistent compatibility patch and into FlashInfer:

• Persist/probe attention kernels during warm-up, with per-kernel locking and a nonblocking busy result. All five compiler entry points bypass their in-process memo in these modes, including tuple-returning decode compilers. Failed persistence is reported rather than returning a misleading placeholder.
• Cache sparse MLA geometry after argument validation and explicit CUDA-device resolution.
• Reuse prepared native D512 single-token launches, with complete pool/sink metadata keys, per-call query/index/output checks, retained storage, and invalidation on replan.
• Load paged-context helpers before CuTe tracing. Pass the page-size field to eliminate their circular dependency on the resource module.

The prepared fast path is restricted to native D512, fixed single-token queries, host scales, caller-provided BF16 output, and no returned/caller-provided LSE or output quantization. Other configurations retain the full run path. Existing kernel math and public signatures are unchanged.

Measured on a B300 PCIe allocation (SM103; the driver reports NVIDIA Graphics Device), BF16, one KV pool, Q length 1. benchmarks/bench_prims_ts_sparse_mla_host.py uses preallocated output, 20 warm-ups, and the median of five sets of 500 calls. Baseline: 9da5a602ae4cd6bd0823fd77abafa65f96136adc.

Rows / heads / top-k	CPU submission µs, before → after	Completed eager wall µs, before → after	Workspace query µs, before → after
4 / 16 / 128	40.73 → 19.04	40.77 → 19.05	544.55 → 5.87
16 / 64 / 256	41.89 → 19.53	41.93 → 19.57	539.84 → 6.08

These are host submission/completed eager measurements, not pure GPU kernel times or model throughput. Reproduce by running the script with PYTHONPATH selecting each checkout, using --rows 4 --heads 16 --topk 128 and --rows 16 --heads 64 --topk 256.

🔍 Related Issues

Extracted from TensorRT-LLM 3rdparty/vendor_patches/flashinfer-prims-ts.patch. The companion consumer PR NVIDIA/TensorRT-LLM#21337 pins a backport onto its existing source revision, avoiding an unrelated full vendor refresh.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).
• This PR requires full-scope CI, not targeted tests (expensive; justify below).
  Reason:

Validated:

• CPU cache lifecycle/locking/device-resolution/import-order regressions: 9 passed.
• B300 cache lifecycle plus attention feature suite: 25 passed, including fresh-process numerical cold/warm cache checks for decode, context, MLA, block-sparse, and Sage attention.
• Final prepared-input validation follow-up: 5 passed on B300.
• All applicable pre-commit checks passed on changed files, including mypy and Ruff. --all-files was not run; the full test suite was not run.
• Additional D576 forced-split, NVFP4, and fused output-quantization cases: 4 passed on B300.
• Pinned-source backport: 9 CPU and 6 GPU tests passed, including paged-context accuracy. The actual TRT-LLM vendored files independently passed 7 CPU and 5 GPU cache/graph tests.

Suggested CI targets: tests/cpu/attention/test_prims_ts_cache.py tests/attention/test_attention_ts_cache_lifecycle.py tests/attention/test_attention_ts_features.py tests/attention/test_attention_ts_sparse_mla.py

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

This uses the normal review path for existing PrimTS code and adds no public API or automatic backend selection. The companion backport retains TRT-LLM's older sparse MLA argument names and scalar ABI. The positional output-quantization launch fix is already on main (see #6774); it is retained only in that backport.

The forced D576 split test clears the geometry memo around its policy override so each requested split remains exercised.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}