[TRTLLM-35237][feat] Tune cute dsl paged MQA logits decode kernel

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -5986,9 +5986,15 @@ def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim,
             max_blocks_per_seq = cute.sym_int()
             num_ctas = cute.sym_int()
 
-            kv_fake = cute.runtime.make_fake_compact_tensor(
+            # KV may come from the indexer K-cache pool view, which is
+            # strided in dim 0 (pool layout interleaves layers:
+            # [num_blocks, num_layers, kvFactor, blockSize]). Declare outer
+            # stride as sym so the actual per-block stride is read at
+            # runtime; innermost stride is fixed to 1 (byte-contig within a
+            # logical block view).
+            kv_fake = cute.runtime.make_fake_tensor(
                 cutlass.Uint8, (sym_num_phys_blocks, block_bytes),
-                stride_order=(1, 0))
+                stride=(cute.sym_int64(), 1))
 
             q_fake = cute.runtime.make_fake_compact_tensor(cutlass.Uint8,
                                                            (N, head_dim, sym_B),
@@ -6770,12 +6776,21 @@ class CuteDSLFP4PagedMQALogitsRunner:
         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,
-                     output_dtype):
+        def _compile(cls,
+                     compute_block_kv,
+                     phys_block_kv,
+                     num_heads,
+                     head_dim,
+                     next_n,
+                     num_sms,
+                     num_epi_subtiles,
+                     epi_dtype,
+                     output_dtype,
+                     remove_online_sf_transpose=False):
             """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, output_dtype)
+                   num_sms, num_epi_subtiles, epi_dtype, output_dtype,
+                   remove_online_sf_transpose)
             if key in cls.kernel_cache:
                 return
 
@@ -6791,9 +6806,15 @@ def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim,
             max_blocks_per_seq = cute.sym_int()
             num_ctas = cute.sym_int()
 
-            kv_fake = cute.runtime.make_fake_compact_tensor(
+            # KV may come from the indexer K-cache pool view, which is
+            # strided in dim 0 (pool layout interleaves layers:
+            # [num_blocks, num_layers, kvFactor, blockSize]). Declare outer
+            # stride as sym so the actual per-block stride is read at
+            # runtime; innermost stride is fixed to 1 (byte-contig within a
+            # logical block view).
+            kv_fake = cute.runtime.make_fake_tensor(
                 cutlass.Uint8, (sym_num_phys_blocks, block_bytes),
-                stride_order=(1, 0))
+                stride=(cute.sym_int64(), 1))
 
             # Q is FP4 packed bytes: head_dim/2 bytes per row
             q_fake = cute.runtime.make_fake_compact_tensor(
@@ -6845,6 +6866,7 @@ def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim,
                 num_epi_subtiles=num_epi_subtiles,
                 epi_dtype=to_cutlass[epi_dtype],
                 output_dtype=to_cutlass[output_dtype],
+                remove_online_sf_transpose=remove_online_sf_transpose,
             )
 
             compiled = cute.compile(
@@ -6879,6 +6901,7 @@ def forward(
             num_epi_subtiles: int = 1,
             epi_dtype: torch.dtype = torch.float32,
             output_dtype: torch.dtype = torch.float32,
+            remove_online_sf_transpose: bool = False,
         ) -> torch.Tensor:
             """Execute FP4 paged MQA logits kernel.
 
@@ -6946,10 +6969,20 @@ def forward(
 
             # 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, output_dtype)
+                   num_epi_subtiles, epi_dtype, output_dtype,
+                   remove_online_sf_transpose)
             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, output_dtype)
+                cls._compile(
+                    compute_block_kv,
+                    phys_block_kv,
+                    H,
+                    D,
+                    next_n,
+                    num_sms,
+                    num_epi_subtiles,
+                    epi_dtype,
+                    output_dtype,
+                    remove_online_sf_transpose=remove_online_sf_transpose)
             compiled = cls.kernel_cache[key]
 
             # TVM FFI: pass raw tensors, no dlpack/stream needed
