[None][feat] Add PrimTS MLA prefill, MTP, and Rubin support

3rdparty/vendor_patches/flashinfer-prims-ts.patch

@@ -256,7 +256,7 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/context.py b/context.py
-index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9592697ad 100644
+index 1cab933986621c42e01ae1f42fe8c23c99e94b04..b41e7b47cca1ed11eba11bddf1b0426c51071681 100644
 --- a/context.py
 +++ b/context.py
 @@ -48,7 +48,7 @@ from ._hardware import get_max_active_clusters
@@ -268,7 +268,7 @@ index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9
 
 
  if TYPE_CHECKING:
-@@ -493,7 +493,7 @@ def _validate_device(device: torch.device) -> int:
+@@ -500,7 +500,7 @@ def _validate_device(device: torch.device) -> int:
      # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8
      # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import.
      if capability == (10, 7):
@@ -278,7 +278,7 @@ index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9
          require_cute_dsl_arch(device_index)
      return device_index
 diff --git a/decode.py b/decode.py
-index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac722fed527 100644
+index 8c331d05a5b55717df3ba4dd1ff55b6ca44c831e..9c0a5e7243ee2d4ad33eed1afd45f9b8aa87f9e7 100644
 --- a/decode.py
 +++ b/decode.py
 @@ -1,3 +1,4 @@
@@ -308,7 +308,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7
 
          require_cute_dsl_arch(device_index)
 
-@@ -4075,7 +4072,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4143,7 +4140,7 @@ class BatchDecodePagedTSWrapper:
          # previous complete plan revision usable.
          self._plan_state = candidate
 
@@ -317,7 +317,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7
      def run(
          self,
          q: torch.Tensor,
-@@ -4257,7 +4254,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4340,7 +4337,7 @@ class BatchDecodePagedTSWrapper:
          )
 
 
@@ -327,7 +327,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py
-index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2fe30a922 100644
+index d29e7145f1da027ec86be7247a64a0a7d730c102..c4b11379ab3076644f4e7bdadf9948049bb23556 100644
 --- a/kernels/fmha_context/fmha_resources.py
 +++ b/kernels/fmha_context/fmha_resources.py
 @@ -1,3 +1,4 @@
@@ -385,7 +385,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2
          tile_idx = _resolve_kv_tile_idx_context(
              stage_info, kv_tile_start, tile_offset=tile_offset
          )
-@@ -5770,3 +5759,14 @@ class GmemOResource(MemoryResource):
+@@ -5942,3 +5931,14 @@ class GmemOResource(MemoryResource):
                  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)
