diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index f90ef8596940..af6360bd32a8 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -723,12 +723,10 @@ def __init__( # GVR emission-assisted decode (opt-in, experimental): the FP4/FP8 # indexer epilogue emits candidates the GVR Top-K consumes (see # gvr_emission / gvr_routing; state lives on the TopK module) - # only the FP4 scoring op accepts emission kwargs self.use_gvr_emission = ( os.environ.get("TRTLLM_GVR_EMISSION", "0") == "1" and decode_top_k_implementation == TopKImplementation.CUTE_DSL_GVR and self.use_cute_dsl_paged_mqa_logits - and self.use_fp4 ) # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM 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 bbaa63c14f00..4597742a6bbf 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -8054,6 +8054,83 @@ def _( from ..cute_dsl_kernels.blackwell.paged_mqa_logits import ( FP4MQALogitsKernel, FP8MQALogitsKernel) + def _emission_block_max(B, next_n, nb_pad, block_max_out, device): + """Allocate/validate the [rows, nb_pad*4] fp32 warp-partial buffer + (4 records per 128-token block).""" + nrec = nb_pad * 4 + if block_max_out is None: + block_max_out = torch.empty((B * next_n, nrec), + device=device, + dtype=torch.float32) + assert (block_max_out.shape == (B * next_n, nrec) + and block_max_out.is_contiguous()) + return block_max_out + + def _emission_seed_buffers(B, next_n, emit_block_meta, seed_thr, + seed_counts_out): + """Validate the seed contract shared by the FP4/FP8 scoring runners. + + Packed (seed_counts_out is None): seed_thr IS the [rows, 8] fp32 + seed row - lines at cols 0..2, counts accumulate as fp32 at cols + 3..5; the caller zeroes cols 3..7 and writes lines each step. + Split: seed_thr [rows, 3] fp32 + caller-zeroed seed_counts_out + [rows, 3] int32. Returns (seed_packed, seed_counts_out). + """ + assert emit_block_meta, "emit_seed_counts requires emit_block_meta" + if seed_counts_out is None: + assert ( + seed_thr is not None and seed_thr.dtype == torch.float32 + and seed_thr.is_cuda and seed_thr.is_contiguous() + and seed_thr.shape == (B * next_n, 8) + ), (f"packed emit_seed_counts requires seed_thr fp32 " + f"[{B * next_n}, 8]; got " + f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}" + ) + return True, seed_thr + assert ( + seed_thr is not None and seed_thr.dtype == torch.float32 + and seed_thr.is_cuda and seed_thr.is_contiguous() + and seed_thr.shape == (B * next_n, 3) + ), (f"emit_seed_counts requires seed_thr fp32 " + f"[{B * next_n}, 3]; got " + f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}" + ) + assert (seed_counts_out.dtype == torch.int32 and seed_counts_out.is_cuda + and seed_counts_out.is_contiguous() + and seed_counts_out.shape == (B * next_n, 3)), ( + "emit_seed_counts requires a caller-zeroed " + "seed_counts_out int32 [B*next_n, 3]") + return False, seed_counts_out + + def _emission_bucketed_cand(B, next_n, accept_cap, cand_out, cand_idx_out, + cand_ctl_out, cand_cur_out): + """Validate the bucketed SoA contract: cand_out fp32 VALUES + [rows, 2*segA+capC], cand_idx_out int32 positions (same width), + cand_ctl_out int32 [rows, 4] {n0, void, n1, n2} (caller-zeroed), + cand_cur_out int32 [rows, 4] cursors (caller-zeroed). + Returns cand_cap = W - 2*accept_cap.""" + assert cand_out is not None, ("emit_cand_bucketed requires cand_out") + W = cand_out.shape[1] + assert (cand_out.dtype == torch.float32 and cand_out.is_cuda + and cand_out.is_contiguous() and cand_out.dim() == 2 + and cand_out.shape[0] == B * next_n and W > 2 * accept_cap), ( + "bucketed requires cand_out fp32 [rows, 2*segA+capC]") + assert (cand_idx_out is not None and cand_idx_out.dtype == torch.int32 + and cand_idx_out.is_cuda and cand_idx_out.is_contiguous() + and cand_idx_out.shape == cand_out.shape), ( + "bucketed requires cand_idx_out int32, same shape") + assert (cand_ctl_out is not None and cand_ctl_out.dtype == torch.int32 + and cand_ctl_out.is_cuda and cand_ctl_out.is_contiguous() + and cand_ctl_out.shape == (B * next_n, 4)), ( + "bucketed requires caller-zeroed cand_ctl_out " + "int32 [rows, 4]") + assert (cand_cur_out is not None and cand_cur_out.dtype == torch.int32 + and cand_cur_out.is_cuda and cand_cur_out.is_contiguous() + and cand_cur_out.shape == (B * next_n, 4)), ( + "bucketed requires caller-zeroed cand_cur_out " + "int32 [rows, 4]") + return W - 2 * accept_cap + class CuteDSLPagedMQALogitsRunner: """Runner for CuTe DSL FP8 Paged MQA Logits kernel (Blackwell SM100). @@ -8064,13 +8141,28 @@ class CuteDSLPagedMQALogitsRunner: kernel_cache = dict() @classmethod - def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim, - next_n, num_sms, num_epi_subtiles, epi_dtype, acc_dtype, - output_dtype): + 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, + emit_block_meta=False, + emit_seed_counts=False, + seed_packed=False, + emit_cand_bucketed=False, + accept_cap=8192, + cand_cap=0): """Compile kernel using fake tensors + TVM FFI.""" key = (compute_block_kv, phys_block_kv, num_heads, head_dim, next_n, num_sms, num_epi_subtiles, epi_dtype, acc_dtype, - output_dtype) + output_dtype, emit_block_meta, emit_seed_counts, seed_packed, + emit_cand_bucketed, accept_cap, cand_cap) if key in cls.kernel_cache: return @@ -8119,6 +8211,58 @@ def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim, (num_ctas, 2), stride_order=(1, 0)) + # Emission fakes (same contracts as the FP4 runner; hit-stats + # and the plain-list variants are FP4-only). + block_max_fake = None + if emit_block_meta: + block_max_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), cute.sym_int()), + stride_order=(1, 0), + assumed_align=16) + seed_thr_fake = None + seed_counts_fake = None + if emit_seed_counts: + if seed_packed: + seed_thr_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 8), + stride_order=(1, 0), + assumed_align=4) + seed_counts_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 8), + stride_order=(1, 0), + assumed_align=4) + else: + seed_thr_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 3), + stride_order=(1, 0), + assumed_align=4) + seed_counts_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 3), + stride_order=(1, 0), + assumed_align=4) + cand_fake = None + cand_idx_fake = None + cand_ctl_fake = None + cand_cur_fake = None + if emit_cand_bucketed: + wtot = 2 * accept_cap + cand_cap + cand_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), wtot), + stride_order=(1, 0), + assumed_align=4) + cand_idx_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), wtot), + stride_order=(1, 0), + assumed_align=4) + cand_ctl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 4), + stride_order=(1, 0), + assumed_align=4) + cand_cur_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int(), 4), + stride_order=(1, 0), + assumed_align=4) + fake_stream = cute.runtime.make_fake_stream( use_tvm_ffi_env_stream=True) @@ -8133,6 +8277,12 @@ def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim, epi_dtype=to_cutlass[epi_dtype], acc_dtype=to_cutlass[acc_dtype], output_dtype=to_cutlass[output_dtype], + emit_block_meta=emit_block_meta, + emit_seed_counts=emit_seed_counts, + seed_packed=seed_packed, + emit_cand_bucketed=emit_cand_bucketed, + accept_cap=accept_cap, + cand_cap=cand_cap, ) compiled = cute.compile( @@ -8147,6 +8297,13 @@ def _compile(cls, compute_block_kv, phys_block_kv, num_heads, head_dim, cutlass.Int32(1), cutlass.Int32(1), fake_stream, + block_max=block_max_fake, + seed_thr=seed_thr_fake, + seed_counts=seed_counts_fake, + cand=cand_fake, + cand_ctl=cand_ctl_fake, + cand_idx_t=cand_idx_fake, + cand_cur=cand_cur_fake, options="--enable-tvm-ffi", ) cls.kernel_cache[key] = compiled @@ -8168,6 +8325,17 @@ def forward( epi_dtype: torch.dtype = torch.float32, acc_dtype: torch.dtype = torch.float32, output_dtype: torch.dtype = torch.float32, + emit_block_meta: bool = False, + block_max_out: Optional[torch.Tensor] = None, + emit_seed_counts: bool = False, + seed_thr: Optional[torch.Tensor] = None, + seed_counts_out: Optional[torch.Tensor] = None, + emit_cand_bucketed: bool = False, + accept_cap: int = 8192, + cand_out: Optional[torch.Tensor] = None, + cand_idx_out: Optional[torch.Tensor] = None, + cand_ctl_out: Optional[torch.Tensor] = None, + cand_cur_out: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Execute FP8 paged MQA logits kernel. @@ -8183,8 +8351,26 @@ def forward( epi_dtype: epilogue compute dtype acc_dtype: MMA accumulator dtype output_dtype: output logits dtype + emit_block_meta: also emit per-128-block metadata for the + fused GVR top-k. ``block_max_out`` + [B*next_n, nb_pad*4] fp32 (4 warp-partial records per + block) is allocated when not supplied. + emit_seed_counts: with ``seed_counts_out`` None, packed + contract: ``seed_thr`` IS the [B*next_n, 8] fp32 seed + row (lines at cols 0..2, counts at 3..5); otherwise + split contract: ``seed_thr`` [B*next_n, 3] fp32 + + caller-zeroed ``seed_counts_out`` [B*next_n, 3] int32. + emit_cand_bucketed: bucketed A/B/C candidate list into + ``cand_out`` / ``cand_idx_out`` [B*next_n, + 2*accept_cap + cand_cap] + ``cand_ctl_out`` + [B*next_n, 4] + ``cand_cur_out`` [B*next_n, 4]. Returns: - logits: [B*next_n, max_context_len] output_dtype + logits [B*next_n, max_context_len]; with emit_block_meta, + the tuple (logits, block_max_out). + + The optional emission tensors are written by the kernel but + cannot be declared in ``mutates_args`` (see the FP4 runner + note); emission is eager / CUDA-graph only. """ B, next_n, H, D = q.shape N = next_n * H @@ -8221,13 +8407,55 @@ def forward( ) logits = logits[:, :max_context_len] + # Emission buffers: written by the kernel but out of + # mutates_args (see the docstring note). nb_pad mirrors the + # logits padding so WG1's odd-num_kv OOB tile lands in padding. + if emit_block_meta: + nb_pad = aligned_max_ctx // compute_block_kv + block_max_out = _emission_block_max(B, next_n, nb_pad, + block_max_out, q.device) + seed_packed = False + if emit_seed_counts: + seed_packed, seed_counts_out = _emission_seed_buffers( + B, next_n, emit_block_meta, seed_thr, seed_counts_out) + else: + seed_thr = None + seed_counts_out = None + cand_cap = 0 + if emit_cand_bucketed: + assert emit_seed_counts, ( + "emit_cand_bucketed requires emit_seed_counts") + cand_cap = _emission_bucketed_cand(B, next_n, accept_cap, + cand_out, cand_idx_out, + cand_ctl_out, cand_cur_out) + else: + cand_out = None + cand_idx_out = None + cand_ctl_out = None + cand_cur_out = None + # Compile if needed (fake tensors, no real data required) key = (compute_block_kv, phys_block_kv, H, D, next_n, num_sms, - num_epi_subtiles, epi_dtype, acc_dtype, output_dtype) + num_epi_subtiles, epi_dtype, acc_dtype, output_dtype, + emit_block_meta, emit_seed_counts, seed_packed, + emit_cand_bucketed, accept_cap, cand_cap) if key not in cls.kernel_cache: - cls._compile(compute_block_kv, phys_block_kv, H, D, next_n, - num_sms, num_epi_subtiles, epi_dtype, acc_dtype, - output_dtype) + cls._compile(compute_block_kv, + phys_block_kv, + H, + D, + next_n, + num_sms, + num_epi_subtiles, + epi_dtype, + acc_dtype, + output_dtype, + emit_block_meta=emit_block_meta, + emit_seed_counts=emit_seed_counts, + seed_packed=seed_packed, + emit_cand_bucketed=emit_cand_bucketed, + accept_cap=accept_cap, + cand_cap=cand_cap) compiled = cls.kernel_cache[key] # FP8 q needs uint8 view to match compile-time dtype @@ -8235,10 +8463,19 @@ def forward( in (torch.float8_e4m3fn, torch.float8_e5m2) else q_3d) # TVM FFI: pass raw tensors, no dlpack/stream needed + if emit_block_meta: + compiled(kv_flat, q_for_ffi, w_2d, logits, block_table, + context_lens, schedule_meta, num_phys_blocks, B, + block_max_out, seed_thr, seed_counts_out, cand_out, + cand_ctl_out, cand_idx_out, cand_cur_out) + return logits, block_max_out compiled(kv_flat, q_for_ffi, w_2d, logits, block_table, - context_lens, schedule_meta, num_phys_blocks, B) + context_lens, schedule_meta, num_phys_blocks, B, None, + None, None, None, None, None, None) return logits + # NOTE: the optional emission tensors ARE written by the kernel but must + # stay out of mutates_args (torch.library IndexErrors on None defaults). @torch.library.custom_op("trtllm::cute_dsl_fp8_paged_mqa_logits", mutates_args=(), device_types="cuda") @@ -8254,6 +8491,13 @@ def cute_dsl_fp8_paged_mqa_logits( epi_dtype: torch.dtype = torch.float32, acc_dtype: torch.dtype = torch.float32, output_dtype: torch.dtype = torch.float32, + block_max_out: Optional[torch.Tensor] = None, + seed_thr: Optional[torch.Tensor] = None, + cand_out: Optional[torch.Tensor] = None, + cand_idx_out: Optional[torch.Tensor] = None, + cand_ctl_out: Optional[torch.Tensor] = None, + cand_cur_out: Optional[torch.Tensor] = None, + accept_cap: int = 8192, ) -> torch.Tensor: if not is_sm_100f(): raise ValueError( @@ -8274,7 +8518,7 @@ def cute_dsl_fp8_paged_mqa_logits( f"epi_dtype={epi_dtype} acc_dtype={acc_dtype} output_dtype={output_dtype}", key="cute_dsl_fp8_paged_mqa_logits_inputs", ) - return CuteDSLPagedMQALogitsRunner.forward( + ret = CuteDSLPagedMQALogitsRunner.forward( q, kv_fused, weights, @@ -8285,7 +8529,20 @@ def cute_dsl_fp8_paged_mqa_logits( num_epi_subtiles=num_epi_subtiles, epi_dtype=epi_dtype, acc_dtype=acc_dtype, - output_dtype=output_dtype) + output_dtype=output_dtype, + emit_block_meta=block_max_out is not None, + block_max_out=block_max_out, + emit_seed_counts=seed_thr is not None, + seed_thr=seed_thr, + emit_cand_bucketed=cand_out is not None, + accept_cap=accept_cap, + cand_out=cand_out, + cand_idx_out=cand_idx_out, + cand_ctl_out=cand_ctl_out, + cand_cur_out=cand_cur_out) + # with emission on the runner returns a tuple; the op returns + # logits only (emission buffers are caller-owned) + return ret[0] if isinstance(ret, tuple) else ret @torch.library.register_fake("trtllm::cute_dsl_fp8_paged_mqa_logits") def _( @@ -8300,6 +8557,13 @@ def _( epi_dtype: torch.dtype = torch.float32, acc_dtype: torch.dtype = torch.float32, output_dtype: torch.dtype = torch.float32, + block_max_out: Optional[torch.Tensor] = None, + seed_thr: Optional[torch.Tensor] = None, + cand_out: Optional[torch.Tensor] = None, + cand_idx_out: Optional[torch.Tensor] = None, + cand_ctl_out: Optional[torch.Tensor] = None, + cand_cur_out: Optional[torch.Tensor] = None, + accept_cap: int = 8192, ) -> torch.Tensor: B = q.shape[0] next_n = q.shape[1] @@ -9419,7 +9683,7 @@ def forward( if emit_block_meta: nb_pad = aligned_max_ctx // compute_block_kv # 4 warp-partial records per block (see FP4MQALogitsKernel). - nrec = nb_pad * 4 + nb_pad * 4 if emit_hit_stats: assert ( hit_bitmap is not None @@ -9443,50 +9707,13 @@ def forward( else: hit_bitmap = None hit_stats_out = None - if block_max_out is None: - block_max_out = torch.empty((B * next_n, nrec), - device=q.device, - dtype=torch.float32) - assert (block_max_out.shape == (B * next_n, nrec) - and block_max_out.is_contiguous()) + block_max_out = _emission_block_max(B, next_n, nb_pad, + block_max_out, q.device) seed_packed = False if emit_seed_counts: - assert emit_block_meta, ( - "emit_seed_counts requires emit_block_meta") - if seed_counts_out is None: - # Packed contract: seed_thr IS the [rows, 8] fp32 seed - # row; lines at cols 0..2, counts accumulate as fp32 at - # cols 3..5. Caller zeroes cols 3..7 and writes lines - # each step. - seed_packed = True - assert ( - seed_thr is not None and seed_thr.dtype == torch.float32 - and seed_thr.is_cuda and seed_thr.is_contiguous() - and seed_thr.shape == (B * next_n, 8) - ), (f"packed emit_seed_counts requires seed_thr fp32 " - f"[{B * next_n}, 8]; got " - f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}" - ) - seed_counts_out = seed_thr - else: - # Legacy split contract: 3 thresholds per row (fp32), - # counts accumulated with red.global.add.s32 into a - # caller-zeroed int32 [rows, 3]. - assert ( - seed_thr is not None and seed_thr.dtype == torch.float32 - and seed_thr.is_cuda and seed_thr.is_contiguous() - and seed_thr.shape == (B * next_n, 3) - ), (f"emit_seed_counts requires seed_thr fp32 " - f"[{B * next_n}, 3]; got " - f"{None if seed_thr is None else (seed_thr.dtype, tuple(seed_thr.shape))}" - ) - assert (seed_counts_out.dtype == torch.int32 - and seed_counts_out.is_cuda - and seed_counts_out.is_contiguous() - and seed_counts_out.shape == (B * next_n, 3)), ( - "emit_seed_counts requires a caller-zeroed " - "seed_counts_out int32 [B*next_n, 3]") + seed_packed, seed_counts_out = _emission_seed_buffers( + B, next_n, emit_block_meta, seed_thr, seed_counts_out) else: seed_thr = None seed_counts_out = None @@ -9515,40 +9742,9 @@ def forward( elif emit_cand_bucketed: assert emit_seed_counts, ( "emit_cand_bucketed requires emit_seed_counts") - # SoA contract: cand_out = fp32 VALUES [rows, 2*segA+capC], - # cand_idx_out = int32 positions (same width), cand_cur_out = - # int32 [rows, 4] cursors (caller-zeroed), cand_ctl_out = - # int32 [rows, 4] {n0, void, n1, n2} (caller-zeroed) - assert cand_out is not None, ( - "emit_cand_bucketed requires cand_out") - W = cand_out.shape[1] - assert ( - cand_out.dtype == torch.float32 and cand_out.is_cuda - and cand_out.is_contiguous() and cand_out.dim() == 2 - and cand_out.shape[0] == B * next_n - and W > 2 * accept_cap), ( - "bucketed requires cand_out fp32 [rows, 2*segA+capC]") - assert (cand_idx_out is not None - and cand_idx_out.dtype == torch.int32 - and cand_idx_out.is_cuda - and cand_idx_out.is_contiguous() - and cand_idx_out.shape == cand_out.shape), ( - "bucketed requires cand_idx_out int32, same shape") - assert (cand_ctl_out is not None - and cand_ctl_out.dtype == torch.int32 - and cand_ctl_out.is_cuda - and cand_ctl_out.is_contiguous() - and cand_ctl_out.shape == (B * next_n, 4)), ( - "bucketed requires caller-zeroed cand_ctl_out " - "int32 [rows, 4]") - assert (cand_cur_out is not None - and cand_cur_out.dtype == torch.int32 - and cand_cur_out.is_cuda - and cand_cur_out.is_contiguous() - and cand_cur_out.shape == (B * next_n, 4)), ( - "bucketed requires caller-zeroed cand_cur_out " - "int32 [rows, 4]") - cand_cap = W - 2 * accept_cap + cand_cap = _emission_bucketed_cand(B, next_n, accept_cap, + cand_out, cand_idx_out, + cand_ctl_out, cand_cur_out) else: cand_out = None cand_ctl_out = None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py index 89f32856d77c..2c289bb18c48 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py @@ -434,7 +434,301 @@ def utccp_required_smem_warp_transpose(smem_ptr) -> None: st_shared_b32(smem_ptr + offset, values[i]) -class FP4MQALogitsKernel: +class _PagedMQAEmissionMixin: + """GVR emission epilogue shared by the FP4/FP8 scoring kernels. + + Consumes the host-validated constexpr attrs next_n / seed_packed / + accept_cap / cand_cap / CAND_WIN; the two producers must stay + byte-identical to the GVR consumer's identities. + """ + + @cute.jit + def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None, cand_ctl=None): + """Warp-redux the lane-local seed counters and fire one lane-0 + red.global.add per (t, threshold). Caller zero-initializes the + count slots each step; cross-CTA totals accumulate atomically. + + seed_packed: mSeedCounts IS the [num_rows, 8] packed seed row - + counts land as fp32 at cols 3..5 (exact to 2^24).""" + next_n = cutlass.const_expr(self.next_n) + base_addr = mSeedCounts.iterator.toint() + for t in cutlass.range_constexpr(next_n): + for j in cutlass.range_constexpr(3): + w_cnt = cute.arch.warp_redux_sync(scnt[t * 3 + j], "add") + if meta_lane == cutlass.Int32(0): + row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + if cutlass.const_expr(self.seed_packed): + addr = base_addr + ( + cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(3 + j) + ) * cutlass.Int64(4) + _red_global_add_f32(addr, cutlass.Float32(w_cnt)) + else: + addr = base_addr + ( + cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j) + ) * cutlass.Int64(4) + _red_global_add_s32(addr, w_cnt) + if cutlass.const_expr(self.emit_cand_bucketed): + if j >= 1: + # consumer contract: ctl = {n0, void, n1, n2} + if meta_lane == cutlass.Int32(0): + row_b = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + ctl_a = cand_ctl.iterator.toint() + ( + cutlass.Int64(row_b) * cutlass.Int64(4) + cutlass.Int64(j + 1) + ) * cutlass.Int64(4) + _red_global_add_s32(ctl_a, w_cnt) + scnt[t * 3 + j] = cutlass.Int32(0) + if cutlass.const_expr(self.seed_packed and spass is not None): + # packed col 6: adaptive-skip pass count (lane0-accumulated) + for t in cutlass.range_constexpr(next_n): + w_bp = cute.arch.warp_redux_sync(spass[t], "add") + if meta_lane == cutlass.Int32(0): + row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + addr = base_addr + ( + cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(6) + ) * cutlass.Int64(4) + _red_global_add_f32(addr, cutlass.Float32(w_bp)) + spass[t] = cutlass.Int32(0) + + @cute.jit + def _flush_cand_window_bucketed(self, mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane): + """Sentinel-fill the unconsumed C-window tail in BOTH SoA columns + (score -inf, idx -1: the consumer pads by score) and invalidate + the window. Segment C sits at base 2*segA in each row.""" + next_n = cutlass.const_expr(self.next_n) + segA_f = cutlass.const_expr(self.accept_cap) + capC_f = cutlass.const_expr(self.cand_cap) + wtot_f = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) + vbase_f = mCand.iterator.toint() + ibase_f = mCandIdx.iterator.toint() + for t in cutlass.range_constexpr(next_n): + if cwleft[t] > cutlass.Int32(0): + row_f = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) + sl_f = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and sl_f < cutlass.Int32(capC_f): + off_f = ( + cutlass.Int64(row_f) * cutlass.Int64(wtot_f) + + cutlass.Int64(2 * segA_f + sl_f) + ) * cutlass.Int64(4) + vp_f = cute.make_ptr( + cutlass.Float32, + vbase_f + off_f, + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp_f, cute.make_layout((1,)))[0] = cutlass.Float32( + _META_NEG_FLT_MAX + ) + ip_f = cute.make_ptr( + cutlass.Int32, + ibase_f + off_f, + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip_f, cute.make_layout((1,)))[0] = cutlass.Int32(-1) + cwbase[t] = cutlass.Int32(0) + cwleft[t] = cutlass.Int32(0) + + @cute.jit + def _emit_cand_bucketed_step( + self, + mCand, + mCandIdx, + mCandCtl, + mCandCur, + q_idx, + t, + next_n, + r_bmax, + sthr, + f32_t, + kv_pos, + meta_valid, + meta_lane, + cwbase, + cwleft, + ): + """Bucketed SoA: A/B EXACT ballot claims + (their prefixes must stay pad-free for + the consumer's prefix math), C keeps the + claim-window; a full segment spills to + the next looser one. Every warp + collective sits at the TOP level of this + warp-uniform bound gate - no collectives + inside nested dynamic branches (DSL).""" + if r_bmax >= sthr[t * 3 + 0]: + segA_k = cutlass.const_expr(self.accept_cap) + capC_k = cutlass.const_expr(self.cand_cap) + wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) + row_k = q_idx * next_n + t + vb_k = mCand.iterator.toint() + cutlass.Int64(row_k) * cutlass.Int64( + wtot_k + ) * cutlass.Int64(4) + ib_k = mCandIdx.iterator.toint() + cutlass.Int64(row_k) * cutlass.Int64( + wtot_k + ) * cutlass.Int64(4) + cur_k = mCandCur.iterator.toint() + cutlass.Int64(row_k) * cutlass.Int64(16) + ctl_k = mCandCtl.iterator.toint() + cutlass.Int64(row_k) * cutlass.Int64(16) + lmk_k = (cutlass.Uint32(1) << cutlass.Uint32(meta_lane)) - cutlass.Uint32(1) + # exclusive class predicates + pA_k = cutlass.Int32(0) + pB_k = cutlass.Int32(0) + pC_k = cutlass.Int32(0) + if meta_valid: + if f32_t >= sthr[t * 3 + 2]: + pA_k = cutlass.Int32(1) + if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0): + pB_k = cutlass.Int32(1) + if ( + f32_t >= sthr[t * 3 + 0] + and pA_k == cutlass.Int32(0) + and pB_k == cutlass.Int32(0) + ): + pC_k = cutlass.Int32(1) + # Band gates on the warp-uniform block max: a block whose max + # sits below a line cannot contain a hit for that band (or a + # spill from a tighter one), so its ballot rounds are skipped + # wholesale. Collectives stay at the top level of warp-uniform + # gates, same safety class as the enclosing t0 gate. + spA_k = cutlass.Int32(0) + spB_k = cutlass.Int32(0) + if r_bmax >= sthr[t * 3 + 1]: + cntA_k = cutlass.Int32(0) + if r_bmax >= sthr[t * 3 + 2]: + # ---- A: exact claim ---- + mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0)) + cntA_k = cutlass.Int32(cute.arch.popc(mA_k)) + offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k)) + baseA_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0): + baseA_k = _atom_global_add_s32(cur_k, cntA_k) + baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0)) + slotA_k = baseA_k + offA_k + if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32(segA_k): + spA_k = cutlass.Int32(1) + if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k): + vp_k = cute.make_ptr( + cutlass.Float32, + vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t + ip_k = cute.make_ptr( + cutlass.Int32, + ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos + # ---- B: exact claim (native + A spill) ---- + pBe_k = cutlass.Int32(0) + if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0): + pBe_k = cutlass.Int32(1) + mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0)) + cntB_k = cutlass.Int32(cute.arch.popc(mB_k)) + offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k)) + baseB_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0): + baseB_k = _atom_global_add_s32(cur_k + cutlass.Int64(4), cntB_k) + baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0)) + slotB_k = baseB_k + offB_k + if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32(segA_k): + spB_k = cutlass.Int32(1) + if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32(segA_k): + vp2_k = cute.make_ptr( + cutlass.Float32, + vb_k + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t + ip2_k = cute.make_ptr( + cutlass.Int32, + ib_k + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos + # n0 += exact placements in A and B + plc_k = ( + cntA_k + - cutlass.Int32( + cute.arch.popc(cute.arch.vote_ballot_sync(spA_k != cutlass.Int32(0))) + ) + ) + ( + cntB_k + - cutlass.Int32( + cute.arch.popc(cute.arch.vote_ballot_sync(spB_k != cutlass.Int32(0))) + ) + ) + if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0): + _atom_global_add_s32(ctl_k, plc_k) + # ---- C: claim window (native + B spill) ---- + pCe_k = cutlass.Int32(0) + if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0): + pCe_k = cutlass.Int32(1) + mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0)) + cntC_k = cutlass.Int32(cute.arch.popc(mC_k)) + offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k)) + if cntC_k > cwleft[t]: + # sentinel-fill the old window tail + # (BOTH columns: the consumer pads + # by score -inf, idx -1) + slo_k = cwbase[t] + meta_lane + if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k): + vpo_k = cute.make_ptr( + cutlass.Float32, + vb_k + cutlass.Int64(2 * segA_k + slo_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = cutlass.Float32( + _META_NEG_FLT_MAX + ) + ipo_k = cute.make_ptr( + cutlass.Int32, + ib_k + cutlass.Int64(2 * segA_k + slo_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = cutlass.Int32(-1) + mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN) + nbC_k = cutlass.Int32(0) + if meta_lane == cutlass.Int32(0): + nbC_k = _atom_global_add_s32(cur_k + cutlass.Int64(8), mC2_k) + _atom_global_add_s32(ctl_k, mC2_k) + if nbC_k + mC2_k > cutlass.Int32(capC_k) and nbC_k <= cutlass.Int32(capC_k): + vdp_k = cute.make_ptr( + cutlass.Int32, + ctl_k + cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vdp_k, cute.make_layout((1,)))[0] = cutlass.Int32(1) + nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0)) + cwbase[t] = nbC_k + cwleft[t] = mC2_k + slotC_k = cwbase[t] + offC_k + if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32(capC_k): + vpc_k = cute.make_ptr( + cutlass.Float32, + vb_k + cutlass.Int64(2 * segA_k + slotC_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t + ipc_k = cute.make_ptr( + cutlass.Int32, + ib_k + cutlass.Int64(2 * segA_k + slotC_k) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos + cwbase[t] = cwbase[t] + cntC_k + cwleft[t] = cwleft[t] - cntC_k + + +class FP4MQALogitsKernel(_PagedMQAEmissionMixin): """FP4 (MXFP4) paged MQA logits kernel for Blackwell (SM100). Each CTA processes a range of (q_idx, kv_split) pairs. @@ -1075,92 +1369,6 @@ def _flush_hit_agg( hacc_sum[t] = cutlass.Float32(0.0) hacc_cnt[t] = cutlass.Int32(0) - @cute.jit - def _flush_seed_counts(self, mSeedCounts, q_idx, scnt, meta_lane, spass=None, cand_ctl=None): - """Warp-redux the lane-local seed counters and fire one lane-0 - red.global.add per (t, threshold). Caller zero-initializes the - count slots each step; cross-CTA totals accumulate atomically. - - seed_packed: mSeedCounts IS the [num_rows, 8] packed seed row - - counts land as fp32 at cols 3..5 (exact to 2^24).""" - next_n = cutlass.const_expr(self.next_n) - base_addr = mSeedCounts.iterator.toint() - for t in cutlass.range_constexpr(next_n): - for j in cutlass.range_constexpr(3): - w_cnt = cute.arch.warp_redux_sync(scnt[t * 3 + j], "add") - if meta_lane == cutlass.Int32(0): - row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) - if cutlass.const_expr(self.seed_packed): - addr = base_addr + ( - cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(3 + j) - ) * cutlass.Int64(4) - _red_global_add_f32(addr, cutlass.Float32(w_cnt)) - else: - addr = base_addr + ( - cutlass.Int64(row) * cutlass.Int64(3) + cutlass.Int64(j) - ) * cutlass.Int64(4) - _red_global_add_s32(addr, w_cnt) - if cutlass.const_expr(self.emit_cand_bucketed): - if j >= 1: - # consumer contract: ctl = {n0, void, n1, n2} - if meta_lane == cutlass.Int32(0): - row_b = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) - ctl_a = cand_ctl.iterator.toint() + ( - cutlass.Int64(row_b) * cutlass.Int64(4) + cutlass.Int64(j + 1) - ) * cutlass.Int64(4) - _red_global_add_s32(ctl_a, w_cnt) - scnt[t * 3 + j] = cutlass.Int32(0) - if cutlass.const_expr(self.seed_packed and spass is not None): - # packed col 6: adaptive-skip pass count (lane0-accumulated) - for t in cutlass.range_constexpr(next_n): - w_bp = cute.arch.warp_redux_sync(spass[t], "add") - if meta_lane == cutlass.Int32(0): - row = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) - addr = base_addr + ( - cutlass.Int64(row) * cutlass.Int64(8) + cutlass.Int64(6) - ) * cutlass.Int64(4) - _red_global_add_f32(addr, cutlass.Float32(w_bp)) - spass[t] = cutlass.Int32(0) - - @cute.jit - def _flush_cand_window_bucketed(self, mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane): - """Sentinel-fill the unconsumed C-window tail in BOTH SoA columns - (score -inf, idx -1: the consumer pads by score) and invalidate - the window. Segment C sits at base 2*segA in each row.""" - next_n = cutlass.const_expr(self.next_n) - segA_f = cutlass.const_expr(self.accept_cap) - capC_f = cutlass.const_expr(self.cand_cap) - wtot_f = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) - vbase_f = mCand.iterator.toint() - ibase_f = mCandIdx.iterator.toint() - for t in cutlass.range_constexpr(next_n): - if cwleft[t] > cutlass.Int32(0): - row_f = q_idx * cutlass.Int32(next_n) + cutlass.Int32(t) - sl_f = cwbase[t] + meta_lane - if meta_lane < cwleft[t] and sl_f < cutlass.Int32(capC_f): - off_f = ( - cutlass.Int64(row_f) * cutlass.Int64(wtot_f) - + cutlass.Int64(2 * segA_f + sl_f) - ) * cutlass.Int64(4) - vp_f = cute.make_ptr( - cutlass.Float32, - vbase_f + off_f, - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vp_f, cute.make_layout((1,)))[0] = cutlass.Float32( - _META_NEG_FLT_MAX - ) - ip_f = cute.make_ptr( - cutlass.Int32, - ibase_f + off_f, - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ip_f, cute.make_layout((1,)))[0] = cutlass.Int32(-1) - cwbase[t] = cutlass.Int32(0) - cwleft[t] = cutlass.Int32(0) - @cute.jit def _flush_cand_window(self, mCand, q_idx, cwbase, cwleft, meta_lane): """Sentinel-fill the unconsumed tail of each per-(warp, t) claim @@ -2644,219 +2852,23 @@ def kernel( cwbase[t] = cwbase[t] + cnt_c cwleft[t] = cwleft[t] - cnt_c if cutlass.const_expr(self.emit_cand_bucketed): - # Bucketed SoA: A/B EXACT ballot claims - # (their prefixes must stay pad-free for - # the consumer's prefix math), C keeps the - # claim-window; a full segment spills to - # the next looser one. Every warp - # collective sits at the TOP level of this - # warp-uniform bound gate - no collectives - # inside nested dynamic branches (DSL). - if r_bmax >= sthr[t * 3 + 0]: - segA_k = cutlass.const_expr(self.accept_cap) - capC_k = cutlass.const_expr(self.cand_cap) - wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) - row_k = q_idx * next_n + t - vb_k = mCand.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) - ib_k = mCandIdx.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) - cur_k = mCandCur.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(16) - ctl_k = mCandCtl.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(16) - lmk_k = ( - cutlass.Uint32(1) << cutlass.Uint32(meta_lane) - ) - cutlass.Uint32(1) - # exclusive class predicates - pA_k = cutlass.Int32(0) - pB_k = cutlass.Int32(0) - pC_k = cutlass.Int32(0) - if meta_valid: - if f32_t >= sthr[t * 3 + 2]: - pA_k = cutlass.Int32(1) - if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0): - pB_k = cutlass.Int32(1) - if ( - f32_t >= sthr[t * 3 + 0] - and pA_k == cutlass.Int32(0) - and pB_k == cutlass.Int32(0) - ): - pC_k = cutlass.Int32(1) - # ---- A: exact claim ---- - mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0)) - cntA_k = cutlass.Int32(cute.arch.popc(mA_k)) - offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k)) - baseA_k = cutlass.Int32(0) - if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0): - baseA_k = _atom_global_add_s32(cur_k, cntA_k) - baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0)) - slotA_k = baseA_k + offA_k - spA_k = cutlass.Int32(0) - if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32( - segA_k - ): - spA_k = cutlass.Int32(1) - if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k): - vp_k = cute.make_ptr( - cutlass.Float32, - vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t - ip_k = cute.make_ptr( - cutlass.Int32, - ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos - # ---- B: exact claim (native + A spill) ---- - pBe_k = cutlass.Int32(0) - if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0): - pBe_k = cutlass.Int32(1) - mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0)) - cntB_k = cutlass.Int32(cute.arch.popc(mB_k)) - offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k)) - baseB_k = cutlass.Int32(0) - if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0): - baseB_k = _atom_global_add_s32( - cur_k + cutlass.Int64(4), cntB_k - ) - baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0)) - slotB_k = baseB_k + offB_k - spB_k = cutlass.Int32(0) - if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32( - segA_k - ): - spB_k = cutlass.Int32(1) - if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32( - segA_k - ): - vp2_k = cute.make_ptr( - cutlass.Float32, - vb_k - + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t - ip2_k = cute.make_ptr( - cutlass.Int32, - ib_k - + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos - # n0 += exact placements in A and B - plc_k = ( - cntA_k - - cutlass.Int32( - cute.arch.popc( - cute.arch.vote_ballot_sync( - spA_k != cutlass.Int32(0) - ) - ) - ) - ) + ( - cntB_k - - cutlass.Int32( - cute.arch.popc( - cute.arch.vote_ballot_sync( - spB_k != cutlass.Int32(0) - ) - ) - ) - ) - if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0): - _atom_global_add_s32(ctl_k, plc_k) - # ---- C: claim window (native + B spill) ---- - pCe_k = cutlass.Int32(0) - if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0): - pCe_k = cutlass.Int32(1) - mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0)) - cntC_k = cutlass.Int32(cute.arch.popc(mC_k)) - offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k)) - if cntC_k > cwleft[t]: - # sentinel-fill the old window tail - # (BOTH columns: the consumer pads - # by score -inf, idx -1) - slo_k = cwbase[t] + meta_lane - if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k): - vpo_k = cute.make_ptr( - cutlass.Float32, - vb_k - + cutlass.Int64(2 * segA_k + slo_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = ( - cutlass.Float32(_META_NEG_FLT_MAX) - ) - ipo_k = cute.make_ptr( - cutlass.Int32, - ib_k - + cutlass.Int64(2 * segA_k + slo_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = ( - cutlass.Int32(-1) - ) - mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN) - nbC_k = cutlass.Int32(0) - if meta_lane == cutlass.Int32(0): - nbC_k = _atom_global_add_s32( - cur_k + cutlass.Int64(8), mC2_k - ) - _atom_global_add_s32(ctl_k, mC2_k) - if nbC_k + mC2_k > cutlass.Int32( - capC_k - ) and nbC_k <= cutlass.Int32(capC_k): - vdp_k = cute.make_ptr( - cutlass.Int32, - ctl_k + cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vdp_k, cute.make_layout((1,)))[ - 0 - ] = cutlass.Int32(1) - nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0)) - cwbase[t] = nbC_k - cwleft[t] = mC2_k - slotC_k = cwbase[t] + offC_k - if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32( - capC_k - ): - vpc_k = cute.make_ptr( - cutlass.Float32, - vb_k - + cutlass.Int64(2 * segA_k + slotC_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t - ipc_k = cute.make_ptr( - cutlass.Int32, - ib_k - + cutlass.Int64(2 * segA_k + slotC_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos - cwbase[t] = cwbase[t] + cntC_k - cwleft[t] = cwleft[t] - cntC_k + self._emit_cand_bucketed_step( + mCand, + mCandIdx, + mCandCtl, + mCandCur, + q_idx, + t, + next_n, + r_bmax, + sthr, + f32_t, + kv_pos, + meta_valid, + meta_lane, + cwbase, + cwleft, + ) if cutlass.const_expr(self.emit_hit_stats): # Lane-local accumulation; keep it branchless # (`if meta_hit` compiles to real divergent @@ -3411,219 +3423,23 @@ def kernel( cwbase[t] = cwbase[t] + cnt_c cwleft[t] = cwleft[t] - cnt_c if cutlass.const_expr(self.emit_cand_bucketed): - # Bucketed SoA: A/B EXACT ballot claims - # (their prefixes must stay pad-free for - # the consumer's prefix math), C keeps the - # claim-window; a full segment spills to - # the next looser one. Every warp - # collective sits at the TOP level of this - # warp-uniform bound gate - no collectives - # inside nested dynamic branches (DSL). - if r_bmax >= sthr[t * 3 + 0]: - segA_k = cutlass.const_expr(self.accept_cap) - capC_k = cutlass.const_expr(self.cand_cap) - wtot_k = cutlass.const_expr(2 * self.accept_cap + self.cand_cap) - row_k = q_idx * next_n + t - vb_k = mCand.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) - ib_k = mCandIdx.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(wtot_k) * cutlass.Int64(4) - cur_k = mCandCur.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(16) - ctl_k = mCandCtl.iterator.toint() + cutlass.Int64( - row_k - ) * cutlass.Int64(16) - lmk_k = ( - cutlass.Uint32(1) << cutlass.Uint32(meta_lane) - ) - cutlass.Uint32(1) - # exclusive class predicates - pA_k = cutlass.Int32(0) - pB_k = cutlass.Int32(0) - pC_k = cutlass.Int32(0) - if meta_valid: - if f32_t >= sthr[t * 3 + 2]: - pA_k = cutlass.Int32(1) - if f32_t >= sthr[t * 3 + 1] and pA_k == cutlass.Int32(0): - pB_k = cutlass.Int32(1) - if ( - f32_t >= sthr[t * 3 + 0] - and pA_k == cutlass.Int32(0) - and pB_k == cutlass.Int32(0) - ): - pC_k = cutlass.Int32(1) - # ---- A: exact claim ---- - mA_k = cute.arch.vote_ballot_sync(pA_k != cutlass.Int32(0)) - cntA_k = cutlass.Int32(cute.arch.popc(mA_k)) - offA_k = cutlass.Int32(cute.arch.popc(mA_k & lmk_k)) - baseA_k = cutlass.Int32(0) - if meta_lane == cutlass.Int32(0) and cntA_k > cutlass.Int32(0): - baseA_k = _atom_global_add_s32(cur_k, cntA_k) - baseA_k = cute.arch.shuffle_sync(baseA_k, cutlass.Int32(0)) - slotA_k = baseA_k + offA_k - spA_k = cutlass.Int32(0) - if pA_k != cutlass.Int32(0) and slotA_k >= cutlass.Int32( - segA_k - ): - spA_k = cutlass.Int32(1) - if pA_k != cutlass.Int32(0) and slotA_k < cutlass.Int32(segA_k): - vp_k = cute.make_ptr( - cutlass.Float32, - vb_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vp_k, cute.make_layout((1,)))[0] = f32_t - ip_k = cute.make_ptr( - cutlass.Int32, - ib_k + cutlass.Int64(slotA_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ip_k, cute.make_layout((1,)))[0] = kv_pos - # ---- B: exact claim (native + A spill) ---- - pBe_k = cutlass.Int32(0) - if pB_k != cutlass.Int32(0) or spA_k != cutlass.Int32(0): - pBe_k = cutlass.Int32(1) - mB_k = cute.arch.vote_ballot_sync(pBe_k != cutlass.Int32(0)) - cntB_k = cutlass.Int32(cute.arch.popc(mB_k)) - offB_k = cutlass.Int32(cute.arch.popc(mB_k & lmk_k)) - baseB_k = cutlass.Int32(0) - if meta_lane == cutlass.Int32(0) and cntB_k > cutlass.Int32(0): - baseB_k = _atom_global_add_s32( - cur_k + cutlass.Int64(4), cntB_k - ) - baseB_k = cute.arch.shuffle_sync(baseB_k, cutlass.Int32(0)) - slotB_k = baseB_k + offB_k - spB_k = cutlass.Int32(0) - if pBe_k != cutlass.Int32(0) and slotB_k >= cutlass.Int32( - segA_k - ): - spB_k = cutlass.Int32(1) - if pBe_k != cutlass.Int32(0) and slotB_k < cutlass.Int32( - segA_k - ): - vp2_k = cute.make_ptr( - cutlass.Float32, - vb_k - + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vp2_k, cute.make_layout((1,)))[0] = f32_t - ip2_k = cute.make_ptr( - cutlass.Int32, - ib_k - + cutlass.Int64(segA_k + slotB_k) * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ip2_k, cute.make_layout((1,)))[0] = kv_pos - # n0 += exact placements in A and B - plc_k = ( - cntA_k - - cutlass.Int32( - cute.arch.popc( - cute.arch.vote_ballot_sync( - spA_k != cutlass.Int32(0) - ) - ) - ) - ) + ( - cntB_k - - cutlass.Int32( - cute.arch.popc( - cute.arch.vote_ballot_sync( - spB_k != cutlass.Int32(0) - ) - ) - ) - ) - if meta_lane == cutlass.Int32(0) and plc_k > cutlass.Int32(0): - _atom_global_add_s32(ctl_k, plc_k) - # ---- C: claim window (native + B spill) ---- - pCe_k = cutlass.Int32(0) - if pC_k != cutlass.Int32(0) or spB_k != cutlass.Int32(0): - pCe_k = cutlass.Int32(1) - mC_k = cute.arch.vote_ballot_sync(pCe_k != cutlass.Int32(0)) - cntC_k = cutlass.Int32(cute.arch.popc(mC_k)) - offC_k = cutlass.Int32(cute.arch.popc(mC_k & lmk_k)) - if cntC_k > cwleft[t]: - # sentinel-fill the old window tail - # (BOTH columns: the consumer pads - # by score -inf, idx -1) - slo_k = cwbase[t] + meta_lane - if meta_lane < cwleft[t] and slo_k < cutlass.Int32(capC_k): - vpo_k = cute.make_ptr( - cutlass.Float32, - vb_k - + cutlass.Int64(2 * segA_k + slo_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vpo_k, cute.make_layout((1,)))[0] = ( - cutlass.Float32(_META_NEG_FLT_MAX) - ) - ipo_k = cute.make_ptr( - cutlass.Int32, - ib_k - + cutlass.Int64(2 * segA_k + slo_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ipo_k, cute.make_layout((1,)))[0] = ( - cutlass.Int32(-1) - ) - mC2_k = cntC_k + cutlass.Int32(self.CAND_WIN) - nbC_k = cutlass.Int32(0) - if meta_lane == cutlass.Int32(0): - nbC_k = _atom_global_add_s32( - cur_k + cutlass.Int64(8), mC2_k - ) - _atom_global_add_s32(ctl_k, mC2_k) - if nbC_k + mC2_k > cutlass.Int32( - capC_k - ) and nbC_k <= cutlass.Int32(capC_k): - vdp_k = cute.make_ptr( - cutlass.Int32, - ctl_k + cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vdp_k, cute.make_layout((1,)))[ - 0 - ] = cutlass.Int32(1) - nbC_k = cute.arch.shuffle_sync(nbC_k, cutlass.Int32(0)) - cwbase[t] = nbC_k - cwleft[t] = mC2_k - slotC_k = cwbase[t] + offC_k - if pCe_k != cutlass.Int32(0) and slotC_k < cutlass.Int32( - capC_k - ): - vpc_k = cute.make_ptr( - cutlass.Float32, - vb_k - + cutlass.Int64(2 * segA_k + slotC_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(vpc_k, cute.make_layout((1,)))[0] = f32_t - ipc_k = cute.make_ptr( - cutlass.Int32, - ib_k - + cutlass.Int64(2 * segA_k + slotC_k) - * cutlass.Int64(4), - cute.AddressSpace.gmem, - assumed_align=4, - ) - cute.make_tensor(ipc_k, cute.make_layout((1,)))[0] = kv_pos - cwbase[t] = cwbase[t] + cntC_k - cwleft[t] = cwleft[t] - cntC_k + self._emit_cand_bucketed_step( + mCand, + mCandIdx, + mCandCtl, + mCandCur, + q_idx, + t, + next_n, + r_bmax, + sthr, + f32_t, + kv_pos, + meta_valid, + meta_lane, + cwbase, + cwleft, + ) if cutlass.const_expr(self.emit_hit_stats): # Lane-local accumulation; keep it branchless # (`if meta_hit` compiles to real divergent 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 722ac73cebd0..d8acc0bb1a0f 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 @@ -69,6 +69,15 @@ from cutlass.cutlass_dsl import dsl_user_op from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +# Emission shares the FP4 kernel's gmem-reduction helpers and the GVR +# sentinel constants (single definition; the two producers must stay +# byte-identical to the consumer's identities). +from .fp4_paged_mqa_logits import ( # noqa: E501 isort: skip + _META_FLT_MAX, + _PagedMQAEmissionMixin, + _META_NEG_FLT_MAX, +) + # CuTe DSL CUDA 13 validates rounding modes as string literals. The string # form is also accepted by older wrappers, so keep it version-independent. _RND_RN = "rn" @@ -177,7 +186,7 @@ def add_f16x2( ) -class FP8MQALogitsKernel: +class FP8MQALogitsKernel(_PagedMQAEmissionMixin): """FP8 paged MQA logits kernel for Blackwell (SM100). Each CTA processes a range of (q_idx, kv_split) pairs. @@ -202,6 +211,12 @@ def __init__( epi_dtype=cutlass.Float32, acc_dtype=cutlass.Float32, output_dtype=cutlass.Float32, + emit_block_meta: bool = False, + emit_seed_counts: bool = False, + seed_packed: bool = False, + emit_cand_bucketed: bool = False, + accept_cap: int = 8192, + cand_cap: int = 5120, ): self.block_kv = block_kv self.phys_block_kv = phys_block_kv @@ -291,6 +306,33 @@ def __init__( self.cluster_shape_mn = (1, 1) self.mma_tiler_mn = (block_kv, self.N) + # GVR emission (same contract as the FP4 kernel; hit-stats and the + # plain-list variants are FP4-only — production wiring never passes + # them). Metadata is computed on the POST-conversion stored logit + # (f32(output_dtype(result * scale))) so it bounds what the GVR + # consumer reads back bit-exactly. + # block_max [num_rows, nb_pad*4] fp32 — warp-partial max per + # (128-token tile, warp); valid positions only (kv_pos < ctx). + # seed row [num_rows, 8] fp32 (seed_packed) — lines at cols + # 0..2, counts accumulated at cols 3..5, skip-pass count col 6. + # candidate list — bucketed SoA segments A/B/C (see the FP4 + # kernel's emit_cand_bucketed contract). + if emit_block_meta and block_kv != 128: + raise ValueError("emission's tile*4+warp record layout assumes block_kv == 128") + self.emit_block_meta = emit_block_meta + if emit_seed_counts and not emit_block_meta: + raise ValueError("emit_seed_counts requires emit_block_meta") + self.emit_seed_counts = emit_seed_counts + if seed_packed and not emit_seed_counts: + raise ValueError("seed_packed requires emit_seed_counts") + self.seed_packed = seed_packed + if emit_cand_bucketed and not emit_seed_counts: + raise ValueError("emit_cand_bucketed requires emit_seed_counts") + self.emit_cand_bucketed = emit_cand_bucketed + self.accept_cap = accept_cap + self.cand_cap = cand_cap + self.CAND_WIN = 8 + def _setup_mma(self, a_dtype, b_dtype, a_major, b_major): self.a_dtype = a_dtype self.b_dtype = b_dtype @@ -365,6 +407,15 @@ def __call__( num_phys_blocks: cutlass.Int32, batch_size: cutlass.Int32, stream: cuda.CUstream, + # emission-only tensors; defaulted so the positional signature + # stays the one callers already use + block_max: cute.Tensor = None, # [num_rows, nb_pad*4] fp32 warp-partials + seed_thr: cute.Tensor = None, # [num_rows, 3] fp32 / [num_rows, 8] packed + seed_counts: cute.Tensor = None, # [num_rows, 3] int32 / [num_rows, 8] fp32 packed + cand: cute.Tensor = None, # bucketed: [num_rows, 2*segA+capC] fp32 SoA values + cand_ctl: cute.Tensor = None, # [num_rows, 4] int32 {n0, void, n1, n2}, zeroed + cand_idx_t: cute.Tensor = None, # bucketed: [num_rows, 2*segA+capC] int32 SoA + cand_cur: cute.Tensor = None, # bucketed: [num_rows, 4] int32 cursors, zeroed ): # Derive KV and Scale views from fused buffer using CuTE ops. # Fused layout per physical block: [KV data (phys_block_kv*head_dim)] [Scales (phys_block_kv*4)] @@ -511,6 +562,13 @@ class SharedStorage: context_lens, schedule_meta, batch_size, + block_max, + seed_thr, + seed_counts, + cand, + cand_ctl, + cand_idx_t, + cand_cur, self.cluster_layout_vmnk, self.a_smem_layout_staged, self.b_smem_layout_staged, @@ -543,6 +601,13 @@ def kernel( mContextLens: cute.Tensor, # [batch_size] mScheduleMeta: cute.Tensor, # [num_sms+1, 2] int32 batch_size: cutlass.Int32, + mBlockMax: cute.Tensor, # [num_rows, nb_pad*4] fp32 (emit_block_meta) + mSeedThr: cute.Tensor, # [num_rows, 3] fp32 / [num_rows, 8] packed + mSeedCounts: cute.Tensor, # [num_rows, 3] int32 / [num_rows, 8] fp32 packed + mCand: cute.Tensor, # bucketed SoA values (emit_cand_bucketed) + mCandCtl: cute.Tensor, # [num_rows, 4] int32 {n0, void, n1, n2} + mCandIdx: cute.Tensor, # bucketed SoA indices + mCandCur: cute.Tensor, # [num_rows, 4] int32 cursors cluster_layout_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, b_smem_layout_staged: cute.ComposedLayout, @@ -1282,9 +1347,32 @@ def kernel( 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 + if cutlass.const_expr(self.emit_block_meta): + # free registers for the meta accumulators; the weight + # cache sits at the spill edge + MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) q_stage_local = cutlass.Int32(0) + if cutlass.const_expr(self.emit_block_meta): + ctx_cur = cutlass.Int32(0) + meta_warp = local_tidx // 32 + meta_lane = local_tidx % 32 + if cutlass.const_expr(self.emit_seed_counts): + sthr = cute.make_fragment(next_n * 3, cutlass.Float32) + scnt = cute.make_fragment(next_n * 3, cutlass.Int32) + spass = cute.make_fragment(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + spass[_i] = cutlass.Int32(0) + for _i in cutlass.range_constexpr(next_n * 3): + sthr[_i] = cutlass.Float32(_META_FLT_MAX) + scnt[_i] = cutlass.Int32(0) + if cutlass.const_expr(self.emit_cand_bucketed): + cwbase = cute.make_fragment(next_n, cutlass.Int32) + cwleft = cute.make_fragment(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + cwbase[_i] = cutlass.Int32(0) + cwleft[_i] = cutlass.Int32(0) while has_work: # fetch_next_task: commit next → current @@ -1306,11 +1394,43 @@ def kernel( w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ (t_i * num_heads + w_j, q_stage_local) ] + if cutlass.const_expr(self.emit_block_meta): + if cutlass.const_expr(self.emit_seed_counts): + # Flush the PREVIOUS request's counters + # before switching context. + if q_idx_old < batch_size: + self._flush_seed_counts( + mSeedCounts, + q_idx_old, + scnt, + meta_lane, + spass=spass, + cand_ctl=mCandCtl, + ) + # (re)load this q's thresholds (a stale + # FLT_MAX default zeroes every counter) + for _t in cutlass.range_constexpr(next_n): + for _j in cutlass.range_constexpr(3): + sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)] + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx_old < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane + ) + # Context len for the meta valid mask + # (kv_pos < ctx_cur): the FP8 epilogue runs + # unconditionally over the last partial KV + # block, so aligned-padding GEMM garbage must + # be kept out of every emitted statistic. + ctx_cur = mContextLens[q_idx] # 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.emit_block_meta): + meta_kv_tile = kv_idx + meta_valid = kv_pos < ctx_cur if cutlass.const_expr(self.remove_kv_wait_in_epilogue): # Skip KV wait, rely on UMMA barrier's @@ -1473,11 +1593,60 @@ def kernel( 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) - ) + stored_t = self.output_dtype(result_t * Float16(scale_val)) else: - mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + stored_t = self.output_dtype(result_t * scale_val) + mLogits[(out_row, kv_pos)] = stored_t + if cutlass.const_expr(self.emit_block_meta): + # Meta reduction on the POST-conversion value so + # block_max bounds what GVR reads back bit-exactly. + f32_t = cutlass.Float32(stored_t) + bmax_v = cutlass.Float32(_META_NEG_FLT_MAX) + if meta_valid: + bmax_v = f32_t + r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax") + # Warp-autonomous store: record index = + # tile*4 + warp; the GVR consumer folds the + # 4 warp-partials per block. + if meta_lane == cutlass.Int32(0): + out_row_m = q_idx * next_n + t + rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp + mBlockMax[(out_row_m, rec_m)] = r_bmax + if cutlass.const_expr(self.emit_seed_counts): + # Seed-count accumulation: branchless 0/1 + # adds on the post-conversion value; the + # valid mask keeps aligned-padding garbage + # out (same contract as block_max). + valid_i1 = cutlass.Int32(meta_valid) + for _j in cutlass.range_constexpr(3): + ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j]) + scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1) + if cutlass.const_expr(self.seed_packed): + # adaptive-skip pass count: one record + # per (tile, warp); r_bmax is warp- + # uniform so lane0 alone accumulates + if meta_lane == cutlass.Int32(0): + spass[t] = spass[t] + cutlass.Int32( + r_bmax >= sthr[t * 3 + 0] + ) + if cutlass.const_expr(self.emit_cand_bucketed): + self._emit_cand_bucketed_step( + mCand, + mCandIdx, + mCandCtl, + mCandCur, + q_idx, + t, + next_n, + r_bmax, + sthr, + f32_t, + kv_pos, + meta_valid, + meta_lane, + cwbase, + cwleft, + ) # Advance: inline fetch_next_task next_kv_idx = kv_idx + NUM_MATH_WG @@ -1491,6 +1660,18 @@ def kernel( # Update while-loop condition has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + # Flush the final request's emission state (WG 0). + if cutlass.const_expr(self.emit_seed_counts): + if q_idx < batch_size: + self._flush_seed_counts( + mSeedCounts, q_idx, scnt, meta_lane, spass=spass, cand_ctl=mCandCtl + ) + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane + ) + # Release last Q stage (WG 0) if q_idx < batch_size: q_pipeline.consumer_release(q_cons_state) @@ -1517,9 +1698,32 @@ def kernel( 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 + if cutlass.const_expr(self.emit_block_meta): + # free registers for the meta accumulators; the weight + # cache sits at the spill edge + MAX_NUM_W_IN_REG = MAX_NUM_W_IN_REG - 8 NUM_W_IN_REG = min(MAX_NUM_W_IN_REG, num_heads) w_cache = cute.make_rmem_tensor(NUM_W_IN_REG * next_n, self.epi_dtype) q_stage_local = cutlass.Int32(0) + if cutlass.const_expr(self.emit_block_meta): + ctx_cur = cutlass.Int32(0) + meta_warp = local_tidx // 32 + meta_lane = local_tidx % 32 + if cutlass.const_expr(self.emit_seed_counts): + sthr = cute.make_fragment(next_n * 3, cutlass.Float32) + scnt = cute.make_fragment(next_n * 3, cutlass.Int32) + spass = cute.make_fragment(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + spass[_i] = cutlass.Int32(0) + for _i in cutlass.range_constexpr(next_n * 3): + sthr[_i] = cutlass.Float32(_META_FLT_MAX) + scnt[_i] = cutlass.Int32(0) + if cutlass.const_expr(self.emit_cand_bucketed): + cwbase = cute.make_fragment(next_n, cutlass.Int32) + cwleft = cute.make_fragment(next_n, cutlass.Int32) + for _i in cutlass.range_constexpr(next_n): + cwbase[_i] = cutlass.Int32(0) + cwleft[_i] = cutlass.Int32(0) while has_work: # fetch_next_task: commit next → current @@ -1541,12 +1745,44 @@ def kernel( w_cache[t_i * NUM_W_IN_REG + w_j] = sW[ (t_i * num_heads + w_j, q_stage_local) ] + if cutlass.const_expr(self.emit_block_meta): + if cutlass.const_expr(self.emit_seed_counts): + # Flush the PREVIOUS request's counters + # before switching context. + if q_idx_old < batch_size: + self._flush_seed_counts( + mSeedCounts, + q_idx_old, + scnt, + meta_lane, + spass=spass, + cand_ctl=mCandCtl, + ) + # (re)load this q's thresholds (a stale + # FLT_MAX default zeroes every counter) + for _t in cutlass.range_constexpr(next_n): + for _j in cutlass.range_constexpr(3): + sthr[_t * 3 + _j] = mSeedThr[(q_idx * next_n + _t, _j)] + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx_old < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx_old, cwbase, cwleft, meta_lane + ) + # Context len for the meta valid mask + # (kv_pos < ctx_cur): the FP8 epilogue runs + # unconditionally over the last partial KV + # block, so aligned-padding GEMM garbage must + # be kept out of every emitted statistic. + ctx_cur = mContextLens[q_idx] # 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.emit_block_meta): + meta_kv_tile = kv_idx_1 + meta_valid = kv_pos < ctx_cur if cutlass.const_expr(self.remove_kv_wait_in_epilogue): umma_pipeline_1.consumer_wait(umma_cons_state_1) @@ -1689,11 +1925,60 @@ def kernel( 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) - ) + stored_t = self.output_dtype(result_t * Float16(scale_val)) else: - mLogits[(out_row, kv_pos)] = self.output_dtype(result_t * scale_val) + stored_t = self.output_dtype(result_t * scale_val) + mLogits[(out_row, kv_pos)] = stored_t + if cutlass.const_expr(self.emit_block_meta): + # Meta reduction on the POST-conversion value so + # block_max bounds what GVR reads back bit-exactly. + f32_t = cutlass.Float32(stored_t) + bmax_v = cutlass.Float32(_META_NEG_FLT_MAX) + if meta_valid: + bmax_v = f32_t + r_bmax = cute.arch.warp_redux_sync(bmax_v, "fmax") + # Warp-autonomous store: record index = + # tile*4 + warp; the GVR consumer folds the + # 4 warp-partials per block. + if meta_lane == cutlass.Int32(0): + out_row_m = q_idx * next_n + t + rec_m = meta_kv_tile * cutlass.Int32(4) + meta_warp + mBlockMax[(out_row_m, rec_m)] = r_bmax + if cutlass.const_expr(self.emit_seed_counts): + # Seed-count accumulation: branchless 0/1 + # adds on the post-conversion value; the + # valid mask keeps aligned-padding garbage + # out (same contract as block_max). + valid_i1 = cutlass.Int32(meta_valid) + for _j in cutlass.range_constexpr(3): + ge_j = cutlass.Int32(f32_t >= sthr[t * 3 + _j]) + scnt[t * 3 + _j] = scnt[t * 3 + _j] + (ge_j & valid_i1) + if cutlass.const_expr(self.seed_packed): + # adaptive-skip pass count: one record + # per (tile, warp); r_bmax is warp- + # uniform so lane0 alone accumulates + if meta_lane == cutlass.Int32(0): + spass[t] = spass[t] + cutlass.Int32( + r_bmax >= sthr[t * 3 + 0] + ) + if cutlass.const_expr(self.emit_cand_bucketed): + self._emit_cand_bucketed_step( + mCand, + mCandIdx, + mCandCtl, + mCandCur, + q_idx, + t, + next_n, + r_bmax, + sthr, + f32_t, + kv_pos, + meta_valid, + meta_lane, + cwbase, + cwleft, + ) # Advance: inline fetch_next_task next_kv_idx = kv_idx + NUM_MATH_WG @@ -1707,6 +1992,18 @@ def kernel( # Update while-loop condition has_work = (next_q_idx != end_q_idx) | (next_kv_idx != end_kv_idx) + # Flush the final request's emission state (WG 1). + if cutlass.const_expr(self.emit_seed_counts): + if q_idx < batch_size: + self._flush_seed_counts( + mSeedCounts, q_idx, scnt, meta_lane, spass=spass, cand_ctl=mCandCtl + ) + if cutlass.const_expr(self.emit_cand_bucketed): + if q_idx < batch_size: + self._flush_cand_window_bucketed( + mCand, mCandIdx, q_idx, cwbase, cwleft, meta_lane + ) + # Release last Q stage (WG 1) if q_idx < batch_size: q_pipeline.consumer_release(q_cons_state) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py index fbaec11736c1..f76a58431be0 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_emission.py @@ -219,8 +219,8 @@ def update_seed_rows(self, num_rows: int, emit_tier: str = "counts") -> None: self.cand_cur[:nc].zero_() def indexer_emit_kwargs(self, emit_tier: str, num_rows: int) -> dict: - """kwargs for CuteDSLFP4PagedMQALogitsRunner.forward covering the - planned emission tier (caller merges into its call).""" + """kwargs for the paged-MQA scoring runners (FP4 / FP8) covering + the planned emission tier (caller merges into its call).""" kw: dict = {} if emit_tier in ("counts", "list"): kw["seed_thr"] = self.seed_row[:num_rows] 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 3ca2d837cbca..d74839e13190 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 @@ -452,6 +452,421 @@ def test_cute_dsl_fp8_paged_mqa_logits_multi_block( ) +# --------------------------------------------------------------------------- +# Emission tests (emit_block_meta / emit_seed_counts / emit_cand_bucketed — +# fused-GVR support). Ports of the FP4 emission tests; hit-stats and the +# plain (non-bucketed) candidate list are FP4-only and not ported. The FP8 +# epilogue multiplies by the per-token dequant scale before the store and +# emission is computed on the post-conversion stored logit, so references +# recomputed from the KERNEL'S OWN logits output remain the right oracle. +# --------------------------------------------------------------------------- + +_FLT_MAX_F32 = torch.finfo(torch.float32).max + +_EMISSION_COMMON = dict( + num_epi_subtiles=1, + epi_dtype=torch.float32, + acc_dtype=torch.float32, + output_dtype=torch.float32, +) + + +def _emission_test_data(batch_size, next_n, avg_ctx, phys_block_kv, fix_length, seed): + """FP8 inputs for the emission tests: ctx == avg_ctx (or randomized + around it) with the logits buffer spanning max_model_len == 2 * avg_ctx, + mirroring the FP4 emission tests (ctx < buffer proves no writes land + past each row's valid region). Q/KV/weights use the same recipes as + ``_generate_test_data``; the block table is a random permutation.""" + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + num_heads, head_dim = 64, 128 + max_model_len = max(avg_ctx * 2, 2048) + device = "cuda" + + if fix_length: + context_lens = torch.full((batch_size,), avg_ctx, dtype=torch.int32, device=device) + else: + lo = max(phys_block_kv, int(0.7 * avg_ctx)) + context_lens = torch.randint( + lo, int(1.3 * avg_ctx) + 1, (batch_size,), dtype=torch.int32, device=device + ).clamp(max=max_model_len) + + num_blocks_per_seq = (context_lens + phys_block_kv - 1) // phys_block_kv + num_total_blocks = int(num_blocks_per_seq.sum().item()) + batch_size * 2 + max_blocks_per_seq = int(num_blocks_per_seq.max().item()) + block_table = torch.zeros((batch_size, max_blocks_per_seq), dtype=torch.int32, device=device) + pool = torch.randperm(num_total_blocks, device=device, dtype=torch.int32) + off = 0 + for i, n_blks in enumerate(num_blocks_per_seq.tolist()): + block_table[i, :n_blks] = pool[off : off + n_blks] + off += n_blks + + q_fp8 = torch.randn(batch_size, next_n, num_heads, head_dim, device=device).to( + torch.float8_e4m3fn + ) + kv_bf16 = torch.randn(num_total_blocks, phys_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) + weights = torch.randn(batch_size * next_n, num_heads, device=device, dtype=torch.float32) + kv_fused = _make_fused_kv(kv_fp8, kv_scale, phys_block_kv, head_dim) + + from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + # `block_kv = 64` — see test_cute_dsl_fp8_paged_mqa_logits for reasoning. + DG_METADATA_BLOCK_KV = 64 + schedule_meta = get_paged_mqa_logits_metadata( + context_lens.unsqueeze(-1), DG_METADATA_BLOCK_KV, num_sms + ) + return q_fp8, kv_fused, weights, context_lens, block_table, schedule_meta, max_model_len + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2, 3, 4]) +# 4224 = 33 blocks of 128 -> odd num_kv exercises WG1's OOB padding tile. +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("fix_length", [True, False]) +def test_cute_dsl_fp8_paged_mqa_logits_block_meta( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + fix_length, +): + """emit_block_meta correctness: block_max recomputed from the KERNEL'S + OWN logits output (the meta contract is defined on what the kernel + stores — the per-token dequant scale is already folded in). + NaN-prefilled buffers prove no writes land outside + [0, num_kv (+1 when odd)) per row.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLPagedMQALogitsRunner + + device = "cuda" + ( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) = _emission_test_data(batch_size, next_n, avg_ctx, phys_block_kv, fix_length, seed=7) + + aligned_max_ctx = ((max_model_len + 255) // 256) * 256 + nb_pad = aligned_max_ctx // 128 + num_rows = batch_size * next_n + + # block_max: 4 warp-partial records per block; consumers fold. NaN + # prefill proves write coverage is exactly [0, written_hi*4) per row. + nan = float("nan") + block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device) + logits, bm = CuteDSLPagedMQALogitsRunner.forward( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + emit_block_meta=True, + block_max_out=block_max, + **_EMISSION_COMMON, + ) + torch.cuda.synchronize() + + lf = logits.float() + for row in range(num_rows): + req = row // next_n + ctx = int(context_lens[req].item()) + num_kv = (ctx + 127) // 128 + tag = f"row={row} req={req} ctx={ctx} next_n={next_n} pbk={phys_block_kv}" + + # Fold the kernel's 4 warp-partials per block (the consumer-side + # contract); per-warp partials themselves depend on the TMEM + # lane->row mapping and are not checked individually. + bm_fold = bm[row].view(nb_pad, 4).amax(-1) + + # block_max reference from the kernel's own stored logits. + padded = torch.full((nb_pad * 128,), -_FLT_MAX_F32, device=device) + padded[:ctx] = lf[row, :ctx] + ref_bmax = padded.view(nb_pad, 128).amax(-1) + torch.testing.assert_close( + bm_fold[:num_kv], + ref_bmax[:num_kv], + atol=0.0, + rtol=0.0, + msg=lambda m, tag=tag: f"block_max mismatch: {tag}\n{m}", + ) + + # Odd num_kv: WG1's OOB tile writes pure identities into block + # slot num_kv (every lane invalid). + written_hi = num_kv + (num_kv % 2) + if written_hi > num_kv: + assert bm_fold[num_kv].item() == -_FLT_MAX_F32, tag + # No stray writes past the padding tile: NaN prefill intact. + assert bm[row, written_hi * 4 :].isnan().all(), f"stray block_max write: {tag}" + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2, 3, 4]) +# 4224 = 33 blocks of 128 -> odd num_kv exercises WG1's OOB padding tile. +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("fix_length", [True, False]) +@pytest.mark.parametrize("packed", [False, True]) +def test_cute_dsl_fp8_paged_mqa_logits_seed_counts( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + fix_length, + packed, +): + """emit_seed_counts exactness: per-row counts of logits >= threshold + recomputed from the KERNEL'S OWN logits output (the count contract is + defined on post-conversion values, same as block_max). Thresholds are + per-row quantiles of the row's own logits so each of the 3 counters + lands in a different regime (loose/mid/tight).""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLPagedMQALogitsRunner + + device = "cuda" + ( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) = _emission_test_data(batch_size, next_n, avg_ctx, phys_block_kv, fix_length, seed=11) + + aligned_max_ctx = ((max_model_len + 255) // 256) * 256 + nb_pad = aligned_max_ctx // 128 + num_rows = batch_size * next_n + base_args = ( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) + + # First pass without seed counts to harvest per-row logits for + # threshold picking (post-conversion value domain). + nan = float("nan") + block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device) + logits0, _ = CuteDSLPagedMQALogitsRunner.forward( + *base_args, + emit_block_meta=True, + block_max_out=block_max, + **_EMISSION_COMMON, + ) + torch.cuda.synchronize() + # clone: the runner returns a persistent arena buffer that the second + # forward overwrites in place (fp32 output makes .float() a no-copy). + lf0 = logits0.float().clone() + + seed_thr = torch.empty((num_rows, 3), dtype=torch.float32, device=device) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + vals = lf0[row, :ctx] + # Loose / mid / tight rungs; ties on exact stored values are the + # point (>= must count them all). + seed_thr[row, 0] = torch.quantile(vals, 0.10) + seed_thr[row, 1] = torch.quantile(vals, 0.90) + seed_thr[row, 2] = torch.quantile(vals, 0.998) + + if packed: + # Packed contract: one [rows, 8] fp32 seed row, lines at cols + # 0..2, counts accumulate as fp32 at cols 3..5 (caller zeroes). + seed_row = torch.zeros((num_rows, 8), dtype=torch.float32, device=device) + seed_row[:, 0:3] = seed_thr + thr_arg, counts_arg = seed_row, None + else: + seed_counts = torch.zeros((num_rows, 3), dtype=torch.int32, device=device) + thr_arg, counts_arg = seed_thr, seed_counts + block_max.fill_(nan) + logits, _ = CuteDSLPagedMQALogitsRunner.forward( + *base_args, + emit_block_meta=True, + block_max_out=block_max, + emit_seed_counts=True, + seed_thr=thr_arg, + seed_counts_out=counts_arg, + **_EMISSION_COMMON, + ) + torch.cuda.synchronize() + + lf = logits.float() + # compare valid prefixes only: past ctx the buffer is unwritten + # allocator garbage and differs run-to-run + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + torch.testing.assert_close(lf[row, :ctx], lf0[row, :ctx], atol=0.0, rtol=0.0) + if packed: + assert torch.equal(seed_row[:, 0:3], seed_thr), "lines clobbered" + # col 6 carries the adaptive-skip pass count (lane0-accumulated + # diagnostic; the top-k consumer reads it when block_max rides + # along), so it is a legitimate output here - bounded by the + # block-max record count. Only col 7 must stay untouched. + assert (seed_row[:, 7] == 0).all(), "stray write past counts" + nrec = block_max.shape[1] + assert ((seed_row[:, 6] >= 0) & (seed_row[:, 6] <= nrec)).all(), ( + "adaptive-skip pass count out of range" + ) + seed_counts = seed_row[:, 3:6].to(torch.int32) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + tag = f"row={row} ctx={ctx} next_n={next_n} pbk={phys_block_kv}" + ref = (lf[row, :ctx].unsqueeze(0) >= seed_thr[row].unsqueeze(1)).sum(-1) + got = seed_counts[row].to(torch.int64) + assert torch.equal(got.cpu(), ref.cpu().to(torch.int64)), ( + f"seed_counts mismatch: {tag} got={got.tolist()} ref={ref.tolist()}" + ) + + +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +@pytest.mark.parametrize("phys_block_kv", [64, 128]) +@pytest.mark.parametrize("cap_mode", ["roomy", "tight"]) +def test_cute_dsl_fp8_paged_mqa_logits_cand_bucketed( + batch_size, + next_n, + avg_ctx, + phys_block_kv, + cap_mode, +): + """emit_cand_bucketed (v5 SoA contract): three fixed segments with + pad-free A/B prefixes and spill-to-looser, ctl {n0, void, n1, n2} + with n1/n2 mirrored from the seed counters, C-window pads carrying + score -inf / idx -1. All invariants recomputed from the kernel's + own logits.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLPagedMQALogitsRunner + + device = "cuda" + ( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) = _emission_test_data(batch_size, next_n, avg_ctx, phys_block_kv, True, seed=23) + + aligned_max_ctx = ((max_model_len + 255) // 256) * 256 + nb_pad = aligned_max_ctx // 128 + num_rows = batch_size * next_n + nan = float("nan") + block_max = torch.full((num_rows, nb_pad * 4), nan, dtype=torch.float32, device=device) + base_args = ( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) + logits0, _ = CuteDSLPagedMQALogitsRunner.forward( + *base_args, + emit_block_meta=True, + block_max_out=block_max, + **_EMISSION_COMMON, + ) + torch.cuda.synchronize() + # clone: arena buffer, see test_cute_dsl_fp8_paged_mqa_logits_seed_counts. + lf0 = logits0.float().clone() + seed_row = torch.zeros((num_rows, 8), dtype=torch.float32, device=device) + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + vals = lf0[row, :ctx] + seed_row[row, 0] = torch.quantile(vals, 0.60) + seed_row[row, 1] = torch.quantile(vals, 0.90) + seed_row[row, 2] = torch.quantile(vals, 0.99) + # segment caps: roomy fits everything; tight forces A/B spill and a + # C-window void + if cap_mode == "roomy": + segA, capC = 2048, 4096 + else: + segA, capC = 32, 128 + W = 2 * segA + capC + cand_vals = torch.full((num_rows, W), nan, dtype=torch.float32, device=device) + cand_idx = torch.full((num_rows, W), -7, dtype=torch.int32, device=device) + cand_ctl = torch.zeros((num_rows, 4), dtype=torch.int32, device=device) + cand_cur = torch.zeros((num_rows, 4), dtype=torch.int32, device=device) + block_max.fill_(nan) + logits, _ = CuteDSLPagedMQALogitsRunner.forward( + *base_args, + emit_block_meta=True, + block_max_out=block_max, + emit_seed_counts=True, + seed_thr=seed_row, + emit_cand_bucketed=True, + accept_cap=segA, + cand_out=cand_vals, + cand_idx_out=cand_idx, + cand_ctl_out=cand_ctl, + cand_cur_out=cand_cur, + **_EMISSION_COMMON, + ) + torch.cuda.synchronize() + lf = logits.float() + for row in range(num_rows): + ctx = int(context_lens[row // next_n].item()) + torch.testing.assert_close(lf[row, :ctx], lf0[row, :ctx], atol=0.0, rtol=0.0) + t0, t1, t2 = (float(seed_row[row, j]) for j in range(3)) + v = lf[row, :ctx] + n0_ref = int((v >= t0).sum()) + n1_ref = int((v >= t1).sum()) + n2_ref = int((v >= t2).sum()) + n0c, voidc, n1c, n2c = (int(cand_ctl[row, j]) for j in range(4)) + tag = f"row={row} caps=({segA},{capC}) refs=({n0_ref},{n1_ref},{n2_ref})" + assert n1c == n1_ref and n2c == n2_ref, f"n1/n2 mismatch {tag} got {n1c},{n2c}" + curA, curB, curC = (int(cand_cur[row, j]) for j in range(3)) + lenA = min(n2_ref, segA) + lenB = min(n1_ref - n2_ref + max(n2_ref - segA, 0), segA) + assert min(curA, segA) >= lenA or curA == n2_ref, f"curA {curA} {tag}" + # A prefix: pad-free, every entry >= t2, positions valid + unique + pa = cand_idx[row, :lenA] + va = cand_vals[row, :lenA] + assert (pa >= 0).all() and (pa < ctx).all(), f"A idx {tag}" + assert (va >= t2).all(), f"A vals {tag}" + got_a = lf[row, pa.long()] + torch.testing.assert_close(got_a, va, atol=0.0, rtol=0.0) + # B prefix: pad-free, [t1, t2) or A-spill (>= t2) + pb = cand_idx[row, segA : segA + lenB] + vb = cand_vals[row, segA : segA + lenB] + assert (pb >= 0).all() and (pb < ctx).all(), f"B idx {tag}" + assert (vb >= t1).all(), f"B vals {tag}" + torch.testing.assert_close(lf[row, pb.long()], vb, atol=0.0, rtol=0.0) + if voidc == 0: + # full coverage: union of live entries == the >= t0 set + lenC = n0c - lenA - lenB + pc = cand_idx[row, 2 * segA : 2 * segA + lenC] + vc = cand_vals[row, 2 * segA : 2 * segA + lenC] + live = pc >= 0 + assert (vc[live] >= t0).all(), f"C vals {tag}" + # pads carry -FLT_MAX (never ranks; the emu uses -inf, the + # kernel the finite sentinel - both satisfy the contract) + assert (vc[~live] <= -3e38).all(), f"C pads {tag}" + allp = torch.cat([pa, pb, pc[live]]) + assert allp.unique().numel() == allp.numel() == n0_ref, ( + f"coverage {tag}: {allp.unique().numel()} vs {n0_ref}" + ) + else: + assert cap_mode == "tight", f"unexpected void {tag}" + if cap_mode == "tight": + assert int(cand_ctl[:, 1].sum()) > 0, "tight caps never voided" + + 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 @@ -828,6 +1243,83 @@ def dsl_f32_fn(data=data): print() +@skip_not_sm100 +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("avg_ctx", [4096, 4224]) +def test_cute_dsl_fp8_paged_mqa_logits_op_emission_surface(batch_size, avg_ctx): + """The op-level emission seam: kwargs produced by GvrEmissionState (the + production wiring) feed torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits. + The op must accept the kwarg names as-is, unwrap the runner tuple, and + produce the deterministic emission outputs (logits, block_max, packed + seed row, ctl counts, cursor totals) bit-identical to the runner path.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLPagedMQALogitsRunner + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k.gvr_emission import GvrEmissionState + + device = "cuda" + next_n, phys_block_kv = 1, 64 + ( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + ) = _emission_test_data(batch_size, next_n, avg_ctx, phys_block_kv, False, seed=11) + num_rows = batch_size * next_n + + def run(op_path: bool): + state = GvrEmissionState( + max_rows=num_rows, top_k=2048, device=torch.device(device), own_prior=False + ) + state.update_seed_rows(num_rows, "list") + kwargs = state.indexer_emit_kwargs("list", num_rows) + kwargs["block_max_out"] = state.ensure_block_max(max_model_len)[:num_rows] + if op_path: + logits = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + **kwargs, + ) + assert isinstance(logits, torch.Tensor), "op must unwrap the runner tuple" + else: + logits, _ = CuteDSLPagedMQALogitsRunner.forward( + q_fp8, + kv_fused, + weights, + context_lens, + block_table, + schedule_meta, + max_model_len, + emit_block_meta=True, + emit_seed_counts=True, + emit_cand_bucketed=True, + seed_thr=kwargs["seed_thr"], + accept_cap=kwargs["accept_cap"], + cand_out=kwargs["cand_out"], + cand_idx_out=kwargs["cand_idx_out"], + cand_ctl_out=kwargs["cand_ctl_out"], + cand_cur_out=kwargs["cand_cur_out"], + block_max_out=kwargs["block_max_out"], + **_EMISSION_COMMON, + ) + torch.cuda.synchronize() + return logits, state + + logits_op, st_op = run(op_path=True) + logits_rn, st_rn = run(op_path=False) + torch.testing.assert_close(logits_op.float(), logits_rn.float(), atol=0.0, rtol=0.0) + torch.testing.assert_close(st_op.block_max, st_rn.block_max, atol=0.0, rtol=0.0, equal_nan=True) + torch.testing.assert_close(st_op.seed_row, st_rn.seed_row, atol=0.0, rtol=0.0) + assert torch.equal(st_op.cand_ctl[:num_rows], st_rn.cand_ctl[:num_rows]) + assert torch.equal(st_op.cand_cur[:num_rows], st_rn.cand_cur[:num_rows]) + + if __name__ == "__main__": import argparse import os