[None][feat] Support PrimsTS SageAttention for VisualGen VSA and SOL

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

@@ -1,13 +1,13 @@
 diff --git a/block_sparse.py b/block_sparse.py
-index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644
+index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe2532a1c7 100644
 --- a/block_sparse.py
 +++ b/block_sparse.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");
-@@ -36,13 +37,7 @@ from typing import Literal
+@@ -37,13 +38,7 @@ from typing import Literal
 
  import torch
 
@@ -22,7 +22,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
 
  from ._block_sparse.common import (
      _num_sparse_pattern_heads,
-@@ -301,7 +296,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
+@@ -320,7 +315,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
          # previously published revision intact and runnable.
          self._plan_state = candidate
 
@@ -31,7 +31,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
      def run(
          self,
          q: torch.Tensor,
-@@ -433,7 +428,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
+@@ -452,7 +447,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase):
          return self._launch_validated_run(state, run_args, run_stream)
 
 
@@ -40,7 +40,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
  def block_sparse_attention(
      q: torch.Tensor,
      k: torch.Tensor,
-@@ -726,9 +721,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
+@@ -757,9 +752,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
              )
          self._plan_state = candidate
 
@@ -51,7 +51,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
      def run(
          self,
          q: torch.Tensor,
-@@ -850,7 +843,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
+@@ -881,7 +874,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase):
          return self._launch_validated_run(state, run_args, run_stream)
 
 
@@ -61,10 +61,10 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/context.py b/context.py
-index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415450369df 100644
+index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d1bcc2020 100644
 --- a/context.py
 +++ b/context.py
-@@ -46,7 +46,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional
+@@ -45,7 +45,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional
 
  import torch
 
@@ -73,7 +73,7 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415
 
 
  if TYPE_CHECKING:
-@@ -488,7 +488,7 @@ def _validate_device(device: torch.device) -> int:
+@@ -490,7 +490,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):
@@ -83,7 +83,7 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415
          require_cute_dsl_arch(device_index)
      return device_index
 diff --git a/decode.py b/decode.py
-index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa6d76a609 100644
+index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329fc4a3c11 100644
 --- a/decode.py
 +++ b/decode.py
 @@ -1,3 +1,4 @@
@@ -113,7 +113,7 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa
 
          require_cute_dsl_arch(device_index)
 
-@@ -4046,7 +4043,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4048,7 +4045,7 @@ class BatchDecodePagedTSWrapper:
          # previous complete plan revision usable.
          self._plan_state = candidate
 
@@ -122,7 +122,7 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa
      def run(
          self,
          q: torch.Tensor,
-@@ -4228,7 +4225,7 @@ class BatchDecodePagedTSWrapper:
+@@ -4230,7 +4227,7 @@ class BatchDecodePagedTSWrapper:
          )
 
 
@@ -132,10 +132,10 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa
      q: torch.Tensor,
      paged_kv_cache: PagedKVCache,
 diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py
-index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb2d16f080 100644
+index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633428e32c3 100644
 --- a/kernels/fmha_context/fmha_resources.py
 +++ b/kernels/fmha_context/fmha_resources.py
-@@ -1200,8 +1200,6 @@ class GmemQKVResource(MemoryResource):
+@@ -1199,8 +1199,6 @@ class GmemQKVResource(MemoryResource):
                  if cutlass.const_expr(self.cfg.has_uniform_varlen):
                      seqlen_k = Int32(self.cfg.uniform_seq_len_k)
                  else:
@@ -144,7 +144,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb
                      seqlen_k = _load_runtime_seq_len_kv(
                          self.seqlens_kv, self.max_seq_len_kv, batch_coord
                      )
-@@ -1224,8 +1222,6 @@ class GmemQKVResource(MemoryResource):
+@@ -1223,8 +1221,6 @@ class GmemQKVResource(MemoryResource):
                      kv_request_begin = batch_coord * Int32(self.block_table_row_stride)
                      kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1)
                  else:
@@ -153,7 +153,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb
                      kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds(
                          Int32(self.block_table_row_stride),
                          self.cfg,
-@@ -1301,13 +1297,9 @@ class GmemQKVResource(MemoryResource):
+@@ -1300,13 +1296,9 @@ class GmemQKVResource(MemoryResource):
              kv_request_begin = batch_coord * Int32(self.block_table_row_stride)
              kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1)
          else:
@@ -167,7 +167,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb
              kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds(
                  Int32(self.block_table_row_stride),
                  self.cfg,
-@@ -1744,8 +1736,6 @@ class SmemPageOffsetsKvResource(MemoryResource):
+@@ -1743,8 +1735,6 @@ class SmemPageOffsetsKvResource(MemoryResource):
          kv_request_begin: Int32,
          kv_page_idx_ub: Int32,
      ) -> None:
