[None][feat] Add the PrimTS native sparse MLA FMHA for DeepSeek-V4 generation and enable it by default on SM100

3rdparty/vendor_sources.lock.yaml

@@ -13,16 +13,16 @@ vendors:
     patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff
     digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac
   flashinfer-prims-ts:
-    url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261012
-    commit: 0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf
+    url: https://github.com/yunruis/flashinfer.git
+    branch: trtllm-prims-ts-dev-20261023
+    commit: aa709061230a58858375c513056e7733fa74294a
     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:1ec340aac4fea9c77b9e9d956b1676c09cff5081053ef1e41f962ed38a6f728c
-    digest: sha256-tree-v1:e4129294b2b26295581cd9cbf1ac5ac4d7eec2ec65f0108626692ddf8a036f96
+    patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07
+    digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

docs/source/release-notes.md

@@ -8,6 +8,11 @@ All published functionality in the Release Notes has been fully tested and verif
 
 ### Key Features and Enhancements
 
+- **DeepSeek-V4 generation on B200/GB200 (SM100) now runs on the PrimTS native sparse MLA FMHA by default.** The new `prims_ts_sparse_mla` FMHA library (a CuTe DSL kernel vendored from FlashInfer) serves the DeepSeek-V4 sparse MLA generation phase ahead of the TRTLLM-Gen path on SM100; it covers BF16 and FP8 KV caches, the fused inverse-RoPE + FP8 output epilogue, speculative decoding (MTP) rows and CUDA graph replay. SM103 (GB300), `fp8_ds_mla` and NVFP4 KV caches keep the TRTLLM-Gen path. Opt out with `TLLM_FMHA_LIBS=-prims_ts_sparse_mla`.
+  - **First-launch compile cost.** The kernels are JIT-compiled by the CuTe DSL and cached on disk under the FlashInfer workspace. The first start on a host (or after a FlashInfer / CuTe DSL upgrade) compiles the kernel variants the deployment's generation row counts select during `LLM()` startup, about 5 to 10 minutes on B200 for DeepSeek-V4; later starts load the cache. Under `trtllm-llmapi-launch` and the default `LLM()` MPI session every rank has its own FlashInfer workspace (`TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=1`), so each rank compiles its own copy; set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` to share one cache between the ranks of a host.
+  - **Warm-up controls.** `TLLM_PRIMTS_WARM_WORKERS` sets the compile subprocesses per rank (default 4, reduced so that the ranks on a host, counted from the model's parallel mapping, start at most 16 together; `0` compiles in-process), `TLLM_PRIMTS_WARM_TIMEOUT` bounds the wait for them (default 1800 s), and `TLLM_PRIMTS_WARM_STRICT=1` turns a worker failure into an error instead of an in-process retry. Each worker imports torch and the CuTe DSL and opens a CUDA context, so lower the worker count on nodes with little host memory.
+  - **`FLASHINFER_DISABLE_JIT=1`** now also refuses to compile these kernels (with the on-disk CuTe DSL cache enabled, its default): startup then fails on a variant missing from the cache instead of falling back to TRTLLM-Gen, so a deployment that sets it must pre-populate the cache (one start without the variable) or opt out of the library.
+
 ### Infrastructure Changes
 
 ### API Changes

tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md

@@ -479,7 +479,11 @@ the selected library.
 it may add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the
 defaults or `TLLM_FMHA_LIBS=fallback` to force the fallback path. Generic
 block-sparse PrimTS remains enabled by default because a dense fallback cannot
-preserve its routing semantics. Delta entries update the
+preserve its routing semantics. The PrimTS sparse MLA library
+(`prims_ts_sparse_mla`) is also enabled by default: it serves only DeepSeek-V4
+generation on SM100 and sits ahead of the TRTLLM-Gen libraries, so DeepSeek-V4
+decode on B200/GB200 runs on it unless `TLLM_FMHA_LIBS=-prims_ts_sparse_mla`
+removes it. Delta entries update the
 default membership and follow canonical registry order, while an exact list
 preserves the user-specified order. Each FMHA library exposes `is_available()`
 for module/static environment checks and `is_supported()` for per-forward
@@ -552,6 +556,10 @@ The FMHA package is split by role:
   With an FP8 KV cache it passes `mla_bmm1_scale` / `mla_bmm2_scale` to the
   kernel as device tensors (element 0 of each), so the scales need no host
   read and stay valid under CUDA graph replay.
+- `fmha/prims_ts_sparse_mla.py` runs DeepSeek-V4 generation on the vendored
+  PrimTS two-source sparse MLA kernel (SM100 only; SM103 stays on TRTLLM-Gen);
+  it is in the default library set and `TLLM_FMHA_LIBS=-prims_ts_sparse_mla`
+  removes it.
 - `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the
   vendored PrimTS contiguous and paged wrappers. Paged generation passes a
   live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it

tensorrt_llm/_torch/attention/backends/fmha/__init__.py

@@ -24,6 +24,7 @@
 from .msa_prefill import MsaPrefillFmha
 from .phased import FmhaParams, PhasedFmha
 from .prims_ts import PrimsTSFmha
+from .prims_ts_sparse_mla import PrimsTSSparseMlaFmha
 from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes
 from .triton_custom_mask import TritonCustomMaskFmha
 
@@ -44,6 +45,7 @@
     "MsaPrefillFmha",
     "PhasedFmha",
     "PrimsTSFmha",
+    "PrimsTSSparseMlaFmha",
     "TritonCustomMaskFmha",
     "get_enabled_fmha_lib_classes",
 ]

tensorrt_llm/_torch/attention/backends/fmha/manager.py

@@ -311,6 +311,20 @@ def __init__(self, attn: TrtllmAttention) -> None:
         for fmha_cls in get_enabled_fmha_lib_classes():
             if fmha_cls.is_available(attn):
                 self.fmha_libs.append(fmha_cls(attn))
+        self._layer_idx = getattr(attn, "layer_idx", None)
+
+    def _announce(self, fmha: Fmha, cache_key: _FmhaCacheKey) -> None:
+        # One line per library per layer is enough to tell which kernel
+        # serves a model; the logger dedups on the key.
+        logger.info_once(
+            f"FMHA layer {self._layer_idx}: {type(fmha).__name__} serves "
+            f"{cache_key.attention_mask_type.name} requests "
+            f"(context_batch<={cache_key.context_batch_size}, "
+            f"generation_batch<={cache_key.generation_batch_size}, "
+            f"generation_seq_len_q<={cache_key.generation_seq_len_q}, "
+            f"output_dtype={cache_key.output_dtype}).",
+            key=f"fmha_serves:{self._layer_idx}:{type(fmha).__name__}",
+        )
 
     def _make_cache_key(
         self,
@@ -422,6 +436,7 @@ def select(
         if fmha is None:
             return None
         self._cache[cache_key] = fmha
+        self._announce(fmha, cache_key)
         if self._cache_sanity_check_enabled:
             self._cache_inputs[cache_key] = _snapshot_fmha_cache_inputs(
                 q, k, v, metadata, forward_args

tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py

@@ -0,0 +1,281 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Parallel precompilation of PrimTS sparse MLA kernel variants.
+
+``PrimsTSSparseMlaFmha._warm_kernel_variants`` plans every kernel variant a
+deployment can select. Each variant is one ``cute.compile`` (tens of seconds),
+and FlashInfer's on-disk CuTe DSL cache serializes the builds of one op family
+behind a file lock, so N ranks warming the same variants finish no sooner than
+one rank compiling them all in sequence. This module fans the variants out to
+worker subprocesses that compile through ``_compile_cache.precompile_only()``
+(``cute.compile`` outside the lock, the artifact committed under it); the
+serving process then loads every variant from disk.
+
+Workers are plain interpreter subprocesses rather than ``multiprocessing``
+children: importing ``tensorrt_llm`` initializes MPI, so a worker must not
+inherit the launcher's MPI environment, and a subprocess gets a fully
+controlled one. Worker ``i`` of ``N`` starts with payload slots ``i, i+N, ...``
+and then goes through the rest of the list, deferring kernels another process
+is compiling (per-kernel lock) and finishing with a blocking pass, so every
+worker ends with every variant on disk whether the workspace is private to the
+rank (TRT-LLM's default) or shared between ranks. Failures are best-effort:
+the caller falls back to compiling in-process.
+
+Host load. Every worker is an interpreter that imports torch and the CuTe DSL,
+opens a CUDA context on its rank's GPU and runs ``cute.compile`` (one to a few
+GB of host memory and one busy core each). Under the default launcher each
+rank owns a private FlashInfer workspace, so the ranks of one host do not
+share compiles and every rank fans out on its own. The default worker count
+therefore follows the number of ranks on the host (from the model's
+``Mapping``, or the launcher's environment): at most 4 per rank and at most
+``_HOST_COMPILE_BUDGET`` (16) workers per host in total, so a TP=8 node runs 2
+workers per rank, not 32 subprocesses at once. Set
+``TLLM_PRIMTS_WARM_WORKERS`` explicitly to trade startup time against host
+memory and CPU on a given node.
+
+Environment:
+    TLLM_PRIMTS_WARM_WORKERS   subprocess count per rank; default 4 capped so
+                               that all ranks on the host start at most 16;
+                               0 or 1 disables.
+    TLLM_PRIMTS_WARM_TIMEOUT   seconds to wait for the workers, default 1800.
+    TLLM_PRIMTS_WARM_STRICT    1 turns a worker failure into an error (tests).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess  # nosec B404 - runs sys.executable workers with a fixed argv, no shell
+import sys
+import tempfile
+import time
+from typing import Any, Sequence
+
+from tensorrt_llm.logger import logger
+
+_WORKER_BOOTSTRAP = """
+import json, os, sys
+_path = os.environ.get("TLLM_PRIMTS_PRECOMPILE_SYS_PATH")
+if _path:
+    sys.path[:] = json.loads(_path)
+from tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile import _main
+sys.exit(_main(sys.argv))
+"""
+
+
+_MAX_DEFAULT_WORKERS = 4
+_HOST_COMPILE_BUDGET = 16  # default workers of all ranks on one host, together
+_LOCAL_RANK_ENV = ("OMPI_COMM_WORLD_LOCAL_SIZE", "RAY_LOCAL_WORLD_SIZE", "SLURM_NTASKS_PER_NODE")
+
+
+def local_rank_count(mapping) -> int:
+    """Ranks of this job on this host, without any MPI call.
+
+    From the model's ``Mapping`` (``min(world_size, gpus_per_node)``) when the
+    caller has one, else from the launcher's environment, else 1. No
+    communicator query: this runs inside the first forward, where a collective
+    would hang if the ranks reached it at different points.
+    """
+    if mapping is not None:
+        world = int(getattr(mapping, "world_size", 0) or 0)
+        per_node = int(getattr(mapping, "gpus_per_node", 0) or 0)
+        if world > 0 and per_node > 0:
+            return max(1, min(world, per_node))
+        if world > 0:
+            return world
+    for name in _LOCAL_RANK_ENV:
+        raw = os.environ.get(name)
+        if raw:
+            try:
+                return max(1, int(raw))
+            except ValueError:
+                continue
+    return 1
+
+
+def default_worker_count(local_ranks: int = 1) -> int:
+    """Per-rank default: 4, lowered so the ranks on the host start at most 16 workers."""
+    return max(1, min(_MAX_DEFAULT_WORKERS, _HOST_COMPILE_BUDGET // max(1, int(local_ranks))))
+
+
+def worker_count(num_payloads: int, *, local_ranks: int = 1) -> int:
+    """Subprocess count from ``TLLM_PRIMTS_WARM_WORKERS``; ``0``/``1`` disables.
+
+    Without the variable the default comes from ``default_worker_count()``
+    for ``local_ranks`` ranks on this host (see the module note on host load).
+    Zero when the on-disk cache or JIT compilation is disabled: a worker could
+    then compile but keep nothing.
+    """
+    if os.environ.get("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "0") == "1" or os.environ.get(
+        "FLASHINFER_DISABLE_JIT"
+    ):
+        return 0
+    raw = os.environ.get("TLLM_PRIMTS_WARM_WORKERS")
+    if raw is None:
+        workers = default_worker_count(local_ranks)
+    else:
+        try:
+            workers = int(raw)
+        except ValueError:
+            workers = default_worker_count(local_ranks)
+            logger.warning(
+                f"Ignoring TLLM_PRIMTS_WARM_WORKERS={raw!r}; expected an integer "
+                f"(using the default of {workers})."
+            )
+    return max(0, min(workers, num_payloads))
+
+
+def strict() -> bool:
+    return os.environ.get("TLLM_PRIMTS_WARM_STRICT", "0") == "1"
+
+
+def encode_payload(device_index: int, rows: int, heads: int, plan_kwargs: dict) -> dict[str, Any]:
+    """JSON-safe description of one ``BatchSparseMLADecodePagedTSWrapper.plan`` call."""
+    kwargs = dict(plan_kwargs)
+    dtype = kwargs.pop("q_data_type")
+    kwargs["q_data_type"] = str(dtype).removeprefix("torch.")
+    return {"device_index": device_index, "rows": rows, "heads": heads, "plan_kwargs": kwargs}
+
+
+def shard_order(payloads: Sequence[Any], slot: int, slots: int) -> list:
+    """Own shard first (``payloads[slot::slots]``), then the rest in list order."""
+    mine = list(payloads[slot::slots])
+    return mine + [p for p in payloads if not any(p is m for m in mine)]
+
+
+def _worker_env() -> dict[str, str]:
+    from tensorrt_llm.llmapi.mpi_session import split_mpi_env
+
+    env, _ = split_mpi_env()
+    env["TLLM_PRIMTS_WARM_WORKERS"] = "0"  # workers never fan out again
+    env.setdefault("TLLM_LOG_LEVEL", "warning")
+    # The on-disk cache key includes the installed CuTe DSL distributions as
+    # enumerated along sys.path; a worker must see the parent's sys.path
+    # exactly (entries, order and duplicates) to write the same meta.json.
+    env["TLLM_PRIMTS_PRECOMPILE_SYS_PATH"] = json.dumps(sys.path)
+    return env
+
+
+def run_precompile_workers(payloads: Sequence[dict[str, Any]], workers: int) -> bool:
+    """Compile ``payloads`` with ``workers`` subprocesses; True iff every worker succeeded.
+
+    Worker output goes to temporary files (a pipe could fill while the worker
+    holds a kernel lock another worker waits on). The wait is bounded by
+    ``TLLM_PRIMTS_WARM_TIMEOUT``; on timeout or any error every worker is
+    killed and the caller compiles in-process.
+    """
+    if workers <= 1 or not payloads:
+        return False
+    timeout_s = float(os.environ.get("TLLM_PRIMTS_WARM_TIMEOUT", "1800"))
+    env = _worker_env()
+    payload_json = json.dumps(list(payloads))
+    procs: list[tuple[int, subprocess.Popen, Any]] = []
+    ok = True
+    try:
+        for slot in range(workers):
+            out = tempfile.TemporaryFile(mode="w+", prefix="prims_ts_precompile_", suffix=".log")
+            proc = subprocess.Popen(
+                [sys.executable, "-c", _WORKER_BOOTSTRAP, payload_json, str(slot), str(workers)],
+                env=env,
+                stdout=out,
+                stderr=subprocess.STDOUT,
+                text=True,
+            )
+            procs.append((slot, proc, out))
+        deadline = time.monotonic() + timeout_s
+        for slot, proc, out in procs:
+            remaining = deadline - time.monotonic()
+            try:
+                proc.wait(timeout=max(remaining, 0.0))
+            except subprocess.TimeoutExpired:
+                ok = False
+                logger.warning(
+                    f"PrimTS precompile worker (slot {slot}) exceeded TLLM_PRIMTS_WARM_TIMEOUT="
+                    f"{timeout_s:g}s; the serving process compiles what is missing."
+                )
+                continue
+            out.seek(0)
+            text = out.read()
+            if proc.returncode != 0:
+                ok = False
+                tail = "\n".join(text.strip().splitlines()[-15:])
+                logger.warning(
+                    f"PrimTS precompile worker (slot {slot}) exited with {proc.returncode}; "
+                    f"the serving process compiles what is missing.\n{tail}"
+                )
+                continue
+            meta_line = next(
+                (
+                    line
+                    for line in text.splitlines()
+                    if line.startswith("[prims-ts precompile] meta=")
+                ),
+                "",
+            )
+            expected = str(_expected_cache_meta())
+            if meta_line and meta_line[len("[prims-ts precompile] meta=") :] != expected:
+                ok = False
+                logger.warning(
+                    f"PrimTS precompile worker (slot {slot}) wrote a different cache key than this "
+                    f"process expects; its kernels will be recompiled.\n worker: {meta_line}\n parent: {expected}"
+                )
+            else:
+                logger.debug(f"PrimTS precompile worker (slot {slot}) finished")
+    finally:
+        for _, proc, out in procs:
+            if proc.poll() is None:
+                proc.kill()
+                proc.wait()
+            out.close()
+    return ok
+
+
+def _expected_cache_meta() -> dict:
+    from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache
+
+    return _compile_cache.expected_cache_meta()
+
+
+def _plan(mla_decode, payload: dict[str, Any], torch) -> None:
+    device = torch.device("cuda", int(payload["device_index"]))
+    torch.cuda.set_device(device)
+    kwargs = dict(payload["plan_kwargs"])
+    kwargs["q_data_type"] = getattr(torch, kwargs["q_data_type"])
+    mla_decode.BatchSparseMLADecodePagedTSWrapper().plan(
+        device, int(payload["rows"]), int(payload["heads"]), **kwargs
+    )
+
+
+def _main(argv: list[str]) -> int:
+    """Worker entry: compile ``payloads[slot::slots]`` first, then everything else."""
+    payloads = json.loads(argv[1])
+    slot, slots = (int(argv[2]), int(argv[3])) if len(argv) > 3 else (0, 1)
+    import torch
+
+    from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache, mla_decode
+
+    deferred = []
+    for payload in shard_order(payloads, slot, slots):
+        try:
+            with _compile_cache.precompile_only(skip_busy=True):
+                _plan(mla_decode, payload, torch)
+        except _compile_cache.KernelBusy:
+            deferred.append(payload)
+            continue
+        print(
+            f"[prims-ts precompile] rows={payload['rows']} extra={payload['plan_kwargs'].get('max_extra_topk')} done",
+            flush=True,
+        )
+    for payload in deferred:
+        with _compile_cache.precompile_only():
+            _plan(mla_decode, payload, torch)
+        extra = payload["plan_kwargs"].get("max_extra_topk")
+        print(
+            f"[prims-ts precompile] rows={payload['rows']} extra={extra} done (waited)", flush=True
+        )
+    print(f"[prims-ts precompile] meta={_compile_cache.expected_cache_meta()}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(_main(sys.argv))