[TRTLLM-34871][feat] Add cute dsl FP8 paged MQA logits decode kernel

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -893,13 +893,15 @@ def prepare(self):
         # Because the fp8_paged_mqa_logits only supports seq_len == 1/2/4 (i.e., max_draft_tokens == 0/1/3) on sm100, and
         # seq_len == 1/2 (i.e., max_draft_tokens == 0/1) on sm90, for other cases, we need to flatten the q tensor and
         # expand the kv_lens and block_table for MTP support.
+        # The CuTe DSL kernel supports arbitrary next_n natively, so it never needs expansion.
         # TODO:
         # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90.
         # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens.
-        self.use_expanded_buffers_for_mtp = (
-            (self.max_draft_tokens > 1 and get_sm_version() == 90)
-            or ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3)
-                and get_sm_version() >= 100))
+        _use_dsl = self.sparse_attention_config.use_cute_dsl_paged_mqa_logits
+        self.use_expanded_buffers_for_mtp = (not _use_dsl and (
+            (self.max_draft_tokens > 1 and get_sm_version() == 90) or
+            ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3)
+             and get_sm_version() >= 100)))
         if self.use_expanded_buffers_for_mtp:
             # Expand kv_lens_cuda (only generation)
             num_tokens = self.num_generations * (1 + self.max_draft_tokens)