@@ -6972,6 +7005,7 @@ def cute_dsl_fp4_paged_mqa_logits(
         num_epi_subtiles: int = 1,
         epi_dtype: torch.dtype = torch.float32,
         output_dtype: torch.dtype = torch.float32,
+        remove_online_sf_transpose: bool = False,
     ) -> torch.Tensor:
         if not is_sm_100f():
             raise ValueError(
@@ -7008,7 +7042,8 @@ def cute_dsl_fp4_paged_mqa_logits(
             max_context_len,
             num_epi_subtiles=num_epi_subtiles,
             epi_dtype=epi_dtype,
-            output_dtype=output_dtype)
+            output_dtype=output_dtype,
+            remove_online_sf_transpose=remove_online_sf_transpose)
 
     @torch.library.register_fake("trtllm::cute_dsl_fp4_paged_mqa_logits")
     def _(
@@ -7023,6 +7058,7 @@ def _(
         num_epi_subtiles: int = 1,
         epi_dtype: torch.dtype = torch.float32,
         output_dtype: torch.dtype = torch.float32,
+        remove_online_sf_transpose: bool = False,
     ) -> torch.Tensor:
         B = q.shape[0]
         next_n = q.shape[1]

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

@@ -364,6 +364,8 @@ def __init__(
         num_epi_subtiles: int = 1,
         epi_dtype=cutlass.Float32,
         output_dtype=cutlass.Float32,
+        remove_online_sf_transpose: bool = False,
+        use_batched_store: bool = True,
     ):
         # Static FP4 invariants — see plan Sanity checklist.
         assert num_heads == 64, "FP4 kernel hardcodes num_heads=64 for TMEM/SMEM budget"
@@ -398,6 +400,16 @@ def __init__(
         self.num_sms = num_sms
         self.num_epi_subtiles = num_epi_subtiles
         self.epi_dtype = epi_dtype
+        # When True, skip the in-kernel SMEM warp_transpose for KV SF; assume
+        # the host has pre-arranged GMEM SF into UTCCP chunk layout. Only valid
+        # for phys_block_kv=128 (1 phys block = 1 UTCCP atom). Q SF transpose
+        # is NOT affected by this flag (deferred to a separate phase).
+        if remove_online_sf_transpose and phys_block_kv != 128:
+            remove_online_sf_transpose = False
+        self.remove_online_sf_transpose = remove_online_sf_transpose
+        # When True, defer per-t STG to register array and emit all STGs in
+        # one contiguous LSU phase after the for-t loop (epilogue micro-opt).
+        self.use_batched_store = use_batched_store
         # epi_bytes covers fp16 and bf16 (FP8 only handled fp16).
         self.epi_bytes = 2 if epi_dtype in (cutlass.Float16, cutlass.BFloat16) else 4
         # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy.
@@ -625,7 +637,6 @@ def __call__(
         #                              [SF   phys_block_kv*4         bytes (= phys_block_kv int32)]
         phys_block_kv = self.phys_block_kv
         half_head_dim = self.head_dim // 2  # FP4 packed bytes per row
-        phys_block_bytes = phys_block_kv * (half_head_dim + 4)
         scale_offset_bytes = phys_block_kv * half_head_dim  # to SF region of each phys block
 
         # Recast the fused buffer to FP4. Each uint8 byte becomes 2 FP4 elements,
@@ -636,23 +647,31 @@ def __call__(
         # type inference and TMA descriptors are correct.
         b = cute.recast_tensor(b, Float4E2M1FN)
 
+        # Read the real per-block stride (bytes) from the input tensor.
+        # When KV is the indexer K-cache pool view, the pool is laid out as
+        # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride =
+        # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes).
+        # Using the input stride keeps both the contiguous test path and
+        # the strided prod path correct.
+        kv_block_stride_bytes = kv_fused.layout.stride[0]
+
         # KV data view: [phys_block_kv, head_dim, num_phys_blocks] FP4 elements.
         # Innermost stride 1 = consecutive FP4 elem = packed pair share a byte.
         # Per-row stride = head_dim FP4 elem = head_dim/2 bytes.
-        # Per-block stride = phys_block_bytes * 2 FP4 elem (data + SF region).
+        # Per-block stride (FP4 elem) = kv_block_stride_bytes * 2 (uint8→FP4 doubles).
         kv_layout = cute.make_layout(
             (phys_block_kv, self.head_dim, num_phys_blocks),
-            stride=(self.head_dim, 1, phys_block_bytes * 2),
+            stride=(self.head_dim, 1, kv_block_stride_bytes * 2),
         )
         a = cute.make_tensor(kv_fp4.iterator, kv_layout)
 
         # SF KV view: int32 (4 UE8M0 packed). Build a uint8 view at the SF
         # offset, then recast to int32.
-        # Layout in bytes: (phys_block_kv * 4, num_phys_blocks) stride (1, phys_block_bytes)
-        # After recast int32: (phys_block_kv, num_phys_blocks) stride (1, phys_block_bytes/4)
+        # Layout in bytes: (phys_block_kv * 4, num_phys_blocks) stride (1, kv_block_stride_bytes)
+        # After recast int32: (phys_block_kv, num_phys_blocks) stride (1, kv_block_stride_bytes/4)
         sf_kv_uint8_layout = cute.make_layout(
             (phys_block_kv * 4, num_phys_blocks),
-            stride=(1, phys_block_bytes),
+            stride=(1, kv_block_stride_bytes),
         )
         sf_kv_uint8 = cute.make_tensor(kv_fused.iterator + scale_offset_bytes, sf_kv_uint8_layout)
         sf_kv = cute.recast_tensor(sf_kv_uint8, cutlass.Int32)
@@ -1568,14 +1587,18 @@ def kernel(
 
                     # Step 5.6: SF KV transpose + UTCCP. block_kv = 128 = 1
                     # UTCCP atom; loop is constexpr-1 but kept for clarity.
-                    sf_kv_atoms = self.block_kv // 128
-                    for atom_idx in cutlass.range_constexpr(sf_kv_atoms):
-                        atom_offset = atom_idx * 128
-                        stage_offset = kv_stage * sSF_KV_0.layout.stride[1]
-                        utccp_required_smem_warp_transpose(
-                            sSF_KV_0.iterator + stage_offset + atom_offset
-                        )
-                    cute.arch.fence_view_async_shared()
+                    # When remove_online_sf_transpose=True, the host has already
+                    # pre-arranged GMEM SF into UTCCP chunk layout, so the
+                    # in-kernel SMEM transpose (and its fence) can be skipped.
+                    if cutlass.const_expr(not self.remove_online_sf_transpose):
+                        sf_kv_atoms = self.block_kv // 128
+                        for atom_idx in cutlass.range_constexpr(sf_kv_atoms):
+                            atom_offset = atom_idx * 128
+                            stage_offset = kv_stage * sSF_KV_0.layout.stride[1]
+                            utccp_required_smem_warp_transpose(
+                                sSF_KV_0.iterator + stage_offset + atom_offset
+                            )
+                        cute.arch.fence_view_async_shared()
                     # int32 SMEM → UE8M0 view for UTCCP atom + chunk layout.
                     sSF_KV_0_ue8m0 = cute.recast_tensor(sSF_KV_0, Float8E8M0FNU)
                     stage_off_kv0_ue8m0 = kv_stage * sSF_KV_0_ue8m0.layout.stride[1]
@@ -1705,14 +1728,18 @@ def kernel(
                     kv_stage_1 = kv_cons_state_umma_1.index
 
                     # Step 5.6: SF KV (group 1) transpose + UTCCP.
-                    sf_kv_atoms_1 = self.block_kv // 128
-                    for atom_idx in cutlass.range_constexpr(sf_kv_atoms_1):
-                        atom_offset = atom_idx * 128
-                        stage_offset = kv_stage_1 * sSF_KV_1.layout.stride[1]
-                        utccp_required_smem_warp_transpose(
-                            sSF_KV_1.iterator + stage_offset + atom_offset
-                        )
-                    cute.arch.fence_view_async_shared()
+                    # When remove_online_sf_transpose=True, the host has already
+                    # pre-arranged GMEM SF into UTCCP chunk layout, so the
+                    # in-kernel SMEM transpose (and its fence) can be skipped.
+                    if cutlass.const_expr(not self.remove_online_sf_transpose):
+                        sf_kv_atoms_1 = self.block_kv // 128
+                        for atom_idx in cutlass.range_constexpr(sf_kv_atoms_1):
+                            atom_offset = atom_idx * 128
+                            stage_offset = kv_stage_1 * sSF_KV_1.layout.stride[1]
+                            utccp_required_smem_warp_transpose(
+                                sSF_KV_1.iterator + stage_offset + atom_offset
+                            )
+                        cute.arch.fence_view_async_shared()
                     sSF_KV_1_ue8m0 = cute.recast_tensor(sSF_KV_1, Float8E8M0FNU)
                     stage_off_kv1_ue8m0 = kv_stage_1 * sSF_KV_1_ue8m0.layout.stride[1]
                     sSF_KV_1_chunk = cute.make_tensor(
@@ -1816,6 +1843,13 @@ def kernel(
                     MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
+                # Batched STG: hold reduced result per t in register; the
+                # actual STG happens once after the for-t loop to land all
+                # STGs in one contiguous LSU phase.
+                if cutlass.const_expr(self.use_batched_store):
+                    result_arr = cute.make_fragment(next_n, self.output_dtype)
+                else:
+                    result_arr = None
                 q_stage_local = cutlass.Int32(0)
 
                 while has_work:
@@ -2022,9 +2056,18 @@ def kernel(
                             result_t = sum_lo + sum_hi
                         else:
                             result_t = s0x + s0y + s1x + s1y
-                        out_row = q_idx * next_n + t
                         # Step 5.7: drop * scale_val (FP4 SF baked into acc).
-                        mLogits[(out_row, kv_pos)] = self.output_dtype(result_t)
+                        if cutlass.const_expr(self.use_batched_store):
+                            result_arr[t] = self.output_dtype(result_t)
+                        else:
+                            out_row = q_idx * next_n + t
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t)
+
+                    if cutlass.const_expr(self.use_batched_store):
+                        # Batched STG: all result_arr[t] → mLogits in one pass.
+                        for t in cutlass.range_constexpr(next_n):
+                            out_row = q_idx * next_n + t
+                            mLogits[(out_row, kv_pos)] = result_arr[t]
 
                     # Advance: inline fetch_next_task
                     next_kv_idx = kv_idx + NUM_MATH_WG
@@ -2063,6 +2106,13 @@ def kernel(
                     MAX_NUM_W_IN_REG = 56 if next_n == 3 else 64
                 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads)
                 w_cache = cute.make_fragment(NUM_W_IN_REG * next_n, self.epi_dtype)
+                # Batched STG: hold reduced result per t in register; the
+                # actual STG happens once after the for-t loop to land all
+                # STGs in one contiguous LSU phase.
+                if cutlass.const_expr(self.use_batched_store):
+                    result_arr = cute.make_fragment(next_n, self.output_dtype)
+                else:
+                    result_arr = None
                 q_stage_local = cutlass.Int32(0)
 
                 while has_work:
@@ -2259,9 +2309,18 @@ def kernel(
                             result_t = sum_lo + sum_hi
                         else:
                             result_t = s0x + s0y + s1x + s1y
-                        out_row = q_idx * next_n + t
                         # Step 5.7: drop * scale_val (FP4 SF baked into acc).
-                        mLogits[(out_row, kv_pos)] = self.output_dtype(result_t)
+                        if cutlass.const_expr(self.use_batched_store):
+                            result_arr[t] = self.output_dtype(result_t)
+                        else:
+                            out_row = q_idx * next_n + t
+                            mLogits[(out_row, kv_pos)] = self.output_dtype(result_t)
+
+                    if cutlass.const_expr(self.use_batched_store):
+                        # Batched STG: all result_arr[t] → mLogits in one pass.
+                        for t in cutlass.range_constexpr(next_n):
+                            out_row = q_idx * next_n + t
+                            mLogits[(out_row, kv_pos)] = result_arr[t]
 
                     # Advance: inline fetch_next_task
                     next_kv_idx = kv_idx + NUM_MATH_WG

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

@@ -372,7 +372,6 @@ def __call__(
         # Derive KV and Scale views from fused buffer using CuTE ops.
         # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)]
         phys_block_kv = self.phys_block_kv
-        phys_block_bytes = phys_block_kv * (self.head_dim + 4)
         scale_offset_elems = phys_block_kv * self.head_dim  # in FP8 elements
 
         # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference)
@@ -382,19 +381,28 @@ def __call__(
         # recast back to FP8 so MMA type inference and TMA descriptors are correct.
         b = cute.recast_tensor(b, cutlass.Float8E4M3FN)
 
+        # Read the real per-block stride (bytes = FP8 elements) from the input.
+        # When KV is the indexer K-cache pool view, the pool is laid out as
+        # [num_blocks, num_layers, kvFactor, blockSize], so dim-0 stride =
+        # num_layers * kvFactor * phys_block_bytes (not phys_block_bytes).
+        # Using the input stride keeps both contiguous test path and the
+        # strided prod path correct. FP8 = 1 byte per element, so the byte
+        # stride is the same as the element stride after recast.
+        kv_block_stride = kv_fused.layout.stride[0]
+
         # KV view: [phys_block_kv, head_dim, num_phys_blocks] FP8
         # Each TMA loads one physical block; multiple TMAs fill a compute tile.
         kv_layout = cute.make_layout(
             (phys_block_kv, self.head_dim, num_phys_blocks),
-            stride=(self.head_dim, 1, phys_block_bytes),
+            stride=(self.head_dim, 1, kv_block_stride),
         )
         a = cute.make_tensor(kv_fp8.iterator, kv_layout)
 
         # Scale view: offset pointer to scale region, recast FP8 → Float32
         # [phys_block_kv, num_phys_blocks] float32 (after recast)
         scale_fp8_layout = cute.make_layout(
             (phys_block_kv * 4, num_phys_blocks),
-            stride=(1, phys_block_bytes),
+            stride=(1, kv_block_stride),
         )
         scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout)
         scales = cute.recast_tensor(scale_fp8, cutlass.Float32)