@@ -176,7 +176,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb
          cfg = self.cfg
          # Context's K/V tile index is kv_tile_start + loop_offset; for V the
          # producer runs one tile ahead in TAIL, so reuse the same expression
-@@ -1836,8 +1826,6 @@ class SmemPageOffsetsKvResource(MemoryResource):
+@@ -1835,8 +1825,6 @@ class SmemPageOffsetsKvResource(MemoryResource):
          tile_offset: cutlass.Constexpr[int] = 0,
      ) -> cutlass.Array:
          """Retain one tile's page IDs after its SMEM window is released."""
@@ -185,7 +185,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb
          tile_idx = _resolve_kv_tile_idx_context(
              stage_info, kv_tile_start, tile_offset=tile_offset
          )
-@@ -5798,3 +5786,14 @@ class GmemOResource(MemoryResource):
+@@ -5770,3 +5758,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)

3rdparty/vendor_sources.lock.yaml

@@ -1,16 +1,16 @@
 schema_version: 1
 vendors:
   flashinfer-prims-ts:
-    url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261012
-    commit: ae312f7cc1723b06fda5168105ade0f328ee47d8
+    url: https://github.com/heyuhhh/flashinfer.git
+    branch: trtllm-prims-ts-dev
+    commit: 0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf
     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:a4360ca506d25545db408d22875e678871970fec94bdc1166dd6108d3b504d1c
-    digest: sha256-tree-v1:1bd57713d2487e94c0fc4b134bed4f7eff275282cb8ef471775e87db503ded03
+    patch_digest: sha256:1ec340aac4fea9c77b9e9d956b1676c09cff5081053ef1e41f962ed38a6f728c
+    digest: sha256-tree-v1:e4129294b2b26295581cd9cbf1ac5ac4d7eec2ec65f0108626692ddf8a036f96
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -903,6 +903,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3
     size_t const sage_k_sfs_buffer_size = sizeof(float) * mNumAttnKVHeads * static_cast<size_t>(k_max_n_blk);
     size_t const sage_v_sfs_buffer_size = sizeof(float) * static_cast<size_t>(v_max_n_blk);
     size_t const sage_k_mean_buffer_size = mSageAttnSmoothK ? sizeof(float) * local_hidden_units_kv : 0;
+    size_t const sage_k_mean_sum_buffer_size = mSageAttnSmoothK ? sizeof(int64_t) * local_hidden_units_kv : 0;
 
     size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_num_tokens;
     size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_kv_tokens;
@@ -946,6 +947,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3
     workspaceSizes.sageKScale = sage_k_sfs_buffer_size;
     workspaceSizes.sageVScale = sage_v_sfs_buffer_size;
     workspaceSizes.sageKMean = sage_k_mean_buffer_size;
+    workspaceSizes.sageKMeanSum = sage_k_mean_sum_buffer_size;
     workspaceSizes.cpWorkspace = cpWorkspaceSize;
     workspaceSizes.fmhaMultiCtasKvScratch = fmha_multi_ctas_kv_scratch_size;
     context_workspace_size = AttentionWorkspaceManager::buildContextLayout(workspaceSizes).totalSize;
@@ -1629,6 +1631,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
     size_t const sage_k_sfs_buffer_size = sizeof(float) * mNumAttnKVHeads * static_cast<size_t>(k_max_n_blk);
     size_t const sage_v_sfs_buffer_size = sizeof(float) * static_cast<size_t>(v_max_n_blk);
     size_t const sage_k_mean_buffer_size = mSageAttnSmoothK ? sizeof(float) * local_hidden_units_kv : 0;
+    size_t const sage_k_mean_sum_buffer_size = mSageAttnSmoothK ? sizeof(int64_t) * local_hidden_units_kv : 0;
 
     size_t const padding_offset_size
         = mEnableContextFMHA ? 0 : sizeof(int) * params.batch_size * params.input_seq_length;
@@ -1675,6 +1678,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
     workspaceSizes.sageKScale = sage_k_sfs_buffer_size;
     workspaceSizes.sageVScale = sage_v_sfs_buffer_size;
     workspaceSizes.sageKMean = sage_k_mean_buffer_size;
+    workspaceSizes.sageKMeanSum = sage_k_mean_sum_buffer_size;
     workspaceSizes.cpWorkspace = cpWorkspaceSize;
     workspaceSizes.fmhaMultiCtasKvScratch = fmha_multi_ctas_kv_scratch_size;
     auto const workspaceLayout = AttentionWorkspaceManager::buildContextLayout(workspaceSizes);