@@ -419,15 +419,10 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a
  from .helpers_common import (
      _TASK_CACHE_TMEM_BASE_OFFSET,
 diff --git a/mla_decode.py b/mla_decode.py
-index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e231c98a52 100644
+index 7774f881518b11f8e3fc0626afbcd757c77ca6cc..3c3a5d2197cfbc33fcedc7d6c2fb21ebaac2a74d 100644
 --- a/mla_decode.py
 +++ b/mla_decode.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");
-@@ -27,11 +28,7 @@ from ._hardware import get_max_active_clusters
+@@ -28,11 +28,7 @@ from ._hardware import get_max_active_clusters
 
  import torch
 
@@ -440,7 +435,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
 
  from .decode import (
      _WorkspaceSection,
-@@ -1047,11 +1044,35 @@ def _make_mla_decode_compile_spec(
+@@ -1155,11 +1151,35 @@ def _make_mla_decode_compile_spec(
      )
 
 
@@ -481,34 +476,16 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
 
      import cutlass
      import cutlass.cute as cute
-@@ -1459,11 +1480,12 @@ def _launch_mla_decode(
-         runtime.sparse_inputs,
-         runtime.bmm1_scale,
-         runtime.bmm2_scale,
--        **(
--            {"output_quant_args": runtime.output_quant_args}
--            if runtime.output_quant_args is not None
--            else {}
--        ),
-+        # Passed positionally and also when None. The kernel was compiled with
-+        # ``output_quant_args`` in its signature (``None`` for unquantized
-+        # output), and a kernel reloaded from the persistent compile cache is
-+        # a positional TVM-FFI entry with no Python defaults: it rejects both
-+        # a missing parameter and a keyword argument.
-+        runtime.output_quant_args,
-     )
-     return runtime.out
-
-@@ -1637,7 +1659,7 @@ class BatchMLADecodePagedTSWrapper:
-             split_kv=int(dict(policy)["split_kv"]),
+@@ -1798,7 +1818,7 @@ class BatchMLADecodePagedTSWrapper:
+             use_kv_bounds=use_kv_bounds,
          )
 
 -    @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch)
 +    @flashinfer_experimental_api
      def run(
          self,
          query: torch.Tensor,
-@@ -1742,7 +1764,7 @@ class BatchMLADecodePagedTSWrapper:
+@@ -1927,7 +1947,7 @@ class BatchMLADecodePagedTSWrapper:
          )
 
 
@@ -517,15 +494,15 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
  def batch_mla_decode_with_paged_kv_cache(
      query: torch.Tensor,
      kv_cache: torch.Tensor,
-@@ -2099,6 +2121,7 @@ def _prepare_sparse_output_quant(
+@@ -2306,6 +2326,7 @@ def _prepare_sparse_output_quant(
      )
 
 
 +@functools.lru_cache(maxsize=1024)
  def _resolve_sparse_mla_plan(
      device,
      batch_size,
-@@ -2119,17 +2142,17 @@ def _resolve_sparse_mla_plan(
+@@ -2326,17 +2347,17 @@ def _resolve_sparse_mla_plan(
  ):
      """Resolve sparse geometry and scratch sizing without allocation/compilation."""
      import cutlass
@@ -547,7 +524,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
          MlaDecodeTs,
      )
 
-@@ -2264,6 +2287,9 @@ def _resolve_sparse_mla_plan(
+@@ -2471,6 +2492,9 @@ def _resolve_sparse_mla_plan(
      return device, profile, spec, sections, byte_end, capacity, output_quant
 
 
@@ -557,7 +534,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
  class BatchSparseMLADecodePagedTSWrapper:
      """Plan native BF16/E4M3 attention, then bind one or two source index lists.
 
-@@ -2297,6 +2323,11 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -2504,6 +2528,11 @@ class BatchSparseMLADecodePagedTSWrapper:
          self._workspace_buffer = workspace_buffer
          self._state = None
          self._scalar_cache = {}
@@ -569,15 +546,15 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
          self.workspace_size_bytes = 0
 
      def plan(
-@@ -2386,6 +2417,7 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -2593,6 +2622,7 @@ class BatchSparseMLADecodePagedTSWrapper:
                  kernel_signature=_mla_kernel_compile_signature(static_kernel),
              )
          )
 +        self._prepared.clear()
          self._state = dict(
              device=device,
              batch=batch_size,
-@@ -2608,6 +2640,38 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -2815,6 +2845,38 @@ class BatchSparseMLADecodePagedTSWrapper:
          state = self._state
          if state is None:
              raise RuntimeError("plan() must be called before run()")
@@ -616,7 +593,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
          _validate_query(
              query,
              packed_query=state["packed"],
-@@ -2781,10 +2845,124 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -2988,10 +3050,126 @@ class BatchSparseMLADecodePagedTSWrapper:
                  if static_scales
                  else state["compiled"],
              )
@@ -735,6 +712,8 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2
 +            bmm1,
 +            bmm2,
 +            None,
++            None,
++            None,
 +        )
 +        return out
 +

3rdparty/vendor_sources.lock.yaml

@@ -14,15 +14,15 @@ vendors:
     digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261012
-    commit: aa709061230a58858375c513056e7733fa74294a
+    branch: prims-ts-mla-prefill--flashinfer
+    commit: 069139a65aad98f3e5c61f049c761e28eff0b14e
     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:5cbcaf2b60029742c55f622eadc84b32f49e574038c417972955706157715141
-    digest: sha256-tree-v1:fa20dc33e52628519dba7f071f5b28a4eb7963768adb6b46eb8000e4fdb805c4
+    patch_digest: sha256:011e0dab0ea60efff540c3ce80a96970f1be9c85f6f5b67522d157c31fa710fa
+    digest: sha256-tree-v1:a5e826af66f2cb96c5c7ce1bb785d5a413e2d2a997d7c5057fdb99080ee8c5db
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp

@@ -105,6 +105,132 @@ void mergeChunkedAttentionForMLAHelper(torch::Tensor& merged_attn, torch::Tensor
 
 } // namespace
 
+// Shared dense-MLA context prologue, without dispatching an attention kernel.
+// Cached/chunked callers pass no latent_cache: their Q/K RoPE and cache append
+// have already run. Output buffers are owned and accounted for by the caller.
+void mlaContextPreprocess(torch::Tensor q, torch::Tensor k, torch::Tensor const& v,
+    torch::optional<torch::Tensor> latent_cache, torch::Tensor const& cu_q_seqlens,
+    torch::Tensor const& sequence_lengths, torch::Tensor const& cos_sin_cache,
+    torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers,
+    torch::Tensor const& host_kv_cache_pool_mapping, torch::optional<torch::Tensor> kv_scale_orig_quant,
+    torch::optional<torch::Tensor> kv_scale_quant_orig, int64_t num_heads, int64_t max_input_seq_len, int64_t layer_idx,
+    int64_t tokens_per_block, int64_t attention_window_size, int64_t quant_mode, double sm_scale, torch::Tensor q_out,
+    torch::Tensor k_out, torch::Tensor v_out, torch::Tensor bmm1_scale, torch::Tensor bmm2_scale)
+{
+    TORCH_CHECK(q.scalar_type() == torch::kBFloat16 || q.scalar_type() == torch::kFloat16);
+    CHECK_INPUT(q, q.scalar_type());
+    CHECK_INPUT(k, q.scalar_type());
+    TORCH_CHECK(q.dim() == 2 && k.dim() == 2 && v.dim() == 2 && num_heads > 0);
+    TORCH_CHECK(q.size(1) == num_heads * 192 && k.size(1) == num_heads * 192);
+    TORCH_CHECK(v.device() == q.device() && k.device() == q.device());
+    TORCH_CHECK(v.scalar_type() == q.scalar_type() && v.size(0) == k.size(0));
+    TORCH_CHECK(v.size(1) == num_heads * 128 && v.stride(1) == 1);
+    CHECK_INPUT(cu_q_seqlens, torch::kInt32);
+    CHECK_INPUT(sequence_lengths, torch::kInt32);
+    CHECK_INPUT(cos_sin_cache, torch::kFloat32);
+    TORCH_CHECK(cu_q_seqlens.device() == q.device() && sequence_lengths.device() == q.device()
+        && cos_sin_cache.device() == q.device());
+    TORCH_CHECK(cu_q_seqlens.dim() == 1 && cu_q_seqlens.numel() > 1);
+    auto const batch_size = cu_q_seqlens.numel() - 1;
+    TORCH_CHECK(sequence_lengths.numel() >= batch_size && max_input_seq_len > 0);
+    TORCH_CHECK(cos_sin_cache.numel() >= attention_window_size * 64);
+    auto const mode = tc::QuantMode(static_cast<uint32_t>(quant_mode));
+    TORCH_CHECK(!mode.hasInt8KvCache() && !mode.hasFp4KvCache());
+    bool const fp8 = mode.hasFp8KvCache();
+    auto const output_dtype = fp8 ? torch::kFloat8_e4m3fn : q.scalar_type();
+    for (auto const& tensor : {q_out, k_out, v_out})
+    {
+        CHECK_INPUT(tensor, output_dtype);
+        TORCH_CHECK(tensor.device() == q.device());
+    }
+    TORCH_CHECK(q_out.sizes() == q.sizes() && k_out.sizes() == k.sizes() && v_out.sizes() == v.sizes());
+    CHECK_INPUT(bmm1_scale, torch::kFloat32);
+    CHECK_INPUT(bmm2_scale, torch::kFloat32);
+    TORCH_CHECK(bmm1_scale.numel() == 2 && bmm2_scale.numel() == 1);
+    TORCH_CHECK(bmm1_scale.device() == q.device() && bmm2_scale.device() == q.device());
+    if (fp8)
+    {
+        // The shared quantizer consumes the kv_b_proj split view of V.
+        TORCH_CHECK(v.stride(0) == num_heads * 256);
+    }
+    else
+    {
+        TORCH_CHECK(q_out.data_ptr() == q.data_ptr() && k_out.data_ptr() == k.data_ptr());
+    }
+    for (auto const& scale : {kv_scale_orig_quant, kv_scale_quant_orig})
+    {
+        if (scale.has_value())
+        {
+            CHECK_INPUT(scale.value(), torch::kFloat32);
+            TORCH_CHECK(scale->device() == q.device() && scale->numel() >= 1);
+        }
+    }
+    if (latent_cache.has_value())
+    {
+        CHECK_INPUT(latent_cache.value(), q.scalar_type());
+        TORCH_CHECK(latent_cache->device() == q.device());
+        TORCH_CHECK(latent_cache->numel() == q.size(0) * 576 && k.size(0) == q.size(0));
+    }
+    auto cache
+        = buildPagedKvCacheBuffers(std::optional(kv_cache_block_offsets), std::optional(host_kv_cache_pool_pointers),
+            std::optional(host_kv_cache_pool_mapping), mode, layer_idx, batch_size, tokens_per_block, 1, 576,
+            attention_window_size, attention_window_size, 1, 0, true, q.element_size())
+              .kvCacheBuffer;
+    auto stream = at::cuda::getCurrentCUDAStream(q.get_device());
+    auto launch = [&](auto type_tag)
+    {
+        using T = decltype(type_tag);
+        tk::MlaParams<T> params{};
+        params.q_buf = static_cast<T*>(q.data_ptr());
+        params.k_buf = static_cast<T*>(k.data_ptr());
+        params.v_buf = static_cast<T const*>(v.data_ptr());
+        params.latent_cache = latent_cache ? static_cast<T const*>(latent_cache->data_ptr()) : nullptr;
+        params.cu_q_seqlens = cu_q_seqlens.data_ptr<int>();
+        params.cache_seq_lens = sequence_lengths.data_ptr<int>();
+        params.cos_sin_cache = static_cast<float2 const*>(cos_sin_cache.data_ptr());
+        params.batch_size = batch_size;
+        params.acc_q_len = q.size(0);
+        params.head_num = num_heads;
+        params.max_input_seq_len = max_input_seq_len;
+        params.meta.kv_lora_rank = 512;
+        params.meta.qk_nope_head_dim = 128;
+        params.meta.qk_rope_head_dim = 64;
+        params.meta.v_head_dim = 128;
+        params.cache_type = tk::cacheTypeFromQuantMode(mode);
+        params.quant_scale_kv = kv_scale_orig_quant ? kv_scale_orig_quant->data_ptr<float>() : nullptr;
+        if (latent_cache)
+        {
+            tk::invokeMLARopeContext(params, cache, stream);
+        }
+        if (fp8)
+        {
+            params.quant_q_buf = q_out.data_ptr();
+            params.quant_k_buf = k_out.data_ptr();
+            params.quant_v_buf = v_out.data_ptr();
+            params.quant_scale_qkv = params.quant_scale_kv;
+            params.dequant_scale_q = kv_scale_quant_orig ? kv_scale_quant_orig->data_ptr<float>() : nullptr;
+            params.dequant_scale_kv = params.dequant_scale_q;
+            params.host_bmm1_scale = sm_scale;
+            params.bmm1_scale = bmm1_scale.data_ptr<float>();
+            params.bmm2_scale = bmm2_scale.data_ptr<float>();
+            tk::invokeMLAContextFp8Quantize(params, k.size(0), stream);
+        }
+    };
+    if (q.scalar_type() == torch::kBFloat16)
+    {
+        launch(__nv_bfloat16{});
+    }
+    else
+    {
+        launch(half{});
+    }
+    if (!fp8)
+    {
+        // Q/K alias the caller's inputs; only V's split projection needs packing.
+        v_out.copy_(v);
+    }
+}
+
 std::vector<torch::Tensor> loadPagedKVCacheForMLA(torch::ScalarType out_dtype, int64_t const num_contexts,
     int64_t const num_ctx_cached_tokens, int64_t const max_ctx_cached_kv_len, torch::Tensor& cu_ctx_cached_kv_lens,
     torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers,
@@ -419,6 +545,13 @@ TRTLLM_NAMESPACE_END
 
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
+    m.def(
+        "mla_context_preprocess(Tensor(a!) q, Tensor(b!) k, Tensor v, Tensor? latent_cache, "
+        "Tensor cu_q_seqlens, Tensor sequence_lengths, Tensor cos_sin_cache, Tensor kv_cache_block_offsets, "
+        "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_scale_orig_quant, "
+        "Tensor? kv_scale_quant_orig, int num_heads, int max_input_seq_len, int layer_idx, int tokens_per_block, "
+        "int attention_window_size, int quant_mode, float sm_scale, Tensor(a!) q_out, Tensor(b!) k_out, "
+        "Tensor(c!) v_out, Tensor(d!) bmm1_scale, Tensor(e!) bmm2_scale) -> ()");
     m.def(
         "load_paged_kv_cache_for_mla("
         "ScalarType out_dtype"
@@ -442,6 +575,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
+    m.impl("mla_context_preprocess", &tensorrt_llm::torch_ext::mlaContextPreprocess);
     m.impl("load_paged_kv_cache_for_mla", &tensorrt_llm::torch_ext::loadPagedKVCacheForMLA);
 }