From 255cd9d3899df0d9c4b0509b63662c583d1a50f9 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 04:42:34 -0700 Subject: [PATCH 01/26] [None][feat] Add CuTE DSL FP8 paged MQA logits kernel for Blackwell SM100 Replace DeepGEMM-based indexer logits with a CuTE DSL kernel on SM100+, gated by `use_cute_dsl_logits` config flag. Includes kernel implementation, PyTorch custom op registration, config plumbing, and unit tests. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 19 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 209 ++ .../blackwell/paged_mqa_logits/__init__.py | 16 + .../paged_mqa_logits/fp8_paged_mqa_logits.py | 2317 +++++++++++++++++ tensorrt_llm/_torch/model_config.py | 3 + tensorrt_llm/llmapi/llm_args.py | 5 + .../test_cute_dsl_fp8_paged_mqa_logits.py | 545 ++++ 7 files changed, 3109 insertions(+), 5 deletions(-) create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py create mode 100644 tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 99a41a19b09c..5ea98f2767f6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1111,6 +1111,10 @@ 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_logits = (getattr(sparse_attention_config, + 'use_cute_dsl_logits', False) + and IS_CUTLASS_DSL_AVAILABLE + and get_sm_version() >= 100) self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5 self._enable_heuristic_topk = ( @@ -1665,11 +1669,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_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) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index f91d331dbe52..3eb037c2e609 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5127,3 +5127,212 @@ 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 \ + FP8MQALogitsDGFullKKernel + + class CuteDSLPagedMQALogitsRunner: + """Runner for CuTe DSL FP8 Paged MQA Logits kernel (Blackwell SM100). + + Caches compiled kernels keyed by static params + (block_kv, num_heads, head_dim, next_n, num_sms). + """ + + kernel_cache = dict() + + @classmethod + def _make_dlpacks(cls, kv_flat, q_3d, w_2d, logits, block_table, + context_lens, schedule_meta): + """Wrap tensors with dynamic shape markers for JIT reuse.""" + from cutlass.cute.runtime import from_dlpack + dl_kv = from_dlpack(kv_flat).mark_compact_shape_dynamic(mode=0) + q_for_dl = q_3d.view( + torch.uint8) if q_3d.dtype in (torch.float8_e4m3fn, + torch.float8_e5m2) else q_3d + dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( + mode=2, stride_order=(2, 0, 1)) + dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic( + mode=1, stride_order=(1, 0)) + dl_logits = from_dlpack(logits).mark_compact_shape_dynamic( + mode=0, stride_order=(0, 1)).mark_compact_shape_dynamic( + mode=1, stride_order=(0, 1)) + dl_bt = from_dlpack(block_table).mark_compact_shape_dynamic( + mode=0, stride_order=(0, 1)).mark_compact_shape_dynamic( + mode=1, stride_order=(0, 1)) + dl_cl = from_dlpack(context_lens).mark_compact_shape_dynamic(mode=0) + dl_sm = from_dlpack(schedule_meta).mark_compact_shape_dynamic( + mode=0) + return dl_kv, dl_q, dl_w, dl_logits, dl_bt, dl_cl, dl_sm + + _TORCH_TO_CUTLASS_DTYPE = { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + } + + @classmethod + def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, + kv_flat, q_3d, w_2d, logits, block_table, context_lens, + schedule_meta, num_phys_blocks, B, stream, + num_epi_subtiles, epi_dtype, acc_dtype, output_dtype): + """Compile kernel using from_dlpack with dynamic shape markers.""" + key = (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 = cls._TORCH_TO_CUTLASS_DTYPE + dl_args = cls._make_dlpacks(kv_flat, q_3d, w_2d, logits, + block_table, context_lens, + schedule_meta) + kernel = FP8MQALogitsDGFullKKernel( + block_kv=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, *dl_args, num_phys_blocks, B, + stream) + cls.kernel_cache[key] = compiled + + @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, 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 + block_kv = kv_fused.shape[1] + 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 = 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] + + # Create stream + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # Compile if needed (uses real tensors for shape marking) + key = (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(block_kv, H, D, next_n, num_sms, kv_flat, q_3d, + w_2d, logits, block_table, context_lens, + schedule_meta, num_phys_blocks, B, stream, + num_epi_subtiles, epi_dtype, acc_dtype, + output_dtype) + compiled = cls.kernel_cache[key] + + # Wrap tensors for runtime call + dl_args = cls._make_dlpacks(kv_flat, q_3d, w_2d, logits, + block_table, context_lens, + schedule_meta) + compiled(*dl_args, num_phys_blocks, B, stream) + torch.cuda.synchronize() + 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: + 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) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py new file mode 100644 index 000000000000..e2190e217e6a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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 FP8MQALogitsDGFullKKernel diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..b1fa2daee12a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -0,0 +1,2317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +DeepGEMM-aligned 2-group kernel with full-K TMA and fused KV layout (SM100). +Supports multi-batch via in-kernel scheduler matching DeepGEMM's PagedMQALogitsScheduler. + +Architecture: + - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA) + - Full-K TMA: 1 TMA per KV block [128, 128], UMMA iterates 4x K=32 + - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2) + - Q reloaded via TMA pipeline when q_idx (batch) changes + - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs + - Weights cached in registers: preloaded once per q_idx change (not per KV block) + - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes) + +Merged KV+Scale pipeline (tma_5, matches DeepGEMM): + - KV data and scales share a single TMA barrier per group + - TMA loads both KV and Scale under one barrier (combined tx_count) + - UMMA waits on merged barrier (for KV GEMM), does NOT release + - Math waits on merged barrier (for scale read), Math releases + - This eliminates the separate scale pipeline overhead + +Fused KV layout (tma_4, matches DeepGEMM): + - KV data and scales stored contiguously per physical block: + [num_phys_blocks, block_kv * (head_dim + 4)] bytes + - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)] + - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic + - Benefit: L2 cache locality — scale data shares cache lines with KV data + +Scheduler (aligned with DeepGEMM's PagedMQALogitsScheduler): + - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups) + - schedule_meta[sm_idx+1] = end boundary for this CTA + - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx) + - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step + - exist_q_idx(qi): checks if qi is within this CTA's assigned range (for Q prefetch) + +Dynamic shape support (tma_8): + - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen + - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq, + num_ctas) are marked dynamic via mark_compact_shape_dynamic + - Allows JIT cache reuse across different batch sizes / sequence lengths + +Epilogue dtype flows (--acc_dtype / --epi_dtype): + + Flow 1: --acc_dtype fp32 --epi_dtype fp16 + Q(FP8) x K(FP8) -> MMA acc(FP32) -> TMEM(FP32) + -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16) + -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) + -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) + Benefits: weights SMEM BW halved, weight regs halved, FP16 FMA + Unchanged: TMEM(FP32), LDTM BW(FP32), acc regs(FP32) + + Flow 2: --acc_dtype fp16 --epi_dtype fp16 + Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b) + -> LDTM -> Reg(FP16) -> ReLU(FP16) + -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) + -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) + Extra benefits over Flow 1: TMEM halved (more umma stages), LDTM BW halved, acc regs halved + Risk: MMA FP16 accumulation over K=128 has precision loss; epilogue sum may overflow FP16 + + --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion. + + Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32 (original FP32 baseline) + +Run scripts: + - Single values: + python paged_mqa_logits_dg_fullk_tma_7_dynamic_improve_v3.py \ + --batch_size 1 --next_n 2 --avg_ctx 4096 --num_sms 148 + - Multiple values: + python paged_mqa_logits_dg_fullk_tma_7_dynamic_improve_v3.py \ + --batch_size 1 32 --next_n 1 2 4 --avg_ctx 256 4096 --num_sms 148 + - Full sweep: python paged_mqa_logits_dg_fullk_tma_7_dynamic_improve_v3.py --sweep + - Default (no args): uses batch_size=[32], next_n=[1], avg_ctx=[32768], num_sms=[148] as before +""" + +from typing import Tuple + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass import Float16, Int32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, nvvm, vector +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait # noqa: F401 + + +@dsl_user_op +def pack_f16x2( + a: Float16, + b: Float16, + *, + loc=None, + ip=None, +) -> Int32: + f16_ty = Float16.mlir_type + i32_ty = Int32.mlir_type + vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc) + v = vector.from_elements( + vec2_f16, + (Float16(a).ir_value(loc=loc, ip=ip), Float16(b).ir_value(loc=loc, ip=ip)), + loc=loc, + ip=ip, + ) + return Int32(llvm.bitcast(i32_ty, v, loc=loc, ip=ip)) + + +@dsl_user_op +def unpack_f16x2( + packed: Int32, + *, + loc=None, + ip=None, +) -> Tuple[Float16, Float16]: + f16_ty = Float16.mlir_type + vec2_f16 = ir.VectorType.get([2], f16_ty, loc=loc) + v = llvm.bitcast(vec2_f16, Int32(packed).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + r0 = Float16(vector.extract(v, dynamic_position=[], static_position=[0], loc=loc, ip=ip)) + r1 = Float16(vector.extract(v, dynamic_position=[], static_position=[1], loc=loc, ip=ip)) + return r0, r1 + + +@dsl_user_op +def fma_f16x2( + a: Int32, + b: Int32, + c: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [ + Int32(a).ir_value(loc=loc, ip=ip), + Int32(b).ir_value(loc=loc, ip=ip), + Int32(c).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f16x2 $0, $1, $2, $3;", + "=r,r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def max_f16x2( + a: Int32, + b: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "max.f16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def add_f16x2( + a: Int32, + b: Int32, + *, + loc=None, + ip=None, +) -> Int32: + i32_ty = Int32.mlir_type + return Int32( + llvm.inline_asm( + i32_ty, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "add.f16x2 $0, $1, $2;", + "=r,r,r", + loc=loc, + ip=ip, + ) + ) + + +class FP8MQALogitsDGFullKKernel: + """ + DG-Aligned 2-group kernel with full-K TMA, multi-batch support. + + Each CTA processes a range of (q_idx, kv_split) pairs. + A split = 2 consecutive KV blocks within a sequence (one per warp group). + Q is shared between warp groups and reloaded when q_idx changes. + """ + + def __init__( + self, + block_kv: int = 128, + num_heads: int = 64, + head_dim: int = 128, + next_n: int = 1, + num_sms: int = 148, + remove_kv_wait_in_epilogue: bool = False, + early_tmem_copy: bool = False, + smem_subpartition_opt: bool = False, + max_kv_pipeline: bool = False, + max_umma_pipeline: bool = False, + num_epi_subtiles: int = 1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, + ): + self.block_kv = block_kv + self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue + self.early_tmem_copy = early_tmem_copy + self.smem_subpartition_opt = smem_subpartition_opt + self.num_heads = num_heads + self.head_dim = head_dim + self.next_n = next_n + self.N = next_n * num_heads + self.num_sms = num_sms + self.num_epi_subtiles = num_epi_subtiles + self.epi_dtype = epi_dtype + self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4 + self.output_dtype = output_dtype + if num_epi_subtiles > 1: + if num_heads % num_epi_subtiles != 0: + raise ValueError("num_heads must be divisible by num_epi_subtiles") + if (num_heads // num_epi_subtiles) % 4 != 0: + raise ValueError( + "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)" + ) + self.num_groups = 2 + + self.num_math_threads = 256 + self.num_specialized_threads = 128 + self.threads_per_cta = 384 + self.num_math_warps = 8 + self.tma_warp_base = 8 + self.umma_warp_base = 10 + + self.num_q_stages = 3 # 3 stages for Q pipelining across batch sequences + + # TMEM: 512 columns total, each group needs N columns per UMMA stage + # max_umma_stages = 512 // (2 * N) + TMEM_COLS = 512 + if max_umma_pipeline: + self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N)) + else: + self.num_umma_stages = 1 + + if max_kv_pipeline: + smem_capacity = utils.get_smem_capacity_in_bytes() + # Reserve ~1 KB for barriers and misc + SMEM_BUDGET = smem_capacity - 1024 + # KV+Scale per stage (×2 groups): + # 2 * (block_kv * head_dim * 1B + block_kv * 4B) + kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4) + # Q+W per stage: N * head_dim * 1B + N * 4B + qw_per_stage = self.N * head_dim + self.N * 4 + qw_total = qw_per_stage * self.num_q_stages + self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage + else: + self.num_kv_stages = 3 + + # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB), + # avoiding sub-bank conflicts with UMMA reading sKV. + # Layout: barriers(~256B) | sKV_0 | sKV_1 | sQ | [pad] | sW | sScales + if self.smem_subpartition_opt: + BOUNDARY = 128 * 1024 + used = ( + 256 + + 2 * (block_kv * head_dim * self.num_kv_stages) + + self.N * head_dim * self.num_q_stages + ) + used = ((used + 127) // 128) * 128 + if used < BOUNDARY: + self.smem_pad_bytes = ((BOUNDARY - used + 1023) // 1024) * 1024 + else: + self.smem_pad_bytes = 0 + else: + self.smem_pad_bytes = 0 + + self.acc_dtype = acc_dtype + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape_mn = (1, 1) + self.mma_tiler_mn = (block_kv, self.N) + + def _setup_mma(self, a_dtype, b_dtype, a_major, b_major): + self.a_dtype = a_dtype + self.b_dtype = b_dtype + self.a_major_mode = a_major + self.b_major_mode = b_major + + self.mma_tiler = (*self.mma_tiler_mn, 1) + tiled_mma = sm100_utils.make_trivial_tiled_mma( + a_dtype, + a_major, + b_major, + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) # 32 + + # Full-K: tile K = head_dim (128), 1 TMA per block + mma_inst_tile_k = self.head_dim // mma_inst_shape_k # 4 + full_k = mma_inst_shape_k * mma_inst_tile_k # 128 + self.mma_tiler = ( + self.mma_tiler_mn[0], + self.mma_tiler_mn[1], + full_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.epi_tile = self.cta_tile_shape_mnk[:2] + + # KV SMEM: 3 stages per group, each stage holds full [128, 128] + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + a_dtype, + self.num_kv_stages, + ) + # Q SMEM: 1 stage, holds full [N, 128] + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + b_dtype, + self.num_q_stages, + ) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + self.num_tmem_alloc_cols_total = ( + self.num_tmem_alloc_cols * self.num_groups * self.num_umma_stages + ) + + return tiled_mma + + @cute.jit + def __call__( + self, + kv_fused: cute.Tensor, # Fused KV: [num_phys_blocks, block_bytes] FP8 + b: cute.Tensor, # Q: [N, head_dim, batch_size] + weights: cute.Tensor, # [N, batch_size] (transposed for TMA) + logits: cute.Tensor, # [batch_size * next_n, max_context_len] + block_table: cute.Tensor, # [batch_size, max_blocks_per_seq] + context_lens: cute.Tensor, # [batch_size] + schedule_meta: cute.Tensor, # [num_sms+1, 2] int32 + num_phys_blocks: cutlass.Int32, + batch_size: cutlass.Int32, + stream: cuda.CUstream, + ): + # Derive KV and Scale views from fused buffer using CuTE ops. + # Fused layout per block: [KV data (block_kv*head_dim)] [Scales (block_kv*4)] + block_bytes = self.block_kv * (self.head_dim + 4) + scale_offset_elems = self.block_kv * self.head_dim # in FP8 elements + + # Recast fused buffer to FP8 (same 1-byte elements, needed for MMA type inference) + kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN) + + # Q (b) was passed as uint8 to work around DLPack's lack of float8 support; + # recast back to FP8 so MMA type inference and TMA descriptors are correct. + b = cute.recast_tensor(b, cutlass.Float8E4M3FN) + + # KV view: [block_kv, head_dim, num_phys_blocks] FP8 + # Pointer is fused base, layout strides: (head_dim, 1, block_bytes) + kv_layout = cute.make_layout( + (self.block_kv, self.head_dim, num_phys_blocks), + stride=(self.head_dim, 1, block_bytes), + ) + a = cute.make_tensor(kv_fp8.iterator, kv_layout) + + # Scale view: offset pointer to scale region, recast FP8 → Float32 + # Step 1: create FP8 tensor at offset scale_offset_elems + scale_fp8_layout = cute.make_layout( + (self.block_kv * 4, num_phys_blocks), + stride=(1, block_bytes), + ) + scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout) + # Step 2: recast from FP8 (1 byte) to Float32 (4 bytes) + scales = cute.recast_tensor(scale_fp8, cutlass.Float32) + + a_dtype = a.element_type + b_dtype = b.element_type + a_major = utils.LayoutEnum.from_tensor(a).mma_major_mode() + b_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + + tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # TMA for KV (A) — full K=128 per load + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA for Q (B) — full K=128, L dim = batch_size + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA for Weights — [N, batch_size], tile [N], L=batch_size + tma_load_op = cpasync.CopyBulkTensorTileG2SOp() + self.w_smem_layout_staged = cute.make_layout((self.N, self.num_q_stages)) + w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0]) + tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom( + tma_load_op, + weights, + w_smem_per_stage, + self.w_smem_layout_staged.shape[:1], + ) + + # TMA for Scales — [block_kv, num_phys_blocks], tile [block_kv], L=num_phys_blocks + self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages)) + s_smem_per_stage = cute.select(self.s_smem_layout_staged, mode=[0]) + tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom( + tma_load_op, + scales, + s_smem_per_stage, + self.s_smem_layout_staged.shape[:1], + ) + + a_copy_size = cute.size_in_bytes(a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout) + w_copy_size = self.N * self.epi_bytes + kv_tma_bytes = a_copy_size * atom_thr_size + scale_tma_bytes = self.block_kv * 4 + # KV + Scale share barrier (like DeepGEMM) + self.num_kv_scale_tma_bytes = kv_tma_bytes + scale_tma_bytes + # Q + Weights share barrier (like DeepGEMM) + self.num_q_tma_bytes = b_copy_size * atom_thr_size + w_copy_size + + num_ctas = self.num_sms + + @cute.struct + class SharedStorage: + kv_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2] + kv_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_kv_stages * 2] + q_mbar: cute.struct.MemRange[cutlass.Int64, self.num_q_stages * 2] + umma_mbar_0: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2] + umma_mbar_1: cute.struct.MemRange[cutlass.Int64, self.num_umma_stages * 2] + tmem_holding_buf: cutlass.Int32 + + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_w, + tma_tensor_w, + tma_atom_s, + tma_tensor_s, + logits, + block_table, + context_lens, + schedule_meta, + batch_size, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.w_smem_layout_staged, + self.s_smem_layout_staged, + self.epi_tile, + SharedStorage, + ).launch( + grid=(1, 1, num_ctas), + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, # KV pool + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, # Q (L dim = batch_size) + tma_atom_w: cute.CopyAtom, + mW_tma: cute.Tensor, # Weights TMA coord tensor [N, batch_size] + tma_atom_s: cute.CopyAtom, + mS_tma: cute.Tensor, # Scales TMA coord tensor [block_kv, num_phys_blocks] + mLogits: cute.Tensor, # [batch_size * next_n, max_context_len] + mBlockTable: cute.Tensor, # [batch_size, max_blocks_per_seq] + mContextLens: cute.Tensor, # [batch_size] + mScheduleMeta: cute.Tensor, # [num_sms+1, 2] int32 + batch_size: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + w_smem_layout_staged: cute.Layout, + s_smem_layout_staged: cute.Layout, + epi_tile: cute.Tile, + SharedStorage: cutlass.Constexpr, + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + # Warp roles (matches DeepGEMM SM100) + warpgroup_idx = warp_idx // 4 + is_math_warp = warp_idx < 8 + is_tma_warp_0 = warp_idx == 8 + is_tma_warp_1 = warp_idx == 9 + is_tma_warp = is_tma_warp_0 | is_tma_warp_1 + is_umma_warp_0 = warp_idx == 10 + is_umma_warp_1 = warp_idx == 11 + is_umma_warp = is_umma_warp_0 | is_umma_warp_1 # noqa: F841 + + # Early schedule metadata load: issue global loads ASAP so their + # ~200-cycle L2 latency overlaps with subsequent prologue setup + # (SMEM alloc, TMA partition, MMA fragment creation, etc.) + NUM_MATH_WG = 2 # kNumMathWarpGroups + sm_idx = bidz + start_q = mScheduleMeta[(sm_idx, 0)] + start_kv_half = mScheduleMeta[(sm_idx, 1)] + end_q_idx = mScheduleMeta[(sm_idx + 1, 0)] + end_kv_half = mScheduleMeta[(sm_idx + 1, 1)] + # Early mContextLens load: overlap ~200-cycle L2 latency with the + # entire prologue setup (pipelines, SMEM alloc, TMA partition, etc.) + current_num_kv = (mContextLens[start_q] + self.block_kv - 1) // self.block_kv + + if is_tma_warp: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_w) + cpasync.prefetch_descriptor(tma_atom_s) + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + block_kv_val = self.block_kv + num_heads = self.num_heads + next_n = self.next_n + num_epi_subtiles = self.num_epi_subtiles + num_q_stages = self.num_q_stages # noqa: F841 + + # === Pipelines === + prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_mcast_a = cute.size(cluster_layout_vmnk.shape[2]) + num_mcast_b = cute.size(cluster_layout_vmnk.shape[1]) + num_tma_prod = num_mcast_a + num_mcast_b - 1 + cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_prod) # noqa: F841 + + # Q pipeline: TMA producer → Math consumer (8 math warps) + # PipelineTmaAsync: consumer_release uses is_signalling_thread + # (lane 0 per warp). 8 math warps × 1 lane-0 = 8 arrives. + q_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 8) + q_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.q_mbar.data_ptr(), + num_stages=self.num_q_stages, + producer_group=prod_group, + consumer_group=q_cons_group, + tx_count=self.num_q_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + q_prod_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_q_stages + ) + # Both Math WGs share the same pipeline state (advance in lockstep) + q_cons_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + # UMMA warps observe Q pipeline (wait only, no release) + # to ensure Q is in SMEM before GEMM. Critical for UMMA warp 1 + # since TMA warp 1 only loads KV1 (not Q). + q_cons_state_umma_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + q_cons_state_umma_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_q_stages + ) + + # Merged KV+Scale pipelines (per-group, 3 stages each) + # Like DeepGEMM: KV data and scales share one barrier. + # TMA loads both under one barrier. Math is consumer (releases). + # UMMA also waits on this barrier (for KV GEMM) but does NOT release. + math_warps_per_group = self.num_math_warps // 2 # 4 warps + kv_cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, math_warps_per_group) + kv_pipeline_0 = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.kv_mbar_0.data_ptr(), + num_stages=self.num_kv_stages, + producer_group=prod_group, + consumer_group=kv_cons_group, + tx_count=self.num_kv_scale_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + kv_pipeline_1 = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.kv_mbar_1.data_ptr(), + num_stages=self.num_kv_stages, + producer_group=prod_group, + consumer_group=kv_cons_group, + tx_count=self.num_kv_scale_tma_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + tidx=tidx, + defer_sync=True, + ) + + kv_prod_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_kv_stages + ) + kv_prod_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_kv_stages + ) + # UMMA consumer states (wait only, no release) + kv_cons_state_umma_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + kv_cons_state_umma_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + # Math consumer states (wait + release) + kv_cons_state_math_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + kv_cons_state_math_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_kv_stages + ) + + # UMMA pipelines (per-group) + math_threads_per_group = self.num_math_threads // 2 + umma_pipeline_0 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.umma_mbar_0.data_ptr(), + num_stages=self.num_umma_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + umma_pipeline_1 = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.umma_mbar_1.data_ptr(), + num_stages=self.num_umma_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, math_threads_per_group), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + umma_prod_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_umma_stages + ) + umma_prod_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_umma_stages + ) + umma_cons_state_0 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_umma_stages + ) + umma_cons_state_1 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_umma_stages + ) + + # TMEM — only Math warps (8×32=256) + UMMA warps (2×32=64) = 320 threads + # TMA warps do NOT participate, so they can start TMA loads earlier. + # Math warp 0 is the allocator (like fp16_gemm_3's epilogue warp 0), + # because math warps are the last TMEM consumers (epilogue reads). + tmem_alloc_num_threads = 320 # 10 warps: warp 0-7 (math) + warp 10-11 (umma) + tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=tmem_alloc_num_threads) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=0, # math warp 0 does alloc+free (last TMEM consumer) + is_two_cta=False, + ) + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # SMEM allocation: per-group KV + shared Q + sKV_0 = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sKV_1 = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sQ = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + # Pad SMEM to push sW/sScales into sub-partition 1 (>= 128KB) + # to avoid sub-bank conflicts with UMMA reading sKV from + # sub-partition 0. + if cutlass.const_expr(self.smem_pad_bytes > 0): + _ = smem.allocate(self.smem_pad_bytes) + # Weights SMEM: [N, num_q_stages], shared Q barrier + sW = smem.allocate_tensor( + element_type=self.epi_dtype, + layout=w_smem_layout_staged, + byte_alignment=128, + ) + # Scales SMEM: [block_kv, num_kv_stages] float32, per group + sScales_0 = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=s_smem_layout_staged, + byte_alignment=128, + ) + sScales_1 = smem.allocate_tensor( + element_type=cutlass.Float32, + layout=s_smem_layout_staged, + byte_alignment=128, + ) + + a_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # Partition KV (A): per-group SMEM targets + gA_mkl = cute.local_tile( + mA_mkl, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + + tAsA_0, tAgA_0 = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sKV_0, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsA_1, tAgA_1 = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sKV_1, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAgA_0 = tAgA_0[(None, 0, None, None)] # [tma, K, L] + tAgA_1 = tAgA_1[(None, 0, None, None)] + + # Partition Q (B): shared SMEM, L dim = batch_size + gB_nkl = cute.local_tile( + mB_nkl, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBgB = tBgB[(None, 0, None, None)] # [tma, K, L] + + # Partition Weights: standalone TMA, [N, batch_size] → [N] per stage + w_cta_layout = cute.make_layout((1,)) + tWsW, tWgW = cpasync.tma_partition( + tma_atom_w, + 0, + w_cta_layout, + cute.group_modes(sW, 0, 1), + cute.group_modes(mW_tma, 0, 1), + ) + + # Partition Scales: standalone TMA, [block_kv, num_phys_blocks] + # tile [block_kv], L=num_phys_blocks. Per-group SMEM targets. + s_cta_layout = cute.make_layout((1,)) + tSsS_0, tSgS_0 = cpasync.tma_partition( + tma_atom_s, + 0, + s_cta_layout, + cute.group_modes(sScales_0, 0, 1), + cute.group_modes(mS_tma, 0, 1), + ) + tSsS_1, tSgS_1 = cpasync.tma_partition( + tma_atom_s, + 0, + s_cta_layout, + cute.group_modes(sScales_1, 0, 1), + cute.group_modes(mS_tma, 0, 1), + ) + + # MMA fragments + tCrA_0 = tiled_mma.make_fragment_A(sKV_0) + tCrA_1 = tiled_mma.make_fragment_A(sKV_1) + tCrB = tiled_mma.make_fragment_B(sQ) # shared + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) # noqa: F841 + + # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim + # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank + us = self.num_umma_stages + cols = self.num_tmem_alloc_cols + acc_shape_staged = cute.append(acc_shape, us) + tCtAcc_fake_staged = tiled_mma.make_fragment_C(acc_shape_staged) + + # TMEM layout info (allocation deferred to UMMA/Math warp branches) + cols_per_group = cols * us * (32 // self.acc_dtype.width) + num_tmem_alloc_cols_total = self.num_tmem_alloc_cols_total + + # Epilogue setup + c_layout = utils.LayoutEnum.ROW_MAJOR + epi_sub_mn = (epi_tile[0], num_heads // num_epi_subtiles) + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + c_layout, + self.acc_dtype, + self.acc_dtype, + epi_sub_mn, + use_2cta_instrs, + ) + + # ===== SCHEDULER: derive values from early-loaded schedule metadata ===== + end_kv_idx = end_kv_half * NUM_MATH_WG + + # Convert start to KV block units (matching DeepGEMM) + current_q_idx = start_q + current_kv_idx = start_kv_half * NUM_MATH_WG + + # ===== COMMON SCHEDULER STATE (before warp branches) ===== + # Each warp role independently maintains its own copy of these + # variables (like DeepGEMM where each role creates its own scheduler). + # Pre-fetch first task (current_num_kv loaded early above for latency hiding) + next_q_idx = current_q_idx + next_kv_idx = current_kv_idx + next_num_kv = current_num_kv + # Sentinel: no previous batch (matches DeepGEMM's q_idx = batch_size) + q_idx = batch_size + # While-loop termination flag (matches DeepGEMM's fetch_next_task pattern). + # True if this CTA has work assigned (start != end in schedule_meta). + has_work = (current_q_idx != end_q_idx) | (current_kv_idx != end_kv_idx) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # ===== WARP-SPECIALIZED EXECUTION ===== + + if is_tma_warp_0: + # TMA warp 0: loads Q (prefetch) + KV for group 0 + # Matches DeepGEMM's TMA warp with kv_group_idx == 0 + cute.arch.warpgroup_reg_dealloc(24) + lane_idx = tidx % 32 + + # Block table prefetch: 32 lanes cache 32 block indices, + # distributed via shuffle. (Matches DeepGEMM L233-244) + cached_blk_idx = cutlass.Int32(0) + kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use + + # Prefetch first Q before loop (like DeepGEMM line 203-204) + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + cute.copy( + tma_atom_b, + tBgB[(None, 0, next_q_idx)], + tBsB[(None, q_prod_state.index)], + tma_bar_ptr=q_bar, + mcast_mask=b_mcast_mask, + ) + cute.copy( + tma_atom_w, + tWgW[(None, next_q_idx)], + tWsW[(None, q_prod_state.index)], + tma_bar_ptr=q_bar, + ) + q_prod_state.advance() + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q prefetch: when batch changes, load Q for NEXT batch + if q_idx != q_idx_old: + kv_blk_ptr = cutlass.Int32(32) # force re-prefetch + prefetch_next = q_idx + 1 + if prefetch_next < end_q_idx: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + cute.copy( + tma_atom_b, + tBgB[(None, 0, prefetch_next)], + tBsB[(None, q_prod_state.index)], + tma_bar_ptr=q_bar, + mcast_mask=b_mcast_mask, + ) + cute.copy( + tma_atom_w, + tWgW[(None, prefetch_next)], + tWsW[(None, q_prod_state.index)], + tma_bar_ptr=q_bar, + ) + q_prod_state.advance() + elif prefetch_next == end_q_idx: + if end_kv_idx > 0: + q_pipeline.producer_acquire(q_prod_state) + q_bar = q_pipeline.producer_get_barrier(q_prod_state) + cute.copy( + tma_atom_b, + tBgB[(None, 0, prefetch_next)], + tBsB[(None, q_prod_state.index)], + tma_bar_ptr=q_bar, + mcast_mask=b_mcast_mask, + ) + cute.copy( + tma_atom_w, + tWgW[(None, prefetch_next)], + tWsW[(None, q_prod_state.index)], + tma_bar_ptr=q_bar, + ) + q_prod_state.advance() + + # Block table prefetch for group 0 (like DeepGEMM L233-241) + # Each lane prefetches block_table[q_idx][kv_idx + lane_i * 2] + if kv_blk_ptr == 32: + kv_blk_ptr = cutlass.Int32(0) + prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG + if prefetch_kv < num_kv: + cached_blk_idx = mBlockTable[(q_idx, prefetch_kv)] + else: + cached_blk_idx = cutlass.Int32(0) + + # Get block index via shuffle (like DeepGEMM L244) + phys_blk = cute.arch.shuffle_sync(cached_blk_idx, kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + + # Load KV + Scale for group 0 (kv_idx + 0) + # Unconditional TMA (like DeepGEMM): OOB kv_idx uses + # phys_blk=0 from block_table guard, writes to aligned + # padding region in logits. Keeps pipeline timing aligned. + kv_pipeline_0.producer_acquire(kv_prod_state_0) + bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0) + cute.copy( + tma_atom_a, + tAgA_0[(None, 0, phys_blk)], + tAsA_0[(None, kv_prod_state_0.index)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_0[(None, phys_blk)], + tSsS_0[(None, kv_prod_state_0.index)], + tma_bar_ptr=bar, + ) + kv_prod_state_0.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_tma_warp_1: + # TMA warp 1: loads KV + Scale for group 1 only + # Matches DeepGEMM's TMA warp with kv_group_idx == 1 + cute.arch.warpgroup_reg_dealloc(24) + lane_idx = tidx % 32 + + # Block table prefetch for group 1 + cached_blk_idx = cutlass.Int32(0) + kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # New q_idx → force block table re-prefetch + if q_idx != q_idx_old: + kv_blk_ptr = cutlass.Int32(32) + + # Block table prefetch for group 1 (like DeepGEMM L233-241) + # Each lane prefetches block_table[q_idx][kv_idx + 1 + lane_i * 2] + if kv_blk_ptr == 32: + kv_blk_ptr = cutlass.Int32(0) + prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG + if prefetch_kv < num_kv: + cached_blk_idx = mBlockTable[(q_idx, prefetch_kv)] + else: + cached_blk_idx = cutlass.Int32(0) + + # Get block index via shuffle (like DeepGEMM L244) + phys_blk = cute.arch.shuffle_sync(cached_blk_idx, kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + + # Load KV + Scale for group 1 (kv_idx + 1) + # Unconditional TMA (like DeepGEMM) + kv_pipeline_1.producer_acquire(kv_prod_state_1) + bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1) + cute.copy( + tma_atom_a, + tAgA_1[(None, 0, phys_blk)], + tAsA_1[(None, kv_prod_state_1.index)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_1[(None, phys_blk)], + tSsS_1[(None, kv_prod_state_1.index)], + tma_bar_ptr=bar, + ) + kv_prod_state_1.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = (mContextLens[next_q_idx] + block_kv_val - 1) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_umma_warp_0: + # UMMA warp for group 0 + # Must wait on Q pipeline: TMA operations with different + # barriers are NOT visibility-ordered even within the same + # warp. KV0 barrier arriving does not guarantee Q SMEM + # writes are visible. + cute.arch.warpgroup_reg_dealloc(24) + + # TMEM: wait for math warp 0's allocation, retrieve pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + if is_leader_cta: + num_k_blocks = cute.size(tCrA_0.shape[2]) + q_stage_0 = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Wait for Q pipeline when batch changes + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_cons_state_umma_0.advance() + q_pipeline.consumer_wait(q_cons_state_umma_0) + q_stage_0 = q_cons_state_umma_0.index + + # Process KV block for group 0 (kv_idx + 0) + # Unconditional UMMA (like DeepGEMM): OOB iterations + # compute on garbage data; results written to aligned + # padding region in logits buffer. + # Wait KV first, then TMEM empty (like DeepGEMM) + kv_pipeline_0.consumer_wait(kv_cons_state_umma_0) + umma_pipeline_0.producer_acquire(umma_prod_state_0) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + kv_stage = kv_cons_state_umma_0.index + tCtAcc_0 = tCtAcc_base_0[(None, None, None, umma_prod_state_0.index)] + for k_block in cutlass.range_constexpr(num_k_blocks): + cute.gemm( + tiled_mma, + tCtAcc_0, + tCrA_0[None, None, k_block, kv_stage], + tCrB[None, None, k_block, q_stage_0], + tCtAcc_0, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + # No consumer_release here — Math WG 0 releases + kv_cons_state_umma_0.advance() + + umma_pipeline_0.producer_commit(umma_prod_state_0) + umma_prod_state_0.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_umma_warp_1: + # UMMA warp for group 1 + # Explicitly waits on Q pipeline — critical because TMA warp 1 + # only loads KV1, not Q. Without this wait, UMMA warp 1 can + # start GEMM before TMA warp 0 finishes loading Q into SMEM. + cute.arch.warpgroup_reg_dealloc(24) + + # TMEM: wait for umma_warp_0's allocation, retrieve pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + if is_leader_cta: + num_k_blocks_1 = cute.size(tCrA_1.shape[2]) + q_stage_1 = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Wait for Q pipeline when batch changes + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_cons_state_umma_1.advance() + q_pipeline.consumer_wait(q_cons_state_umma_1) + q_stage_1 = q_cons_state_umma_1.index + + # Process KV block for group 1 (kv_idx + 1) + # Unconditional UMMA (like DeepGEMM) + # Wait KV first, then TMEM empty (like DeepGEMM) + kv_pipeline_1.consumer_wait(kv_cons_state_umma_1) + umma_pipeline_1.producer_acquire(umma_prod_state_1) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + kv_stage_1 = kv_cons_state_umma_1.index + tCtAcc_1 = tCtAcc_base_1[(None, None, None, umma_prod_state_1.index)] + for k_block in cutlass.range_constexpr(num_k_blocks_1): + cute.gemm( + tiled_mma, + tCtAcc_1, + tCrA_1[None, None, k_block, kv_stage_1], + tCrB[None, None, k_block, q_stage_1], + tCtAcc_1, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + # No consumer_release here — Math WG 1 releases + kv_cons_state_umma_1.advance() + + umma_pipeline_1.producer_commit(umma_prod_state_1) + umma_prod_state_1.advance() + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + elif is_math_warp: + cute.arch.warpgroup_reg_alloc(240) + + # TMEM: math warp 0 is the allocator; all math warps wait + retrieve + tmem.allocate(num_tmem_alloc_cols_total) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base_0 = cute.make_tensor(tmem_ptr, tCtAcc_fake_staged.layout) + tCtAcc_base_1 = cute.make_tensor(tmem_ptr + cols_per_group, tCtAcc_fake_staged.layout) + + local_tidx = tidx % 128 + cC = cute.make_identity_tensor(epi_sub_mn) + + if warpgroup_idx == 0: + # Math WG 0: process group 0 + # Reference setup (stage 0) for m_coord + # flat_divide by sub-tile to get sub-tile partitions + tAcc_0_ref = tCtAcc_base_0[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_0_ref_epi = cute.flat_divide(tAcc_0_ref, epi_sub_mn) + tiled_copy_ref_0 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_0_ref_epi[(None, None, 0, 0)] + ) + thr_copy_ref_0 = tiled_copy_ref_0.get_slice(local_tidx) + tTR_cC = thr_copy_ref_0.partition_D(cC) + m_coord = tTR_cC[0][0] + + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + + # Weight register cache: only first NUM_W_IN_REG + # per next_n slot (like DeepGEMM min(52, kNumHeads)). + # Remaining weights read from SMEM in epilogue. + # FP16 weights use half the regs, so we can fit + # all heads for next_n <= 3. + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + MAX_NUM_W_IN_REG = 64 if next_n <= 3 else 48 + else: + MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 + 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) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 0 (kv_idx + 0) + # Unconditional Math (like DeepGEMM): OOB results + # written to aligned padding region in logits buffer. + kv_pos = kv_idx * block_kv_val + m_coord + + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + # Skip KV wait, rely on UMMA barrier's + # transitive visibility. + umma_pipeline_0.consumer_wait(umma_cons_state_0) + else: + # Default: wait KV first to overlap lds with + # UMMA computation. + kv_pipeline_0.consumer_wait(kv_cons_state_math_0) + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + umma_pipeline_0.consumer_wait(umma_cons_state_0) + + # --- TMEM sub-tile setup --- + # flat_divide accumulator by sub-tile shape; + # partition once, then loop over sub-tiles. + tCtAcc_c0 = tCtAcc_base_0[(None, None, None, umma_cons_state_0.index)] + tAcc_c0 = tCtAcc_c0[((None, None), 0, 0)] + tAcc_c0_epi = cute.flat_divide(tAcc_c0, epi_sub_mn) + tc_0 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c0_epi[(None, None, 0, 0)]) + tr_0 = tc_0.get_slice(local_tidx) + tTR_0 = tr_0.partition_S(tAcc_c0_epi) + + # --- First sub-tile LDTM + KV release --- + if cutlass.const_expr(self.early_tmem_copy): + # Issue first sub-tile LDTM early + cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc) + # Scale LDS + KV release fill latency + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + cute.arch.fence_view_async_tmem_load() + else: + # Default: scale LDS + KV release first + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_0 = kv_cons_state_math_0.index + scale_val = sScales_0[(m_coord, sc_stage_0)] + kv_pipeline_0.consumer_release(kv_cons_state_math_0) + kv_cons_state_math_0.advance() + cute.copy(tc_0, tTR_0[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # --- Sub-tile compute loop --- + # Each sub-tile: LDTM.xN → fence → load → + # ReLU+FMA. Breaks FMA chain (16→4 per chunk) + # and interleaves LDTM with FP32 compute to + # reduce ShadowPipeThrottle. + # Sub-tiles are within each t-slot (num_heads + # // num_epi_subtiles wide). flat_divide yields + # next_n * num_epi_subtiles sub-tiles total; + # global index = t * num_epi_subtiles + i. + subtile_n = num_heads // num_epi_subtiles + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + packed_zero = pack_f16x2(Float16(0.0), Float16(0.0)) + for t in cutlass.range_constexpr(next_n): + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps0 = packed_zero + ps1 = packed_zero + else: + s0x = cutlass.Float32(0.0) + s0y = cutlass.Float32(0.0) + s1x = cutlass.Float32(0.0) + s1y = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(num_epi_subtiles): + # LDTM for sub-tiles 1..N-1 + # (sub-tile 0 handled above) + if t > 0 or i > 0: + cute.copy( + tc_0, + tTR_0[(None, None, None, 0, t * num_epi_subtiles + i)], + tTR_rAcc, + ) + cute.arch.fence_view_async_tmem_load() + # Release UMMA after last LDTM+fence + if t == next_n - 1 and i == num_epi_subtiles - 1: + umma_pipeline_0.consumer_release(umma_cons_state_0) + umma_cons_state_0.advance() + acc_vec = tTR_rAcc.load() + # Reg-path: weights from registers + reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n)) + for h in cutlass.range_constexpr(0, reg_h_end, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + r0 = t * NUM_W_IN_REG + h_g + pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1]) + pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3]) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + r0 = t * NUM_W_IN_REG + h_g + w0 = w_cache[r0] + w1 = w_cache[r0 + 1] + w2 = w_cache[r0 + 2] + w3 = w_cache[r0 + 3] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=nvvm.RoundingModeKind.RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=nvvm.RoundingModeKind.RN + ) + # SMEM-path: weights from shared mem + smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n) + for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + pw01 = pack_f16x2( + sW[(t * num_heads + h_g, q_stage_local)], + sW[(t * num_heads + h_g + 1, q_stage_local)], + ) + pw23 = pack_f16x2( + sW[(t * num_heads + h_g + 2, q_stage_local)], + sW[(t * num_heads + h_g + 3, q_stage_local)], + ) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + w0 = sW[(t * num_heads + h_g, q_stage_local)] + w1 = sW[(t * num_heads + h_g + 1, q_stage_local)] + w2 = sW[(t * num_heads + h_g + 2, q_stage_local)] + w3 = sW[(t * num_heads + h_g + 3, q_stage_local)] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=nvvm.RoundingModeKind.RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=nvvm.RoundingModeKind.RN + ) + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps_sum = add_f16x2(ps0, ps1) + sum_lo, sum_hi = unpack_f16x2(ps_sum) + result_t = sum_lo + sum_hi + else: + result_t = s0x + s0y + s1x + s1y + out_row = q_idx * next_n + t + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + mLogits[(out_row, kv_pos)] = self.output_dtype( + result_t * Float16(scale_val) + ) + else: + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 0) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + + else: + # Math WG 1: process group 1 + tAcc_1_ref = tCtAcc_base_1[(None, None, None, 0)][((None, None), 0, 0)] + tAcc_1_ref_epi = cute.flat_divide(tAcc_1_ref, epi_sub_mn) + tiled_copy_ref_1 = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_1_ref_epi[(None, None, 0, 0)] + ) + thr_copy_ref_1 = tiled_copy_ref_1.get_slice(local_tidx) + tTR_cC = thr_copy_ref_1.partition_D(cC) + m_coord = tTR_cC[0][0] + + tTR_rAcc = cute.make_fragment_like(tTR_cC, self.acc_dtype) + + # Weight register cache: only first NUM_W_IN_REG + # per next_n slot (like DeepGEMM min(52, kNumHeads)). + # FP16 weights use half the regs, so we can fit + # all heads for next_n <= 3. + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + MAX_NUM_W_IN_REG = 64 if next_n <= 3 else 48 + else: + MAX_NUM_W_IN_REG = 64 if next_n == 1 else 40 if next_n >= 4 else 52 + 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) + q_stage_local = cutlass.Int32(0) + + while has_work: + # fetch_next_task: commit next → current + q_idx_old = q_idx + q_idx = next_q_idx + kv_idx = next_kv_idx + num_kv = next_num_kv + + # Q pipeline consumer: wait for Q+Weights SMEM + if q_idx != q_idx_old: + if q_idx_old < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + q_pipeline.consumer_wait(q_cons_state) + q_stage_local = q_cons_state.index + # Preload first NUM_W_IN_REG weights per slot + for t_i in cutlass.range_constexpr(next_n): + for w_j in cutlass.range_constexpr(NUM_W_IN_REG): + w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ + (t_i * num_heads + w_j, q_stage_local) + ] + + # Process KV block for group 1 (kv_idx + 1) + # Unconditional Math (like DeepGEMM) + kv_idx_1 = kv_idx + 1 + + kv_pos = kv_idx_1 * block_kv_val + m_coord + + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + umma_pipeline_1.consumer_wait(umma_cons_state_1) + else: + kv_pipeline_1.consumer_wait(kv_cons_state_math_1) + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + umma_pipeline_1.consumer_wait(umma_cons_state_1) + + # --- TMEM sub-tile setup (WG1) --- + tCtAcc_c1 = tCtAcc_base_1[(None, None, None, umma_cons_state_1.index)] + tAcc_c1 = tCtAcc_c1[((None, None), 0, 0)] + tAcc_c1_epi = cute.flat_divide(tAcc_c1, epi_sub_mn) + tc_1 = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_c1_epi[(None, None, 0, 0)]) + tr_1 = tc_1.get_slice(local_tidx) + tTR_1 = tr_1.partition_S(tAcc_c1_epi) + + # --- First sub-tile LDTM + KV release (WG1) --- + if cutlass.const_expr(self.early_tmem_copy): + cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc) + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + cute.arch.fence_view_async_tmem_load() + else: + if cutlass.const_expr(self.remove_kv_wait_in_epilogue): + sc_stage_1 = kv_cons_state_math_1.index + scale_val = sScales_1[(m_coord, sc_stage_1)] + kv_pipeline_1.consumer_release(kv_cons_state_math_1) + kv_cons_state_math_1.advance() + cute.copy(tc_1, tTR_1[(None, None, None, 0, 0)], tTR_rAcc) + cute.arch.fence_view_async_tmem_load() + + # --- Sub-tile compute loop (WG1) --- + subtile_n = num_heads // num_epi_subtiles + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + packed_zero = pack_f16x2(Float16(0.0), Float16(0.0)) + for t in cutlass.range_constexpr(next_n): + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps0 = packed_zero + ps1 = packed_zero + else: + s0x = cutlass.Float32(0.0) + s0y = cutlass.Float32(0.0) + s1x = cutlass.Float32(0.0) + s1y = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(num_epi_subtiles): + if t > 0 or i > 0: + cute.copy( + tc_1, + tTR_1[(None, None, None, 0, t * num_epi_subtiles + i)], + tTR_rAcc, + ) + cute.arch.fence_view_async_tmem_load() + if t == next_n - 1 and i == num_epi_subtiles - 1: + umma_pipeline_1.consumer_release(umma_cons_state_1) + umma_cons_state_1.advance() + acc_vec = tTR_rAcc.load() + # Reg-path + reg_h_end = min(subtile_n, max(0, NUM_W_IN_REG - i * subtile_n)) + for h in cutlass.range_constexpr(0, reg_h_end, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + r0 = t * NUM_W_IN_REG + h_g + pw01 = pack_f16x2(w_cache[r0], w_cache[r0 + 1]) + pw23 = pack_f16x2(w_cache[r0 + 2], w_cache[r0 + 3]) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + r0 = t * NUM_W_IN_REG + h_g + w0 = w_cache[r0] + w1 = w_cache[r0 + 1] + w2 = w_cache[r0 + 2] + w3 = w_cache[r0 + 3] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=nvvm.RoundingModeKind.RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=nvvm.RoundingModeKind.RN + ) + # SMEM-path + smem_h_start = max(0, NUM_W_IN_REG - i * subtile_n) + for h in cutlass.range_constexpr(smem_h_start, subtile_n, 4): + n0 = h + h_g = i * subtile_n + h + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + pa01 = pack_f16x2( + Float16(acc_vec[n0]), Float16(acc_vec[n0 + 1]) + ) + pa23 = pack_f16x2( + Float16(acc_vec[n0 + 2]), Float16(acc_vec[n0 + 3]) + ) + pa01 = max_f16x2(pa01, packed_zero) + pa23 = max_f16x2(pa23, packed_zero) + pw01 = pack_f16x2( + sW[(t * num_heads + h_g, q_stage_local)], + sW[(t * num_heads + h_g + 1, q_stage_local)], + ) + pw23 = pack_f16x2( + sW[(t * num_heads + h_g + 2, q_stage_local)], + sW[(t * num_heads + h_g + 3, q_stage_local)], + ) + ps0 = fma_f16x2(pa01, pw01, ps0) + ps1 = fma_f16x2(pa23, pw23, ps1) + else: + a0 = cutlass.max(acc_vec[n0], cutlass.Float32(0.0)) + a1 = cutlass.max(acc_vec[n0 + 1], cutlass.Float32(0.0)) + a2 = cutlass.max(acc_vec[n0 + 2], cutlass.Float32(0.0)) + a3 = cutlass.max(acc_vec[n0 + 3], cutlass.Float32(0.0)) + w0 = sW[(t * num_heads + h_g, q_stage_local)] + w1 = sW[(t * num_heads + h_g + 1, q_stage_local)] + w2 = sW[(t * num_heads + h_g + 2, q_stage_local)] + w3 = sW[(t * num_heads + h_g + 3, q_stage_local)] + s0x, s0y = cute.arch.fma_packed_f32x2( + (a0, a1), (w0, w1), (s0x, s0y), rnd=nvvm.RoundingModeKind.RN + ) + s1x, s1y = cute.arch.fma_packed_f32x2( + (a2, a3), (w2, w3), (s1x, s1y), rnd=nvvm.RoundingModeKind.RN + ) + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + ps_sum = add_f16x2(ps0, ps1) + sum_lo, sum_hi = unpack_f16x2(ps_sum) + result_t = sum_lo + sum_hi + else: + result_t = s0x + s0y + s1x + s1y + out_row = q_idx * next_n + t + if cutlass.const_expr(self.epi_dtype == cutlass.Float16): + mLogits[(out_row, kv_pos)] = self.output_dtype( + result_t * Float16(scale_val) + ) + else: + mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + + # Advance: inline fetch_next_task + next_kv_idx = kv_idx + NUM_MATH_WG + if next_kv_idx >= num_kv: + next_q_idx = q_idx + 1 + next_kv_idx = 0 + if next_q_idx < batch_size: + next_num_kv = ( + mContextLens[next_q_idx] + block_kv_val - 1 + ) // block_kv_val + # Update while-loop condition + has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + + # Release last Q stage (WG 1) + if q_idx < batch_size: + q_pipeline.consumer_release(q_cons_state) + q_cons_state.advance() + + # TMEM dealloc: math warps are allocator + last consumer + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + else: + cute.arch.warpgroup_reg_dealloc(24) + + +def cdiv(a, b): + return (a + b - 1) // b + + +def compute_schedule_metadata(context_lens, block_kv, num_ctas): + """Compute schedule metadata: [num_ctas+1, 2] int32. + + Each row stores (q_idx, kv_idx / kNumMathWarpGroups) marking CTA boundaries. + Matches DeepGEMM's PagedMQALogitsScheduler metadata format: + - schedule[i] = start boundary for CTA i + - schedule[i+1] = end boundary for CTA i (= start of CTA i+1) + - schedule[num_ctas] = past-the-end sentinel (batch_size, 0) + + The kernel multiplies the second column by kNumMathWarpGroups (=2) to get + the actual kv_idx in units of KV blocks. + """ + batch_size = context_lens.shape[0] + splits_per_seq = [] + total_splits = 0 + for b in range(batch_size): + ctx = context_lens[b].item() + num_kv = cdiv(ctx, block_kv) + ns = cdiv(num_kv, 2) + splits_per_seq.append(ns) + total_splits += ns + + # Balanced distribution + q_div = total_splits // num_ctas + r_mod = total_splits % num_ctas + + schedule = torch.zeros((num_ctas + 1, 2), dtype=torch.int32) + + # For each CTA boundary, find (q_idx, kv_half_idx) + cum = 0 + seq_idx = 0 + seq_offset = 0 + for i in range(num_ctas + 1): + target = i * q_div + min(i, r_mod) + while seq_idx < batch_size and cum + (splits_per_seq[seq_idx] - seq_offset) <= target: + cum += splits_per_seq[seq_idx] - seq_offset + seq_idx += 1 + seq_offset = 0 + if seq_idx >= batch_size: + break + if seq_idx >= batch_size: + # Past-the-end: (batch_size, 0) — matches DeepGEMM's end sentinel. + # When the scheduler wraps past the last batch, it reaches + # (batch_size, 0), and the end check (q==end_q and kv==end_kv) + # correctly terminates. + schedule[i] = torch.tensor([batch_size, 0], dtype=torch.int32) + else: + local_split = target - cum + seq_offset + schedule[i] = torch.tensor([seq_idx, local_split], dtype=torch.int32) + + return schedule + + +def make_fused_kv(kv_cache_fp8, kv_cache_scales, block_kv, head_dim): + """Create fused KV tensor from separate KV and scale tensors. + + Output shape matches DeepGEMM: [num_phys_blocks, block_kv, 1, per_token_size] uint8 + where per_token_size = head_dim + 4. + + Per token: [KV (head_dim bytes)] [Scale (4 bytes)] + """ + num_phys_blocks = kv_cache_fp8.shape[0] + per_token_size = head_dim + 4 + block_bytes = block_kv * per_token_size + scale_offset = block_kv * head_dim + + fused = torch.zeros( + num_phys_blocks, + block_bytes, + dtype=torch.uint8, + device=kv_cache_fp8.device, + ) + for blk in range(num_phys_blocks): + fused[blk, :scale_offset] = kv_cache_fp8[blk].view(torch.uint8).reshape(-1) + fused[blk, scale_offset:] = kv_cache_scales[blk].view(torch.uint8).reshape(-1) + return fused.view(num_phys_blocks, block_kv, 1, per_token_size) + + +def fused_kv_views(kv_fused, block_kv, head_dim): + """Create KV (FP8) and Scale (float32) views from per-block fused buffer. + + Both views share the same underlying memory for L2 cache locality. + + Args: + kv_fused: [num_phys_blocks, block_kv * per_token_size] uint8 + Returns: + kv_pool: [block_kv, head_dim, num_phys_blocks] FP8 (strided view) + scales: [block_kv, num_phys_blocks] float32 (strided view) + """ + num_phys_blocks = kv_fused.shape[0] + per_token_size = head_dim + 4 + block_bytes = block_kv * per_token_size + scale_offset = block_kv * head_dim + + fused_flat = kv_fused.reshape(-1) + + # KV view: [block_kv, head_dim, num_phys_blocks] FP8 + # Within each block, KV data is contiguous [block_kv, head_dim] + # Element [m, k, l] → byte: l * block_bytes + m * head_dim + k + kv_pool = torch.as_strided( + fused_flat.view(torch.float8_e4m3fn), + size=(block_kv, head_dim, num_phys_blocks), + stride=(head_dim, 1, block_bytes), + ) + + # Scale view: [block_kv, num_phys_blocks] float32 + # Within each block, scales start at byte offset scale_offset + # and are contiguous [block_kv] float32. + # Element [m, l] → byte: l * block_bytes + scale_offset + m * 4 + # From float32 base at byte offset scale_offset: + # float32 index [m, l] → offset: l * (block_bytes / 4) + m + scale_base = fused_flat[scale_offset:].view(torch.float32) + scales = torch.as_strided( + scale_base, + size=(block_kv, num_phys_blocks), + stride=(1, block_bytes // 4), + ) + + return kv_pool, scales + + +def _make_dynamic_dlpacks( + kv_fused, q_3d, w_2d, logits, block_table_gpu, context_lens_gpu, schedule_meta_gpu +): + """Wrap tensors with dynamic shape markers for JIT reuse. + + Static dims (model constants): block_kv, head_dim, N, per_token, next_n + Dynamic dims (vary per call): batch_size, num_phys_blocks, max_model_len, + max_blocks_per_seq, num_ctas + """ + dl_kv = from_dlpack(kv_fused).mark_compact_shape_dynamic(mode=0) # [?phys, blk, 1, pt] + # DLPack does not support float8 types; view as uint8 (same 1-byte layout), + # then recast back to Float8E4M3FN inside the kernel. + q_for_dl = ( + q_3d.view(torch.uint8) if q_3d.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) else q_3d + ) + # q_3d is [N, D, B] with stride (D, 1, N*D) after permute(1,2,0); + # use dim_order() to get the correct stride_order for the non-contiguous layout. + # print("limin: q_3d.shape", q_3d.shape) + # print("limin: q_3d.stride()", q_3d.stride()) + # print("limin: q_3d.dim_order()", q_3d.dim_order()) + # dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( + # mode=2, stride_order=q_3d.dim_order()) # [N, D, ?B] + # dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic( + # mode=1, stride_order=w_2d.dim_order()) # [N, ?B] + # dl_logits = from_dlpack(logits).mark_compact_shape_dynamic( + # mode=0, stride_order=logits.dim_order()).mark_compact_shape_dynamic( + # mode=1, stride_order=logits.dim_order()) # [?B, ?ctx] + # dl_bt = from_dlpack(block_table_gpu).mark_compact_shape_dynamic( + # mode=0, stride_order=block_table_gpu.dim_order()).mark_compact_shape_dynamic( + # mode=1, stride_order=block_table_gpu.dim_order()) # [?B, ?blks] + # dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( + # mode=2) # [N, D, ?B] + # dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic( + # mode=1) # [N, ?B] + # dl_logits = from_dlpack(logits).mark_compact_shape_dynamic( + # mode=0).mark_compact_shape_dynamic( + # mode=1) # [?B, ?ctx] + # dl_bt = from_dlpack(block_table_gpu).mark_compact_shape_dynamic( + # mode=0).mark_compact_shape_dynamic( + # mode=1) # [?B, ?blks] + dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( + mode=2, stride_order=(2, 0, 1) + ) # [N, D, ?B] + dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic(mode=1, stride_order=(1, 0)) # [N, ?B] + dl_logits = ( + from_dlpack(logits) + .mark_compact_shape_dynamic(mode=0, stride_order=(0, 1)) + .mark_compact_shape_dynamic(mode=1, stride_order=(0, 1)) + ) # [?B, ?ctx] + dl_bt = ( + from_dlpack(block_table_gpu) + .mark_compact_shape_dynamic(mode=0, stride_order=(0, 1)) + .mark_compact_shape_dynamic(mode=1, stride_order=(0, 1)) + ) # [?B, ?blks] + dl_cl = from_dlpack(context_lens_gpu).mark_compact_shape_dynamic(mode=0) # [?B] + dl_sm = from_dlpack(schedule_meta_gpu).mark_compact_shape_dynamic(mode=0) # [?ctas, 2] + return dl_kv, dl_q, dl_w, dl_logits, dl_bt, dl_cl, dl_sm + + +def _prepare_inputs( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + max_model_len, + block_kv, + num_sms=148, + epi_dtype=None, + output_dtype=None, +): + """Prepare all host/device tensors and compute schedule metadata.""" + B, next_n, H, D = q_fp8.shape + N = next_n * H + num_phys_blocks = kv_fused.shape[0] + + q_3d = q_fp8.reshape(B, N, D).permute(1, 2, 0) + if epi_dtype is not None and epi_dtype == cutlass.Float16: + w_2d = weights.reshape(B, N).half().t() + else: + w_2d = weights.reshape(B, N).t() + block_table_gpu = block_table.to(device=q_fp8.device, dtype=torch.int32) + context_lens_gpu = context_lens.to(device=q_fp8.device, dtype=torch.int32) + + _TORCH_DTYPE = { + cutlass.Float32: torch.float32, + cutlass.Float16: torch.float16, + cutlass.BFloat16: torch.bfloat16, + } + logits_torch_dtype = _TORCH_DTYPE.get(output_dtype, torch.float32) + + # Align logits columns to SPLIT_KV = block_kv * NUM_MATH_WG (like + # DeepGEMM's aligned_max_context_len). OOB unconditional stores from + # the kernel write into this padding region safely. + SPLIT_KV = block_kv * 2 # NUM_MATH_WG = 2 + aligned_max_ctx = ((max_model_len + SPLIT_KV - 1) // SPLIT_KV) * SPLIT_KV + logits = torch.full( + (B * next_n, aligned_max_ctx), + float("-inf"), + device="cuda", + dtype=logits_torch_dtype, + ) + logits = logits[:, :max_model_len] + + # Grid size = num_sms (like DeepGEMM). Empty SMs auto-skip via + # while-loop boundary check (start == end in schedule_meta). + num_ctas = num_sms + schedule_meta = compute_schedule_metadata(context_lens, block_kv, num_ctas) + schedule_meta_gpu = schedule_meta.to(device=q_fp8.device) + + return ( + kv_fused, + q_3d, + w_2d, + logits, + block_table_gpu, + context_lens_gpu, + schedule_meta_gpu, + num_phys_blocks, + B, + ) + + +# Kernel cache: keyed by static params (block_kv, num_heads, head_dim, next_n, num_sms, remove_kv_wait) +_compiled_cache = {} + + +def _get_or_compile_kernel( + block_kv, + num_heads, + head_dim, + next_n, + num_sms, + kv_fused, + q_3d, + w_2d, + logits, + block_table_gpu, + context_lens_gpu, + schedule_meta_gpu, + num_phys_blocks, + B, + stream, + remove_kv_wait_in_epilogue=False, + early_tmem_copy=False, + smem_subpartition_opt=False, + max_kv_pipeline=False, + max_umma_pipeline=False, + num_epi_subtiles=1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, +): + """Return a compiled kernel, compiling only on first call per static config.""" + cache_key = ( + block_kv, + num_heads, + head_dim, + next_n, + num_sms, + remove_kv_wait_in_epilogue, + early_tmem_copy, + smem_subpartition_opt, + max_kv_pipeline, + max_umma_pipeline, + num_epi_subtiles, + epi_dtype, + acc_dtype, + output_dtype, + ) + if cache_key not in _compiled_cache: + kernel = FP8MQALogitsDGFullKKernel( + block_kv=block_kv, + num_heads=num_heads, + head_dim=head_dim, + next_n=next_n, + num_sms=num_sms, + remove_kv_wait_in_epilogue=remove_kv_wait_in_epilogue, + early_tmem_copy=early_tmem_copy, + smem_subpartition_opt=smem_subpartition_opt, + max_kv_pipeline=max_kv_pipeline, + max_umma_pipeline=max_umma_pipeline, + num_epi_subtiles=num_epi_subtiles, + epi_dtype=epi_dtype, + acc_dtype=acc_dtype, + output_dtype=output_dtype, + ) + dl_args = _make_dynamic_dlpacks( + kv_fused, q_3d, w_2d, logits, block_table_gpu, context_lens_gpu, schedule_meta_gpu + ) + compiled = cute.compile(kernel, *dl_args, num_phys_blocks, B, stream) + _compiled_cache[cache_key] = compiled + print( + f" [compile] {cache_key} kv_stages={kernel.num_kv_stages} umma_stages={kernel.num_umma_stages}" + ) + return _compiled_cache[cache_key] + + +def dsl_fp8_paged_mqa_logits_dg_fullk( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + max_model_len, + block_kv, + num_sms=148, + remove_kv_wait_in_epilogue=False, + early_tmem_copy=False, + smem_subpartition_opt=False, + max_kv_pipeline=False, + max_umma_pipeline=False, + num_epi_subtiles=1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, +): + """DG-FullK 2-group kernel with fused KV layout. Supports multi-batch. + + Args: + kv_fused: [num_phys_blocks, block_kv, 1, head_dim + 4] uint8 + """ + B, next_n, H, D = q_fp8.shape + + (kv_f, q_3d, w_2d, logits, bt_gpu, cl_gpu, sm_gpu, num_phys_blocks, B) = _prepare_inputs( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + max_model_len, + block_kv, + num_sms, + epi_dtype=epi_dtype, + output_dtype=output_dtype, + ) + + torch_stream = torch.cuda.Stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + compiled = _get_or_compile_kernel( + block_kv, + H, + D, + next_n, + num_sms, + kv_f, + q_3d, + w_2d, + logits, + bt_gpu, + cl_gpu, + sm_gpu, + num_phys_blocks, + B, + stream, + remove_kv_wait_in_epilogue=remove_kv_wait_in_epilogue, + early_tmem_copy=early_tmem_copy, + smem_subpartition_opt=smem_subpartition_opt, + max_kv_pipeline=max_kv_pipeline, + max_umma_pipeline=max_umma_pipeline, + num_epi_subtiles=num_epi_subtiles, + epi_dtype=epi_dtype, + acc_dtype=acc_dtype, + output_dtype=output_dtype, + ) + + dl_args = _make_dynamic_dlpacks(kv_f, q_3d, w_2d, logits, bt_gpu, cl_gpu, sm_gpu) + compiled( + *dl_args, + num_phys_blocks, + B, + stream, + ) + torch.cuda.synchronize() + + return logits + + +def run_test( + batch_size_list=None, + next_n_list=None, + avg_ctx_list=None, + num_sms=148, + remove_kv_wait_in_epilogue=False, + early_tmem_copy=False, + smem_subpartition_opt=False, + max_kv_pipeline=False, + max_umma_pipeline=False, + num_epi_subtiles=1, + epi_dtype=cutlass.Float32, + acc_dtype=cutlass.Float32, + output_dtype=cutlass.Float32, +): + """Test DG-FullK kernel against reference.""" + import sys + import time + + sys.path.insert(0, ".") + from paged_mqa_logits_helpers import calc_diff, generate_test_data, ref_fp8_paged_mqa_logits + + if batch_size_list is None: + batch_size_list = [32] + if next_n_list is None: + next_n_list = [1] + if avg_ctx_list is None: + avg_ctx_list = [32768] + + opt_str = "" + if remove_kv_wait_in_epilogue: + opt_str += " +remove_kv_wait" + if early_tmem_copy: + opt_str += " +early_tmem_copy" + if smem_subpartition_opt: + opt_str += " +smem_subpart" + if max_kv_pipeline: + opt_str += " +max_kv_pipeline" + if max_umma_pipeline: + opt_str += " +max_umma_pipeline" + print(f"=== DG-FullK (2-Group + Full-K TMA + Fused KV + Dynamic Shapes{opt_str}) Tests ===") + t0 = time.time() + n_passed = 0 + n_total = 0 + for test_batch in batch_size_list: + for test_next_n in next_n_list: + for avg_ctx in avg_ctx_list: + data = generate_test_data( + batch_size=test_batch, + next_n=test_next_n, + num_heads=64, + head_dim=128, + block_kv=128, + avg_context_len=avg_ctx, + max_model_len=max(avg_ctx * 2, 2048), + device="cuda", + ) + kv_fused = make_fused_kv( + data["kv_cache"], + data["kv_cache_scales"], + data["block_kv"], + 128, + ) + fk_logits = dsl_fp8_paged_mqa_logits_dg_fullk( + data["q"], + kv_fused, + data["weights"], + data["context_lens"], + data["block_table"], + data["max_model_len"], + data["block_kv"], + num_sms=num_sms, + remove_kv_wait_in_epilogue=remove_kv_wait_in_epilogue, + early_tmem_copy=early_tmem_copy, + smem_subpartition_opt=smem_subpartition_opt, + max_kv_pipeline=max_kv_pipeline, + max_umma_pipeline=max_umma_pipeline, + num_epi_subtiles=num_epi_subtiles, + epi_dtype=epi_dtype, + acc_dtype=acc_dtype, + output_dtype=output_dtype, + ) + ref_logits = ref_fp8_paged_mqa_logits( + data["q"], + data["kv_cache"], + data["kv_cache_scales"], + data["weights"], + data["context_lens"], + data["block_table"], + data["max_model_len"], + data["block_kv"], + ) + + B_test = data["batch_size"] + mask = torch.zeros_like(ref_logits, dtype=torch.bool) + for b in range(B_test): + ctx = data["context_lens"][b].item() + for t in range(test_next_n): + row = b * test_next_n + t + q_pos = ctx - test_next_n + t + mask[row, : q_pos + 1] = True + + diff = calc_diff( + fk_logits.float().masked_fill(~mask, 0), + ref_logits.masked_fill(~mask, 0), + ) + total_blks = sum(cdiv(data["context_lens"][b].item(), 128) for b in range(B_test)) + n_total += 1 + passed = diff < 1e-3 + if passed: + n_passed += 1 + status = "PASSED" if passed else "FAILED" + print( + f" B={test_batch}, next_n={test_next_n}, " + f"avg_ctx={avg_ctx}, nblk={total_blks}, num_sms={num_sms}: " + f"diff={diff:.2e} {status}" + ) + elapsed = time.time() - t0 + print(f"\n{n_passed}/{n_total} passed in {elapsed:.1f}s ({len(_compiled_cache)} compilations)") + + +def parse_args(): + import argparse + + parser = argparse.ArgumentParser(description="DG-FullK paged MQA logits kernel test") + parser.add_argument( + "--batch_size", + type=int, + nargs="+", + default=None, + help="batch size(s), e.g. --batch_size 1 32", + ) + parser.add_argument( + "--next_n", type=int, nargs="+", default=None, help="next_n value(s), e.g. --next_n 1 2 4" + ) + parser.add_argument( + "--avg_ctx", + type=int, + nargs="+", + default=None, + help="avg context len(s), e.g. --avg_ctx 256 4096", + ) + parser.add_argument( + "--num_sms", type=int, default=148, help="number of SMs for scheduling (default: 148)" + ) + parser.add_argument( + "--sweep", action="store_true", help="run full sweep over predefined ranges" + ) + parser.add_argument( + "--remove_kv_wait", + action="store_true", + help="remove KV barrier wait in epilogue (epilogue-bound opt)", + ) + parser.add_argument( + "--early_tmem_copy", + action="store_true", + help="issue LDTM early to hide latency behind scale LDS + KV release", + ) + parser.add_argument( + "--smem_subpartition_opt", + action="store_true", + help="pad SMEM to put sW/sScales in sub-partition 1, avoid UMMA conflicts", + ) + parser.add_argument( + "--max_kv_pipeline", + action="store_true", + help="maximize KV pipeline stages to fill SMEM budget", + ) + parser.add_argument( + "--max_umma_pipeline", + action="store_true", + help="maximize UMMA pipeline stages to fill TMEM budget", + ) + parser.add_argument( + "--num_epi_subtiles", type=int, default=1, help="number of epilogue sub-tiles (default: 1)" + ) + parser.add_argument( + "--epi_dtype", + type=str, + default="fp32", + choices=["fp32", "fp16"], + help="epilogue dtype (default: fp32)", + ) + parser.add_argument( + "--acc_dtype", + type=str, + default="fp32", + choices=["fp32", "fp16"], + help="accumulator dtype (default: fp32)", + ) + parser.add_argument( + "--output_dtype", + type=str, + default="fp32", + choices=["fp32", "fp16", "bf16"], + help="output logits dtype (default: fp32)", + ) + return parser.parse_args() + + +if __name__ == "__main__": + _DTYPE_MAP = { + "fp32": cutlass.Float32, + "fp16": cutlass.Float16, + "bf16": cutlass.BFloat16, + } + args = parse_args() + if args.sweep: + batch_size_override = [1, 32] + next_n_override = [1, 2, 3, 4] + avg_ctx_override = [256, 1024, 4096, 8192, 16384, 32768] + else: + batch_size_override = args.batch_size + next_n_override = args.next_n + avg_ctx_override = args.avg_ctx + run_test( + batch_size_list=batch_size_override, + next_n_list=next_n_override, + avg_ctx_list=avg_ctx_override, + num_sms=args.num_sms, + remove_kv_wait_in_epilogue=args.remove_kv_wait, + early_tmem_copy=args.early_tmem_copy, + smem_subpartition_opt=args.smem_subpartition_opt, + max_kv_pipeline=args.max_kv_pipeline, + max_umma_pipeline=args.max_umma_pipeline, + num_epi_subtiles=args.num_epi_subtiles, + epi_dtype=_DTYPE_MAP[args.epi_dtype], + acc_dtype=_DTYPE_MAP[args.acc_dtype], + output_dtype=_DTYPE_MAP[args.output_dtype], + ) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index f9988fc9442d..ac868d83c3b7 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -523,6 +523,7 @@ def from_pretrained(cls, indexer_max_chunk_size = sparse_attention_config.indexer_max_chunk_size skip_indexer_for_short_seqs = sparse_attention_config.skip_indexer_for_short_seqs use_cute_dsl_topk = sparse_attention_config.use_cute_dsl_topk + use_cute_dsl_logits = sparse_attention_config.use_cute_dsl_logits q_split_threshold = sparse_attention_config.q_split_threshold enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk else: @@ -532,6 +533,7 @@ def from_pretrained(cls, indexer_max_chunk_size = None skip_indexer_for_short_seqs = True use_cute_dsl_topk = False + use_cute_dsl_logits = False q_split_threshold = 8192 enable_heuristic_topk = False kwargs[ @@ -543,6 +545,7 @@ def from_pretrained(cls, skip_indexer_for_short_seqs= skip_indexer_for_short_seqs, use_cute_dsl_topk=use_cute_dsl_topk, + use_cute_dsl_logits=use_cute_dsl_logits, q_split_threshold=q_split_threshold, indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f6e7639d9644..e745b9b344ee 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -328,6 +328,11 @@ class DeepSeekSparseAttentionConfig(BaseSparseAttentionConfig): description= "Whether to use CuTE DSL top-k kernel instead of the CUDA C++ indexer_topk_decode." ) + use_cute_dsl_logits: bool = Field( + default=False, + description= + "Whether to use CuTE DSL paged MQA logits kernel on SM100 instead of C++ DeepGEMM." + ) q_split_threshold: int = Field( default=8192, description= diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py new file mode 100644 index 000000000000..26d34b0d3510 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -0,0 +1,545 @@ +# 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. +""" +Test CuTe DSL fp8_paged_mqa_logits kernel against C++ DeepGEMM reference. +""" + +import random + +import pytest +import torch + +try: + from utils.util import skip_pre_blackwell +except ModuleNotFoundError: + skip_pre_blackwell = pytest.mark.skipif(False, reason="") + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + +def has_deep_gemm(): + try: + from tensorrt_llm import deep_gemm + + return deep_gemm is not None + except Exception: + return False + + +def _ceil_to_ue8m0(x: torch.Tensor): + return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) + + +def _calc_diff(x: torch.Tensor, y: torch.Tensor): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + if denominator == 0: + return 0.0 + sim = 2 * (x * y).sum() / denominator + return (1 - sim).item() + + +def _ref_fp8_paged_mqa_logits( + q_fp8, kv_fp8, kv_scales, weights, context_lens, block_table, max_model_len, block_kv +): + """Pure PyTorch reference for fp8_paged_mqa_logits. + + Args: + q_fp8: [B, next_n, H, D] float8_e4m3fn + kv_fp8: [num_blocks, block_kv, D] float8_e4m3fn + kv_scales: [num_blocks, block_kv] float32 + weights: [B*next_n, H] float32 + context_lens: [B] int32 + block_table: [B, max_blocks] int32 + max_model_len: int + block_kv: int + + Returns: + logits: [B*next_n, max_model_len] float32 + """ + B, next_n, H, D = q_fp8.shape + device = q_fp8.device + + logits = torch.full( + (B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32 + ) + + q_f32 = q_fp8.float() + + for b in range(B): + ctx_len = context_lens[b].item() + q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device) + + w = weights[b * next_n : (b + 1) * next_n, :] + + for blk_idx in range((ctx_len + block_kv - 1) // block_kv): + phys_blk = block_table[b, blk_idx].item() + + k_f32 = kv_fp8[phys_blk].float() + scales = kv_scales[phys_blk] + + k_positions = torch.arange(blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device) + + mask = (k_positions[None, :] < ctx_len) & (k_positions[None, :] <= q_positions[:, None]) + + qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T) # [H, next_n, block_kv] + qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device)) + qk = torch.relu(qk) + + weighted = (w.T[:, :, None] * qk).sum(dim=0) # [next_n, block_kv] + weighted = weighted * scales[None, :] + + start_pos = blk_idx * block_kv + end_pos = start_pos + block_kv + logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where( + mask, weighted, torch.tensor(float("-inf"), device=device) + ) + + return logits + + +def _make_fused_kv(kv_fp8, kv_scales, block_kv, head_dim): + """Create fused KV in packed-by-type layout matching DeepGEMM/DSL kernel. + + Per block: [all FP8 bytes (block_kv * head_dim)] [all scale bytes (block_kv * 4)] + Viewed as [num_blocks, block_kv, 1, head_dim + 4]. + """ + num_phys_blocks = kv_fp8.shape[0] + per_token_size = head_dim + 4 + block_bytes = block_kv * per_token_size + scale_offset = block_kv * head_dim + + fused = torch.zeros(num_phys_blocks, block_bytes, dtype=torch.uint8, device=kv_fp8.device) + for blk in range(num_phys_blocks): + fused[blk, :scale_offset] = kv_fp8[blk].view(torch.uint8).reshape(-1) + fused[blk, scale_offset:] = ( + kv_scales[blk].float().contiguous().view(torch.uint8).reshape(-1) + ) + return fused.view(num_phys_blocks, block_kv, 1, per_token_size) + + +def _generate_test_data( + batch_size, next_n, num_heads, head_dim, block_kv, avg_context_len, max_model_len, device="cuda" +): + """Generate random test data for fp8 paged MQA logits.""" + context_lens = torch.randint( + max(block_kv, int(0.7 * avg_context_len)), + int(1.3 * avg_context_len) + 1, + (batch_size,), + dtype=torch.int32, + device="cpu", + ) + context_lens = context_lens.clamp(max=max_model_len) + + max_blocks_per_seq = (max_model_len + block_kv - 1) // block_kv + total_blocks = ((context_lens + block_kv - 1) // block_kv).sum().item() + num_phys_blocks = total_blocks + batch_size * 2 + + block_table = torch.full((batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device) + blk_offset = 0 + for i in range(batch_size): + n_blks = (context_lens[i].item() + block_kv - 1) // block_kv + block_table[i, :n_blks] = torch.arange( + blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device + ) + blk_offset += n_blks + + q_bf16 = torch.randn(batch_size, next_n, num_heads, head_dim, device=device) + q_fp8 = q_bf16.to(torch.float8_e4m3fn) + + kv_bf16 = torch.randn(num_phys_blocks, block_kv, head_dim, device=device) + kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4) + kv_scale = _ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1) + kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + + kv_fused = _make_fused_kv(kv_fp8, kv_scale, block_kv, head_dim) + + weights = torch.randn(batch_size * next_n, num_heads, device=device, dtype=torch.float32) + + return { + "q_fp8": q_fp8, + "kv_fp8": kv_fp8, + "kv_scales": kv_scale, + "kv_fused": kv_fused, + "weights": weights, + "context_lens": context_lens.to(device), + "block_table": block_table, + "max_model_len": max_model_len, + "block_kv": block_kv, + "num_phys_blocks": num_phys_blocks, + } + + +skip_if_unsupported = pytest.mark.skipif( + not (has_deep_gemm() and IS_CUTLASS_DSL_AVAILABLE), reason="Requires DeepGEMM and CuTe DSL" +) + + +@skip_if_unsupported +@skip_pre_blackwell +@pytest.mark.parametrize("batch_size", [1, 4, 32]) +@pytest.mark.parametrize("next_n", [1, 2, 3, 4]) +@pytest.mark.parametrize("avg_ctx", [256, 4096, 32768]) +@pytest.mark.parametrize("output_dtype", [torch.float32, torch.float16]) +def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype): + """Compare CuTe DSL kernel output against reference. + + Uses C++ DeepGEMM as reference when available (next_n in {1,2,4}), + falls back to pure PyTorch reference otherwise (e.g. next_n=3). + Tests both fp32 and fp16 epi/acc/output paths. + """ + torch.manual_seed(42) + random.seed(42) + + num_heads = 64 + head_dim = 128 + block_kv = 128 + max_model_len = max(avg_ctx * 2, 2048) + + data = _generate_test_data( + batch_size, next_n, num_heads, head_dim, block_kv, avg_ctx, max_model_len + ) + + from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + + # DSL kernel always uses full num_sms as grid size. + dsl_schedule_meta = get_paged_mqa_logits_metadata(data["context_lens"], block_kv, num_sms) + + # Reference: try C++ DeepGEMM first (fp32 only), fall back to PyTorch ref. + ref_logits = None + if output_dtype == torch.float32: + try: + from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits + + num_kv_multicast = 2 if next_n == 4 else 1 + num_clusters = num_sms // num_kv_multicast + dg_schedule_meta = get_paged_mqa_logits_metadata( + data["context_lens"], block_kv, num_clusters + ) + ref_logits = fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fused"], + data["weights"], + data["context_lens"], + data["block_table"], + dg_schedule_meta, + max_model_len, + ) + except RuntimeError: + pass + + if ref_logits is None: + ref_logits = _ref_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fp8"], + data["kv_scales"], + data["weights"], + data["context_lens"], + data["block_table"], + max_model_len, + block_kv, + ) + + # CuTe DSL kernel + dsl_logits = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fused"], + data["weights"], + data["context_lens"], + data["block_table"], + dsl_schedule_meta, + max_model_len, + epi_dtype=output_dtype, + acc_dtype=output_dtype, + output_dtype=output_dtype, + ) + + assert dsl_logits.dtype == output_dtype + + # Mask invalid positions + B = batch_size + positions = torch.arange(max_model_len, device="cuda").unsqueeze(0) + row_indices = torch.arange(B * next_n, device="cuda") // next_n + next_n_offset = torch.arange(B * next_n, device="cuda") % next_n + end_pos = data["context_lens"][row_indices] - next_n + next_n_offset + mask = positions <= end_pos.unsqueeze(1) + + dsl_masked = dsl_logits.float().masked_fill(~mask, 0) + ref_masked = ref_logits.float().masked_fill(~mask, 0) + finite = torch.isfinite(dsl_masked) & torch.isfinite(ref_masked) + diff = _calc_diff(dsl_masked.masked_fill(~finite, 0), ref_masked.masked_fill(~finite, 0)) + + tol = 5e-3 if output_dtype == torch.float16 else 1e-3 + assert diff < tol, ( + f"Accuracy check failed: diff={diff:.2e}, " + f"B={batch_size}, next_n={next_n}, avg_ctx={avg_ctx}, " + f"dtype={output_dtype}" + ) + + +def _profile_kernel_us(fn, num_warmup=10, num_iterations=30): + """Profile CUDA kernel time in microseconds using torch.profiler.""" + from torch.profiler import ProfilerActivity, profile + + for _ in range(num_warmup): + fn() + torch.cuda.synchronize() + + with profile(activities=[ProfilerActivity.CUDA], record_shapes=False) as prof: + for _ in range(num_iterations): + fn() + torch.cuda.synchronize() + + total_cuda_us = 0 + for evt in prof.events(): + if evt.device_type == torch.autograd.DeviceType.CUDA: + # for fp16 dtype, we use .half() to convert weights to fp16 dtype currently. + # so we need to skip the vectorized_elementwise_kernel event. + if "vectorized_elementwise_kernel" in evt.name: + continue + total_cuda_us += evt.device_time_total + return total_cuda_us / num_iterations + + +def _generate_bench_data( + batch_size, context_len, next_n, num_heads=64, head_dim=128, block_kv=128, device="cuda" +): + """Generate benchmark data with uniform context length.""" + torch.manual_seed(42) + num_blocks_per_seq = (context_len + block_kv - 1) // block_kv + total_blocks = batch_size * num_blocks_per_seq + + # fix-length workload: all sequences have the same context length. + context_lens = torch.full((batch_size,), context_len, dtype=torch.int32, device=device) + block_table = torch.arange(total_blocks, dtype=torch.int32, device=device).reshape( + batch_size, num_blocks_per_seq + ) + + q_fp8 = torch.randn( + batch_size, next_n, num_heads, head_dim, device=device, dtype=torch.bfloat16 + ).to(torch.float8_e4m3fn) + weights = torch.randn(batch_size * next_n, num_heads, device=device, dtype=torch.float32) + + kv_fp8 = torch.randn(total_blocks, block_kv, head_dim, device=device, dtype=torch.bfloat16).to( + torch.float8_e4m3fn + ) + kv_scales = ( + torch.rand(total_blocks, block_kv, device=device, dtype=torch.float32) * 0.01 + 0.001 + ) + + kv_fused = _make_fused_kv(kv_fp8, kv_scales, block_kv, head_dim) + + return { + "q_fp8": q_fp8, + "kv_fused": kv_fused, + "weights": weights, + "context_lens": context_lens, + "block_table": block_table, + "max_model_len": context_len, + "total_blocks": total_blocks, + } + + +def benchmark_fp8_paged_mqa_logits( + batch_sizes, + next_ns, + context_lens, + num_warmup=10, + num_iterations=30, + max_mem_gb=30, + output_dtype=torch.float32, + num_epi_subtiles=1, +): + """Benchmark CuTe DSL vs C++ DeepGEMM kernel time.""" + from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + + num_heads = 64 + head_dim = 128 + block_kv = 128 + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + + dtype_str = str(output_dtype).split(".")[-1] + print(f"output_dtype={dtype_str} num_epi_subtiles={num_epi_subtiles}") + is_non_default = output_dtype != torch.float32 or num_epi_subtiles != 1 + hdr = ( + f"{'batch':>5s} {'ctx':>7s} {'next_n':>6s} {'nblk':>7s} | " + f"{'DSL(us)':>8s} {'DG(fp32,us)':>12s} {'DG/DSL':>7s}" + ) + if is_non_default: + hdr += f" {'DSL(fp32,us)':>13s} {'DSL(fp32)/DSL':>13s}" + print(hdr) + print("-" * len(hdr)) + + for next_n in next_ns: + for context_len in context_lens: + for batch_size in batch_sizes: + mem_gb = ( + batch_size + * ((context_len + block_kv - 1) // block_kv) + * block_kv + * head_dim + / 1e9 + ) + nblk = batch_size * ((context_len + block_kv - 1) // block_kv) + if mem_gb > max_mem_gb: + print( + f"{batch_size:5d} {context_len:7d} {next_n:6d} " + f"{nblk:7d} | SKIP ({mem_gb:.1f}GB)" + ) + continue + + data = _generate_bench_data( + batch_size, context_len, next_n, num_heads, head_dim, block_kv + ) + + dsl_schedule_meta = get_paged_mqa_logits_metadata( + data["context_lens"], block_kv, num_sms + ) + + def dsl_fn(data=data): + torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fused"], + data["weights"], + data["context_lens"], + data["block_table"], + dsl_schedule_meta, + data["max_model_len"], + num_epi_subtiles=num_epi_subtiles, + epi_dtype=output_dtype, + acc_dtype=output_dtype, + output_dtype=output_dtype, + ) + + dsl_us = _profile_kernel_us(dsl_fn, num_warmup, num_iterations) + + dg_us = None + try: + from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits + + num_kv_multicast = 2 if next_n == 4 else 1 + num_clusters = num_sms // num_kv_multicast + dg_schedule_meta = get_paged_mqa_logits_metadata( + data["context_lens"], block_kv, num_clusters + ) + + def dg_fn(data=data): + fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fused"], + data["weights"], + data["context_lens"], + data["block_table"], + dg_schedule_meta, + data["max_model_len"], + ) + + dg_us = _profile_kernel_us(dg_fn, num_warmup, num_iterations) + except RuntimeError: + pass + + dsl_f32_us = None + if is_non_default: + + def dsl_f32_fn(data=data): + torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fused"], + data["weights"], + data["context_lens"], + data["block_table"], + dsl_schedule_meta, + data["max_model_len"], + ) + + dsl_f32_us = _profile_kernel_us(dsl_f32_fn, num_warmup, num_iterations) + + ratio_str = f"{dg_us / dsl_us:6.3f}x" if dg_us else " N/A " + dg_str = f"{dg_us:11.1f}" if dg_us else " N/A" + line = ( + f"{batch_size:5d} {context_len:7d} {next_n:6d} " + f"{nblk:7d} | {dsl_us:7.1f} {dg_str} {ratio_str}" + ) + if is_non_default: + f32_str = f"{dsl_f32_us:12.1f}" if dsl_f32_us else " N/A" + f32_ratio = f"{dsl_f32_us / dsl_us:12.3f}x" if dsl_f32_us else " N/A " + line += f" {f32_str} {f32_ratio}" + print(line) + + del data + torch.cuda.empty_cache() + print() + + +if __name__ == "__main__": + import argparse + import os + import sys + + sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir, os.pardir, os.pardir)) + + parser = argparse.ArgumentParser(description="Benchmark CuTe DSL fp8_paged_mqa_logits kernel") + parser.add_argument( + "--batch_size", + type=int, + nargs="+", + default=[1, 32, 128], + help="batch sizes (default: 1 32 128)", + ) + parser.add_argument( + "--next_n", type=int, nargs="+", default=[1, 2, 4], help="next_n values (default: 1 2 4)" + ) + parser.add_argument( + "--context_len", + type=int, + nargs="+", + default=[4096, 32768, 131072], + help="context lengths (default: 4096 32768 131072)", + ) + parser.add_argument("--warmup", type=int, default=10, help="warmup iterations (default: 10)") + parser.add_argument("--repeat", type=int, default=30, help="profiling iterations (default: 30)") + # TODO: do we need this argument? + parser.add_argument( + "--max_mem_gb", type=float, default=30, help="max KV cache memory in GB (default: 30)" + ) + parser.add_argument( + "--output_dtype", + type=str, + default="float32", + choices=["float32", "float16"], + help="output dtype (default: float32)", + ) + parser.add_argument( + "--num_epi_subtiles", + type=int, + default=1, + choices=[1, 2, 4], + help="epilogue sub-tile count (default: 1)", + ) + args = parser.parse_args() + + dtype_map = {"float32": torch.float32, "float16": torch.float16} + benchmark_fp8_paged_mqa_logits( + batch_sizes=args.batch_size, + next_ns=args.next_n, + context_lens=args.context_len, + num_warmup=args.warmup, + num_iterations=args.repeat, + max_mem_gb=args.max_mem_gb, + output_dtype=dtype_map[args.output_dtype], + num_epi_subtiles=args.num_epi_subtiles, + ) From c32db8fdb88e66c45402c81d79f3a764e6d1a2c2 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 04:43:55 -0700 Subject: [PATCH 02/26] [None][fix] Fix docstring script filename in fp8_paged_mqa_logits.py Replace stale development script name with the actual module filename. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index b1fa2daee12a..b53ea175a438 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -64,12 +64,12 @@ Run scripts: - Single values: - python paged_mqa_logits_dg_fullk_tma_7_dynamic_improve_v3.py \ + python fp8_paged_mqa_logits.py \ --batch_size 1 --next_n 2 --avg_ctx 4096 --num_sms 148 - Multiple values: - python paged_mqa_logits_dg_fullk_tma_7_dynamic_improve_v3.py \ + python fp8_paged_mqa_logits.py \ --batch_size 1 32 --next_n 1 2 4 --avg_ctx 256 4096 --num_sms 148 - - Full sweep: python paged_mqa_logits_dg_fullk_tma_7_dynamic_improve_v3.py --sweep + - Full sweep: python fp8_paged_mqa_logits.py --sweep - Default (no args): uses batch_size=[32], next_n=[1], avg_ctx=[32768], num_sms=[148] as before """ From b235eb13e8f2a4ba8b85f7b0b8938d8576a15381 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 04:45:04 -0700 Subject: [PATCH 03/26] [None][fix] Update copyright year in paged_mqa_logits __init__.py Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py index e2190e217e6a..82ca27da8e58 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); From 69d7170830aec1268eb64d0f0e46d884943fddcf Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 04:48:33 -0700 Subject: [PATCH 04/26] [None][fix] Remove unused max_mem_gb argument from benchmark Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../test_cute_dsl_fp8_paged_mqa_logits.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 26d34b0d3510..052163aeee77 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -360,7 +360,6 @@ def benchmark_fp8_paged_mqa_logits( context_lens, num_warmup=10, num_iterations=30, - max_mem_gb=30, output_dtype=torch.float32, num_epi_subtiles=1, ): @@ -387,20 +386,7 @@ def benchmark_fp8_paged_mqa_logits( for next_n in next_ns: for context_len in context_lens: for batch_size in batch_sizes: - mem_gb = ( - batch_size - * ((context_len + block_kv - 1) // block_kv) - * block_kv - * head_dim - / 1e9 - ) nblk = batch_size * ((context_len + block_kv - 1) // block_kv) - if mem_gb > max_mem_gb: - print( - f"{batch_size:5d} {context_len:7d} {next_n:6d} " - f"{nblk:7d} | SKIP ({mem_gb:.1f}GB)" - ) - continue data = _generate_bench_data( batch_size, context_len, next_n, num_heads, head_dim, block_kv @@ -512,10 +498,6 @@ def dsl_f32_fn(data=data): ) parser.add_argument("--warmup", type=int, default=10, help="warmup iterations (default: 10)") parser.add_argument("--repeat", type=int, default=30, help="profiling iterations (default: 30)") - # TODO: do we need this argument? - parser.add_argument( - "--max_mem_gb", type=float, default=30, help="max KV cache memory in GB (default: 30)" - ) parser.add_argument( "--output_dtype", type=str, @@ -539,7 +521,6 @@ def dsl_f32_fn(data=data): context_lens=args.context_len, num_warmup=args.warmup, num_iterations=args.repeat, - max_mem_gb=args.max_mem_gb, output_dtype=dtype_map[args.output_dtype], num_epi_subtiles=args.num_epi_subtiles, ) From a2b9b43254df1dbe7e53c0b0fdf348e0a7d749b6 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 21:41:12 -0700 Subject: [PATCH 05/26] [None][fix] Fix stream handling, add arch guard, and rename kernel class - Use current_stream() instead of creating a new stream to avoid data races - Remove unnecessary torch.cuda.synchronize() - Add is_sm_100f() check in custom op to fail fast on unsupported GPUs - Rename FP8MQALogitsDGFullKKernel to FP8MQALogitsKernel - Add __all__ to paged_mqa_logits __init__.py - Fix test skip logic to precisely match SM 100/103 only Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 14 ++++++++------ .../blackwell/paged_mqa_logits/__init__.py | 6 +++++- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 4 ++-- .../sparse/test_cute_dsl_fp8_paged_mqa_logits.py | 13 +++++++------ 4 files changed, 22 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 3eb037c2e609..b773196a5ca2 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5131,8 +5131,7 @@ def warmup_cute_dsl_indexer_topk( # ------------------------------------------------------------------ # # CuTE DSL FP8 Paged MQA Logits (Blackwell SM100) # # ------------------------------------------------------------------ # - from ..cute_dsl_kernels.blackwell.paged_mqa_logits import \ - FP8MQALogitsDGFullKKernel + from ..cute_dsl_kernels.blackwell.paged_mqa_logits import FP8MQALogitsKernel class CuteDSLPagedMQALogitsRunner: """Runner for CuTe DSL FP8 Paged MQA Logits kernel (Blackwell SM100). @@ -5187,7 +5186,7 @@ def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, dl_args = cls._make_dlpacks(kv_flat, q_3d, w_2d, logits, block_table, context_lens, schedule_meta) - kernel = FP8MQALogitsDGFullKKernel( + kernel = FP8MQALogitsKernel( block_kv=block_kv, num_heads=num_heads, head_dim=head_dim, @@ -5264,8 +5263,8 @@ def forward( ) logits = logits[:, :max_context_len] - # Create stream - torch_stream = torch.cuda.Stream() + # Get current stream + torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) # Compile if needed (uses real tensors for shape marking) @@ -5284,7 +5283,6 @@ def forward( block_table, context_lens, schedule_meta) compiled(*dl_args, num_phys_blocks, B, stream) - torch.cuda.synchronize() return logits @torch.library.custom_op("trtllm::cute_dsl_fp8_paged_mqa_logits", @@ -5303,6 +5301,10 @@ def cute_dsl_fp8_paged_mqa_logits( 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.") return CuteDSLPagedMQALogitsRunner.forward( q, kv_fused, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py index 82ca27da8e58..ab0c83ef090e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py @@ -13,4 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .fp8_paged_mqa_logits import FP8MQALogitsDGFullKKernel +from .fp8_paged_mqa_logits import FP8MQALogitsKernel + +__all__ = [ + "FP8MQALogitsKernel", +] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index b53ea175a438..e8b74c0e325b 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -194,7 +194,7 @@ def add_f16x2( ) -class FP8MQALogitsDGFullKKernel: +class FP8MQALogitsKernel: """ DG-Aligned 2-group kernel with full-K TMA, multi-batch support. @@ -1977,7 +1977,7 @@ def _get_or_compile_kernel( output_dtype, ) if cache_key not in _compiled_cache: - kernel = FP8MQALogitsDGFullKKernel( + kernel = FP8MQALogitsKernel( block_kv=block_kv, num_heads=num_heads, head_dim=head_dim, diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 052163aeee77..50516f5112d6 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -21,12 +21,13 @@ import pytest import torch -try: - from utils.util import skip_pre_blackwell -except ModuleNotFoundError: - skip_pre_blackwell = pytest.mark.skipif(False, reason="") - from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._utils import get_sm_version + +skip_not_sm100 = pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason=f"CuTe DSL FP8 Paged MQA Logits only supports SM 100/103, got SM {get_sm_version()}", +) def has_deep_gemm(): @@ -188,7 +189,7 @@ def _generate_test_data( @skip_if_unsupported -@skip_pre_blackwell +@skip_not_sm100 @pytest.mark.parametrize("batch_size", [1, 4, 32]) @pytest.mark.parametrize("next_n", [1, 2, 3, 4]) @pytest.mark.parametrize("avg_ctx", [256, 4096, 32768]) From a2b0de777c40468fae590d9bdcfa255403ebefa7 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 21:51:03 -0700 Subject: [PATCH 06/26] [None][fix] Ensure CuTE DSL op registration when only logits kernel is enabled Broaden the cute_dsl_custom_ops import guard to also trigger when use_cute_dsl_logits is True, so the custom op is registered even when use_cute_dsl_topk is False. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 5ea98f2767f6..60a7fc7295ee 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1121,13 +1121,15 @@ def __init__(self, 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_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+).""" From fe637c8d352227b1bbdaec09fe556f2de8563f5a Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 21:56:34 -0700 Subject: [PATCH 07/26] [None][refactor] Rename use_cute_dsl_logits to use_cute_dsl_paged_mqa_logits Align config field name with the op name cute_dsl_fp8_paged_mqa_logits for clarity. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 12 ++++++------ tensorrt_llm/_torch/model_config.py | 7 ++++--- tensorrt_llm/llmapi/llm_args.py | 2 +- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 60a7fc7295ee..d691fbdb2db9 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1111,10 +1111,10 @@ 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_logits = (getattr(sparse_attention_config, - 'use_cute_dsl_logits', False) - and IS_CUTLASS_DSL_AVAILABLE - and get_sm_version() >= 100) + self.use_cute_dsl_paged_mqa_logits = (getattr( + sparse_attention_config, 'use_cute_dsl_paged_mqa_logits', False) + and IS_CUTLASS_DSL_AVAILABLE + and get_sm_version() >= 100) self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5 self._enable_heuristic_topk = ( @@ -1122,7 +1122,7 @@ def __init__(self, and get_sm_version() >= 100) if (self.use_cute_dsl_topk - or self.use_cute_dsl_logits) and layer_idx == 0: + or self.use_cute_dsl_paged_mqa_logits) and layer_idx == 0: from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops if self.use_cute_dsl_topk: @@ -1671,7 +1671,7 @@ def sparse_attn_indexer( k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers( self.layer_idx) - if self.use_cute_dsl_logits: + 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) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index ac868d83c3b7..28520f3c3e39 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -523,7 +523,7 @@ def from_pretrained(cls, indexer_max_chunk_size = sparse_attention_config.indexer_max_chunk_size skip_indexer_for_short_seqs = sparse_attention_config.skip_indexer_for_short_seqs use_cute_dsl_topk = sparse_attention_config.use_cute_dsl_topk - use_cute_dsl_logits = sparse_attention_config.use_cute_dsl_logits + use_cute_dsl_paged_mqa_logits = sparse_attention_config.use_cute_dsl_paged_mqa_logits q_split_threshold = sparse_attention_config.q_split_threshold enable_heuristic_topk = sparse_attention_config.enable_heuristic_topk else: @@ -533,7 +533,7 @@ def from_pretrained(cls, indexer_max_chunk_size = None skip_indexer_for_short_seqs = True use_cute_dsl_topk = False - use_cute_dsl_logits = False + use_cute_dsl_paged_mqa_logits = False q_split_threshold = 8192 enable_heuristic_topk = False kwargs[ @@ -545,7 +545,8 @@ def from_pretrained(cls, skip_indexer_for_short_seqs= skip_indexer_for_short_seqs, use_cute_dsl_topk=use_cute_dsl_topk, - use_cute_dsl_logits=use_cute_dsl_logits, + use_cute_dsl_paged_mqa_logits= + use_cute_dsl_paged_mqa_logits, q_split_threshold=q_split_threshold, indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e745b9b344ee..365a3874cadd 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -328,7 +328,7 @@ class DeepSeekSparseAttentionConfig(BaseSparseAttentionConfig): description= "Whether to use CuTE DSL top-k kernel instead of the CUDA C++ indexer_topk_decode." ) - use_cute_dsl_logits: bool = Field( + use_cute_dsl_paged_mqa_logits: bool = Field( default=False, description= "Whether to use CuTE DSL paged MQA logits kernel on SM100 instead of C++ DeepGEMM." From 8f9b0e15ebf4037a72b729b8a473a712227bf364 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 22:00:24 -0700 Subject: [PATCH 08/26] [None][fix] Remove redundant SM version check from DSA logits config The custom op already validates SM 100/103 via is_sm_100f(). Also replace getattr with direct attribute access to match use_cute_dsl_topk style. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/sparse/dsa.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index d691fbdb2db9..99f206ac8f33 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1111,10 +1111,9 @@ 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 = (getattr( - sparse_attention_config, 'use_cute_dsl_paged_mqa_logits', False) - and IS_CUTLASS_DSL_AVAILABLE - and get_sm_version() >= 100) + 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 = ( From cf418abde1cce274725bca0b9949d769ee02d405 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Mon, 20 Apr 2026 22:11:58 -0700 Subject: [PATCH 09/26] [None][fix] Clean up DeepGEMM/DG-FullK references in docstrings and user-facing strings Remove DeepGEMM and DG-FullK references from module docstring, class docstring, function docstrings, print statements, and argparse descriptions. Inline implementation comments retained as design cross-references. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 48 ++++++------------- 1 file changed, 14 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index e8b74c0e325b..0d14a77cbe5c 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -1,40 +1,36 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ -DeepGEMM-aligned 2-group kernel with full-K TMA and fused KV layout (SM100). -Supports multi-batch via in-kernel scheduler matching DeepGEMM's PagedMQALogitsScheduler. +CuTe DSL FP8 paged MQA logits kernel for Blackwell (SM100). Architecture: - 384 threads: 256 math (2 WGs) + 128 specialized (2 TMA + 2 UMMA) - - Full-K TMA: 1 TMA per KV block [128, 128], UMMA iterates 4x K=32 + - 1 TMA per KV block [128, 128], UMMA iterates 4x K=32 - 2 warp groups process 2 KV blocks per iteration (kNumMathWarpGroups=2) - Q reloaded via TMA pipeline when q_idx (batch) changes - Persistent kernel: CTAs iterate through assigned (q_idx, kv_idx) pairs - Weights cached in registers: preloaded once per q_idx change (not per KV block) - KV Scales loaded via TMA to SMEM (separate pipeline per group, Math consumes) -Merged KV+Scale pipeline (tma_5, matches DeepGEMM): +Merged KV+Scale pipeline: - KV data and scales share a single TMA barrier per group - TMA loads both KV and Scale under one barrier (combined tx_count) - UMMA waits on merged barrier (for KV GEMM), does NOT release - Math waits on merged barrier (for scale read), Math releases - - This eliminates the separate scale pipeline overhead -Fused KV layout (tma_4, matches DeepGEMM): +Fused KV layout: - KV data and scales stored contiguously per physical block: [num_phys_blocks, block_kv * (head_dim + 4)] bytes - Per block: [KV_all_tokens (block_kv * head_dim bytes)] [Scales (block_kv * 4 bytes)] - KV and Scale views are derived inside __call__ using CuTE pointer arithmetic - - Benefit: L2 cache locality — scale data shares cache lines with KV data -Scheduler (aligned with DeepGEMM's PagedMQALogitsScheduler): +Scheduler: - schedule_meta[sm_idx] = (start_q_idx, start_kv_idx / kNumMathWarpGroups) - schedule_meta[sm_idx+1] = end boundary for this CTA - fetch_next_task pattern: each warp role independently advances (q_idx, kv_idx) - kv_idx in units of KV blocks, advances by kNumMathWarpGroups=2 per step - - exist_q_idx(qi): checks if qi is within this CTA's assigned range (for Q prefetch) -Dynamic shape support (tma_8): +Dynamic shape support: - Model-constant dims (block_kv, head_dim, N, per_token) remain static for codegen - Runtime-varying dims (batch_size, num_phys_blocks, max_ctx, max_blocks_per_seq, num_ctas) are marked dynamic via mark_compact_shape_dynamic @@ -47,30 +43,15 @@ -> LDTM -> Reg(FP32) -> cvt FP16 -> ReLU(FP16) -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) - Benefits: weights SMEM BW halved, weight regs halved, FP16 FMA - Unchanged: TMEM(FP32), LDTM BW(FP32), acc regs(FP32) Flow 2: --acc_dtype fp16 --epi_dtype fp16 Q(FP8) x K(FP8) -> MMA acc(FP16) -> TMEM(FP16, pack_16b) -> LDTM -> Reg(FP16) -> ReLU(FP16) -> FMA(fma.rn.f16x2) with weights(FP16 from SMEM) -> partial sum(FP16) -> x scale(FP32->FP16) -> cvt output_dtype -> store logits(output_dtype) - Extra benefits over Flow 1: TMEM halved (more umma stages), LDTM BW halved, acc regs halved - Risk: MMA FP16 accumulation over K=128 has precision loss; epilogue sum may overflow FP16 --output_dtype: fp32 (default), fp16, bf16. Controls logits tensor dtype and final store conversion. - - Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32 (original FP32 baseline) - -Run scripts: - - Single values: - python fp8_paged_mqa_logits.py \ - --batch_size 1 --next_n 2 --avg_ctx 4096 --num_sms 148 - - Multiple values: - python fp8_paged_mqa_logits.py \ - --batch_size 1 32 --next_n 1 2 4 --avg_ctx 256 4096 --num_sms 148 - - Full sweep: python fp8_paged_mqa_logits.py --sweep - - Default (no args): uses batch_size=[32], next_n=[1], avg_ctx=[32768], num_sms=[148] as before + Default: --acc_dtype fp32 --epi_dtype fp32 --output_dtype fp32 """ from typing import Tuple @@ -195,8 +176,7 @@ def add_f16x2( class FP8MQALogitsKernel: - """ - DG-Aligned 2-group kernel with full-K TMA, multi-batch support. + """FP8 paged MQA logits kernel for Blackwell (SM100). Each CTA processes a range of (q_idx, kv_split) pairs. A split = 2 consecutive KV blocks within a sequence (one per warp group). @@ -1689,7 +1669,7 @@ def compute_schedule_metadata(context_lens, block_kv, num_ctas): """Compute schedule metadata: [num_ctas+1, 2] int32. Each row stores (q_idx, kv_idx / kNumMathWarpGroups) marking CTA boundaries. - Matches DeepGEMM's PagedMQALogitsScheduler metadata format: + Metadata format: - schedule[i] = start boundary for CTA i - schedule[i+1] = end boundary for CTA i (= start of CTA i+1) - schedule[num_ctas] = past-the-end sentinel (batch_size, 0) @@ -1741,7 +1721,7 @@ def compute_schedule_metadata(context_lens, block_kv, num_ctas): def make_fused_kv(kv_cache_fp8, kv_cache_scales, block_kv, head_dim): """Create fused KV tensor from separate KV and scale tensors. - Output shape matches DeepGEMM: [num_phys_blocks, block_kv, 1, per_token_size] uint8 + Output shape: [num_phys_blocks, block_kv, 1, per_token_size] uint8 where per_token_size = head_dim + 4. Per token: [KV (head_dim bytes)] [Scale (4 bytes)] @@ -2023,7 +2003,7 @@ def dsl_fp8_paged_mqa_logits_dg_fullk( acc_dtype=cutlass.Float32, output_dtype=cutlass.Float32, ): - """DG-FullK 2-group kernel with fused KV layout. Supports multi-batch. + """Run FP8 paged MQA logits kernel. Args: kv_fused: [num_phys_blocks, block_kv, 1, head_dim + 4] uint8 @@ -2100,7 +2080,7 @@ def run_test( acc_dtype=cutlass.Float32, output_dtype=cutlass.Float32, ): - """Test DG-FullK kernel against reference.""" + """Test FP8 paged MQA logits kernel against reference.""" import sys import time @@ -2125,7 +2105,7 @@ def run_test( opt_str += " +max_kv_pipeline" if max_umma_pipeline: opt_str += " +max_umma_pipeline" - print(f"=== DG-FullK (2-Group + Full-K TMA + Fused KV + Dynamic Shapes{opt_str}) Tests ===") + print(f"=== FP8 Paged MQA Logits (Fused KV + Dynamic Shapes{opt_str}) Tests ===") t0 = time.time() n_passed = 0 n_total = 0 @@ -2209,7 +2189,7 @@ def run_test( def parse_args(): import argparse - parser = argparse.ArgumentParser(description="DG-FullK paged MQA logits kernel test") + parser = argparse.ArgumentParser(description="FP8 paged MQA logits kernel test") parser.add_argument( "--batch_size", type=int, From b1260fbce1dc6528a7139095b9907f70a5c12475 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Apr 2026 08:43:59 +0000 Subject: [PATCH 10/26] [None][refactor] Migrate MQA logits runner to fake tensor + TVM FFI Replace from_dlpack + mark_compact_shape_dynamic compile pattern with make_fake_compact_tensor + TVM FFI in CuteDSLPagedMQALogitsRunner. This eliminates real tensor data at compile time, removes per-call dlpack wrapping and manual CUDA stream passing at runtime. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 115 +++++++++++------- 1 file changed, 68 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index b773196a5ca2..4caa5ce76372 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5142,30 +5142,6 @@ class CuteDSLPagedMQALogitsRunner: kernel_cache = dict() - @classmethod - def _make_dlpacks(cls, kv_flat, q_3d, w_2d, logits, block_table, - context_lens, schedule_meta): - """Wrap tensors with dynamic shape markers for JIT reuse.""" - from cutlass.cute.runtime import from_dlpack - dl_kv = from_dlpack(kv_flat).mark_compact_shape_dynamic(mode=0) - q_for_dl = q_3d.view( - torch.uint8) if q_3d.dtype in (torch.float8_e4m3fn, - torch.float8_e5m2) else q_3d - dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( - mode=2, stride_order=(2, 0, 1)) - dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic( - mode=1, stride_order=(1, 0)) - dl_logits = from_dlpack(logits).mark_compact_shape_dynamic( - mode=0, stride_order=(0, 1)).mark_compact_shape_dynamic( - mode=1, stride_order=(0, 1)) - dl_bt = from_dlpack(block_table).mark_compact_shape_dynamic( - mode=0, stride_order=(0, 1)).mark_compact_shape_dynamic( - mode=1, stride_order=(0, 1)) - dl_cl = from_dlpack(context_lens).mark_compact_shape_dynamic(mode=0) - dl_sm = from_dlpack(schedule_meta).mark_compact_shape_dynamic( - mode=0) - return dl_kv, dl_q, dl_w, dl_logits, dl_bt, dl_cl, dl_sm - _TORCH_TO_CUTLASS_DTYPE = { torch.float16: cutlass.Float16, torch.bfloat16: cutlass.BFloat16, @@ -5174,18 +5150,55 @@ def _make_dlpacks(cls, kv_flat, q_3d, w_2d, logits, block_table, @classmethod def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, - kv_flat, q_3d, w_2d, logits, block_table, context_lens, - schedule_meta, num_phys_blocks, B, stream, num_epi_subtiles, epi_dtype, acc_dtype, output_dtype): - """Compile kernel using from_dlpack with dynamic shape markers.""" + """Compile kernel using fake tensors + TVM FFI.""" key = (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 = cls._TORCH_TO_CUTLASS_DTYPE - dl_args = cls._make_dlpacks(kv_flat, q_3d, w_2d, logits, - block_table, context_lens, - schedule_meta) + N = next_n * num_heads + block_bytes = 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=block_kv, num_heads=num_heads, @@ -5197,8 +5210,21 @@ def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, acc_dtype=to_cutlass[acc_dtype], output_dtype=to_cutlass[output_dtype], ) - compiled = cute.compile(kernel, *dl_args, num_phys_blocks, B, - stream) + + 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 @classmethod @@ -5263,26 +5289,21 @@ def forward( ) logits = logits[:, :max_context_len] - # Get current stream - torch_stream = torch.cuda.current_stream() - stream = cuda.CUstream(torch_stream.cuda_stream) - - # Compile if needed (uses real tensors for shape marking) + # Compile if needed (fake tensors, no real data required) key = (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(block_kv, H, D, next_n, num_sms, kv_flat, q_3d, - w_2d, logits, block_table, context_lens, - schedule_meta, num_phys_blocks, B, stream, - num_epi_subtiles, epi_dtype, acc_dtype, - output_dtype) + cls._compile(block_kv, H, D, next_n, num_sms, num_epi_subtiles, + epi_dtype, acc_dtype, output_dtype) compiled = cls.kernel_cache[key] - # Wrap tensors for runtime call - dl_args = cls._make_dlpacks(kv_flat, q_3d, w_2d, logits, - block_table, context_lens, - schedule_meta) - compiled(*dl_args, num_phys_blocks, B, stream) + # 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", From 4ef3d13596e8a99839a98ee22396d487d6e673f8 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Apr 2026 08:51:20 -0700 Subject: [PATCH 11/26] [None][fix] Remove commented-out debug code from fp8_paged_mqa_logits Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 0d14a77cbe5c..125c8a1bf9ab 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -1801,31 +1801,6 @@ def _make_dynamic_dlpacks( q_for_dl = ( q_3d.view(torch.uint8) if q_3d.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) else q_3d ) - # q_3d is [N, D, B] with stride (D, 1, N*D) after permute(1,2,0); - # use dim_order() to get the correct stride_order for the non-contiguous layout. - # print("limin: q_3d.shape", q_3d.shape) - # print("limin: q_3d.stride()", q_3d.stride()) - # print("limin: q_3d.dim_order()", q_3d.dim_order()) - # dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( - # mode=2, stride_order=q_3d.dim_order()) # [N, D, ?B] - # dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic( - # mode=1, stride_order=w_2d.dim_order()) # [N, ?B] - # dl_logits = from_dlpack(logits).mark_compact_shape_dynamic( - # mode=0, stride_order=logits.dim_order()).mark_compact_shape_dynamic( - # mode=1, stride_order=logits.dim_order()) # [?B, ?ctx] - # dl_bt = from_dlpack(block_table_gpu).mark_compact_shape_dynamic( - # mode=0, stride_order=block_table_gpu.dim_order()).mark_compact_shape_dynamic( - # mode=1, stride_order=block_table_gpu.dim_order()) # [?B, ?blks] - # dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( - # mode=2) # [N, D, ?B] - # dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic( - # mode=1) # [N, ?B] - # dl_logits = from_dlpack(logits).mark_compact_shape_dynamic( - # mode=0).mark_compact_shape_dynamic( - # mode=1) # [?B, ?ctx] - # dl_bt = from_dlpack(block_table_gpu).mark_compact_shape_dynamic( - # mode=0).mark_compact_shape_dynamic( - # mode=1) # [?B, ?blks] dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( mode=2, stride_order=(2, 0, 1) ) # [N, D, ?B] From 670cd1670801ed52820d9b35afa62b147429566a Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Apr 2026 09:08:47 -0700 Subject: [PATCH 12/26] [None][fix] Replace compile print with logger.debug in MQA logits Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 3 +++ .../blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py | 3 --- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 4caa5ce76372..a9ef4c18e272 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5226,6 +5226,9 @@ def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, 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( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 125c8a1bf9ab..51a6da1ee20d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -1953,9 +1953,6 @@ def _get_or_compile_kernel( ) compiled = cute.compile(kernel, *dl_args, num_phys_blocks, B, stream) _compiled_cache[cache_key] = compiled - print( - f" [compile] {cache_key} kv_stages={kernel.num_kv_stages} umma_stages={kernel.num_umma_stages}" - ) return _compiled_cache[cache_key] From 0a1e93d9242d880118205adca3a4e2e35e3d434c Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 21 Apr 2026 20:27:02 -0700 Subject: [PATCH 13/26] [None][fix] Add dtype validation to cute_dsl_fp8_paged_mqa_logits wrapper Add _check_fp8_paged_mqa_logits_dtypes helper that validates all tensor and compute dtypes upfront, collecting all errors into a single ValueError so callers see every mismatch at once rather than one at a time. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index a9ef4c18e272..6801f07757e2 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5133,6 +5133,33 @@ def warmup_cute_dsl_indexer_topk( # ------------------------------------------------------------------ # 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). @@ -5329,6 +5356,9 @@ def cute_dsl_fp8_paged_mqa_logits( 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, From fabc9026e1460873fb558e3e91985e38c52fa04c Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 22 Apr 2026 06:34:16 +0000 Subject: [PATCH 14/26] [None][test] Improve fp16 accuracy test and benchmark for MQA logits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add epi_dtype param to PyTorch ref so epilogue matches kernel precision - Add use_int_data mode for fp16 tests to isolate bugs from FP16 rounding - Tighten tolerances: fp32 atol 5e-3→5e-5, fp16 atol 1e-2→1e-3 with rtol - Use torch.testing.assert_close instead of cosine similarity - Add varlen benchmark support for mixed-length serving workloads - Clean up comments and remove stale parametrize line Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../test_cute_dsl_fp8_paged_mqa_logits.py | 198 ++++++++++++++---- 1 file changed, 162 insertions(+), 36 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 50516f5112d6..a93a5ddffc3d 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -53,7 +53,15 @@ def _calc_diff(x: torch.Tensor, y: torch.Tensor): def _ref_fp8_paged_mqa_logits( - q_fp8, kv_fp8, kv_scales, weights, context_lens, block_table, max_model_len, block_kv + q_fp8, + kv_fp8, + kv_scales, + weights, + context_lens, + block_table, + max_model_len, + block_kv, + epi_dtype=torch.float32, ): """Pure PyTorch reference for fp8_paged_mqa_logits. @@ -66,16 +74,16 @@ def _ref_fp8_paged_mqa_logits( block_table: [B, max_blocks] int32 max_model_len: int block_kv: int + epi_dtype: epilogue dtype — GEMM stays fp32, weighted sum + scale + use this dtype (torch.float32 or torch.float16) Returns: - logits: [B*next_n, max_model_len] float32 + logits: [B*next_n, max_model_len] epi_dtype """ B, next_n, H, D = q_fp8.shape device = q_fp8.device - logits = torch.full( - (B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32 - ) + logits = torch.full((B * next_n, max_model_len), float("-inf"), device=device, dtype=epi_dtype) q_f32 = q_fp8.float() @@ -83,29 +91,32 @@ def _ref_fp8_paged_mqa_logits( ctx_len = context_lens[b].item() q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device) - w = weights[b * next_n : (b + 1) * next_n, :] + w = weights[b * next_n : (b + 1) * next_n, :].to(epi_dtype) for blk_idx in range((ctx_len + block_kv - 1) // block_kv): phys_blk = block_table[b, blk_idx].item() k_f32 = kv_fp8[phys_blk].float() - scales = kv_scales[phys_blk] + scales = kv_scales[phys_blk].to(epi_dtype) k_positions = torch.arange(blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device) mask = (k_positions[None, :] < ctx_len) & (k_positions[None, :] <= q_positions[:, None]) + # GEMM in fp32 qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T) # [H, next_n, block_kv] qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device)) qk = torch.relu(qk) + # Epilogue in epi_dtype + qk = qk.to(epi_dtype) weighted = (w.T[:, :, None] * qk).sum(dim=0) # [next_n, block_kv] weighted = weighted * scales[None, :] start_pos = blk_idx * block_kv end_pos = start_pos + block_kv logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where( - mask, weighted, torch.tensor(float("-inf"), device=device) + mask, weighted, torch.tensor(float("-inf"), device=device, dtype=epi_dtype) ) return logits @@ -132,9 +143,23 @@ def _make_fused_kv(kv_fp8, kv_scales, block_kv, head_dim): def _generate_test_data( - batch_size, next_n, num_heads, head_dim, block_kv, avg_context_len, max_model_len, device="cuda" + batch_size, + next_n, + num_heads, + head_dim, + block_kv, + avg_context_len, + max_model_len, + device="cuda", + use_int_data=False, ): - """Generate random test data for fp8 paged MQA logits.""" + """Generate test data for fp8 paged MQA logits. + + Args: + use_int_data: When True, use small random integers ([-3, 3]) for Q/KV + and integer weights so that GEMM accumulation is exact across + FP8/FP16/FP32. Useful for isolating kernel bugs from precision. + """ context_lens = torch.randint( max(block_kv, int(0.7 * avg_context_len)), int(1.3 * avg_context_len) + 1, @@ -157,17 +182,43 @@ def _generate_test_data( ) blk_offset += n_blks - q_bf16 = torch.randn(batch_size, next_n, num_heads, head_dim, device=device) - q_fp8 = q_bf16.to(torch.float8_e4m3fn) + if use_int_data: + q_fp8 = torch.randint( + -3, + 4, + (batch_size, next_n, num_heads, head_dim), + device=device, + dtype=torch.float32, + ).to(torch.float8_e4m3fn) + + kv_fp8 = torch.randint( + -3, + 4, + (num_phys_blocks, block_kv, head_dim), + device=device, + dtype=torch.float32, + ).to(torch.float8_e4m3fn) + kv_scale = torch.ones(num_phys_blocks, block_kv, device=device, dtype=torch.float32) + + weights = torch.randint( + -3, + 4, + (batch_size * next_n, num_heads), + device=device, + dtype=torch.float32, + ) + else: + q_bf16 = torch.randn(batch_size, next_n, num_heads, head_dim, device=device) + q_fp8 = q_bf16.to(torch.float8_e4m3fn) - kv_bf16 = torch.randn(num_phys_blocks, block_kv, head_dim, device=device) - kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4) - kv_scale = _ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1) - kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + kv_bf16 = torch.randn(num_phys_blocks, block_kv, head_dim, device=device) + kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4) + kv_scale = _ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1) + kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) - kv_fused = _make_fused_kv(kv_fp8, kv_scale, block_kv, head_dim) + weights = torch.randn(batch_size * next_n, num_heads, device=device, dtype=torch.float32) - weights = torch.randn(batch_size * next_n, num_heads, device=device, dtype=torch.float32) + kv_fused = _make_fused_kv(kv_fp8, kv_scale, block_kv, head_dim) return { "q_fp8": q_fp8, @@ -210,7 +261,14 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype max_model_len = max(avg_ctx * 2, 2048) data = _generate_test_data( - batch_size, next_n, num_heads, head_dim, block_kv, avg_ctx, max_model_len + batch_size, + next_n, + num_heads, + head_dim, + block_kv, + avg_ctx, + max_model_len, + use_int_data=(output_dtype == torch.float16), ) from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -220,7 +278,8 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype # DSL kernel always uses full num_sms as grid size. dsl_schedule_meta = get_paged_mqa_logits_metadata(data["context_lens"], block_kv, num_sms) - # Reference: try C++ DeepGEMM first (fp32 only), fall back to PyTorch ref. + # Reference: C++ DeepGEMM is fp32-only and doesn't support next_n=3, + # so only used for fp32 + next_n ∈ {1,2,4}. All other cases use PyTorch ref. ref_logits = None if output_dtype == torch.float32: try: @@ -253,6 +312,7 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype data["block_table"], max_model_len, block_kv, + epi_dtype=output_dtype, ) # CuTe DSL kernel @@ -282,13 +342,34 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype dsl_masked = dsl_logits.float().masked_fill(~mask, 0) ref_masked = ref_logits.float().masked_fill(~mask, 0) finite = torch.isfinite(dsl_masked) & torch.isfinite(ref_masked) - diff = _calc_diff(dsl_masked.masked_fill(~finite, 0), ref_masked.masked_fill(~finite, 0)) + dsl_clean = dsl_masked.masked_fill(~finite, 0) + ref_clean = ref_masked.masked_fill(~finite, 0) + + # Element-wise check on the valid (finite + in-context) region. + # Kernel is deterministic (disjoint CTA writes, no atomics), so every + # element must be within elem_atol. + elem_atol = 1e-3 if output_dtype == torch.float16 else 5e-5 + elem_rtol = 1e-3 if output_dtype == torch.float16 else 1e-5 + + # Debug probe: print max/mean abs error for CI failure diagnosis. + valid = mask & finite + elem_abs = (dsl_clean - ref_clean).abs()[valid] + if elem_abs.numel() > 0: + print( + f"[acc-probe] B={batch_size} next_n={next_n} avg_ctx={avg_ctx} " + f"dtype={output_dtype} -> " + f"max_abs={elem_abs.max().item():.3e} " + f"mean_abs={elem_abs.mean().item():.3e}" + ) - tol = 5e-3 if output_dtype == torch.float16 else 1e-3 - assert diff < tol, ( - f"Accuracy check failed: diff={diff:.2e}, " - f"B={batch_size}, next_n={next_n}, avg_ctx={avg_ctx}, " - f"dtype={output_dtype}" + torch.testing.assert_close( + dsl_clean, + ref_clean, + atol=elem_atol, + rtol=elem_rtol, + msg=lambda m: ( + f"{m}\nB={batch_size}, next_n={next_n}, avg_ctx={avg_ctx}, dtype={output_dtype}" + ), ) @@ -317,18 +398,48 @@ def _profile_kernel_us(fn, num_warmup=10, num_iterations=30): def _generate_bench_data( - batch_size, context_len, next_n, num_heads=64, head_dim=128, block_kv=128, device="cuda" + batch_size, + context_len, + next_n, + num_heads=64, + head_dim=128, + block_kv=128, + varlen=False, + device="cuda", ): - """Generate benchmark data with uniform context length.""" + """Generate benchmark data. + + ``context_len`` is treated as the max length. When varlen=False, all + sequences use this exact length. When varlen=True, per-sequence lengths + are drawn uniformly from [min(2048, max), max] to mimic real mixed-batch + serving workloads. + """ torch.manual_seed(42) num_blocks_per_seq = (context_len + block_kv - 1) // block_kv - total_blocks = batch_size * num_blocks_per_seq - # fix-length workload: all sequences have the same context length. - context_lens = torch.full((batch_size,), context_len, dtype=torch.int32, device=device) - block_table = torch.arange(total_blocks, dtype=torch.int32, device=device).reshape( - batch_size, num_blocks_per_seq - ) + if varlen: + lo = min(2048, context_len) + context_lens = torch.randint( + lo, context_len + 1, (batch_size,), dtype=torch.int32, device=device + ) + total_blocks = ((context_lens + block_kv - 1) // block_kv).sum().item() + block_table = torch.zeros( + (batch_size, num_blocks_per_seq), dtype=torch.int32, device=device + ) + cursor = 0 + for i in range(batch_size): + n_blks = (context_lens[i].item() + block_kv - 1) // block_kv + block_table[i, :n_blks] = torch.arange( + cursor, cursor + n_blks, dtype=torch.int32, device=device + ) + cursor += n_blks + else: + total_blocks = batch_size * num_blocks_per_seq + # fix-length workload: all sequences have the same context length. + context_lens = torch.full((batch_size,), context_len, dtype=torch.int32, device=device) + block_table = torch.arange(total_blocks, dtype=torch.int32, device=device).reshape( + batch_size, num_blocks_per_seq + ) q_fp8 = torch.randn( batch_size, next_n, num_heads, head_dim, device=device, dtype=torch.bfloat16 @@ -363,6 +474,7 @@ def benchmark_fp8_paged_mqa_logits( num_iterations=30, output_dtype=torch.float32, num_epi_subtiles=1, + varlen=False, ): """Benchmark CuTe DSL vs C++ DeepGEMM kernel time.""" from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -373,7 +485,8 @@ def benchmark_fp8_paged_mqa_logits( num_sms = torch.cuda.get_device_properties(0).multi_processor_count dtype_str = str(output_dtype).split(".")[-1] - print(f"output_dtype={dtype_str} num_epi_subtiles={num_epi_subtiles}") + mode_str = "varlen" if varlen else "fix-len" + print(f"output_dtype={dtype_str} num_epi_subtiles={num_epi_subtiles} mode={mode_str}") is_non_default = output_dtype != torch.float32 or num_epi_subtiles != 1 hdr = ( f"{'batch':>5s} {'ctx':>7s} {'next_n':>6s} {'nblk':>7s} | " @@ -390,7 +503,13 @@ def benchmark_fp8_paged_mqa_logits( nblk = batch_size * ((context_len + block_kv - 1) // block_kv) data = _generate_bench_data( - batch_size, context_len, next_n, num_heads, head_dim, block_kv + batch_size, + context_len, + next_n, + num_heads, + head_dim, + block_kv, + varlen=varlen, ) dsl_schedule_meta = get_paged_mqa_logits_metadata( @@ -513,6 +632,12 @@ def dsl_f32_fn(data=data): choices=[1, 2, 4], help="epilogue sub-tile count (default: 1)", ) + parser.add_argument( + "--varlen", + action="store_true", + help="use varlen workload (per-seq lengths in [min(2048,max), max]); " + "default is fix-length where all sequences use --context_len", + ) args = parser.parse_args() dtype_map = {"float32": torch.float32, "float16": torch.float16} @@ -524,4 +649,5 @@ def dsl_f32_fn(data=data): num_iterations=args.repeat, output_dtype=dtype_map[args.output_dtype], num_epi_subtiles=args.num_epi_subtiles, + varlen=args.varlen, ) From aeb32d86091a965a3e90a3657747f31e597f7957 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 22 Apr 2026 01:49:04 -0700 Subject: [PATCH 15/26] [None][cleanup] Remove dead standalone test code and deduplicate dtype map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove broken run_test/parse_args/__main__ block and its helper functions (make_fused_kv, fused_kv_views, _prepare_inputs, _make_dynamic_dlpacks, _get_or_compile_kernel, dsl_fp8_paged_mqa_logits_dg_fullk) from fp8_paged_mqa_logits.py — all had no callers after the missing paged_mqa_logits_helpers module made them non-functional. Accuracy validation is covered by the unit tests. Reuse existing module-level _TORCH_TO_CUTLASS_DTYPE in CuteDSLPagedMQALogitsRunner instead of duplicating it. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 8 +- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 552 ------------------ 2 files changed, 1 insertion(+), 559 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 6801f07757e2..49422d2c5d4c 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5169,12 +5169,6 @@ class CuteDSLPagedMQALogitsRunner: kernel_cache = dict() - _TORCH_TO_CUTLASS_DTYPE = { - torch.float16: cutlass.Float16, - torch.bfloat16: cutlass.BFloat16, - torch.float32: cutlass.Float32, - } - @classmethod def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, num_epi_subtiles, epi_dtype, acc_dtype, output_dtype): @@ -5184,7 +5178,7 @@ def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, if key in cls.kernel_cache: return - to_cutlass = cls._TORCH_TO_CUTLASS_DTYPE + to_cutlass = _TORCH_TO_CUTLASS_DTYPE N = next_n * num_heads block_bytes = block_kv * (head_dim + 4) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 51a6da1ee20d..3518a2e74c96 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -67,7 +67,6 @@ from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, nvvm, vector from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.runtime import from_dlpack from cutlass.cutlass_dsl import dsl_user_op from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait # noqa: F401 @@ -1716,554 +1715,3 @@ def compute_schedule_metadata(context_lens, block_kv, num_ctas): schedule[i] = torch.tensor([seq_idx, local_split], dtype=torch.int32) return schedule - - -def make_fused_kv(kv_cache_fp8, kv_cache_scales, block_kv, head_dim): - """Create fused KV tensor from separate KV and scale tensors. - - Output shape: [num_phys_blocks, block_kv, 1, per_token_size] uint8 - where per_token_size = head_dim + 4. - - Per token: [KV (head_dim bytes)] [Scale (4 bytes)] - """ - num_phys_blocks = kv_cache_fp8.shape[0] - per_token_size = head_dim + 4 - block_bytes = block_kv * per_token_size - scale_offset = block_kv * head_dim - - fused = torch.zeros( - num_phys_blocks, - block_bytes, - dtype=torch.uint8, - device=kv_cache_fp8.device, - ) - for blk in range(num_phys_blocks): - fused[blk, :scale_offset] = kv_cache_fp8[blk].view(torch.uint8).reshape(-1) - fused[blk, scale_offset:] = kv_cache_scales[blk].view(torch.uint8).reshape(-1) - return fused.view(num_phys_blocks, block_kv, 1, per_token_size) - - -def fused_kv_views(kv_fused, block_kv, head_dim): - """Create KV (FP8) and Scale (float32) views from per-block fused buffer. - - Both views share the same underlying memory for L2 cache locality. - - Args: - kv_fused: [num_phys_blocks, block_kv * per_token_size] uint8 - Returns: - kv_pool: [block_kv, head_dim, num_phys_blocks] FP8 (strided view) - scales: [block_kv, num_phys_blocks] float32 (strided view) - """ - num_phys_blocks = kv_fused.shape[0] - per_token_size = head_dim + 4 - block_bytes = block_kv * per_token_size - scale_offset = block_kv * head_dim - - fused_flat = kv_fused.reshape(-1) - - # KV view: [block_kv, head_dim, num_phys_blocks] FP8 - # Within each block, KV data is contiguous [block_kv, head_dim] - # Element [m, k, l] → byte: l * block_bytes + m * head_dim + k - kv_pool = torch.as_strided( - fused_flat.view(torch.float8_e4m3fn), - size=(block_kv, head_dim, num_phys_blocks), - stride=(head_dim, 1, block_bytes), - ) - - # Scale view: [block_kv, num_phys_blocks] float32 - # Within each block, scales start at byte offset scale_offset - # and are contiguous [block_kv] float32. - # Element [m, l] → byte: l * block_bytes + scale_offset + m * 4 - # From float32 base at byte offset scale_offset: - # float32 index [m, l] → offset: l * (block_bytes / 4) + m - scale_base = fused_flat[scale_offset:].view(torch.float32) - scales = torch.as_strided( - scale_base, - size=(block_kv, num_phys_blocks), - stride=(1, block_bytes // 4), - ) - - return kv_pool, scales - - -def _make_dynamic_dlpacks( - kv_fused, q_3d, w_2d, logits, block_table_gpu, context_lens_gpu, schedule_meta_gpu -): - """Wrap tensors with dynamic shape markers for JIT reuse. - - Static dims (model constants): block_kv, head_dim, N, per_token, next_n - Dynamic dims (vary per call): batch_size, num_phys_blocks, max_model_len, - max_blocks_per_seq, num_ctas - """ - dl_kv = from_dlpack(kv_fused).mark_compact_shape_dynamic(mode=0) # [?phys, blk, 1, pt] - # DLPack does not support float8 types; view as uint8 (same 1-byte layout), - # then recast back to Float8E4M3FN inside the kernel. - q_for_dl = ( - q_3d.view(torch.uint8) if q_3d.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) else q_3d - ) - dl_q = from_dlpack(q_for_dl).mark_compact_shape_dynamic( - mode=2, stride_order=(2, 0, 1) - ) # [N, D, ?B] - dl_w = from_dlpack(w_2d).mark_compact_shape_dynamic(mode=1, stride_order=(1, 0)) # [N, ?B] - dl_logits = ( - from_dlpack(logits) - .mark_compact_shape_dynamic(mode=0, stride_order=(0, 1)) - .mark_compact_shape_dynamic(mode=1, stride_order=(0, 1)) - ) # [?B, ?ctx] - dl_bt = ( - from_dlpack(block_table_gpu) - .mark_compact_shape_dynamic(mode=0, stride_order=(0, 1)) - .mark_compact_shape_dynamic(mode=1, stride_order=(0, 1)) - ) # [?B, ?blks] - dl_cl = from_dlpack(context_lens_gpu).mark_compact_shape_dynamic(mode=0) # [?B] - dl_sm = from_dlpack(schedule_meta_gpu).mark_compact_shape_dynamic(mode=0) # [?ctas, 2] - return dl_kv, dl_q, dl_w, dl_logits, dl_bt, dl_cl, dl_sm - - -def _prepare_inputs( - q_fp8, - kv_fused, - weights, - context_lens, - block_table, - max_model_len, - block_kv, - num_sms=148, - epi_dtype=None, - output_dtype=None, -): - """Prepare all host/device tensors and compute schedule metadata.""" - B, next_n, H, D = q_fp8.shape - N = next_n * H - num_phys_blocks = kv_fused.shape[0] - - q_3d = q_fp8.reshape(B, N, D).permute(1, 2, 0) - if epi_dtype is not None and epi_dtype == cutlass.Float16: - w_2d = weights.reshape(B, N).half().t() - else: - w_2d = weights.reshape(B, N).t() - block_table_gpu = block_table.to(device=q_fp8.device, dtype=torch.int32) - context_lens_gpu = context_lens.to(device=q_fp8.device, dtype=torch.int32) - - _TORCH_DTYPE = { - cutlass.Float32: torch.float32, - cutlass.Float16: torch.float16, - cutlass.BFloat16: torch.bfloat16, - } - logits_torch_dtype = _TORCH_DTYPE.get(output_dtype, torch.float32) - - # Align logits columns to SPLIT_KV = block_kv * NUM_MATH_WG (like - # DeepGEMM's aligned_max_context_len). OOB unconditional stores from - # the kernel write into this padding region safely. - SPLIT_KV = block_kv * 2 # NUM_MATH_WG = 2 - aligned_max_ctx = ((max_model_len + SPLIT_KV - 1) // SPLIT_KV) * SPLIT_KV - logits = torch.full( - (B * next_n, aligned_max_ctx), - float("-inf"), - device="cuda", - dtype=logits_torch_dtype, - ) - logits = logits[:, :max_model_len] - - # Grid size = num_sms (like DeepGEMM). Empty SMs auto-skip via - # while-loop boundary check (start == end in schedule_meta). - num_ctas = num_sms - schedule_meta = compute_schedule_metadata(context_lens, block_kv, num_ctas) - schedule_meta_gpu = schedule_meta.to(device=q_fp8.device) - - return ( - kv_fused, - q_3d, - w_2d, - logits, - block_table_gpu, - context_lens_gpu, - schedule_meta_gpu, - num_phys_blocks, - B, - ) - - -# Kernel cache: keyed by static params (block_kv, num_heads, head_dim, next_n, num_sms, remove_kv_wait) -_compiled_cache = {} - - -def _get_or_compile_kernel( - block_kv, - num_heads, - head_dim, - next_n, - num_sms, - kv_fused, - q_3d, - w_2d, - logits, - block_table_gpu, - context_lens_gpu, - schedule_meta_gpu, - num_phys_blocks, - B, - stream, - remove_kv_wait_in_epilogue=False, - early_tmem_copy=False, - smem_subpartition_opt=False, - max_kv_pipeline=False, - max_umma_pipeline=False, - num_epi_subtiles=1, - epi_dtype=cutlass.Float32, - acc_dtype=cutlass.Float32, - output_dtype=cutlass.Float32, -): - """Return a compiled kernel, compiling only on first call per static config.""" - cache_key = ( - block_kv, - num_heads, - head_dim, - next_n, - num_sms, - remove_kv_wait_in_epilogue, - early_tmem_copy, - smem_subpartition_opt, - max_kv_pipeline, - max_umma_pipeline, - num_epi_subtiles, - epi_dtype, - acc_dtype, - output_dtype, - ) - if cache_key not in _compiled_cache: - kernel = FP8MQALogitsKernel( - block_kv=block_kv, - num_heads=num_heads, - head_dim=head_dim, - next_n=next_n, - num_sms=num_sms, - remove_kv_wait_in_epilogue=remove_kv_wait_in_epilogue, - early_tmem_copy=early_tmem_copy, - smem_subpartition_opt=smem_subpartition_opt, - max_kv_pipeline=max_kv_pipeline, - max_umma_pipeline=max_umma_pipeline, - num_epi_subtiles=num_epi_subtiles, - epi_dtype=epi_dtype, - acc_dtype=acc_dtype, - output_dtype=output_dtype, - ) - dl_args = _make_dynamic_dlpacks( - kv_fused, q_3d, w_2d, logits, block_table_gpu, context_lens_gpu, schedule_meta_gpu - ) - compiled = cute.compile(kernel, *dl_args, num_phys_blocks, B, stream) - _compiled_cache[cache_key] = compiled - return _compiled_cache[cache_key] - - -def dsl_fp8_paged_mqa_logits_dg_fullk( - q_fp8, - kv_fused, - weights, - context_lens, - block_table, - max_model_len, - block_kv, - num_sms=148, - remove_kv_wait_in_epilogue=False, - early_tmem_copy=False, - smem_subpartition_opt=False, - max_kv_pipeline=False, - max_umma_pipeline=False, - num_epi_subtiles=1, - epi_dtype=cutlass.Float32, - acc_dtype=cutlass.Float32, - output_dtype=cutlass.Float32, -): - """Run FP8 paged MQA logits kernel. - - Args: - kv_fused: [num_phys_blocks, block_kv, 1, head_dim + 4] uint8 - """ - B, next_n, H, D = q_fp8.shape - - (kv_f, q_3d, w_2d, logits, bt_gpu, cl_gpu, sm_gpu, num_phys_blocks, B) = _prepare_inputs( - q_fp8, - kv_fused, - weights, - context_lens, - block_table, - max_model_len, - block_kv, - num_sms, - epi_dtype=epi_dtype, - output_dtype=output_dtype, - ) - - torch_stream = torch.cuda.Stream() - stream = cuda.CUstream(torch_stream.cuda_stream) - - compiled = _get_or_compile_kernel( - block_kv, - H, - D, - next_n, - num_sms, - kv_f, - q_3d, - w_2d, - logits, - bt_gpu, - cl_gpu, - sm_gpu, - num_phys_blocks, - B, - stream, - remove_kv_wait_in_epilogue=remove_kv_wait_in_epilogue, - early_tmem_copy=early_tmem_copy, - smem_subpartition_opt=smem_subpartition_opt, - max_kv_pipeline=max_kv_pipeline, - max_umma_pipeline=max_umma_pipeline, - num_epi_subtiles=num_epi_subtiles, - epi_dtype=epi_dtype, - acc_dtype=acc_dtype, - output_dtype=output_dtype, - ) - - dl_args = _make_dynamic_dlpacks(kv_f, q_3d, w_2d, logits, bt_gpu, cl_gpu, sm_gpu) - compiled( - *dl_args, - num_phys_blocks, - B, - stream, - ) - torch.cuda.synchronize() - - return logits - - -def run_test( - batch_size_list=None, - next_n_list=None, - avg_ctx_list=None, - num_sms=148, - remove_kv_wait_in_epilogue=False, - early_tmem_copy=False, - smem_subpartition_opt=False, - max_kv_pipeline=False, - max_umma_pipeline=False, - num_epi_subtiles=1, - epi_dtype=cutlass.Float32, - acc_dtype=cutlass.Float32, - output_dtype=cutlass.Float32, -): - """Test FP8 paged MQA logits kernel against reference.""" - import sys - import time - - sys.path.insert(0, ".") - from paged_mqa_logits_helpers import calc_diff, generate_test_data, ref_fp8_paged_mqa_logits - - if batch_size_list is None: - batch_size_list = [32] - if next_n_list is None: - next_n_list = [1] - if avg_ctx_list is None: - avg_ctx_list = [32768] - - opt_str = "" - if remove_kv_wait_in_epilogue: - opt_str += " +remove_kv_wait" - if early_tmem_copy: - opt_str += " +early_tmem_copy" - if smem_subpartition_opt: - opt_str += " +smem_subpart" - if max_kv_pipeline: - opt_str += " +max_kv_pipeline" - if max_umma_pipeline: - opt_str += " +max_umma_pipeline" - print(f"=== FP8 Paged MQA Logits (Fused KV + Dynamic Shapes{opt_str}) Tests ===") - t0 = time.time() - n_passed = 0 - n_total = 0 - for test_batch in batch_size_list: - for test_next_n in next_n_list: - for avg_ctx in avg_ctx_list: - data = generate_test_data( - batch_size=test_batch, - next_n=test_next_n, - num_heads=64, - head_dim=128, - block_kv=128, - avg_context_len=avg_ctx, - max_model_len=max(avg_ctx * 2, 2048), - device="cuda", - ) - kv_fused = make_fused_kv( - data["kv_cache"], - data["kv_cache_scales"], - data["block_kv"], - 128, - ) - fk_logits = dsl_fp8_paged_mqa_logits_dg_fullk( - data["q"], - kv_fused, - data["weights"], - data["context_lens"], - data["block_table"], - data["max_model_len"], - data["block_kv"], - num_sms=num_sms, - remove_kv_wait_in_epilogue=remove_kv_wait_in_epilogue, - early_tmem_copy=early_tmem_copy, - smem_subpartition_opt=smem_subpartition_opt, - max_kv_pipeline=max_kv_pipeline, - max_umma_pipeline=max_umma_pipeline, - num_epi_subtiles=num_epi_subtiles, - epi_dtype=epi_dtype, - acc_dtype=acc_dtype, - output_dtype=output_dtype, - ) - ref_logits = ref_fp8_paged_mqa_logits( - data["q"], - data["kv_cache"], - data["kv_cache_scales"], - data["weights"], - data["context_lens"], - data["block_table"], - data["max_model_len"], - data["block_kv"], - ) - - B_test = data["batch_size"] - mask = torch.zeros_like(ref_logits, dtype=torch.bool) - for b in range(B_test): - ctx = data["context_lens"][b].item() - for t in range(test_next_n): - row = b * test_next_n + t - q_pos = ctx - test_next_n + t - mask[row, : q_pos + 1] = True - - diff = calc_diff( - fk_logits.float().masked_fill(~mask, 0), - ref_logits.masked_fill(~mask, 0), - ) - total_blks = sum(cdiv(data["context_lens"][b].item(), 128) for b in range(B_test)) - n_total += 1 - passed = diff < 1e-3 - if passed: - n_passed += 1 - status = "PASSED" if passed else "FAILED" - print( - f" B={test_batch}, next_n={test_next_n}, " - f"avg_ctx={avg_ctx}, nblk={total_blks}, num_sms={num_sms}: " - f"diff={diff:.2e} {status}" - ) - elapsed = time.time() - t0 - print(f"\n{n_passed}/{n_total} passed in {elapsed:.1f}s ({len(_compiled_cache)} compilations)") - - -def parse_args(): - import argparse - - parser = argparse.ArgumentParser(description="FP8 paged MQA logits kernel test") - parser.add_argument( - "--batch_size", - type=int, - nargs="+", - default=None, - help="batch size(s), e.g. --batch_size 1 32", - ) - parser.add_argument( - "--next_n", type=int, nargs="+", default=None, help="next_n value(s), e.g. --next_n 1 2 4" - ) - parser.add_argument( - "--avg_ctx", - type=int, - nargs="+", - default=None, - help="avg context len(s), e.g. --avg_ctx 256 4096", - ) - parser.add_argument( - "--num_sms", type=int, default=148, help="number of SMs for scheduling (default: 148)" - ) - parser.add_argument( - "--sweep", action="store_true", help="run full sweep over predefined ranges" - ) - parser.add_argument( - "--remove_kv_wait", - action="store_true", - help="remove KV barrier wait in epilogue (epilogue-bound opt)", - ) - parser.add_argument( - "--early_tmem_copy", - action="store_true", - help="issue LDTM early to hide latency behind scale LDS + KV release", - ) - parser.add_argument( - "--smem_subpartition_opt", - action="store_true", - help="pad SMEM to put sW/sScales in sub-partition 1, avoid UMMA conflicts", - ) - parser.add_argument( - "--max_kv_pipeline", - action="store_true", - help="maximize KV pipeline stages to fill SMEM budget", - ) - parser.add_argument( - "--max_umma_pipeline", - action="store_true", - help="maximize UMMA pipeline stages to fill TMEM budget", - ) - parser.add_argument( - "--num_epi_subtiles", type=int, default=1, help="number of epilogue sub-tiles (default: 1)" - ) - parser.add_argument( - "--epi_dtype", - type=str, - default="fp32", - choices=["fp32", "fp16"], - help="epilogue dtype (default: fp32)", - ) - parser.add_argument( - "--acc_dtype", - type=str, - default="fp32", - choices=["fp32", "fp16"], - help="accumulator dtype (default: fp32)", - ) - parser.add_argument( - "--output_dtype", - type=str, - default="fp32", - choices=["fp32", "fp16", "bf16"], - help="output logits dtype (default: fp32)", - ) - return parser.parse_args() - - -if __name__ == "__main__": - _DTYPE_MAP = { - "fp32": cutlass.Float32, - "fp16": cutlass.Float16, - "bf16": cutlass.BFloat16, - } - args = parse_args() - if args.sweep: - batch_size_override = [1, 32] - next_n_override = [1, 2, 3, 4] - avg_ctx_override = [256, 1024, 4096, 8192, 16384, 32768] - else: - batch_size_override = args.batch_size - next_n_override = args.next_n - avg_ctx_override = args.avg_ctx - run_test( - batch_size_list=batch_size_override, - next_n_list=next_n_override, - avg_ctx_list=avg_ctx_override, - num_sms=args.num_sms, - remove_kv_wait_in_epilogue=args.remove_kv_wait, - early_tmem_copy=args.early_tmem_copy, - smem_subpartition_opt=args.smem_subpartition_opt, - max_kv_pipeline=args.max_kv_pipeline, - max_umma_pipeline=args.max_umma_pipeline, - num_epi_subtiles=args.num_epi_subtiles, - epi_dtype=_DTYPE_MAP[args.epi_dtype], - acc_dtype=_DTYPE_MAP[args.acc_dtype], - output_dtype=_DTYPE_MAP[args.output_dtype], - ) From d0fa1cb7d3ac5e3f5aed3f8f61ebb28bde1011e9 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 22 Apr 2026 02:14:23 -0700 Subject: [PATCH 16/26] [None][fix] Validate num_heads divisible by 4 regardless of num_epi_subtiles Move the FMA unroll granularity check outside the num_epi_subtiles > 1 guard so it also applies when num_epi_subtiles == 1. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 3518a2e74c96..c20275234268 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -212,13 +212,12 @@ def __init__( self.epi_dtype = epi_dtype self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4 self.output_dtype = output_dtype - if num_epi_subtiles > 1: - if num_heads % num_epi_subtiles != 0: - raise ValueError("num_heads must be divisible by num_epi_subtiles") - if (num_heads // num_epi_subtiles) % 4 != 0: - raise ValueError( - "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)" - ) + if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0: + raise ValueError("num_heads must be divisible by num_epi_subtiles") + if (num_heads // num_epi_subtiles) % 4 != 0: + raise ValueError( + "num_heads // num_epi_subtiles must be divisible by 4 (FMA unroll granularity)" + ) self.num_groups = 2 self.num_math_threads = 256 From 38ec767afbe17ccb7debe6894128f2eb9681cb9e Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 22 Apr 2026 02:44:26 -0700 Subject: [PATCH 17/26] [None][cleanup] Remove unused variables and unnecessary noqa comments Remove is_umma_warp, num_q_stages (local copy), cons_group, num_tma_prod, num_mcast_a/b, and tCtAcc_fake (in __call__) which were assigned but never referenced. Remove unnecessary noqa F401 on pipeline_init_arrive/pipeline_init_wait which are actively used. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index c20275234268..0d010c2eef69 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -68,7 +68,7 @@ from cutlass._mlir.dialects import llvm, nvvm, vector from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cutlass_dsl import dsl_user_op -from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait # noqa: F401 +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait @dsl_user_op @@ -526,7 +526,6 @@ def kernel( is_tma_warp = is_tma_warp_0 | is_tma_warp_1 is_umma_warp_0 = warp_idx == 10 is_umma_warp_1 = warp_idx == 11 - is_umma_warp = is_umma_warp_0 | is_umma_warp_1 # noqa: F841 # Early schedule metadata load: issue global loads ASAP so their # ~200-cycle L2 latency overlaps with subsequent prologue setup @@ -554,14 +553,9 @@ def kernel( num_heads = self.num_heads next_n = self.next_n num_epi_subtiles = self.num_epi_subtiles - num_q_stages = self.num_q_stages # noqa: F841 # === Pipelines === prod_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_mcast_a = cute.size(cluster_layout_vmnk.shape[2]) - num_mcast_b = cute.size(cluster_layout_vmnk.shape[1]) - num_tma_prod = num_mcast_a + num_mcast_b - 1 - cons_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_prod) # noqa: F841 # Q pipeline: TMA producer → Math consumer (8 math warps) # PipelineTmaAsync: consumer_release uses is_signalling_thread @@ -815,7 +809,6 @@ def kernel( tCrA_1 = tiled_mma.make_fragment_A(sKV_1) tCrB = tiled_mma.make_fragment_B(sQ) # shared acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) # noqa: F841 # Staged acc (fp16_gemm_3 pattern): append UMMA stage dim # shape: (*acc_shape, STAGE) — dynamic index on last dim reduces rank From 719b333196fdc38ad1624a6b4ea338c4fa6db57f Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Wed, 22 Apr 2026 14:26:27 +0000 Subject: [PATCH 18/26] [None][fix] Add missing use_cute_dsl_paged_mqa_logits to test mock config Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- tests/unittest/_torch/attention/sparse/test_dsa_indexer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py index 514c836be2ff..47bb9b8cf430 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py @@ -74,6 +74,7 @@ def __init__(self, index_head_dim, index_n_heads, index_topk): self.index_topk = index_topk self.prompt_budget = 1024 self.use_cute_dsl_topk = False + self.use_cute_dsl_paged_mqa_logits = False self.enable_heuristic_topk = False sparse_attn_config = SparseAttentionConfig( From 1a157f4e9ab67ca6997d726c523d5b579c03d46e Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Thu, 23 Apr 2026 03:37:28 +0000 Subject: [PATCH 19/26] [None][fix] Fix TMA SMEM alignment for fp16 epilogue and support DSL arbitrary next_n Fix cudaErrorMisalignedAddress when num_heads=32 with fp16 output by padding sW stage stride to 128-byte boundary for TMA bulk copy. Update qw_per_stage SMEM budget to use padded stride and actual epi_bytes. In dsa.py, skip MTP batch expansion for DSL kernel (supports arbitrary next_n natively) and use correct scheduler_metadata buffer selection. Expand test coverage with num_heads and fix_length parametrization, pure PyTorch reference, and DG expanded benchmark for next_n=3. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 15 +- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 14 +- .../test_cute_dsl_fp8_paged_mqa_logits.py | 231 ++++++++++++------ .../attention/sparse/test_dsa_indexer.py | 39 ++- 4 files changed, 211 insertions(+), 88 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 99f206ac8f33..140d05d38375 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -895,13 +895,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) @@ -1650,7 +1652,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 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 0d010c2eef69..fb4373a9dfb4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -211,6 +211,11 @@ def __init__( self.num_epi_subtiles = num_epi_subtiles self.epi_dtype = epi_dtype self.epi_bytes = 2 if epi_dtype == cutlass.Float16 else 4 + # sW stage stride padded to 128-byte SMEM alignment for TMA bulk copy. + # Without padding, e.g. fp16 + N=32 gives 64B per stage, so stage 1 + # at +64 would be misaligned (TMA requires 128-byte aligned SMEM dest). + w_stage_bytes = self.N * self.epi_bytes + self.w_stage_stride = ((w_stage_bytes + 127) // 128 * 128) // self.epi_bytes self.output_dtype = output_dtype if num_epi_subtiles > 1 and num_heads % num_epi_subtiles != 0: raise ValueError("num_heads must be divisible by num_epi_subtiles") @@ -244,8 +249,8 @@ def __init__( # KV+Scale per stage (×2 groups): # 2 * (block_kv * head_dim * 1B + block_kv * 4B) kv_scale_per_stage = 2 * (block_kv * head_dim + block_kv * 4) - # Q+W per stage: N * head_dim * 1B + N * 4B - qw_per_stage = self.N * head_dim + self.N * 4 + # Q+W per stage: Q is N * head_dim * 1B, W uses padded stride + qw_per_stage = self.N * head_dim + self.w_stage_stride * self.epi_bytes qw_total = qw_per_stage * self.num_q_stages self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage else: @@ -413,7 +418,10 @@ def __call__( # TMA for Weights — [N, batch_size], tile [N], L=batch_size tma_load_op = cpasync.CopyBulkTensorTileG2SOp() - self.w_smem_layout_staged = cute.make_layout((self.N, self.num_q_stages)) + self.w_smem_layout_staged = cute.make_layout( + (self.N, self.num_q_stages), + stride=(1, self.w_stage_stride), + ) w_smem_per_stage = cute.select(self.w_smem_layout_staged, mode=[0]) tma_atom_w, tma_tensor_w = cpasync.make_tiled_tma_atom( tma_load_op, diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index a93a5ddffc3d..212810fa0361 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -16,8 +16,6 @@ Test CuTe DSL fp8_paged_mqa_logits kernel against C++ DeepGEMM reference. """ -import random - import pytest import torch @@ -43,15 +41,6 @@ def _ceil_to_ue8m0(x: torch.Tensor): return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) -def _calc_diff(x: torch.Tensor, y: torch.Tensor): - x, y = x.double(), y.double() - denominator = (x * x + y * y).sum() - if denominator == 0: - return 0.0 - sim = 2 * (x * y).sum() / denominator - return (1 - sim).item() - - def _ref_fp8_paged_mqa_logits( q_fp8, kv_fp8, @@ -152,6 +141,7 @@ def _generate_test_data( max_model_len, device="cuda", use_int_data=False, + fix_length=True, ): """Generate test data for fp8 paged MQA logits. @@ -160,14 +150,17 @@ def _generate_test_data( and integer weights so that GEMM accumulation is exact across FP8/FP16/FP32. Useful for isolating kernel bugs from precision. """ - context_lens = torch.randint( - max(block_kv, int(0.7 * avg_context_len)), - int(1.3 * avg_context_len) + 1, - (batch_size,), - dtype=torch.int32, - device="cpu", - ) - context_lens = context_lens.clamp(max=max_model_len) + if fix_length: + context_lens = torch.full((batch_size,), max_model_len, dtype=torch.int32, device="cpu") + else: + context_lens = torch.randint( + max(block_kv, int(0.7 * avg_context_len)), + int(1.3 * avg_context_len) + 1, + (batch_size,), + dtype=torch.int32, + device="cpu", + ) + context_lens = context_lens.clamp(max=max_model_len) max_blocks_per_seq = (max_model_len + block_kv - 1) // block_kv total_blocks = ((context_lens + block_kv - 1) // block_kv).sum().item() @@ -239,13 +232,16 @@ def _generate_test_data( ) -@skip_if_unsupported @skip_not_sm100 @pytest.mark.parametrize("batch_size", [1, 4, 32]) @pytest.mark.parametrize("next_n", [1, 2, 3, 4]) +@pytest.mark.parametrize("num_heads", [64]) @pytest.mark.parametrize("avg_ctx", [256, 4096, 32768]) @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float16]) -def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype): +@pytest.mark.parametrize("fix_length", [True, False]) +def test_cute_dsl_fp8_paged_mqa_logits( + batch_size, next_n, num_heads, avg_ctx, output_dtype, fix_length +): """Compare CuTe DSL kernel output against reference. Uses C++ DeepGEMM as reference when available (next_n in {1,2,4}), @@ -253,9 +249,8 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype Tests both fp32 and fp16 epi/acc/output paths. """ torch.manual_seed(42) - random.seed(42) + torch.cuda.manual_seed(42) - num_heads = 64 head_dim = 128 block_kv = 128 max_model_len = max(avg_ctx * 2, 2048) @@ -269,6 +264,7 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype avg_ctx, max_model_len, use_int_data=(output_dtype == torch.float16), + fix_length=fix_length, ) from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -278,42 +274,17 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype # DSL kernel always uses full num_sms as grid size. dsl_schedule_meta = get_paged_mqa_logits_metadata(data["context_lens"], block_kv, num_sms) - # Reference: C++ DeepGEMM is fp32-only and doesn't support next_n=3, - # so only used for fp32 + next_n ∈ {1,2,4}. All other cases use PyTorch ref. - ref_logits = None - if output_dtype == torch.float32: - try: - from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits - - num_kv_multicast = 2 if next_n == 4 else 1 - num_clusters = num_sms // num_kv_multicast - dg_schedule_meta = get_paged_mqa_logits_metadata( - data["context_lens"], block_kv, num_clusters - ) - ref_logits = fp8_paged_mqa_logits( - data["q_fp8"], - data["kv_fused"], - data["weights"], - data["context_lens"], - data["block_table"], - dg_schedule_meta, - max_model_len, - ) - except RuntimeError: - pass - - if ref_logits is None: - ref_logits = _ref_fp8_paged_mqa_logits( - data["q_fp8"], - data["kv_fp8"], - data["kv_scales"], - data["weights"], - data["context_lens"], - data["block_table"], - max_model_len, - block_kv, - epi_dtype=output_dtype, - ) + ref_logits = _ref_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fp8"], + data["kv_scales"], + data["weights"], + data["context_lens"], + data["block_table"], + max_model_len, + block_kv, + epi_dtype=output_dtype, + ) # CuTe DSL kernel dsl_logits = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( @@ -373,6 +344,112 @@ def test_cute_dsl_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype ) +# @skip_if_unsupported +# @skip_not_sm100 +# @pytest.mark.parametrize("batch_size", [1, 4, 32]) +# @pytest.mark.parametrize("next_n", [1, 2, 4]) +# @pytest.mark.parametrize("avg_ctx", [256, 4096, 32768]) +# @pytest.mark.parametrize("output_dtype", [torch.float32]) +# @pytest.mark.parametrize("fix_length", [True, False]) +# def test_deepgemm_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype, fix_length): +# """Compare DeepGEMM kernel output against reference. +# """ +# from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits + +# torch.manual_seed(42) +# torch.cuda.manual_seed(42) + +# num_heads = 64 +# head_dim = 128 +# block_kv = 128 +# max_model_len = max(avg_ctx * 2, 2048) + +# data = _generate_test_data( +# batch_size, +# next_n, +# num_heads, +# head_dim, +# block_kv, +# avg_ctx, +# max_model_len, +# use_int_data=(output_dtype == torch.float16), +# fix_length=fix_length, +# ) + +# from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + +# num_sms = torch.cuda.get_device_properties(0).multi_processor_count + +# # DeepGEMM kernel always uses num_clusters as grid size. +# num_kv_multicast = 2 if next_n == 4 else 1 +# num_clusters = num_sms // num_kv_multicast +# dg_schedule_meta = get_paged_mqa_logits_metadata( +# data["context_lens"], block_kv, num_clusters +# ) + +# dsl_logits = fp8_paged_mqa_logits( +# data["q_fp8"], +# data["kv_fused"], +# data["weights"], +# data["context_lens"], +# data["block_table"], +# dg_schedule_meta, +# max_model_len, +# ) +# ref_logits = _ref_fp8_paged_mqa_logits( +# data["q_fp8"], +# data["kv_fp8"], +# data["kv_scales"], +# data["weights"], +# data["context_lens"], +# data["block_table"], +# max_model_len, +# block_kv, +# epi_dtype=output_dtype, +# ) + +# # Mask invalid positions +# B = batch_size +# positions = torch.arange(max_model_len, device="cuda").unsqueeze(0) +# row_indices = torch.arange(B * next_n, device="cuda") // next_n +# next_n_offset = torch.arange(B * next_n, device="cuda") % next_n +# end_pos = data["context_lens"][row_indices] - next_n + next_n_offset +# mask = positions <= end_pos.unsqueeze(1) + +# dsl_masked = dsl_logits.float().masked_fill(~mask, 0) +# ref_masked = ref_logits.float().masked_fill(~mask, 0) +# finite = torch.isfinite(dsl_masked) & torch.isfinite(ref_masked) +# dsl_clean = dsl_masked.masked_fill(~finite, 0) +# ref_clean = ref_masked.masked_fill(~finite, 0) + +# # Element-wise check on the valid (finite + in-context) region. +# # Kernel is deterministic (disjoint CTA writes, no atomics), so every +# # element must be within elem_atol. +# elem_atol = 1e-3 if output_dtype == torch.float16 else 5e-5 +# elem_rtol = 1e-3 if output_dtype == torch.float16 else 1e-5 + +# # Debug probe: print max/mean abs error for CI failure diagnosis. +# valid = mask & finite +# elem_abs = (dsl_clean - ref_clean).abs()[valid] +# if elem_abs.numel() > 0: +# print( +# f"[acc-probe] B={batch_size} next_n={next_n} avg_ctx={avg_ctx} " +# f"dtype={output_dtype} -> " +# f"max_abs={elem_abs.max().item():.3e} " +# f"mean_abs={elem_abs.mean().item():.3e}" +# ) + +# torch.testing.assert_close( +# dsl_clean, +# ref_clean, +# atol=elem_atol, +# rtol=elem_rtol, +# msg=lambda m: ( +# f"{m}\nB={batch_size}, next_n={next_n}, avg_ctx={avg_ctx}, dtype={output_dtype}" +# ), +# ) + + def _profile_kernel_us(fn, num_warmup=10, num_iterations=30): """Profile CUDA kernel time in microseconds using torch.profiler.""" from torch.profiler import ProfilerActivity, profile @@ -415,6 +492,7 @@ def _generate_bench_data( serving workloads. """ torch.manual_seed(42) + torch.cuda.manual_seed(42) num_blocks_per_seq = (context_len + block_kv - 1) // block_kv if varlen: @@ -537,21 +615,38 @@ def dsl_fn(data=data): try: from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits - num_kv_multicast = 2 if next_n == 4 else 1 + # DeepGEMM doesn't support next_n=3 natively; expand to + # batch=B*next_n, next_n=1 (same approach as production MTP + # expand path in dsa.py). + dg_next_n = next_n + dg_data = data + if next_n == 3: + exp_bs = batch_size * next_n + dg_data = { + "q_fp8": data["q_fp8"].reshape(exp_bs, 1, num_heads, head_dim), + "kv_fused": data["kv_fused"], + "weights": data["weights"].reshape(exp_bs, num_heads), + "context_lens": data["context_lens"].repeat_interleave(next_n), + "block_table": data["block_table"].repeat_interleave(next_n, dim=0), + "max_model_len": data["max_model_len"], + } + dg_next_n = 1 + + num_kv_multicast = 2 if dg_next_n == 4 else 1 num_clusters = num_sms // num_kv_multicast dg_schedule_meta = get_paged_mqa_logits_metadata( - data["context_lens"], block_kv, num_clusters + dg_data["context_lens"], block_kv, num_clusters ) - def dg_fn(data=data): + def dg_fn(dg_data=dg_data): fp8_paged_mqa_logits( - data["q_fp8"], - data["kv_fused"], - data["weights"], - data["context_lens"], - data["block_table"], + dg_data["q_fp8"], + dg_data["kv_fused"], + dg_data["weights"], + dg_data["context_lens"], + dg_data["block_table"], dg_schedule_meta, - data["max_model_len"], + dg_data["max_model_len"], ) dg_us = _profile_kernel_us(dg_fn, num_warmup, num_iterations) diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py index 47bb9b8cf430..3b5cb37ee856 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py @@ -400,7 +400,8 @@ def _create_mock_metadata(request_ids, max_draft_tokens=0, enable_context_mla_with_cached_kv=False, index_topk=2048, - enable_indexer_skip=False): + enable_indexer_skip=False, + use_cute_dsl_paged_mqa_logits=False): """Helper to create mock metadata for testing.""" class MockKVCacheParams: @@ -537,10 +538,12 @@ def __init__(self): self.runtime_features = RuntimeFeatures() # Add expanded buffers for MTP support + # DSL kernel supports arbitrary next_n natively, so it never needs expansion. 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)) + not use_cute_dsl_paged_mqa_logits + 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))) self.kv_lens_expanded_cuda = torch.zeros( (self.num_seqs * (1 + self.max_draft_tokens), ), device='cuda', @@ -918,7 +921,8 @@ def test_fp8_k_cache_roundtrip(): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper @pytest.mark.parametrize("batch_size,next_n", [(4, 1), (2, 2), (4, 3), (4, 4)]) -def test_indexer_decode_with_paged_kv_cache(batch_size, next_n): +@pytest.mark.parametrize("backend", ["deepgemm", "dsl"]) +def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): """ Test FP8 paged KV cache with two-phase workflow and variable context lengths. @@ -933,9 +937,11 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n): torch.manual_seed(123) random.seed(123) - # Test parameters - heads, head_dim = 32, 128 - block_size = 64 + use_dsl = backend == "dsl" + # Note, DSL kernel requires tokens_per_block=128. + # Will remove this restriction in the future. + heads, head_dim = (32, 128) + block_size = 128 if use_dsl else 64 avg_context_len = 2048 num_gen_tokens = next_n # Number of tokens to generate per sequence max_model_len = 4096 @@ -1005,6 +1011,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n): num_ctx_tokens=total_context_tokens, num_tokens=total_context_tokens, max_draft_tokens=next_n - 1, + use_cute_dsl_paged_mqa_logits=use_dsl, ) Indexer.prepare(metadata_context) @@ -1031,6 +1038,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n): num_ctx_tokens=0, num_tokens=batch_size * num_gen_tokens, max_draft_tokens=next_n - 1, + use_cute_dsl_paged_mqa_logits=use_dsl, ) Indexer.prepare(metadata_gen) @@ -1049,7 +1057,9 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n): q_fp8 = q_fp8 context_lens = metadata_gen.kv_lens_cuda_runtime[0:batch_size] block_table = metadata_gen.indexer_k_cache_block_offsets[0:batch_size] - if q_fp8.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_fp8.shape[1] == 4 and not use_dsl: scheduler_metadata_buffer = metadata_gen.scheduler_metadata_buffer_mtp3 else: scheduler_metadata_buffer = metadata_gen.scheduler_metadata_buffer @@ -1060,9 +1070,14 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n): block_table = metadata_gen.block_table_expanded[:num_tokens] scheduler_metadata_buffer = metadata_gen.scheduler_metadata_buffer_expanded - logits = fp8_paged_mqa_logits(q_fp8, kv_cache_fp8_pool, weights, - context_lens, block_table, - scheduler_metadata_buffer, max_model_len) + if use_dsl: + logits = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + q_fp8, kv_cache_fp8_pool, weights, context_lens, block_table, + scheduler_metadata_buffer, max_model_len) + else: + logits = fp8_paged_mqa_logits(q_fp8, kv_cache_fp8_pool, weights, + context_lens, block_table, + scheduler_metadata_buffer, max_model_len) print(f"✓ Kernel output shape: {logits.shape}") # Reference: Reconstruct BF16 cache from original values From 66810313ac36d2de18b96f4085afcee4a16c1509 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Thu, 23 Apr 2026 07:16:35 +0000 Subject: [PATCH 20/26] [None][fix] Fix OOB read in zero-work CTA and move test seed into helper Clamp start_q before reading mContextLens to prevent out-of-bounds access when the scheduler assigns batch_size as a sentinel for CTAs with no work. Move torch.manual_seed into _generate_test_data for reproducibility. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py | 6 +++++- .../attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py | 5 ++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index fb4373a9dfb4..ec5f2f7b51bb 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -546,7 +546,11 @@ def kernel( end_kv_half = mScheduleMeta[(sm_idx + 1, 1)] # Early mContextLens load: overlap ~200-cycle L2 latency with the # entire prologue setup (pipelines, SMEM alloc, TMA partition, etc.) - current_num_kv = (mContextLens[start_q] + self.block_kv - 1) // self.block_kv + # Clamp to avoid OOB when start_q == batch_size (zero-work CTA sentinel). + # Note: zero-work CTAs get a stale current_num_kv (from the last batch + # element), but it is never used because has_work will be False. + start_q_clamped = min(start_q, batch_size - 1) + current_num_kv = (mContextLens[start_q_clamped] + self.block_kv - 1) // self.block_kv if is_tma_warp: cpasync.prefetch_descriptor(tma_atom_a) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 212810fa0361..06d3cac6b5a6 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -150,6 +150,8 @@ def _generate_test_data( and integer weights so that GEMM accumulation is exact across FP8/FP16/FP32. Useful for isolating kernel bugs from precision. """ + torch.manual_seed(42) + torch.cuda.manual_seed(42) if fix_length: context_lens = torch.full((batch_size,), max_model_len, dtype=torch.int32, device="cpu") else: @@ -248,9 +250,6 @@ def test_cute_dsl_fp8_paged_mqa_logits( falls back to pure PyTorch reference otherwise (e.g. next_n=3). Tests both fp32 and fp16 epi/acc/output paths. """ - torch.manual_seed(42) - torch.cuda.manual_seed(42) - head_dim = 128 block_kv = 128 max_model_len = max(avg_ctx * 2, 2048) From 1c688318b07ee501e8a2a69e308d3f293228afda Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Thu, 23 Apr 2026 13:55:02 +0000 Subject: [PATCH 21/26] [None][feat] Support multi-block TMA for phys_block_kv < 128 in DSL paged MQA logits Decouple compute_block_kv (always 128) from phys_block_kv (physical page size). When phys_block_kv < 128, the kernel issues num_blocks_per_mma TMA copies per compute tile. Removes the tokens_per_block=128 constraint for the DSL kernel path. Signed-off-by: Mindy Li Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 33 ++- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 257 ++++++++++-------- .../test_cute_dsl_fp8_paged_mqa_logits.py | 139 +++++++++- .../attention/sparse/test_dsa_indexer.py | 4 +- 4 files changed, 301 insertions(+), 132 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 49422d2c5d4c..9e0e7ff9a137 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5164,23 +5164,25 @@ class CuteDSLPagedMQALogitsRunner: """Runner for CuTe DSL FP8 Paged MQA Logits kernel (Blackwell SM100). Caches compiled kernels keyed by static params - (block_kv, num_heads, head_dim, next_n, num_sms). + (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n, num_sms). """ kernel_cache = dict() @classmethod - def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, - num_epi_subtiles, epi_dtype, acc_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, acc_dtype, + output_dtype): """Compile kernel using fake tensors + TVM FFI.""" - key = (block_kv, num_heads, head_dim, next_n, num_sms, - num_epi_subtiles, epi_dtype, acc_dtype, output_dtype) + 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 = block_kv * (head_dim + 4) + block_bytes = phys_block_kv * (head_dim + 4) sym_num_phys_blocks = cute.sym_int() sym_B = cute.sym_int() @@ -5221,7 +5223,8 @@ def _compile(cls, block_kv, num_heads, head_dim, next_n, num_sms, use_tvm_ffi_env_stream=True) kernel = FP8MQALogitsKernel( - block_kv=block_kv, + block_kv=compute_block_kv, + phys_block_kv=phys_block_kv, num_heads=num_heads, head_dim=head_dim, next_n=next_n, @@ -5270,7 +5273,7 @@ def forward( Args: q: [B, next_n, H, D] FP8 - kv_fused: [num_blocks, block_kv, 1, D+4] uint8 + 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 @@ -5285,7 +5288,8 @@ def forward( """ B, next_n, H, D = q.shape N = next_n * H - block_kv = kv_fused.shape[1] + phys_block_kv = kv_fused.shape[1] + compute_block_kv = 128 num_phys_blocks = kv_fused.shape[0] num_sms = _get_num_sms() @@ -5303,7 +5307,7 @@ def forward( kv_flat = kv_fused.reshape(num_phys_blocks, -1) # Allocate output with alignment padding - SPLIT_KV = block_kv * 2 # NUM_MATH_WG = 2 + 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( @@ -5314,11 +5318,12 @@ def forward( logits = logits[:, :max_context_len] # Compile if needed (fake tensors, no real data required) - key = (block_kv, H, D, next_n, num_sms, num_epi_subtiles, epi_dtype, - acc_dtype, output_dtype) + 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(block_kv, H, D, next_n, num_sms, num_epi_subtiles, - epi_dtype, acc_dtype, output_dtype) + 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 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index ec5f2f7b51bb..f47e3a754fae 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -185,6 +185,7 @@ class FP8MQALogitsKernel: def __init__( self, block_kv: int = 128, + phys_block_kv: int = 128, num_heads: int = 64, head_dim: int = 128, next_n: int = 1, @@ -200,6 +201,14 @@ def __init__( output_dtype=cutlass.Float32, ): self.block_kv = block_kv + self.phys_block_kv = phys_block_kv + self.num_blocks_per_mma = block_kv // phys_block_kv + assert block_kv % phys_block_kv == 0, ( + f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}" + ) + assert self.num_blocks_per_mma <= 4, ( + f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4" + ) self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue self.early_tmem_copy = early_tmem_copy self.smem_subpartition_opt = smem_subpartition_opt @@ -355,9 +364,10 @@ def __call__( stream: cuda.CUstream, ): # Derive KV and Scale views from fused buffer using CuTE ops. - # Fused layout per block: [KV data (block_kv*head_dim)] [Scales (block_kv*4)] - block_bytes = self.block_kv * (self.head_dim + 4) - scale_offset_elems = self.block_kv * self.head_dim # in FP8 elements + # 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) kv_fp8 = cute.recast_tensor(kv_fused, cutlass.Float8E4M3FN) @@ -366,22 +376,21 @@ def __call__( # recast back to FP8 so MMA type inference and TMA descriptors are correct. b = cute.recast_tensor(b, cutlass.Float8E4M3FN) - # KV view: [block_kv, head_dim, num_phys_blocks] FP8 - # Pointer is fused base, layout strides: (head_dim, 1, block_bytes) + # 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( - (self.block_kv, self.head_dim, num_phys_blocks), - stride=(self.head_dim, 1, block_bytes), + (phys_block_kv, self.head_dim, num_phys_blocks), + stride=(self.head_dim, 1, phys_block_bytes), ) a = cute.make_tensor(kv_fp8.iterator, kv_layout) # Scale view: offset pointer to scale region, recast FP8 → Float32 - # Step 1: create FP8 tensor at offset scale_offset_elems + # [phys_block_kv, num_phys_blocks] float32 (after recast) scale_fp8_layout = cute.make_layout( - (self.block_kv * 4, num_phys_blocks), - stride=(1, block_bytes), + (phys_block_kv * 4, num_phys_blocks), + stride=(1, phys_block_bytes), ) scale_fp8 = cute.make_tensor(kv_fp8.iterator + scale_offset_elems, scale_fp8_layout) - # Step 2: recast from FP8 (1 byte) to Float32 (4 bytes) scales = cute.recast_tensor(scale_fp8, cutlass.Float32) a_dtype = a.element_type @@ -392,19 +401,31 @@ def __call__( tiled_mma = self._setup_mma(a_dtype, b_dtype, a_major, b_major) atom_thr_size = cute.size(tiled_mma.thr_id.shape) - # TMA for KV (A) — full K=128 per load - a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) - a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) - tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( - a_op, + # TMA for KV (A) — fmha_decode_paged pattern. + # Build a TMA SMEM layout via tiled_divide on the full compute-tile + # layout, then select to drop trivial K dim. Atom uses mode [0] as + # single-tile SMEM layout and (phys, head) as cta_tiler. + tma_load_op = cpasync.CopyBulkTensorTileG2SOp() + self.a_tma_view_layout = sm100_utils.make_smem_layout( + tcgen05.OperandMajorMode.K, + (self.block_kv, self.head_dim), + a_dtype, + self.num_kv_stages, + ) + self.a_tma_view_layout = cute.tiled_divide( + self.a_tma_view_layout, (self.phys_block_kv, self.head_dim) + ) + # ((tile_M, tile_K), rest_M, rest_K, stages) → drop trivial rest_K + self.a_tma_view_layout = cute.select(self.a_tma_view_layout, mode=[0, 1, 3]) + # ((tile_M, tile_K), rest_M=num_sub_blocks, stages) + tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom( + tma_load_op, a, - a_smem_layout, - self.mma_tiler, - tiled_mma, - self.cluster_layout_vmnk.shape, + self.a_tma_view_layout[0], # atom SMEM = single-tile (mode 0) + (self.phys_block_kv, self.head_dim), ) - # TMA for Q (B) — full K=128, L dim = batch_size + # TMA for Q (B) — full K=128, L dim = batch_size (unchanged) b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( @@ -417,7 +438,6 @@ def __call__( ) # TMA for Weights — [N, batch_size], tile [N], L=batch_size - tma_load_op = cpasync.CopyBulkTensorTileG2SOp() self.w_smem_layout_staged = cute.make_layout( (self.N, self.num_q_stages), stride=(1, self.w_stage_stride), @@ -430,23 +450,27 @@ def __call__( self.w_smem_layout_staged.shape[:1], ) - # TMA for Scales — [block_kv, num_phys_blocks], tile [block_kv], L=num_phys_blocks + # TMA for Scales — [phys_block_kv, num_phys_blocks], tile [phys_block_kv] + # SMEM holds compute_block_kv scales per stage; filled by + # num_blocks_per_mma sub-block TMAs at consecutive offsets. self.s_smem_layout_staged = cute.make_layout((self.block_kv, self.num_kv_stages)) - s_smem_per_stage = cute.select(self.s_smem_layout_staged, mode=[0]) + s_smem_per_subblock = cute.make_layout((phys_block_kv,)) tma_atom_s, tma_tensor_s = cpasync.make_tiled_tma_atom( tma_load_op, scales, - s_smem_per_stage, - self.s_smem_layout_staged.shape[:1], + s_smem_per_subblock, + (phys_block_kv,), ) - a_copy_size = cute.size_in_bytes(a_dtype, a_smem_layout) b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout) w_copy_size = self.N * self.epi_bytes - kv_tma_bytes = a_copy_size * atom_thr_size - scale_tma_bytes = self.block_kv * 4 - # KV + Scale share barrier (like DeepGEMM) - self.num_kv_scale_tma_bytes = kv_tma_bytes + scale_tma_bytes + # Per sub-block: phys_block_kv * head_dim (KV) + phys_block_kv * 4 (scales) + kv_tma_bytes_per_subblock = phys_block_kv * self.head_dim + scale_tma_bytes_per_subblock = phys_block_kv * 4 + # Total per compute tile = num_blocks_per_mma sub-blocks + self.num_kv_scale_tma_bytes = self.num_blocks_per_mma * ( + kv_tma_bytes_per_subblock + scale_tma_bytes_per_subblock + ) # Q + Weights share barrier (like DeepGEMM) self.num_q_tma_bytes = b_copy_size * atom_thr_size + w_copy_size @@ -481,6 +505,7 @@ class SharedStorage: self.b_smem_layout_staged, self.w_smem_layout_staged, self.s_smem_layout_staged, + self.a_tma_view_layout, self.epi_tile, SharedStorage, ).launch( @@ -512,6 +537,7 @@ def kernel( b_smem_layout_staged: cute.ComposedLayout, w_smem_layout_staged: cute.Layout, s_smem_layout_staged: cute.Layout, + a_tma_view_layout: cute.ComposedLayout, epi_tile: cute.Tile, SharedStorage: cutlass.Constexpr, ): @@ -539,6 +565,7 @@ def kernel( # ~200-cycle L2 latency overlaps with subsequent prologue setup # (SMEM alloc, TMA partition, MMA fragment creation, etc.) NUM_MATH_WG = 2 # kNumMathWarpGroups + NUM_BLOCKS_PER_MMA = self.num_blocks_per_mma sm_idx = bidz start_q = mScheduleMeta[(sm_idx, 0)] start_kv_half = mScheduleMeta[(sm_idx, 1)] @@ -744,32 +771,32 @@ def kernel( cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 ) - # Partition KV (A): per-group SMEM targets - gA_mkl = cute.local_tile( + # Partition KV (A): fmha_decode_paged pattern. + # SMEM view is ((tile), num_sub_blocks, stages) — built in __call__. + # Use .outer (plain layout); swizzle is captured by sKV_0's iterator. + # GMEM: local_tile by (phys, head), then group first 2 modes into tile. + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + sKV_0_for_tma = cute.make_tensor(sKV_0.iterator, a_tma_view_layout.outer) + sKV_1_for_tma = cute.make_tensor(sKV_1.iterator, a_tma_view_layout.outer) + gA = cute.local_tile( mA_mkl, - cute.slice_(self.mma_tiler, (None, 0, None)), - (None, None, None), + (self.phys_block_kv, self.head_dim), + coord=(None, None, None), ) - thr_mma = tiled_mma.get_slice(mma_tile_coord_v) - tCgA = thr_mma.partition_A(gA_mkl) - a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) - tAsA_0, tAgA_0 = cpasync.tma_partition( tma_atom_a, - block_in_cluster_coord_vmnk[2], - a_cta_layout, - cute.group_modes(sKV_0, 0, 3), - cute.group_modes(tCgA, 0, 3), + 0, + cute.make_layout(1), + sKV_0_for_tma, + cute.group_modes(gA, 0, 2), ) tAsA_1, tAgA_1 = cpasync.tma_partition( tma_atom_a, - block_in_cluster_coord_vmnk[2], - a_cta_layout, - cute.group_modes(sKV_1, 0, 3), - cute.group_modes(tCgA, 0, 3), + 0, + cute.make_layout(1), + sKV_1_for_tma, + cute.group_modes(gA, 0, 2), ) - tAgA_0 = tAgA_0[(None, 0, None, None)] # [tma, K, L] - tAgA_1 = tAgA_1[(None, 0, None, None)] # Partition Q (B): shared SMEM, L dim = batch_size gB_nkl = cute.local_tile( @@ -798,22 +825,29 @@ def kernel( cute.group_modes(mW_tma, 0, 1), ) - # Partition Scales: standalone TMA, [block_kv, num_phys_blocks] - # tile [block_kv], L=num_phys_blocks. Per-group SMEM targets. - s_cta_layout = cute.make_layout((1,)) + # Partition Scales: explicit sub_blocks + stages dims. + # Layout (phys_block_kv, num_sub, stages) K-major with custom strides. + s_tma_view_layout = cute.make_layout( + (self.phys_block_kv, self.num_blocks_per_mma, self.num_kv_stages), + stride=(1, self.phys_block_kv, self.block_kv), + ) + sScales_0_for_tma = cute.make_tensor(sScales_0.iterator, s_tma_view_layout) + sScales_1_for_tma = cute.make_tensor(sScales_1.iterator, s_tma_view_layout) + # GMEM: local_tile by phys to match atom's tile size + gS = cute.local_tile(mS_tma, (self.phys_block_kv,), coord=(None, None)) tSsS_0, tSgS_0 = cpasync.tma_partition( tma_atom_s, 0, - s_cta_layout, - cute.group_modes(sScales_0, 0, 1), - cute.group_modes(mS_tma, 0, 1), + cute.make_layout(1), + sScales_0_for_tma, + gS, ) tSsS_1, tSgS_1 = cpasync.tma_partition( tma_atom_s, 0, - s_cta_layout, - cute.group_modes(sScales_1, 0, 1), - cute.group_modes(mS_tma, 0, 1), + cute.make_layout(1), + sScales_1_for_tma, + gS, ) # MMA fragments @@ -875,9 +909,10 @@ def kernel( cute.arch.warpgroup_reg_dealloc(24) lane_idx = tidx % 32 - # Block table prefetch: 32 lanes cache 32 block indices, - # distributed via shuffle. (Matches DeepGEMM L233-244) - cached_blk_idx = cutlass.Int32(0) + # Block table prefetch: 32 lanes cache block indices, + # distributed via shuffle. Each lane holds num_blocks_per_mma + # physical block indices per compute tile. (DeepGEMM L233-244) + cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)] kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use # Prefetch first Q before loop (like DeepGEMM line 203-204) @@ -945,39 +980,40 @@ def kernel( ) q_prod_state.advance() - # Block table prefetch for group 0 (like DeepGEMM L233-241) - # Each lane prefetches block_table[q_idx][kv_idx + lane_i * 2] + # Block table prefetch for group 0 (like DeepGEMM L233-241). + # Each lane loads num_blocks_per_mma physical block indices + # for one compute tile (kv_idx counts compute tiles). if kv_blk_ptr == 32: kv_blk_ptr = cutlass.Int32(0) prefetch_kv = kv_idx + lane_idx * NUM_MATH_WG if prefetch_kv < num_kv: - cached_blk_idx = mBlockTable[(q_idx, prefetch_kv)] + base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = mBlockTable[(q_idx, base_phys + i)] else: - cached_blk_idx = cutlass.Int32(0) - - # Get block index via shuffle (like DeepGEMM L244) - phys_blk = cute.arch.shuffle_sync(cached_blk_idx, kv_blk_ptr) - kv_blk_ptr = kv_blk_ptr + 1 + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = cutlass.Int32(0) - # Load KV + Scale for group 0 (kv_idx + 0) - # Unconditional TMA (like DeepGEMM): OOB kv_idx uses - # phys_blk=0 from block_table guard, writes to aligned - # padding region in logits. Keeps pipeline timing aligned. + # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile. kv_pipeline_0.producer_acquire(kv_prod_state_0) bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0) - cute.copy( - tma_atom_a, - tAgA_0[(None, 0, phys_blk)], - tAsA_0[(None, kv_prod_state_0.index)], - tma_bar_ptr=bar, - mcast_mask=a_mcast_mask, - ) - cute.copy( - tma_atom_s, - tSgS_0[(None, phys_blk)], - tSsS_0[(None, kv_prod_state_0.index)], - tma_bar_ptr=bar, - ) + stage = kv_prod_state_0.index + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blk_i = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + cute.copy( + tma_atom_a, + tAgA_0[(None, 0, 0, phys_blk_i)], + tAsA_0[(None, i, stage)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_0[(None, 0, phys_blk_i)], + tSsS_0[(None, i, stage)], + tma_bar_ptr=bar, + ) + kv_blk_ptr = kv_blk_ptr + 1 kv_prod_state_0.advance() # Advance: inline fetch_next_task @@ -997,7 +1033,7 @@ def kernel( lane_idx = tidx % 32 # Block table prefetch for group 1 - cached_blk_idx = cutlass.Int32(0) + cached_blks = [cutlass.Int32(0) for _ in range(NUM_BLOCKS_PER_MMA)] kv_blk_ptr = cutlass.Int32(32) # force prefetch on first use while has_work: @@ -1011,37 +1047,38 @@ def kernel( if q_idx != q_idx_old: kv_blk_ptr = cutlass.Int32(32) - # Block table prefetch for group 1 (like DeepGEMM L233-241) - # Each lane prefetches block_table[q_idx][kv_idx + 1 + lane_i * 2] + # Block table prefetch for group 1 (like DeepGEMM L233-241). if kv_blk_ptr == 32: kv_blk_ptr = cutlass.Int32(0) prefetch_kv = kv_idx + 1 + lane_idx * NUM_MATH_WG if prefetch_kv < num_kv: - cached_blk_idx = mBlockTable[(q_idx, prefetch_kv)] + base_phys = prefetch_kv * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = mBlockTable[(q_idx, base_phys + i)] else: - cached_blk_idx = cutlass.Int32(0) - - # Get block index via shuffle (like DeepGEMM L244) - phys_blk = cute.arch.shuffle_sync(cached_blk_idx, kv_blk_ptr) - kv_blk_ptr = kv_blk_ptr + 1 + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + cached_blks[i] = cutlass.Int32(0) - # Load KV + Scale for group 1 (kv_idx + 1) - # Unconditional TMA (like DeepGEMM) + # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile. kv_pipeline_1.producer_acquire(kv_prod_state_1) bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1) - cute.copy( - tma_atom_a, - tAgA_1[(None, 0, phys_blk)], - tAsA_1[(None, kv_prod_state_1.index)], - tma_bar_ptr=bar, - mcast_mask=a_mcast_mask, - ) - cute.copy( - tma_atom_s, - tSgS_1[(None, phys_blk)], - tSsS_1[(None, kv_prod_state_1.index)], - tma_bar_ptr=bar, - ) + stage = kv_prod_state_1.index + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blk_i = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + cute.copy( + tma_atom_a, + tAgA_1[(None, 0, 0, phys_blk_i)], + tAsA_1[(None, i, stage)], + tma_bar_ptr=bar, + mcast_mask=a_mcast_mask, + ) + cute.copy( + tma_atom_s, + tSgS_1[(None, 0, phys_blk_i)], + tSsS_1[(None, i, stage)], + tma_bar_ptr=bar, + ) + kv_blk_ptr = kv_blk_ptr + 1 kv_prod_state_1.advance() # Advance: inline fetch_next_task diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 06d3cac6b5a6..042f386a26d2 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -343,6 +343,109 @@ def test_cute_dsl_fp8_paged_mqa_logits( ) +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("num_heads", [64]) +@pytest.mark.parametrize("avg_ctx", [256, 4096]) +@pytest.mark.parametrize("phys_block_kv", [32, 64]) +def test_cute_dsl_fp8_paged_mqa_logits_multi_block( + batch_size, next_n, num_heads, avg_ctx, phys_block_kv +): + """Test multi-block TMA: physical block < compute tile (128). + + When phys_block_kv < 128, the kernel issues num_blocks_per_mma + separate TMA copies per compute tile to fill the 128-token SMEM. + """ + head_dim = 128 + compute_block_kv = 128 + max_model_len = max(avg_ctx * 2, 2048) + output_dtype = torch.float32 + + data = _generate_test_data( + batch_size, + next_n, + num_heads, + head_dim, + phys_block_kv, + avg_ctx, + max_model_len, + fix_length=True, + ) + + from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + + dsl_schedule_meta = get_paged_mqa_logits_metadata( + data["context_lens"], compute_block_kv, num_sms + ) + + ref_logits = _ref_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fp8"], + data["kv_scales"], + data["weights"], + data["context_lens"], + data["block_table"], + max_model_len, + phys_block_kv, + epi_dtype=output_dtype, + ) + + dsl_logits = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + data["q_fp8"], + data["kv_fused"], + data["weights"], + data["context_lens"], + data["block_table"], + dsl_schedule_meta, + max_model_len, + epi_dtype=output_dtype, + acc_dtype=output_dtype, + output_dtype=output_dtype, + ) + + assert dsl_logits.dtype == output_dtype + + B = batch_size + positions = torch.arange(max_model_len, device="cuda").unsqueeze(0) + row_indices = torch.arange(B * next_n, device="cuda") // next_n + next_n_offset = torch.arange(B * next_n, device="cuda") % next_n + end_pos = data["context_lens"][row_indices] - next_n + next_n_offset + mask = positions <= end_pos.unsqueeze(1) + + dsl_masked = dsl_logits.float().masked_fill(~mask, 0) + ref_masked = ref_logits.float().masked_fill(~mask, 0) + finite = torch.isfinite(dsl_masked) & torch.isfinite(ref_masked) + dsl_clean = dsl_masked.masked_fill(~finite, 0) + ref_clean = ref_masked.masked_fill(~finite, 0) + + elem_atol = 5e-5 + elem_rtol = 1e-5 + + valid = mask & finite + elem_abs = (dsl_clean - ref_clean).abs()[valid] + if elem_abs.numel() > 0: + print( + f"[multi-block] B={batch_size} next_n={next_n} avg_ctx={avg_ctx} " + f"phys_block_kv={phys_block_kv} -> " + f"max_abs={elem_abs.max().item():.3e} " + f"mean_abs={elem_abs.mean().item():.3e}" + ) + + torch.testing.assert_close( + dsl_clean, + ref_clean, + atol=elem_atol, + rtol=elem_rtol, + msg=lambda m: ( + f"{m}\nB={batch_size}, next_n={next_n}, avg_ctx={avg_ctx}, " + f"phys_block_kv={phys_block_kv}" + ), + ) + + # @skip_if_unsupported # @skip_not_sm100 # @pytest.mark.parametrize("batch_size", [1, 4, 32]) @@ -552,18 +655,31 @@ def benchmark_fp8_paged_mqa_logits( output_dtype=torch.float32, num_epi_subtiles=1, varlen=False, + block_kv=128, ): - """Benchmark CuTe DSL vs C++ DeepGEMM kernel time.""" + """Benchmark CuTe DSL vs C++ DeepGEMM kernel time. + + Args: + block_kv: physical block size (tokens per page). DSL scheduler always + uses compute_block_kv=128; when block_kv < 128, the DSL kernel + issues num_blocks_per_mma TMA copies per compute tile. + """ from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata num_heads = 64 head_dim = 128 - block_kv = 128 + compute_block_kv = 128 # DSL scheduler / compute tile (always 128 on SM100) + assert compute_block_kv % block_kv == 0, ( + f"compute_block_kv={compute_block_kv} must be divisible by block_kv={block_kv}" + ) num_sms = torch.cuda.get_device_properties(0).multi_processor_count dtype_str = str(output_dtype).split(".")[-1] mode_str = "varlen" if varlen else "fix-len" - print(f"output_dtype={dtype_str} num_epi_subtiles={num_epi_subtiles} mode={mode_str}") + print( + f"output_dtype={dtype_str} num_epi_subtiles={num_epi_subtiles} " + f"mode={mode_str} block_kv={block_kv}" + ) is_non_default = output_dtype != torch.float32 or num_epi_subtiles != 1 hdr = ( f"{'batch':>5s} {'ctx':>7s} {'next_n':>6s} {'nblk':>7s} | " @@ -589,8 +705,9 @@ def benchmark_fp8_paged_mqa_logits( varlen=varlen, ) + # DSL scheduler counts compute tiles (always 128), not pages. dsl_schedule_meta = get_paged_mqa_logits_metadata( - data["context_lens"], block_kv, num_sms + data["context_lens"], compute_block_kv, num_sms ) def dsl_fn(data=data): @@ -633,8 +750,10 @@ def dsl_fn(data=data): num_kv_multicast = 2 if dg_next_n == 4 else 1 num_clusters = num_sms // num_kv_multicast + # block_kv arg ignored by DG (scheduler uses compute tile + # internally); differs from DSL only in num_clusters. dg_schedule_meta = get_paged_mqa_logits_metadata( - dg_data["context_lens"], block_kv, num_clusters + dg_data["context_lens"], compute_block_kv, num_clusters ) def dg_fn(dg_data=dg_data): @@ -732,6 +851,15 @@ def dsl_f32_fn(data=data): help="use varlen workload (per-seq lengths in [min(2048,max), max]); " "default is fix-length where all sequences use --context_len", ) + parser.add_argument( + "--block_kv", + type=int, + default=128, + choices=[32, 64, 128], + help="physical block size / tokens per page (default: 128). " + "DSL compute tile is always 128; when block_kv<128, DSL issues " + "num_blocks_per_mma=128/block_kv TMA copies per compute tile.", + ) args = parser.parse_args() dtype_map = {"float32": torch.float32, "float16": torch.float16} @@ -744,4 +872,5 @@ def dsl_f32_fn(data=data): output_dtype=dtype_map[args.output_dtype], num_epi_subtiles=args.num_epi_subtiles, varlen=args.varlen, + block_kv=args.block_kv, ) diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py index 3b5cb37ee856..4f1462bb1c15 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py @@ -938,10 +938,8 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): random.seed(123) use_dsl = backend == "dsl" - # Note, DSL kernel requires tokens_per_block=128. - # Will remove this restriction in the future. heads, head_dim = (32, 128) - block_size = 128 if use_dsl else 64 + block_size = 64 avg_context_len = 2048 num_gen_tokens = next_n # Number of tokens to generate per sequence max_model_len = 4096 From 862260c5da893a2f6bcfe5bc3bbfbda3a13a302b Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Thu, 23 Apr 2026 14:54:13 +0000 Subject: [PATCH 22/26] [None][fix] Move shuffle before barrier acquire to match DeepGEMM scheduling Hoist shuffle_sync calls before producer_acquire in both TMA warpgroups. The previous ordering placed shuffle after barrier acquire, causing a dependency chain that regressed next_n=3 by 5-7%. SASS now matches baseline exactly. Signed-off-by: Mindy Li Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index f47e3a754fae..860404d7572f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -994,26 +994,30 @@ def kernel( for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): cached_blks[i] = cutlass.Int32(0) + # Get block indices via shuffle before barrier (like DeepGEMM L244) + phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + # Load KV + Scale for group 0: num_blocks_per_mma TMAs per tile. kv_pipeline_0.producer_acquire(kv_prod_state_0) bar = kv_pipeline_0.producer_get_barrier(kv_prod_state_0) stage = kv_prod_state_0.index for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): - phys_blk_i = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) cute.copy( tma_atom_a, - tAgA_0[(None, 0, 0, phys_blk_i)], + tAgA_0[(None, 0, 0, phys_blks[i])], tAsA_0[(None, i, stage)], tma_bar_ptr=bar, mcast_mask=a_mcast_mask, ) cute.copy( tma_atom_s, - tSgS_0[(None, 0, phys_blk_i)], + tSgS_0[(None, 0, phys_blks[i])], tSsS_0[(None, i, stage)], tma_bar_ptr=bar, ) - kv_blk_ptr = kv_blk_ptr + 1 kv_prod_state_0.advance() # Advance: inline fetch_next_task @@ -1059,26 +1063,30 @@ def kernel( for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): cached_blks[i] = cutlass.Int32(0) + # Get block indices via shuffle before barrier (like DeepGEMM L244) + phys_blks = [cutlass.Int32(0)] * NUM_BLOCKS_PER_MMA + for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): + phys_blks[i] = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) + kv_blk_ptr = kv_blk_ptr + 1 + # Load KV + Scale for group 1: num_blocks_per_mma TMAs per tile. kv_pipeline_1.producer_acquire(kv_prod_state_1) bar = kv_pipeline_1.producer_get_barrier(kv_prod_state_1) stage = kv_prod_state_1.index for i in cutlass.range_constexpr(NUM_BLOCKS_PER_MMA): - phys_blk_i = cute.arch.shuffle_sync(cached_blks[i], kv_blk_ptr) cute.copy( tma_atom_a, - tAgA_1[(None, 0, 0, phys_blk_i)], + tAgA_1[(None, 0, 0, phys_blks[i])], tAsA_1[(None, i, stage)], tma_bar_ptr=bar, mcast_mask=a_mcast_mask, ) cute.copy( tma_atom_s, - tSgS_1[(None, 0, phys_blk_i)], + tSgS_1[(None, 0, phys_blks[i])], tSsS_1[(None, i, stage)], tma_bar_ptr=bar, ) - kv_blk_ptr = kv_blk_ptr + 1 kv_prod_state_1.advance() # Advance: inline fetch_next_task From 0331897c21f1a20ce5a88198af496ea679ec95ca Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Thu, 23 Apr 2026 15:25:19 +0000 Subject: [PATCH 23/26] [None][cleanup] Remove commented-out test_deepgemm_fp8_paged_mqa_logits Signed-off-by: Mindy Li Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../test_cute_dsl_fp8_paged_mqa_logits.py | 106 ------------------ 1 file changed, 106 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 042f386a26d2..b35d04b7efb7 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -446,112 +446,6 @@ def test_cute_dsl_fp8_paged_mqa_logits_multi_block( ) -# @skip_if_unsupported -# @skip_not_sm100 -# @pytest.mark.parametrize("batch_size", [1, 4, 32]) -# @pytest.mark.parametrize("next_n", [1, 2, 4]) -# @pytest.mark.parametrize("avg_ctx", [256, 4096, 32768]) -# @pytest.mark.parametrize("output_dtype", [torch.float32]) -# @pytest.mark.parametrize("fix_length", [True, False]) -# def test_deepgemm_fp8_paged_mqa_logits(batch_size, next_n, avg_ctx, output_dtype, fix_length): -# """Compare DeepGEMM kernel output against reference. -# """ -# from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits - -# torch.manual_seed(42) -# torch.cuda.manual_seed(42) - -# num_heads = 64 -# head_dim = 128 -# block_kv = 128 -# max_model_len = max(avg_ctx * 2, 2048) - -# data = _generate_test_data( -# batch_size, -# next_n, -# num_heads, -# head_dim, -# block_kv, -# avg_ctx, -# max_model_len, -# use_int_data=(output_dtype == torch.float16), -# fix_length=fix_length, -# ) - -# from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata - -# num_sms = torch.cuda.get_device_properties(0).multi_processor_count - -# # DeepGEMM kernel always uses num_clusters as grid size. -# num_kv_multicast = 2 if next_n == 4 else 1 -# num_clusters = num_sms // num_kv_multicast -# dg_schedule_meta = get_paged_mqa_logits_metadata( -# data["context_lens"], block_kv, num_clusters -# ) - -# dsl_logits = fp8_paged_mqa_logits( -# data["q_fp8"], -# data["kv_fused"], -# data["weights"], -# data["context_lens"], -# data["block_table"], -# dg_schedule_meta, -# max_model_len, -# ) -# ref_logits = _ref_fp8_paged_mqa_logits( -# data["q_fp8"], -# data["kv_fp8"], -# data["kv_scales"], -# data["weights"], -# data["context_lens"], -# data["block_table"], -# max_model_len, -# block_kv, -# epi_dtype=output_dtype, -# ) - -# # Mask invalid positions -# B = batch_size -# positions = torch.arange(max_model_len, device="cuda").unsqueeze(0) -# row_indices = torch.arange(B * next_n, device="cuda") // next_n -# next_n_offset = torch.arange(B * next_n, device="cuda") % next_n -# end_pos = data["context_lens"][row_indices] - next_n + next_n_offset -# mask = positions <= end_pos.unsqueeze(1) - -# dsl_masked = dsl_logits.float().masked_fill(~mask, 0) -# ref_masked = ref_logits.float().masked_fill(~mask, 0) -# finite = torch.isfinite(dsl_masked) & torch.isfinite(ref_masked) -# dsl_clean = dsl_masked.masked_fill(~finite, 0) -# ref_clean = ref_masked.masked_fill(~finite, 0) - -# # Element-wise check on the valid (finite + in-context) region. -# # Kernel is deterministic (disjoint CTA writes, no atomics), so every -# # element must be within elem_atol. -# elem_atol = 1e-3 if output_dtype == torch.float16 else 5e-5 -# elem_rtol = 1e-3 if output_dtype == torch.float16 else 1e-5 - -# # Debug probe: print max/mean abs error for CI failure diagnosis. -# valid = mask & finite -# elem_abs = (dsl_clean - ref_clean).abs()[valid] -# if elem_abs.numel() > 0: -# print( -# f"[acc-probe] B={batch_size} next_n={next_n} avg_ctx={avg_ctx} " -# f"dtype={output_dtype} -> " -# f"max_abs={elem_abs.max().item():.3e} " -# f"mean_abs={elem_abs.mean().item():.3e}" -# ) - -# torch.testing.assert_close( -# dsl_clean, -# ref_clean, -# atol=elem_atol, -# rtol=elem_rtol, -# msg=lambda m: ( -# f"{m}\nB={batch_size}, next_n={next_n}, avg_ctx={avg_ctx}, dtype={output_dtype}" -# ), -# ) - - def _profile_kernel_us(fn, num_warmup=10, num_iterations=30): """Profile CUDA kernel time in microseconds using torch.profiler.""" from torch.profiler import ProfilerActivity, profile From 26effccdf7ef9ade3eae506fe6cab36f41a1523e Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 28 Apr 2026 05:50:34 +0000 Subject: [PATCH 24/26] [None][cleanup] Remove unused helpers in DSL paged MQA logits kernel and test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop dead code flagged in review: - compute_schedule_metadata + cdiv (and now-unused torch import) from the kernel module — never called. - skip_if_unsupported decorator + has_deep_gemm helper + IS_CUTLASS_DSL_AVAILABLE import from the test — defined but never applied. SM100/103 gate already ensures DeepGEMM and CuTe DSL availability. Also refresh the test docstring: the reference is pure PyTorch, not DeepGEMM. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../paged_mqa_logits/fp8_paged_mqa_logits.py | 58 ------------------- .../test_cute_dsl_fp8_paged_mqa_logits.py | 19 +----- 2 files changed, 1 insertion(+), 76 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py index 860404d7572f..e94073112209 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py @@ -62,7 +62,6 @@ import cutlass.pipeline as pipeline import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils -import torch from cutlass import Float16, Int32 from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, nvvm, vector @@ -1707,60 +1706,3 @@ def kernel( else: cute.arch.warpgroup_reg_dealloc(24) - - -def cdiv(a, b): - return (a + b - 1) // b - - -def compute_schedule_metadata(context_lens, block_kv, num_ctas): - """Compute schedule metadata: [num_ctas+1, 2] int32. - - Each row stores (q_idx, kv_idx / kNumMathWarpGroups) marking CTA boundaries. - Metadata format: - - schedule[i] = start boundary for CTA i - - schedule[i+1] = end boundary for CTA i (= start of CTA i+1) - - schedule[num_ctas] = past-the-end sentinel (batch_size, 0) - - The kernel multiplies the second column by kNumMathWarpGroups (=2) to get - the actual kv_idx in units of KV blocks. - """ - batch_size = context_lens.shape[0] - splits_per_seq = [] - total_splits = 0 - for b in range(batch_size): - ctx = context_lens[b].item() - num_kv = cdiv(ctx, block_kv) - ns = cdiv(num_kv, 2) - splits_per_seq.append(ns) - total_splits += ns - - # Balanced distribution - q_div = total_splits // num_ctas - r_mod = total_splits % num_ctas - - schedule = torch.zeros((num_ctas + 1, 2), dtype=torch.int32) - - # For each CTA boundary, find (q_idx, kv_half_idx) - cum = 0 - seq_idx = 0 - seq_offset = 0 - for i in range(num_ctas + 1): - target = i * q_div + min(i, r_mod) - while seq_idx < batch_size and cum + (splits_per_seq[seq_idx] - seq_offset) <= target: - cum += splits_per_seq[seq_idx] - seq_offset - seq_idx += 1 - seq_offset = 0 - if seq_idx >= batch_size: - break - if seq_idx >= batch_size: - # Past-the-end: (batch_size, 0) — matches DeepGEMM's end sentinel. - # When the scheduler wraps past the last batch, it reaches - # (batch_size, 0), and the end check (q==end_q and kv==end_kv) - # correctly terminates. - schedule[i] = torch.tensor([batch_size, 0], dtype=torch.int32) - else: - local_split = target - cum + seq_offset - schedule[i] = torch.tensor([seq_idx, local_split], dtype=torch.int32) - - return schedule diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index b35d04b7efb7..d592bfb4a78e 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -19,7 +19,6 @@ import pytest import torch -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._utils import get_sm_version skip_not_sm100 = pytest.mark.skipif( @@ -28,15 +27,6 @@ ) -def has_deep_gemm(): - try: - from tensorrt_llm import deep_gemm - - return deep_gemm is not None - except Exception: - return False - - def _ceil_to_ue8m0(x: torch.Tensor): return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -229,11 +219,6 @@ def _generate_test_data( } -skip_if_unsupported = pytest.mark.skipif( - not (has_deep_gemm() and IS_CUTLASS_DSL_AVAILABLE), reason="Requires DeepGEMM and CuTe DSL" -) - - @skip_not_sm100 @pytest.mark.parametrize("batch_size", [1, 4, 32]) @pytest.mark.parametrize("next_n", [1, 2, 3, 4]) @@ -244,10 +229,8 @@ def _generate_test_data( def test_cute_dsl_fp8_paged_mqa_logits( batch_size, next_n, num_heads, avg_ctx, output_dtype, fix_length ): - """Compare CuTe DSL kernel output against reference. + """Compare CuTe DSL kernel output against a pure PyTorch reference. - Uses C++ DeepGEMM as reference when available (next_n in {1,2,4}), - falls back to pure PyTorch reference otherwise (e.g. next_n=3). Tests both fp32 and fp16 epi/acc/output paths. """ head_dim = 128 From f34fb16e13a81e86bdfd4e1d3d734ebefe0ea715 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Tue, 28 Apr 2026 20:20:26 -0700 Subject: [PATCH 25/26] [None][fix] Skip DSL backend on non-SM100 archs in indexer decode test The DSL variant of test_indexer_decode_with_paged_kv_cache calls torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits, which raises ValueError on SM != 100/103. Mark only the "dsl" parameter with a skipif so the "deepgemm" backend still runs on Hopper. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../_torch/attention/sparse/test_dsa_indexer.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py index 4f1462bb1c15..e8d2ef34b479 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py @@ -921,7 +921,16 @@ def test_fp8_k_cache_roundtrip(): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper @pytest.mark.parametrize("batch_size,next_n", [(4, 1), (2, 2), (4, 3), (4, 4)]) -@pytest.mark.parametrize("backend", ["deepgemm", "dsl"]) +@pytest.mark.parametrize("backend", [ + "deepgemm", + pytest.param( + "dsl", + marks=pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason= + f"CuTe DSL FP8 Paged MQA Logits only supports SM 100/103, got SM {get_sm_version()}", + )), +]) def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): """ Test FP8 paged KV cache with two-phase workflow and variable context lengths. From 4bcbd79a5d9024f70c9206b9095171555c7a4c88 Mon Sep 17 00:00:00 2001 From: Mindy Li <11663212+limin2021@users.noreply.github.com> Date: Fri, 8 May 2026 05:18:20 +0000 Subject: [PATCH 26/26] [None][test] Fix DSL indexer scheduler buffer aliasing and extend FP8 MQA next_n coverage test_dsa_indexer: separate scheduler_metadata_buffer for DSL backend (kNumNextNAtoms=1) to avoid the next_n>1 alias used by DeepGEMM. Fixes test_indexer_decode_with_paged_kv_cache across {deepgemm, dsl} x {(4,1),(2,2),(4,3),(4,4)}. test_cute_dsl_fp8_paged_mqa_logits: extend multi_block next_n coverage to {1,2,3,4}; drop now-obsolete next_n==3 expansion in bench (newer DeepGEMM supports it natively); update --block_kv default from 128 to 64 to match the new {32,64} assertion in fp8_paged_mqa_logits. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> --- .../test_cute_dsl_fp8_paged_mqa_logits.py | 46 ++++++------------- .../attention/sparse/test_dsa_indexer.py | 11 ++++- 2 files changed, 24 insertions(+), 33 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py index 3aafe9eeee47..9e5461bc275f 100644 --- a/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py @@ -346,7 +346,7 @@ def test_cute_dsl_fp8_paged_mqa_logits( @skip_not_sm100 @pytest.mark.parametrize("batch_size", [1, 4]) -@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("next_n", [1, 2, 3, 4]) @pytest.mark.parametrize("num_heads", [64]) @pytest.mark.parametrize("avg_ctx", [256, 4096]) @pytest.mark.parametrize("phys_block_kv", [32, 64]) @@ -638,36 +638,18 @@ def dsl_fn(data=data): try: from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits - # DeepGEMM doesn't support next_n=3 natively; expand to - # batch=B*next_n, next_n=1 (same approach as production MTP - # expand path in dsa.py). - dg_next_n = next_n - dg_data = data - if next_n == 3: - exp_bs = batch_size * next_n - dg_data = { - "q_fp8": data["q_fp8"].reshape(exp_bs, 1, num_heads, head_dim), - "kv_fused": data["kv_fused"], - "weights": data["weights"].reshape(exp_bs, num_heads), - "context_lens": data["context_lens"].repeat_interleave(next_n), - "block_table": data["block_table"].repeat_interleave(next_n, dim=0), - "max_model_len": data["max_model_len"], - } - dg_next_n = 1 - # SM100 always uses num_kv_multicast=1 in upgraded DeepGEMM # (cluster(2,1,1) for next_n=4 was removed). Atom-split is # encoded in metadata via num_next_n_atoms which the wrapper - # derives from context_lens.size(1). + # derives from context_lens.size(1). DG natively supports + # next_n in {1,2,3,4}. num_clusters = num_sms - # 2D context_lens shape (B, dg_next_n): for dg_next_n>1 the - # wrapper computes `num_next_n_atoms = dg_next_n / next_n_atom_size` + # 2D context_lens shape (B, next_n): for next_n>1 the wrapper + # computes `num_next_n_atoms = next_n / next_n_atom_size` # which DG's compute kernel expects. All next_n positions # of a batch share the same KV length here (broadcast via # expand) — TRT-LLM does the same in production. - dg_ctx_2d = ( - dg_data["context_lens"].unsqueeze(-1).expand(-1, dg_next_n).contiguous() - ) + dg_ctx_2d = data["context_lens"].unsqueeze(-1).expand(-1, next_n).contiguous() # `block_kv = 64` for the same reason as the DSL path: # metadata SPLIT_KV = block_kv * 4 must equal DG compute # kernel's hardcoded SPLIT_KV = 256 (apis/attention.hpp:353). @@ -677,15 +659,15 @@ def dsl_fn(data=data): dg_ctx_2d, DG_METADATA_BLOCK_KV, num_clusters ) - def dg_fn(dg_data=dg_data, dg_ctx_2d=dg_ctx_2d): + def dg_fn(data=data, dg_ctx_2d=dg_ctx_2d): fp8_paged_mqa_logits( - dg_data["q_fp8"], - dg_data["kv_fused"], - dg_data["weights"], + data["q_fp8"], + data["kv_fused"], + data["weights"], dg_ctx_2d, - dg_data["block_table"], + data["block_table"], dg_schedule_meta, - dg_data["max_model_len"], + data["max_model_len"], ) dg_us = _profile_kernel_us(dg_fn, num_warmup, num_iterations) @@ -775,9 +757,9 @@ def dsl_f32_fn(data=data): parser.add_argument( "--block_kv", type=int, - default=128, + default=64, choices=[32, 64, 128], - help="physical block size / tokens per page (default: 128). " + help="physical block size / tokens per page (default: 64). " "DSL compute tile is always 128; when block_kv<128, DSL issues " "num_blocks_per_mma=128/block_kv TMA copies per compute tile.", ) diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py index 4b917f2af7ff..ab5f3a78a52d 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py @@ -463,7 +463,16 @@ def __init__(self): self.scheduler_metadata_buffer = torch.zeros((self.num_sms + 1, 2), device='cuda', dtype=torch.int32) - self.scheduler_metadata_buffer_full_next_n = self.scheduler_metadata_buffer + # DSL needs the next_n=1 schedule preserved (kNumNextNAtoms=1), + # so allocate a separate buffer for the full-next_n schedule. + # DeepGEMM expects the full-next_n schedule in + # `scheduler_metadata_buffer` itself (the alias makes + # `Indexer.prepare()`'s second populate overwrite the first). + if use_cute_dsl_paged_mqa_logits: + self.scheduler_metadata_buffer_full_next_n = torch.zeros( + (self.num_sms + 1, 2), device='cuda', dtype=torch.int32) + else: + self.scheduler_metadata_buffer_full_next_n = self.scheduler_metadata_buffer # Pre-allocated 2D kv_lens buffer for the DeepGEMM 2D context_lens API. self.kv_lens_cuda_2d = torch.zeros( (self.num_seqs, 1 + self.max_draft_tokens),