@@ -1959,17 +1963,12 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
             // otherwise surface as an empty scale buffer deep inside invokeSageQuant().
             TLLM_CHECK_WITH_INFO(q_max_n_blk > 0 && k_max_n_blk > 0,
                 "No SageQuant kernel on sm_%d for block sizes (q, k, v) = (%d, %d, %d) with qk_int8=%s. SM90 "
-                "requires (2, 16, 1) together with INT8 Q/K; SM100 requires q/k block sizes of 1, 4 or 16.",
+                "requires (2, 16, 1) together with INT8 Q/K; other architectures require q/k block sizes of 1, 4, 16, "
+                "or 32.",
                 mSM, mSageAttnNumEltsPerBlkQ, mSageAttnNumEltsPerBlkK, mSageAttnNumEltsPerBlkV,
                 mSageAttnQkInt8 ? "true" : "false");
             TLLM_CHECK_WITH_INFO(!params.kv_scale_quant_orig,
                 "SageAttention disregards the configured params.kv_scale_quant_orig, invalidating the result.");
-            // Reduction buffers must be initialized prior to invokeSageQuant().
-            check_cuda_error(cudaMemsetAsync(workspaceViews.sageVScale, 0, sage_v_sfs_buffer_size, stream));
-            if (mSageAttnSmoothK)
-            {
-                check_cuda_error(cudaMemsetAsync(workspaceViews.sageKMean, 0, sage_k_mean_buffer_size, stream));
-            }
 
             // Common params for sageQuant
             tc::SageQuantParams sageQuantParams{};
@@ -1984,6 +1983,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
             sageQuantParams.ptrVQuant = workspaceViews.fp8VBuf;
             sageQuantParams.ptrVScale = workspaceViews.sageVScale;
             sageQuantParams.ptrKForMean = params.k_ptr;
+            sageQuantParams.ptrKMeanSum = workspaceViews.sageKMeanSum;
             sageQuantParams.ptrKMean = workspaceViews.sageKMean;
             sageQuantParams.smCount = mMultiProcessorCount;
             sageQuantParams.stream = stream;

cpp/tensorrt_llm/common/attentionWorkspace.h

@@ -66,6 +66,7 @@ struct AttentionContextWorkspaceSizes
     size_t sageKScale{};
     size_t sageVScale{};
     size_t sageKMean{};
+    size_t sageKMeanSum{};
     size_t cpWorkspace{};
     size_t fmhaMultiCtasKvScratch{};
 };
@@ -98,6 +99,7 @@ struct AttentionContextWorkspaceLayout
     WorkspaceSlice sageKScale{};
     WorkspaceSlice sageVScale{};
     WorkspaceSlice sageKMean{};
+    WorkspaceSlice sageKMeanSum{};
     WorkspaceSlice cpWorkspace{};
     WorkspaceSlice fmhaMultiCtasKvScratch{};
     size_t totalSize{};
@@ -132,6 +134,7 @@ struct AttentionContextWorkspaceViews
     float* sageKScale{};
     float* sageVScale{};
     float* sageKMean{};
+    int64_t* sageKMeanSum{};
     T* gatherInBuffer{};
     T* gatherOutBuffer{};
     int* cuCpPartialSeqlens{};
@@ -258,6 +261,7 @@ class AttentionWorkspaceManager
         layout.sageKScale = nextSlice(offset, sizes.sageKScale, alignment);
         layout.sageVScale = nextSlice(offset, sizes.sageVScale, alignment);
         layout.sageKMean = nextSlice(offset, sizes.sageKMean, alignment);
+        layout.sageKMeanSum = nextSlice(offset, sizes.sageKMeanSum, alignment);
         layout.cpWorkspace = nextSlice(offset, sizes.cpWorkspace, alignment);
         layout.fmhaMultiCtasKvScratch = nextSlice(offset, sizes.fmhaMultiCtasKvScratch, alignment);
         layout.totalSize = offset;
@@ -296,6 +300,7 @@ class AttentionWorkspaceManager
         views.sageKScale = ptr<float>(workspace, layout.sageKScale);
         views.sageVScale = ptr<float>(workspace, layout.sageVScale);
         views.sageKMean = ptr<float>(workspace, layout.sageKMean);
+        views.sageKMeanSum = ptr<int64_t>(workspace, layout.sageKMeanSum);
 
         views.gatherInBuffer = ptr<T>(workspace, layout.cpWorkspace);
         if (views.gatherInBuffer != nullptr)