@@ -1109,19 +1111,24 @@ def __init__(self,
         self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
         self.use_cute_dsl_topk = (sparse_attention_config.use_cute_dsl_topk
                                   and IS_CUTLASS_DSL_AVAILABLE)
+        self.use_cute_dsl_paged_mqa_logits = (
+            sparse_attention_config.use_cute_dsl_paged_mqa_logits
+            and IS_CUTLASS_DSL_AVAILABLE)
         self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5
 
         self._enable_heuristic_topk = (
             sparse_attention_config.enable_heuristic_topk
             and get_sm_version() >= 100)
 
-        if self.use_cute_dsl_topk and layer_idx == 0:
+        if (self.use_cute_dsl_topk
+                or self.use_cute_dsl_paged_mqa_logits) and layer_idx == 0:
             from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops
 
-            # the dtype of topk input tensor, which is float32 now.
-            # Note, need to update it if the dtype of topk input tensor is changed.
-            cute_dsl_custom_ops.warmup_cute_dsl_indexer_topk(
-                dtype=torch.float32, top_k=self.index_topk)
+            if self.use_cute_dsl_topk:
+                # the dtype of topk input tensor, which is float32 now.
+                # Note, need to update it if the dtype of topk input tensor is changed.
+                cute_dsl_custom_ops.warmup_cute_dsl_indexer_topk(
+                    dtype=torch.float32, top_k=self.index_topk)
 
     def post_load_weights(self):
         """Fuse wk + weights_proj into single FP32 weight for cuBLAS GEMM (TF32 on Ampere+)."""
@@ -1643,7 +1650,10 @@ def sparse_attn_indexer(
                     num_contexts:num_contexts + num_generations]
                 block_table = metadata.indexer_k_cache_block_offsets[
                     num_contexts:num_contexts + num_generations]
-                if q_decode.shape[1] == 4:
+                # DeepGEMM uses cluster(2,1,1) for next_n=4, requiring halved num_sms metadata.
+                # DSL kernel always uses cluster(1,1,1), so it always uses the full num_sms buffer.
+                if q_decode.shape[
+                        1] == 4 and not self.use_cute_dsl_paged_mqa_logits:
                     scheduler_metadata_buffer = metadata.scheduler_metadata_buffer_mtp3
                 else:
                     scheduler_metadata_buffer = metadata.scheduler_metadata_buffer
@@ -1663,11 +1673,16 @@ def sparse_attn_indexer(
             k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(
                 self.layer_idx)
 
-            logits_decode = fp8_paged_mqa_logits(q_decode, k_cache,
-                                                 weights_decode, context_lens,
-                                                 block_table,
-                                                 scheduler_metadata_buffer,
-                                                 max_seq_len)
+            if self.use_cute_dsl_paged_mqa_logits:
+                logits_decode = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits(
+                    q_decode, k_cache, weights_decode, context_lens,
+                    block_table, scheduler_metadata_buffer, max_seq_len)
+            else:
+                logits_decode = fp8_paged_mqa_logits(q_decode, k_cache,
+                                                     weights_decode,
+                                                     context_lens, block_table,
+                                                     scheduler_metadata_buffer,
+                                                     max_seq_len)
 
             if use_custom_topk:
                 # Kernel expects kv_lens (total cache length), not seq_lens (new tokens)

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -5127,3 +5127,267 @@ def warmup_cute_dsl_indexer_topk(
             f"Warmed up CuTE DSL indexer top-k kernels: dtype={dtype}, "
             f"SingleCTA bucketed_num_cols=[2^{min_seq_len_log2}..2^{max_seq_len_log2}], "
             f"{multi_cta_info}, top_k={top_k}, next_n={next_n}")
+
+    # ------------------------------------------------------------------ #
+    #  CuTE DSL FP8 Paged MQA Logits (Blackwell SM100)                   #
+    # ------------------------------------------------------------------ #
+    from ..cute_dsl_kernels.blackwell.paged_mqa_logits import FP8MQALogitsKernel
+
+    def _check_fp8_paged_mqa_logits_dtypes(q, kv_fused, weights, context_lens,
+                                           block_table, schedule_meta,
+                                           epi_dtype, acc_dtype, output_dtype):
+        errs = []
+        if q.dtype != torch.float8_e4m3fn:
+            errs.append(f"q must be float8_e4m3fn, got {q.dtype}")
+        if kv_fused.dtype != torch.uint8:
+            errs.append(f"kv_fused must be uint8, got {kv_fused.dtype}")
+        # TODO: update to (torch.float32, torch.float16) once fp16 weights
+        # are validated end-to-end and the in-kernel .half() conversion is removed.
+        if weights.dtype != torch.float32:
+            errs.append(f"weights must be float32, got {weights.dtype}")
+        if context_lens.dtype != torch.int32:
+            errs.append(f"context_lens must be int32, got {context_lens.dtype}")
+        if block_table.dtype != torch.int32:
+            errs.append(f"block_table must be int32, got {block_table.dtype}")
+        if schedule_meta.dtype != torch.int32:
+            errs.append(
+                f"schedule_meta must be int32, got {schedule_meta.dtype}")
+        for name, dt in [("epi_dtype", epi_dtype), ("acc_dtype", acc_dtype),
+                         ("output_dtype", output_dtype)]:
+            if dt not in (torch.float16, torch.float32):
+                errs.append(f"{name} must be float16 or float32, got {dt}")
+        if errs:
+            raise ValueError("FP8 Paged MQA Logits dtype errors:\n  " +
+                             "\n  ".join(errs))
+
+    class CuteDSLPagedMQALogitsRunner:
+        """Runner for CuTe DSL FP8 Paged MQA Logits kernel (Blackwell SM100).
+
+        Caches compiled kernels keyed by static params
+        (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n, num_sms).
+        """
+
+        kernel_cache = dict()
+
+        @classmethod
+        def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim,
+                     next_n, num_sms, num_epi_subtiles, epi_dtype, acc_dtype,
+                     output_dtype):
+            """Compile kernel using fake tensors + TVM FFI."""
+            key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n,
+                   num_sms, num_epi_subtiles, epi_dtype, acc_dtype,
+                   output_dtype)
+            if key in cls.kernel_cache:
+                return
+
+            to_cutlass = _TORCH_TO_CUTLASS_DTYPE
+            N = next_n * num_heads
+            block_bytes = phys_block_kv * (head_dim + 4)
+
+            sym_num_phys_blocks = cute.sym_int()
+            sym_B = cute.sym_int()
+            max_ctx = cute.sym_int()
+            max_blocks_per_seq = cute.sym_int()
+            num_ctas = cute.sym_int()
+
+            kv_fake = cute.runtime.make_fake_compact_tensor(
+                cutlass.Uint8, (sym_num_phys_blocks, block_bytes),
+                stride_order=(1, 0))
+
+            q_fake = cute.runtime.make_fake_compact_tensor(cutlass.Uint8,
+                                                           (N, head_dim, sym_B),
+                                                           stride_order=(1, 0,
+                                                                         2))
+
+            w_dtype = (cutlass.Float16
+                       if epi_dtype == torch.float16 else to_cutlass[epi_dtype])
+            w_fake = cute.runtime.make_fake_compact_tensor(w_dtype, (N, sym_B),
+                                                           stride_order=(0, 1))
+
+            logits_fake = cute.runtime.make_fake_tensor(
+                to_cutlass[output_dtype], (cute.sym_int(), max_ctx),
+                stride=(cute.sym_int64(), 1))
+
+            bt_fake = cute.runtime.make_fake_compact_tensor(
+                cutlass.Int32, (sym_B, max_blocks_per_seq), stride_order=(1, 0))
+
+            cl_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32,
+                                                            (sym_B, ),
+                                                            stride_order=(0, ))
+
+            sm_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int32,
+                                                            (num_ctas, 2),
+                                                            stride_order=(1, 0))
+
+            fake_stream = cute.runtime.make_fake_stream(
+                use_tvm_ffi_env_stream=True)
+
+            kernel = FP8MQALogitsKernel(
+                block_kv=compute_block_kv,
+                phys_block_kv=phys_block_kv,
+                num_heads=num_heads,
+                head_dim=head_dim,
+                next_n=next_n,
+                num_sms=num_sms,
+                num_epi_subtiles=num_epi_subtiles,
+                epi_dtype=to_cutlass[epi_dtype],
+                acc_dtype=to_cutlass[acc_dtype],
+                output_dtype=to_cutlass[output_dtype],
+            )
+
+            compiled = cute.compile(
+                kernel,
+                kv_fake,
+                q_fake,
+                w_fake,
+                logits_fake,
+                bt_fake,
+                cl_fake,
+                sm_fake,
+                cutlass.Int32(1),
+                cutlass.Int32(1),
+                fake_stream,
+                options="--enable-tvm-ffi",
+            )
+            cls.kernel_cache[key] = compiled
+            logger.debug(f"[compile cute_dsl fp8_paged_mqa_logits] {key}"
+                         f" kv_stages={kernel.num_kv_stages}"
+                         f" umma_stages={kernel.num_umma_stages}")
+
+        @classmethod
+        def forward(
+            cls,
+            q: torch.Tensor,
+            kv_fused: torch.Tensor,
+            weights: torch.Tensor,
+            context_lens: torch.Tensor,
+            block_table: torch.Tensor,
+            schedule_meta: torch.Tensor,
+            max_context_len: int,
+            num_epi_subtiles: int = 1,
+            epi_dtype: torch.dtype = torch.float32,
+            acc_dtype: torch.dtype = torch.float32,
+            output_dtype: torch.dtype = torch.float32,
+        ) -> torch.Tensor:
+            """Execute FP8 paged MQA logits kernel.
+
+            Args:
+                q: [B, next_n, H, D] FP8
+                kv_fused: [num_blocks, phys_block_kv, 1, D+4] uint8
+                weights: [B*next_n, H] float32
+                context_lens: [B] int32
+                block_table: [B, max_blocks] int32
+                schedule_meta: [num_sms+1, 2] int32
+                max_context_len: int
+                num_epi_subtiles: epilogue sub-tile count (1, 2, or 4)
+                epi_dtype: epilogue compute dtype
+                acc_dtype: MMA accumulator dtype
+                output_dtype: output logits dtype
+            Returns:
+                logits: [B*next_n, max_context_len] output_dtype
+            """
+            B, next_n, H, D = q.shape
+            N = next_n * H
+            phys_block_kv = kv_fused.shape[1]
+            compute_block_kv = 128
+            num_phys_blocks = kv_fused.shape[0]
+            num_sms = _get_num_sms()
+
+            # Reshape Q: [B, next_n, H, D] -> [B, N, D] -> [N, D, B]
+            q_3d = q.reshape(B, N, D).permute(1, 2, 0)
+
+            # Reshape weights: [B*next_n, H] -> [B, N] -> [N, B]
+            if epi_dtype == torch.float16:
+                # TODO: move type conversion to weight loading
+                w_2d = weights.reshape(B, N).half().t()
+            else:
+                w_2d = weights.reshape(B, N).t()
+
+            # Flatten fused KV to [num_phys_blocks, block_bytes]
+            kv_flat = kv_fused.reshape(num_phys_blocks, -1)
+
+            # Allocate output with alignment padding
+            SPLIT_KV = compute_block_kv * 2  # NUM_MATH_WG = 2
+            aligned_max_ctx = (
+                (max_context_len + SPLIT_KV - 1) // SPLIT_KV) * SPLIT_KV
+            logits = torch.empty(
+                (B * next_n, aligned_max_ctx),
+                device=q.device,
+                dtype=output_dtype,
+            )
+            logits = logits[:, :max_context_len]
+
+            # Compile if needed (fake tensors, no real data required)
+            key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms,
+                   num_epi_subtiles, epi_dtype, acc_dtype, output_dtype)
+            if key not in cls.kernel_cache:
+                cls._compile(compute_block_kv, phys_block_kv, H, D, next_n,
+                             num_sms, num_epi_subtiles, epi_dtype, acc_dtype,
+                             output_dtype)
+            compiled = cls.kernel_cache[key]
+
+            # FP8 q needs uint8 view to match compile-time dtype
+            q_for_ffi = (q_3d.view(torch.uint8) if q_3d.dtype
+                         in (torch.float8_e4m3fn, torch.float8_e5m2) else q_3d)
+
+            # TVM FFI: pass raw tensors, no dlpack/stream needed
+            compiled(kv_flat, q_for_ffi, w_2d, logits, block_table,
+                     context_lens, schedule_meta, num_phys_blocks, B)
+            return logits
+
+    @torch.library.custom_op("trtllm::cute_dsl_fp8_paged_mqa_logits",
+                             mutates_args=(),
+                             device_types="cuda")
+    def cute_dsl_fp8_paged_mqa_logits(
+        q: torch.Tensor,
+        kv_fused: torch.Tensor,
+        weights: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        schedule_meta: torch.Tensor,
+        max_context_len: int,
+        num_epi_subtiles: int = 1,
+        epi_dtype: torch.dtype = torch.float32,
+        acc_dtype: torch.dtype = torch.float32,
+        output_dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        if not is_sm_100f():
+            raise ValueError(
+                f"CuteDSL: SM version {get_sm_version()} is not supported. "
+                f"CuteDSL FP8 Paged MQA Logits only supports SM 100 family.")
+        _check_fp8_paged_mqa_logits_dtypes(q, kv_fused, weights, context_lens,
+                                           block_table, schedule_meta,
+                                           epi_dtype, acc_dtype, output_dtype)
+        return CuteDSLPagedMQALogitsRunner.forward(
+            q,
+            kv_fused,
+            weights,
+            context_lens,
+            block_table,
+            schedule_meta,
+            max_context_len,
+            num_epi_subtiles=num_epi_subtiles,
+            epi_dtype=epi_dtype,
+            acc_dtype=acc_dtype,
+            output_dtype=output_dtype)
+
+    @torch.library.register_fake("trtllm::cute_dsl_fp8_paged_mqa_logits")
+    def _(
+        q: torch.Tensor,
+        kv_fused: torch.Tensor,
+        weights: torch.Tensor,
+        context_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        schedule_meta: torch.Tensor,
+        max_context_len: int,
+        num_epi_subtiles: int = 1,
+        epi_dtype: torch.dtype = torch.float32,
+        acc_dtype: torch.dtype = torch.float32,
+        output_dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        B = q.shape[0]
+        next_n = q.shape[1]
+        return torch.empty(B * next_n,
+                           max_context_len,
+                           dtype=output_dtype,
+                           device=q.device)

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .fp8_paged_mqa_logits import FP8MQALogitsKernel
+
+__all__ = [
+    "FP8MQALogitsKernel",
+]