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

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

@@ -256,7 +256,7 @@ index 4769d6b6a8c6009b22abf1850fb7f2df68216f8e..5ac63fcca05d8668e4367d0843c4c6d5
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/context.py b/context.py
-index 1c84d982574d7edf31e62eb4e8ebf8575b0b674d..9feef94adb43ba3e9f608952ddbd40ad570f906f 100644
+index fae5f042194abe87b87ff7d74eef3582670c3f99..da0089bd82e09c402d76e991a462c5d0ee6ade24 100644
 --- a/context.py
 +++ b/context.py
 @@ -61,7 +61,7 @@ from .vc_attention import (
@@ -268,7 +268,7 @@ index 1c84d982574d7edf31e62eb4e8ebf8575b0b674d..9feef94adb43ba3e9f608952ddbd40ad
  from .decode import (
      _validate_native_hnd_tma_strides,
      prepare_lse,
-@@ -577,7 +577,7 @@ def _validate_device(device: torch.device) -> int:
+@@ -617,7 +617,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 1c84d982574d7edf31e62eb4e8ebf8575b0b674d..9feef94adb43ba3e9f608952ddbd40ad
          require_cute_dsl_arch(device_index)
      return device_index
 diff --git a/decode.py b/decode.py
-index d8b009efc5a9224b4c81e3b7fd51bdb27ae1cc7c..1cdad28f48e3d54758735b3dc905335aa4b29886 100644
+index 0e6aeb152ae95a28c02579e4e10c1fd844345f11..8de8b4eed84f227f0b6d42b155591b1a256aea2f 100644
 --- a/decode.py
 +++ b/decode.py
 @@ -1,3 +1,4 @@
@@ -314,7 +314,7 @@ index d8b009efc5a9224b4c81e3b7fd51bdb27ae1cc7c..1cdad28f48e3d54758735b3dc905335a
 
          require_cute_dsl_arch(device_index)
 
-@@ -4457,7 +4454,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4574,7 +4571,7 @@ class BatchDecodePagedTSWrapper:
          # previous complete plan revision usable.
          self._plan_state = candidate
 
@@ -323,7 +323,7 @@ index d8b009efc5a9224b4c81e3b7fd51bdb27ae1cc7c..1cdad28f48e3d54758735b3dc905335a
      def run(
          self,
          q: torch.Tensor,
-@@ -4668,7 +4665,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4799,7 +4796,7 @@ class BatchDecodePagedTSWrapper:
          return (result, lse) if return_lse else result
 
 
@@ -333,7 +333,7 @@ index d8b009efc5a9224b4c81e3b7fd51bdb27ae1cc7c..1cdad28f48e3d54758735b3dc905335a
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py
-index 17df005fcdb7bdd757e2dcd957b294d83b7cc41b..1b25383975ce5d5038748e57517819fd8317c22b 100644
+index 6163ed60278cb31643b09734a8d5126647b645d5..7f40b5094ae1887c23632144909bb36b53c92673 100644
 --- a/kernels/fmha_context/fmha_resources.py
 +++ b/kernels/fmha_context/fmha_resources.py
 @@ -1,3 +1,4 @@
@@ -373,7 +373,7 @@ index 17df005fcdb7bdd757e2dcd957b294d83b7cc41b..1b25383975ce5d5038748e57517819fd
              kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds(
                  Int32(self.block_table_row_stride),
                  self.cfg,
-@@ -7822,3 +7815,13 @@ class GmemOResource(MemoryResource):
+@@ -7984,3 +7977,13 @@ 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,10 +419,10 @@ index 8b2e8a43e5553ed33b67603273e0cc1fbd3585ae..fd3f735ede1f723f66d76f45a6490921
  from .gather import _address, gather4, select_gather_map
  from ...fmha_decode.fmha_decode_resources.tmem_transformed_kv import (
 diff --git a/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py b/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py
-index 0f953c00ed9a42f2eba63de8923bd3a21951245e..fd825a70072c597511d37230198964e2788a2af7 100644
+index b86948ba6a37455033daed88d7520c87270f1cc6..397da9b69057ce1854eb7da335cccbfe6f2a68fd 100644
 --- a/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py
 +++ b/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py
-@@ -353,17 +353,17 @@ def _run_parallel_gmem_reduction_g1_shared_stats(
+@@ -364,17 +364,17 @@ def _run_parallel_gmem_reduction_g1_shared_stats(
                      cutlass.const_expr(lane_slot_i == 0)
                      or Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) < split_limit
                  ):
@@ -444,15 +444,10 @@ index 0f953c00ed9a42f2eba63de8923bd3a21951245e..fd825a70072c597511d37230198964e2
 
      prims.barrier_cta_sync(
 diff --git a/mla_decode.py b/mla_decode.py
-index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7bf4e50862 100644
+index 31dab1f3d41aafbf049e6bb7cac87017a65d47e4..0c9a860440877abe71d0fc8c35da9b0d9543b1fa 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");
-@@ -40,12 +41,8 @@ from .kernels.mla_decode.helpers.nvfp4_kv import (
+@@ -41,12 +41,8 @@ from .kernels.mla_decode.helpers.nvfp4_kv import (
      validate_nvfp4_kv_formats,
  )
 
@@ -466,7 +461,7 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
 
 
  from .decode import (
-@@ -1196,8 +1193,34 @@ def _make_mla_decode_compile_spec(
+@@ -1308,8 +1304,34 @@ def _make_mla_decode_compile_spec(
  _cached_mla_decode_compile_spec = functools.cache(_make_mla_decode_compile_spec)
 
 
@@ -503,16 +498,16 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
      compile_spec: _MLADecodeCompileSpec,
  ):
      """Compile and cache one batch-dynamic TS MLA topology."""
-@@ -1897,7 +1920,7 @@ class BatchMLADecodePagedTSWrapper:
-             softmax_rescale_threshold_log2=softmax_rescale_threshold_log2,
+@@ -2067,7 +2089,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,
-@@ -2013,7 +2036,7 @@ class BatchMLADecodePagedTSWrapper:
+@@ -2209,7 +2231,7 @@ class BatchMLADecodePagedTSWrapper:
          return (result, lse) if return_lse else result
 
 
@@ -521,15 +516,15 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
  def batch_mla_decode_with_paged_kv_cache(
      query: torch.Tensor,
      kv_cache: torch.Tensor,
-@@ -2379,6 +2402,7 @@ def _prepare_sparse_output_quant(
+@@ -2588,6 +2610,7 @@ def _prepare_sparse_output_quant(
      )
 
 
 +@functools.lru_cache(maxsize=1024)
  def _resolve_sparse_mla_plan(
      device,
      batch_size,
-@@ -2411,17 +2435,17 @@ def _resolve_sparse_mla_plan(
+@@ -2620,17 +2643,17 @@ def _resolve_sparse_mla_plan(
      affect geometry or workspace sizing.
      """
      import cutlass
@@ -551,7 +546,7 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
          MlaDecodeTs,
      )
 
-@@ -2809,6 +2833,9 @@ def resolve_sparse_d576_reload_plan(
+@@ -3018,6 +3041,9 @@ def resolve_sparse_d576_reload_plan(
      return device, profile, spec, sections, byte_end, capacity, None
 
 
@@ -561,23 +556,23 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
  class BatchSparseMLADecodePagedTSWrapper:
      """Plan native BF16/E4M3 attention, then bind one or two source index lists.
 
-@@ -2859,6 +2886,7 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -3068,6 +3094,7 @@ class BatchSparseMLADecodePagedTSWrapper:
      def __init__(self, workspace_buffer=None):
          self._workspace_buffer = workspace_buffer
          self._state = None
 +        self._prepared: dict[tuple, dict] = {}
          self.workspace_size_bytes = 0
 
      def plan(
-@@ -3041,6 +3069,7 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -3250,6 +3277,7 @@ class BatchSparseMLADecodePagedTSWrapper:
                  kernel_signature=_mla_kernel_compile_signature(kernel),
              )
          compiled = _get_compiled_mla_decode(spec)
 +        self._prepared.clear()
          self._state = dict(
              reload_d576=reload_d576,
              device=device,
-@@ -3276,6 +3305,49 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -3485,6 +3513,49 @@ class BatchSparseMLADecodePagedTSWrapper:
          state = self._state
          if state is None:
              raise RuntimeError("plan() must be called before run()")
@@ -627,7 +622,7 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
          _validate_query(
              query,
              packed_query=state["packed_query"],
-@@ -3344,15 +3416,6 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -3553,15 +3624,6 @@ class BatchSparseMLADecodePagedTSWrapper:
              raise ValueError("attention_sinks must be contiguous CUDA FP32[H]")
          if validate and torch.isnan(attention_sinks).any().item():
              raise ValueError("attention_sinks must not contain NaN")
@@ -643,7 +638,7 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
          if state["kv_cache_format"] != "fp8" and (
              isinstance(kv_cache_scale, torch.Tensor) or kv_cache_scale != 1.0
          ):
-@@ -3479,10 +3542,127 @@ class BatchSparseMLADecodePagedTSWrapper:
+@@ -3688,10 +3750,129 @@ class BatchSparseMLADecodePagedTSWrapper:
                  compiled=state["compiled"],
                  max_q_len=1,
              )
@@ -765,6 +760,8 @@ index 94547ed86f7365df3638215daa0762ff1344741f..d0916ceffe0839ff9dd5c26fcf380a7b
 +            bmm2,
 +            None,
 +            1,
++            None,
++            None,
 +        )
 +        return out
 +

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-mla-prefill--flashinfer
+    commit: 26b197614300b549eb103fc59d736c3a0e4b3fa9
     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:61470f1c11d8155475ccd791074df1936b783895b03e76a73a29e8d76fe21f06
+    digest: sha256-tree-v1:32178a44da603dad9423890c27c4c14f68b89b84c0e75e3b6020afeb0733742a
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,18 +26,19 @@
 namespace
 {
 
-template <typename T>
+template <typename T, int HeadSize>
 struct MergeSoftmaxTraits
 {
-    static constexpr int kQKNopeSize = 128;
-    static constexpr int kHeadSize = kQKNopeSize;
+    static constexpr int kHeadSize = HeadSize;
 
     static constexpr int kBytesPerElem = sizeof(T);
     static constexpr int kBytesPerLoad = 16;
     static constexpr int kElemPerThread = kBytesPerLoad / sizeof(T);
     static_assert((kHeadSize * kBytesPerElem) % kBytesPerLoad == 0,
         "kHeadSize * kBytesPerElem must be multiple of kBytesPerLoad (16Bytes)");
     static constexpr int kVecPerHead = (kHeadSize * kBytesPerElem) / kBytesPerLoad;
+    // All readers of aliased softmax stats must stay in the same warp as their writer.
+    static_assert(kVecPerHead <= 32, "In-place softmax stats require at most one warp per head");
     static constexpr int kTokenPerBlock
         = std::is_same_v<T, float> ? 4 : 8; // for each block, we fetch 8 token for fp16, 4 tokens for fp32.
     static constexpr int kNumThreads = kVecPerHead * kTokenPerBlock;
@@ -136,14 +137,14 @@ inline __device__ void dequantCopy(
     }
 }
 
-// merged_attn [q_total_len, H=128, D=128] (T)
+// merged_attn [q_total_len, num_heads, HeadSize] (T)
 // merged_softmax_sum [q_total_len, H, 2] (float, max/sum)
-template <typename T>
+template <typename T, int HeadSize>
 __global__ void mergeAttnWithSoftmaxKernel(T* merged_attn, float2* merged_softmax_stats, T const* pre_attn,
     float2 const* pre_softmax_stats, T const* curr_attn, float2 const* curr_softmax_stats, int64_t const* cu_q_seq_len,
     int64_t const* merge_op, int const num_heads, int const head_size)
 {
-    using KT = MergeSoftmaxTraits<T>;
+    using KT = MergeSoftmaxTraits<T, HeadSize>;
     int const batch_idx = static_cast<int>(blockIdx.y);
     int const head_idx = static_cast<int>(blockIdx.z);
 
@@ -221,6 +222,20 @@ __global__ void mergeAttnWithSoftmaxKernel(T* merged_attn, float2* merged_softma
     }
 }
 
+template <typename T, int HeadSize>
+void launchMergeAttnWithSoftmax(T* merged_attn, float* merged_softmax_stats, T const* pre_attn,
+    float const* pre_softmax_stats, T const* curr_attn, float const* curr_softmax_stats, int const batch_size,
+    int64_t const* cu_q_seq_len, int max_q_seq_len, int64_t const* merge_op, int const num_heads, cudaStream_t stream)
+{
+    using KT = MergeSoftmaxTraits<T, HeadSize>;
+    dim3 grid(static_cast<int>(tensorrt_llm::common::divUp(max_q_seq_len, KT::kTokenPerBlock)), batch_size, num_heads);
+    dim3 block(KT::kNumThreads);
+
+    mergeAttnWithSoftmaxKernel<T, HeadSize><<<grid, block, 0, stream>>>(merged_attn,
+        reinterpret_cast<float2*>(merged_softmax_stats), pre_attn, reinterpret_cast<float2 const*>(pre_softmax_stats),
+        curr_attn, reinterpret_cast<float2 const*>(curr_softmax_stats), cu_q_seq_len, merge_op, num_heads, HeadSize);
+}
+
 // kv_output {total_chunk_token=b*chunk_size, h=1, d_lora}
 // k_pe_output {total_chunk_token, h=1, d_rope}
 template <typename T, typename TCache>
@@ -296,7 +311,7 @@ TRTLLM_NAMESPACE_BEGIN
 namespace kernels
 {
 
-// merged_attn [q_total_len, H=128, D=128] (T)
+// merged_attn [q_total_len, num_heads, head_size] (T)
 // merged_softmax_sum [q_total_len, H, 2] (float), the first part is the max value for each
 // row of P = QK^T, the second part is the softmax sum
 // if merge_op[b] == 0, we just skip this batch, if merge_op[b] == 1, we merge the pre-attn and curr-attn, if
@@ -308,15 +323,23 @@ void invokeMergeAttnWithSoftmax(T* merged_attn, float* merged_softmax_stats, T c
     int64_t const* cu_q_seq_len, int max_q_seq_len, int64_t const* merge_op, int const num_heads, int const head_size,
     cudaStream_t stream)
 {
-    using KT = MergeSoftmaxTraits<T>;
-    TLLM_CHECK_WITH_INFO(head_size == KT::kHeadSize, "head dim should be equal to %d", KT::kHeadSize);
-
-    dim3 grid(static_cast<int>(tensorrt_llm::common::divUp(max_q_seq_len, KT::kTokenPerBlock)), batch_size, num_heads);
-    dim3 block(KT::kNumThreads);
-
-    mergeAttnWithSoftmaxKernel<T><<<grid, block, 0, stream>>>(merged_attn,
-        reinterpret_cast<float2*>(merged_softmax_stats), pre_attn, reinterpret_cast<float2 const*>(pre_softmax_stats),
-        curr_attn, reinterpret_cast<float2 const*>(curr_softmax_stats), cu_q_seq_len, merge_op, num_heads, head_size);
+    constexpr int kHeadSize128 = 128;
+    constexpr int kHeadSize256 = 256;
+    if (head_size == kHeadSize128)
+    {
+        launchMergeAttnWithSoftmax<T, kHeadSize128>(merged_attn, merged_softmax_stats, pre_attn, pre_softmax_stats,
+            curr_attn, curr_softmax_stats, batch_size, cu_q_seq_len, max_q_seq_len, merge_op, num_heads, stream);
+    }
+    else if constexpr (!std::is_same_v<T, float>)
+    {
+        TLLM_CHECK_WITH_INFO(head_size == kHeadSize256, "head dim should be 128 or 256 for FP16/BF16");
+        launchMergeAttnWithSoftmax<T, kHeadSize256>(merged_attn, merged_softmax_stats, pre_attn, pre_softmax_stats,
+            curr_attn, curr_softmax_stats, batch_size, cu_q_seq_len, max_q_seq_len, merge_op, num_heads, stream);
+    }
+    else
+    {
+        TLLM_CHECK_WITH_INFO(false, "head dim should be equal to 128 for FP32");
+    }
 }
 
 // load single chunk kv from kv_cache for each request