diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index 9574f2cfbf9c..140e9c90f94d 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -1,5 +1,240 @@ +diff --git a/python/fmha_sm100/api.py b/python/fmha_sm100/api.py +index b02f747..d9d0a41 100644 +--- a/python/fmha_sm100/api.py ++++ b/python/fmha_sm100/api.py +@@ -532,9 +532,8 @@ def _fmha_sm100_plan( + + if kv_block_num > 0 and (sparse_kernel_mode == 'prefill' or (sparse_kernel_mode == 'auto' and max_qo_len_orig > _prefill_qlen_threshold(True))): + # print("Nv-Prefill") +- qo_segment_lens = qo_segment_lens.to(device) +- kv_segment_lens = kv_segment_lens.to(device) +- qo_offset = qo_offset.to(device) ++ # sparse_fmha_plan reads these lengths on the host and stages its own ++ # device copies, so hand them through unchanged. + return sparse_fmha_plan(qo_segment_lens=qo_segment_lens, kv_segment_lens=kv_segment_lens, + num_qo_heads=num_qo_heads, causal=causal, qo_offset=qo_offset, + num_kv_splits=num_kv_splits, page_size=page_size, output_maxscore=output_maxscore, +@@ -941,6 +940,62 @@ def _fmha_sm100( + return out, max_score + + ++def _mixed_batch_split(qo_segment_lens, split_threshold, *, enable, sparse): ++ """Row index splitting a mixed prefill/decode batch, or 0 for no split. ++ ++ Returns ``split`` such that rows ``[0, split)`` and ``[split, batch)`` are ++ each uniformly prefill or uniformly decode. Both halves must be contiguous ++ row ranges because that is all the run path can slice. ++ ++ Returns 0 when the batch is uniform, when the two classes interleave, or ++ when splitting is disabled. ++ ++ Both orderings are recognised: ++ ++ * **Decode-first** (prefill rows form a suffix) is this module's original ++ convention and is honoured for every plan flavour, so existing callers ++ keep their exact behaviour. ++ * **Prefill-first** (prefill rows form a prefix), which is what ++ TensorRT-LLM produces, is honoured only for ``sparse`` plans. The dense ++ flavour already handles a mixed batch in one call, whereas the sparse ++ MM-SA-Nv prefill kernel has no mixed path and would otherwise drag every ++ decode row through it. Confining it to sparse plans also leaves the proxy ++ and dense plans byte-identical to their pre-split behaviour, which ++ matters because the proxy's max-score output drives top-k block ++ selection. ++ """ ++ batch_size = int(qo_segment_lens.shape[0]) ++ if not enable or batch_size < 2: ++ return 0 ++ is_prefill = qo_segment_lens > split_threshold ++ n_prefill = int(is_prefill.sum()) ++ if n_prefill == 0 or n_prefill == batch_size: ++ return 0 ++ if bool(is_prefill[batch_size - n_prefill:].all()): ++ return batch_size - n_prefill ++ if sparse and bool(is_prefill[:n_prefill].all()): ++ return n_prefill ++ # Interleaved: neither half would be uniform, so plan the batch as one. ++ return 0 ++ ++ ++def _cache_split_boundaries(sub_lo): ++ """Precompute the boundaries the run path slices a split batch at. ++ ++ Both are properties of the low sub-plan alone, so they are constant for the ++ whole step. Reading them at plan time keeps the device-to-host read off the ++ per-layer path, where a sparse model would otherwise repeat it once per ++ layer. ++ """ ++ offsets = sub_lo.get("qo_segment_offsets") ++ if offsets is not None: ++ sub_lo["_split_nnz"] = int(offsets[-1].item()) // sub_lo.get("pack_factor", 1) ++ if "kv_page_indptr" in sub_lo: ++ sub_lo["_split_kv_pages"] = int(sub_lo["kv_page_indptr"][-1].item()) ++ elif "total_rows" in sub_lo: ++ sub_lo["_split_kv_pages"] = int(sub_lo["total_rows"]) ++ ++ + def fmha_sm100_plan( + qo_segment_lens: torch.Tensor, + kv_segment_lens: torch.Tensor, +@@ -973,9 +1028,11 @@ def fmha_sm100_plan( + ``kv_segment_lens - qo_segment_lens`` for bottom-right causal masking. + A tensor must have shape ``[batch_size]``. + split_prefill_decode : bool, optional +- If True, a mixed batch ordered as decode requests followed by prefill +- requests is split into two sub-plans. The original order must already +- group short decode sequences before long prefill sequences. ++ If True, a mixed batch whose decode and prefill requests form two ++ contiguous groups is split into two sub-plans. Either group may come ++ first; see ``_mixed_batch_split`` for which orderings apply to which ++ plan flavours. A batch that interleaves the two classes is planned as ++ one. + **kwargs + Planner options forwarded to ``_fmha_sm100_plan``. Common options are + ``num_kv_heads``, ``num_kv_splits``, ``page_size``, +@@ -985,8 +1042,13 @@ def fmha_sm100_plan( + Returns + ------- + tuple +- ``(has_mixed_prefill, split, batch_size, decode_plan, prefill_plan)``. +- Pass this tuple unchanged as ``plan_info`` to ``fmha_sm100``. ++ ``(has_mixed_prefill, split, batch_size, sub_plan_lo, sub_plan_hi)``, ++ where ``sub_plan_lo`` describes batch rows ``[0, split)`` and ++ ``sub_plan_hi`` describes rows ``[split, batch_size)``. Which of the ++ two is the prefill half depends on the caller's ordering, so consumers ++ must test a sub-plan's own tags rather than its position. ++ ``sub_plan_hi`` is None when the batch was not split. Pass this tuple ++ unchanged as ``plan_info`` to ``fmha_sm100``. + """ + + # assert qo_segment_lens.device.type == 'cpu' \ +@@ -999,27 +1061,19 @@ def fmha_sm100_plan( + qo_offset = torch.full_like(qo_segment_lens, qo_offset) + + batch_size = qo_segment_lens.shape[0] +- has_mixed_prefill = False +- qmax = qo_segment_lens.max().item() + sparse = kwargs.get("kv_block_num", -1) > 0 + split_threshold = _prefill_qlen_threshold(sparse) +- if split_prefill_decode and qmax > split_threshold: +- split = (qo_segment_lens > split_threshold).nonzero(as_tuple=False)[0, 0].item() +- has_mixed_prefill = split > 0 +- if has_mixed_prefill: +- # print(f"Split into 2 parts at index {split}") +- decode_qo_segment_lens = qo_segment_lens[:split] +- decode_kv_segment_lens = kv_segment_lens[:split] +- decode_qo_offset = qo_offset[:split] +- decode = _fmha_sm100_plan(decode_qo_segment_lens, decode_kv_segment_lens, *args, +- qo_offset=decode_qo_offset, **kwargs) +- decode = {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in decode.items()} +- prefill_qo_segment_lens = qo_segment_lens[split:] +- prefill_kv_segment_lens = kv_segment_lens[split:] +- prefill_qo_offset = qo_offset[split:] +- prefill = _fmha_sm100_plan(prefill_qo_segment_lens, prefill_kv_segment_lens, *args, +- qo_offset=prefill_qo_offset, **kwargs) +- return (True, split, batch_size, decode, prefill) ++ split = _mixed_batch_split( ++ qo_segment_lens, split_threshold, enable=split_prefill_decode, sparse=sparse ++ ) ++ if split > 0: ++ sub_lo = _fmha_sm100_plan(qo_segment_lens[:split], kv_segment_lens[:split], *args, ++ qo_offset=qo_offset[:split], **kwargs) ++ sub_lo = {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in sub_lo.items()} ++ sub_hi = _fmha_sm100_plan(qo_segment_lens[split:], kv_segment_lens[split:], *args, ++ qo_offset=qo_offset[split:], **kwargs) ++ _cache_split_boundaries(sub_lo) ++ return (True, split, batch_size, sub_lo, sub_hi) + else: + plan = _fmha_sm100_plan(qo_segment_lens, kv_segment_lens, *args, + qo_offset=qo_offset, **kwargs) +@@ -1081,16 +1135,23 @@ def fmha_sm100( + ------- + tuple[torch.Tensor | None, torch.Tensor | None] + ``(out, max_score)``. Either item may be ``None`` if the corresponding +- output was disabled. When both decode and prefill sub-plans are used, +- outputs are concatenated back into the original batch order. ++ output was disabled. When the plan was split into two sub-plans, each ++ half writes its own rows of ``out`` directly; if no ``out`` was supplied ++ the halves are concatenated back into the original batch order instead. + """ + has_mixed_prefill, split, batch_size, decode, prefill = plan_info + if not has_mixed_prefill: + return _fmha_sm100(q, k, v, decode, out=out, max_score=max_score, kv_indices=kv_indices,kv_block_indexes=kv_block_indexes, q_offset_override=q_offset_override, **kwargs) + else: +- +- decode_pack = decode.get("pack_factor", 1) +- decode_nnz = decode["qo_segment_offsets"][-1].item() // decode_pack ++ # The `decode` / `prefill` names are historical: these are really the ++ # sub-plans for batch rows [0, split) and [split, batch_size), and which ++ # one holds the prefill requests depends on the caller's ordering. None ++ # of the slicing below needs to know, because every boundary is derived ++ # from the low sub-plan's own offsets. ++ decode_nnz = decode.get("_split_nnz") ++ if decode_nnz is None: ++ decode_pack = decode.get("pack_factor", 1) ++ decode_nnz = decode["qo_segment_offsets"][-1].item() // decode_pack + is_paged = kv_indices is not None + nnz_qo = q.shape[0] + num_qo_heads = q.shape[1] +@@ -1101,10 +1162,12 @@ def fmha_sm100( + if is_paged: + k_decode, v_decode = k, v + k_prefill, v_prefill = k, v +- if "kv_page_indptr" in decode: +- kv_page_split = decode["kv_page_indptr"][-1].item() +- else: +- kv_page_split = decode["total_rows"] ++ kv_page_split = decode.get("_split_kv_pages") ++ if kv_page_split is None: ++ if "kv_page_indptr" in decode: ++ kv_page_split = decode["kv_page_indptr"][-1].item() ++ else: ++ kv_page_split = decode["total_rows"] + decode_kv_indices = kv_indices[:kv_page_split] + prefill_kv_indices = kv_indices[kv_page_split:] + else: +@@ -1129,27 +1192,35 @@ def fmha_sm100( + decode_qo_offset = None + prefill_qo_offset = None + ++ # Let each half write straight into its own rows of the caller's buffer. ++ # Row views of a contiguous [nnz, H, D] output are themselves ++ # contiguous, which is what the kernels require, so this removes a ++ # full-size concatenate plus copy from every call. ++ write_in_place = out is not None and out.is_contiguous() ++ + # ---- Run kernels ---- + decode_out, decode_ms = _fmha_sm100( + q_decode, k_decode, v_decode, decode, +- out=None, max_score=None, ++ out=out[:decode_nnz] if write_in_place else None, max_score=None, + kv_indices=decode_kv_indices, kv_block_indexes=decode_block_idx, + q_offset_override=decode_qo_offset, + **kwargs) + prefill_out, prefill_ms = _fmha_sm100( + q_prefill, k_prefill, v_prefill, prefill, +- out=None, max_score=None, ++ out=out[decode_nnz:] if write_in_place else None, max_score=None, + kv_indices=prefill_kv_indices, kv_block_indexes=prefill_block_idx, + q_offset_override=prefill_qo_offset, + **kwargs) + + # ---- Merge out ---- +- if decode_out is not None and prefill_out is not None: ++ if write_in_place: ++ combined_out = out ++ elif decode_out is not None and prefill_out is not None: + combined_out = torch.cat([decode_out, prefill_out], dim=0) + else: + combined_out = None + +- if out is not None and combined_out is not None: ++ if out is not None and combined_out is not None and combined_out is not out: + out.copy_(combined_out) + + # ---- Merge max_score ---- diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py -index d72b17a..eca5b8c 100644 +index d72b17a..e09014f 100644 --- a/python/fmha_sm100/cute/interface.py +++ b/python/fmha_sm100/cute/interface.py @@ -136,6 +136,35 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int): @@ -38,7 +273,26 @@ index d72b17a..eca5b8c 100644 def _validate_cu_seqlens( cu_seqlens: torch.Tensor, *, -@@ -736,10 +765,21 @@ def sparse_atten_func( +@@ -622,6 +651,7 @@ def sparse_atten_func( + usable_SM_count: int = -1, + qk_dtype: Optional[torch.dtype] = None, + pv_dtype: Optional[torch.dtype] = None, ++ out: Optional[torch.Tensor] = None, + ): + """Run SM100 CSR block-sparse varlen attention. + +@@ -691,6 +721,10 @@ def sparse_atten_func( + pv_dtype : torch.dtype, optional + Compile-time MMA operand dtype for PV. Defaults to V storage dtype, + except supported FP8 K/V cache staging modes. ++ out : torch.Tensor, optional ++ Destination for the attention output, written in place by the reduction ++ instead of being allocated here. Must be contiguous BF16 of shape ++ ``[total_q, Hq, 128]`` on the same device as ``q``. + + Returns + ------- +@@ -736,10 +770,21 @@ def sparse_atten_func( max_seqlen_q = int(max_seqlen_q) max_seqlen_k = int(max_seqlen_k) @@ -62,6 +316,48 @@ index d72b17a..eca5b8c 100644 k2q_row_ptr.contiguous(), k2q_q_indices.contiguous(), int(topK), +@@ -762,6 +807,7 @@ def sparse_atten_func( + int(max_seqlen_k), + qk_dtype, + pv_dtype, ++ out=out, + ) + + +@@ -1449,6 +1495,7 @@ def _sparse_atten_csr_varlen_forward( + max_seqlen_k: int, + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, ++ out: Optional[torch.Tensor] = None, + ): + total_q, head_q, dim = q.shape + if head_q % head_kv != 0: +@@ -1478,7 +1525,24 @@ def _sparse_atten_csr_varlen_forward( + if kernel_return_temperature_lse + else None + ) +- O_out = torch.empty(total_q, head_q, dim, dtype=torch.bfloat16, device=q.device) ++ if out is None: ++ O_out = torch.empty(total_q, head_q, dim, dtype=torch.bfloat16, device=q.device) ++ else: ++ # `combine` writes its output with a plain contiguous layout, so a ++ # caller that already owns a matching buffer can have it written ++ # directly instead of copying the result afterwards. ++ if tuple(out.shape) != (total_q, head_q, dim): ++ raise ValueError( ++ f"out must have shape {(total_q, head_q, dim)}, " ++ f"got {tuple(out.shape)}" ++ ) ++ if out.dtype != torch.bfloat16: ++ raise TypeError(f"out must be torch.bfloat16, got {out.dtype}") ++ if out.device != q.device: ++ raise ValueError("out must be on the same device as q") ++ if not out.is_contiguous(): ++ raise ValueError("out must be contiguous") ++ O_out = out + LSE_out = torch.empty(total_q, head_q, dtype=torch.float32, device=q.device) + LSE_temperature_out = ( + torch.empty_like(LSE_out) if kernel_return_temperature_lse else None diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cute/test_sparse_atten.py index 21c777e..b5f078b 100644 --- a/python/fmha_sm100/cute/test_sparse_atten.py @@ -224,10 +520,65 @@ index 21c777e..b5f078b 100644 @pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("batch", [3]) diff --git a/python/fmha_sm100/sparse_fmha_adapter.py b/python/fmha_sm100/sparse_fmha_adapter.py -index 306b416..fd24e7e 100644 +index 306b416..a564d47 100644 --- a/python/fmha_sm100/sparse_fmha_adapter.py +++ b/python/fmha_sm100/sparse_fmha_adapter.py -@@ -228,23 +228,83 @@ def _build_page_table( +@@ -129,11 +129,19 @@ def sparse_fmha_plan( + batch = qo_segment_lens.shape[0] + gpu_device = torch.device('cuda') + ++ # Stage the host lengths to the device once, asynchronously, and read every ++ # scalar below from the host copy. Planning runs on the critical path of ++ # input preparation, where a blocking copy or a device .item() would drain ++ # the CUDA queue. ++ qo_lens_gpu = qo_segment_lens.to(torch.int32).to(gpu_device, non_blocking=True) ++ kv_lens_host = kv_segment_lens.to(torch.int32) ++ kv_lens_gpu = kv_lens_host.to(gpu_device, non_blocking=True) ++ + cu_seqlens_q = torch.zeros(batch + 1, dtype=torch.int32, device=gpu_device) +- cu_seqlens_q[1:] = torch.cumsum(qo_segment_lens.to(torch.int32).to(gpu_device), dim=0) ++ cu_seqlens_q[1:] = torch.cumsum(qo_lens_gpu, dim=0) + + cu_seqlens_k = torch.zeros(batch + 1, dtype=torch.int32, device=gpu_device) +- cu_seqlens_k[1:] = torch.cumsum(kv_segment_lens.to(torch.int32).to(gpu_device), dim=0) ++ cu_seqlens_k[1:] = torch.cumsum(kv_lens_gpu, dim=0) + + max_seqlen_q = int(qo_segment_lens.max().item()) + max_seqlen_k = int(kv_segment_lens.max().item()) +@@ -144,11 +152,12 @@ def sparse_fmha_plan( + total_rows = sum((int(kl) + blk_kv - 1) // blk_kv for kl in kv_lens) + + if qo_offset is not None: +- seqused_k = (qo_segment_lens + qo_offset).to(torch.int32).to(gpu_device) ++ seqused_k = (qo_segment_lens + qo_offset).to(torch.int32).to(gpu_device, non_blocking=True) + else: +- seqused_k = kv_segment_lens.to(torch.int32).to(gpu_device) ++ seqused_k = kv_lens_gpu.clone() + +- total_q = int(cu_seqlens_q[-1].item()) ++ # Equals cu_seqlens_q[-1], read on the host. ++ total_q = int(qo_segment_lens.sum()) + qhead_per_kv = num_qo_heads // num_kv_heads if num_kv_heads > 0 else 1 + + target_q_per_cta = SPARSE_SCHEDULE_MODEL.balanced_target_q_per_cta( +@@ -169,11 +178,14 @@ def sparse_fmha_plan( + ) + + return { +- "qo_segment_lens": qo_segment_lens, ++ "qo_segment_lens": qo_lens_gpu, + "cu_seqlens_q": cu_seqlens_q, + "qo_segment_offsets": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, +- "kv_segment_lens": kv_segment_lens.to(torch.int32).to(gpu_device), ++ # Host tensor: only the page-table builder consumes it, and it reads ++ # per-request page counts on the host. Kernels take the lengths from ++ # cu_seqlens_k and seqused_k. ++ "kv_segment_lens": kv_lens_host, + "seqused_k": seqused_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, +@@ -228,23 +240,83 @@ def _build_page_table( page_size: int, batch: int, ) -> torch.Tensor: @@ -326,7 +677,7 @@ index 306b416..fd24e7e 100644 return page_table -@@ -351,8 +411,8 @@ def sparse_fmha( +@@ -351,8 +423,8 @@ def sparse_fmha( if is_paged: if kv_indices is not None: @@ -337,3 +688,33 @@ index 306b416..fd24e7e 100644 ) # build_k2q_csr(return_schedule=True) builds schedule using hardware SM count internally +@@ -382,6 +454,17 @@ def sparse_fmha( + + softmax_scale = sm_scale if sm_scale is not None else q.shape[-1] ** -0.5 + ++ # Hand the caller's buffer to the kernel when its layout already matches ++ # what `combine` produces, so the result does not have to be copied into ++ # place afterwards. Anything else falls back to allocate-then-copy. ++ writable_out = ( ++ out is not None ++ and out.is_contiguous() ++ and out.dtype == torch.bfloat16 ++ and out.device == q.device ++ and tuple(out.shape) == tuple(q.shape) ++ ) ++ + # print(q.shape, k.shape) + result = sparse_atten_func( + q, k, v, +@@ -398,9 +481,10 @@ def sparse_fmha( + seqused_k=seqused_k, + schedule=schedule, + usable_SM_count=usable_SM_count, ++ out=out if writable_out else None, + ) + +- if out is not None: ++ if out is not None and not writable_out: + out.copy_(result) + return out, None + return result, None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index 89a7595c52c9..0d8abed29703 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -30,7 +30,12 @@ import torch from tensorrt_llm._torch.disaggregation.resource.page import MapperKind -from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor +from tensorrt_llm._utils import ( + TensorWrapper, + binding_to_torch_dtype, + convert_to_torch_tensor, + prefer_pinned, +) from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode @@ -455,20 +460,22 @@ def get_block_ids_per_seq(self, request_ids): Drops the base's final ``i // num_local_layers`` step (paired with the base ``index_scales`` multiplication that's also bypassed here). Pads with ``0`` to preserve shape. + + The rows are written through a numpy view of a single zero-filled, + pinned result, so the attention metadata builders ship it to the device + in one asynchronous copy. """ block_ids_per_seq = self.get_batch_cache_indices(request_ids) - block_ids_per_seq_tensors = [ - torch.tensor( - [i if i != BAD_PAGE_INDEX else 0 for i in sublist], - dtype=torch.int, - ) - for sublist in block_ids_per_seq - ] - padded_tensor = torch.nn.utils.rnn.pad_sequence( - block_ids_per_seq_tensors, - batch_first=True, - padding_value=0, + batch = len(block_ids_per_seq) + max_blocks = max((len(block_ids) for block_ids in block_ids_per_seq), default=0) + padded_tensor = torch.zeros( + (batch, max_blocks), dtype=torch.int32, pin_memory=prefer_pinned() ) + rows = padded_tensor.numpy() + for row, block_ids in zip(rows, block_ids_per_seq): + row[: len(block_ids)] = block_ids + # BAD_PAGE_INDEX marks padding, which this tensor reports as 0. + rows[rows == BAD_PAGE_INDEX] = 0 return padded_tensor diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py index ea5ad1f7ee92..94c2c7be7974 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py @@ -11,10 +11,12 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Tuple import torch +from tensorrt_llm._utils import async_tensor_h2d, maybe_pin_memory + from ..params import SparseMetadataParams, SparseParams if TYPE_CHECKING: @@ -194,6 +196,15 @@ def write_kv_slots( cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) +class PagedKvSlotMapping(NamedTuple): + """One step's paged-cache slot mapping (see build_paged_kv_slot_mapping).""" + + req_to_token: torch.Tensor + slot_ids: torch.Tensor + out_cache_loc: torch.Tensor + block_ids_cpu: torch.Tensor + + def build_paged_kv_slot_mapping( *, kv_cache_manager, @@ -201,12 +212,11 @@ def build_paged_kv_slot_mapping( qo_lens_cpu: torch.Tensor, qo_offset_cpu: torch.Tensor, device: torch.device, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> PagedKvSlotMapping: """Build the backend-neutral paged-cache slot mapping. - Returns (req_to_token, slot_ids, out_cache_loc), derived only from the paged - KV cache manager and the per-request query geometry, with no dependency on - any backend-specific metadata. + Derived only from the paged KV cache manager and the per-request query + geometry, with no dependency on any backend-specific metadata. req_to_token is the [batch, max_kv_len] int32 map from (request, position) to a global slot id, expanded from get_block_ids_per_seq with @@ -215,23 +225,28 @@ def build_paged_kv_slot_mapping( lists the per-new-token slot ids in flattened query order: request b contributes positions qo_offset[b] through qo_offset[b] + qo_lens[b] - 1. That one formula covers prefill (qo_offset is the prefix length) and decode - (qo_offset is kv_len - 1 with qo_len 1). - - The req_to_token reads that build out_cache_loc sync the host, so call this - only from prepare(), never from the forward path. + (qo_offset is kv_len - 1 with qo_len 1). block_ids_cpu is the host block-id + table every field above derives from, returned so backends can build their + own page-indexed views without a second manager query or a device round + trip. + + Both out_cache_loc and req_to_token are computed from the host block ids and + staged with non-blocking copies, so this neither syncs the host nor reads + device memory back. It still allocates per step, so call it from prepare(), + never from the forward path. """ tokens_per_block = int(kv_cache_manager.tokens_per_block) # block_ids_per_seq is a [batch, max_blocks_per_seq] tensor; row b holds the # block ids assigned to request_ids[b] in order. KVCacheManagerV2 maps # padded BAD_PAGE_INDEX entries to zero, and the live ranges selected below # never address those padded positions. - block_ids = kv_cache_manager.get_block_ids_per_seq(list(request_ids)) + block_ids = maybe_pin_memory(kv_cache_manager.get_block_ids_per_seq(list(request_ids))) batch = int(qo_lens_cpu.shape[0]) max_blocks = int(block_ids.shape[1]) max_kv_len = max_blocks * tokens_per_block # Expand block ids -> per-token slot ids. - block_ids_dev = block_ids.to(device).to(torch.int64) + block_ids_dev = block_ids.to(device, non_blocking=True).to(torch.int64) within_block = torch.arange(tokens_per_block, device=device, dtype=torch.int64) # Outer product per batch entry: [batch, max_blocks, tokens_per_block] slot_grid = block_ids_dev.unsqueeze(-1) * tokens_per_block + within_block @@ -239,22 +254,37 @@ def build_paged_kv_slot_mapping( slot_ids = torch.arange(batch, device=device, dtype=torch.int32) # out_cache_loc: per-new-token slot ids, in flattened query-token order. - req_to_token_cpu = req_to_token.to("cpu") - qo_lens_list = qo_lens_cpu.to(torch.long).tolist() - qo_offset_list = qo_offset_cpu.to(torch.long).tolist() - out_cache_loc_list: List[int] = [] - for b in range(batch): - start = int(qo_offset_list[b]) - for offset in range(int(qo_lens_list[b])): - out_cache_loc_list.append(int(req_to_token_cpu[b, start + offset].item())) - out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device) - return req_to_token, slot_ids, out_cache_loc + # Expanding the per-request lengths on the host reproduces the same slot + # ids as indexing req_to_token, without copying that [batch, max_kv_len] + # grid back from the device. + qo = qo_lens_cpu.to(torch.long) + total_q = int(qo.sum()) + if total_q == 0: + out_cache_loc_cpu = torch.empty(0, dtype=torch.int32) + else: + row = torch.repeat_interleave(torch.arange(batch, dtype=torch.long), qo) + starts = torch.cumsum(qo, 0) - qo + pos = qo_offset_cpu.to(torch.long)[row] + ( + torch.arange(total_q, dtype=torch.long) - starts[row] + ) + # A zero-length CUDA-graph padding row offsets to -1. Its slot is a + # placeholder no forward reads, so keep it in the row instead of + # indexing off the table. + pos.clamp_(min=0, max=max(max_kv_len - 1, 0)) + block_col = torch.div(pos, tokens_per_block, rounding_mode="floor") + out_cache_loc_cpu = ( + block_ids[row, block_col].to(torch.long) * tokens_per_block + + (pos - block_col * tokens_per_block) + ).to(torch.int32) + out_cache_loc = async_tensor_h2d(out_cache_loc_cpu, torch.int32, device) + return PagedKvSlotMapping(req_to_token, slot_ids, out_cache_loc, block_ids) __all__ = [ "MiniMaxM3SparseConfig", "MiniMaxM3SparseMetadataParams", "MiniMaxM3SparseParams", + "PagedKvSlotMapping", "build_paged_kv_slot_mapping", "write_kv_slots", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index 66871d0c8f7a..5bb62bc47632 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -32,6 +32,7 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType from .common import ( @@ -61,26 +62,6 @@ def _cache_device(meta) -> torch.device: return torch.device(f"cuda:{torch.cuda.current_device()}") -def _stage_sparse_plan_kv_lens_host(plan: tuple, kv_lens_cpu: torch.Tensor) -> None: - """Give the sparse-prefill sub-plan a host copy of its kv_segment_lens. - - sparse_fmha._build_page_table runs once per sparse layer and reads - plan["kv_segment_lens"] with .tolist(), which blocks on a D2H copy while - that tensor lives on the device. The page table still builds on the device - from kv_indices, so the host copy adds no work. Only the sparse-prefill - sub-plan (MM-SA-Nv) qualifies: decode and dense plans need their lengths on - the device for the kernel. - """ - has_mixed, split = plan[0], plan[1] - # A non-mixed batch has one sparse sub-plan (plan[3]); a mixed batch puts - # the sparse prefill rows in plan[4], after the split decode rows. - sparse_dict = plan[4] if has_mixed else plan[3] - if sparse_dict is None or not sparse_dict.get("MM-SA-Nv"): - return - lens = kv_lens_cpu[split:] if has_mixed else kv_lens_cpu - sparse_dict["kv_segment_lens"] = lens.to(torch.int32).contiguous() - - def _worst_case_proxy_max_k_tiles( fmha_sm100, *, @@ -285,12 +266,19 @@ def __post_init__(self) -> None: @property def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]: - """Per-request query length (host int32), from the base seq_lens.""" + """Per-request query length (host int32), from the base seq_lens. + + Pinned where pinning helps, as with the other two length properties: + the planners stage them to the device with non-blocking copies, which + degrade to a synchronous staging copy from pageable memory. + """ seq_lens = self.seq_lens if seq_lens is None: return None out = seq_lens[: self.num_seqs] - return out if out.dtype == torch.int32 else out.to(torch.int32) + if out.dtype != torch.int32: + out = out.to(torch.int32) + return maybe_pin_memory(out) @property def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: @@ -299,7 +287,9 @@ def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: if self.seq_lens is None or kv_lens is None: return None out = kv_lens[: self.num_seqs] - return out if out.dtype == torch.int32 else out.to(torch.int32) + if out.dtype != torch.int32: + out = out.to(torch.int32) + return maybe_pin_memory(out) @property def msa_qo_offset_cpu(self) -> Optional[torch.Tensor]: @@ -308,7 +298,7 @@ def msa_qo_offset_cpu(self) -> Optional[torch.Tensor]: kv = self.msa_kv_lens_cpu if qo is None or kv is None: return None - return kv - qo + return maybe_pin_memory(kv - qo) @property def msa_decode_proxy_plan(self) -> Optional[tuple]: @@ -605,7 +595,6 @@ def _build_step_plans(self) -> None: self._msa_eager_proxy_plan = proxy_plan self._msa_eager_gqa_plan = gqa_plan self._msa_eager_dense_plan = dense_plan - _stage_sparse_plan_kv_lens_host(gqa_plan, kv_lens_cpu) # Stage the valid-block count to the device once for the whole step # (see _msa_eager_n_valid_blocks). n_valid_host = per_token_valid_blocks( @@ -699,14 +688,17 @@ def _build_msa_fields(self) -> None: # fine: forwards read only the persistent buffers filled below. # qo_offset is the prefix length, so one build covers prefill # (num_cached) and decode (kv_len - 1 with qo_len 1). - req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + mapping = build_paged_kv_slot_mapping( kv_cache_manager=kv_cache_manager, request_ids=request_ids, qo_lens_cpu=qo_lens_cpu, qo_offset_cpu=qo_offset_cpu, device=cache_device, ) - kv_indices = build_kv_page_indices(req_to_token, slot_ids, kv_lens_cpu, page_size) + out_cache_loc = mapping.out_cache_loc + # The page table comes from the same host block ids the mapping was + # built from, so it costs no device work. + kv_indices = build_kv_page_indices(mapping.block_ids_cpu, kv_lens_cpu, page_size) total_new_tokens = int(out_cache_loc.shape[0]) total_pages = int(kv_indices.shape[0]) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index 98cbdd9feb9d..4314f3cd70d0 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -11,6 +11,8 @@ import torch +from tensorrt_llm._utils import maybe_pin_memory + from .common import write_kv_slots # fmha_sm100 ships only head_dim 128 variants and the MiniMax-M3 checkpoint @@ -108,36 +110,36 @@ def write_msa_main_kv( def build_kv_page_indices( - req_to_token: torch.Tensor, - slot_ids: torch.Tensor, + block_ids_cpu: torch.Tensor, kv_lens_cpu: torch.Tensor, page_size: int, ) -> torch.Tensor: - """Build the flattened per-request page table fmha_sm100 consumes. - - Returns int32 global page ids concatenated per request. A request's - pages come from the first slot of each page in its req_to_token row. - Page ids are global and non-contiguous in production, so they are not - clamped to a per-request bound. + """Build the flattened per-request page table fmha_sm100 consumes, on the host. + + Returns int32 global page ids concatenated per request, request b + contributing the first ceil(kv_len / page_size) entries of its + block_ids_cpu row. A request with kv_len <= 0 contributes nothing, matching + the page count the plan derives from the same lengths. Page ids are global + and non-contiguous in production, so they are not clamped to a per-request + bound. + + block_ids_cpu is the [batch, max_blocks] host table + build_paged_kv_slot_mapping obtains from the cache manager. Both use the + manager's tokens_per_block as the page size, so + req_to_token[b, p * page_size] // page_size equals block_ids_cpu[b, p] and + the page table needs no device work. The result is pinned where that helps, + so the caller stages it with one asynchronous copy. """ - device = req_to_token.device - req_rows = req_to_token.index_select(0, slot_ids.to(torch.long)).to(torch.long) - batch = int(req_rows.shape[0]) - kv_lens_list = kv_lens_cpu.to(torch.long).tolist() - - page_lists = [] - for b in range(batch): - kv_len = int(kv_lens_list[b]) - if kv_len <= 0: - continue - num_pages = (kv_len + page_size - 1) // page_size - page_starts = torch.arange(num_pages, device=device, dtype=torch.long) * page_size - page_ids = req_rows[b].gather(0, page_starts) // page_size - page_lists.append(page_ids.to(torch.int32)) - - if page_lists: - return torch.cat(page_lists, dim=0) - return torch.empty(0, dtype=torch.int32, device=device) + pages = (kv_lens_cpu.to(torch.long) + (page_size - 1)) // page_size + pages.clamp_(min=0) + total_pages = int(pages.sum()) + if total_pages == 0: + return torch.empty(0, dtype=torch.int32) + batch = int(pages.shape[0]) + row = torch.repeat_interleave(torch.arange(batch, dtype=torch.long), pages) + starts = torch.cumsum(pages, 0) - pages + col = torch.arange(total_pages, dtype=torch.long) - starts[row] + return maybe_pin_memory(block_ids_cpu[row, col].to(torch.int32)) def per_token_valid_blocks( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py index be4b8ae24143..696f0f68055c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py @@ -324,7 +324,7 @@ def _build_runtime_metadata_fresh( prefix_lens_dev = prefix_lens.to(device) if prefix_lens.device != device else prefix_lens qo_lens_cpu = torch.tensor([int(x) for x in extend_seq_lens_cpu], dtype=torch.int32) qo_offset_cpu = prefix_lens.detach().to(device="cpu", dtype=torch.int32) - req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + req_to_token, slot_ids, out_cache_loc, _ = build_paged_kv_slot_mapping( kv_cache_manager=kv_cache_manager, request_ids=request_ids, qo_lens_cpu=qo_lens_cpu, @@ -351,7 +351,7 @@ def _build_runtime_metadata_fresh( # Decode: the new token sits at position seq_lens[b] - 1. qo_lens_cpu = torch.ones(batch, dtype=torch.int32) qo_offset_cpu = seq_lens_cpu.detach().to(device="cpu", dtype=torch.int32) - 1 - req_to_token, slot_ids, out_cache_loc = build_paged_kv_slot_mapping( + req_to_token, slot_ids, out_cache_loc, _ = build_paged_kv_slot_mapping( kv_cache_manager=kv_cache_manager, request_ids=request_ids, qo_lens_cpu=qo_lens_cpu, diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index a6a543b99724..201050818253 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -608,3 +608,98 @@ def test_msa_proxy_max_score_strided_index_k_matches_packed( assert not index_k_strided.is_contiguous() assert index_k_strided.stride(0) == coalescing_scale * page_size * head_dim assert torch.equal(strided_scores, packed_scores) + + +def _expand_slot_rows(block_ids: torch.Tensor, tokens_per_block: int) -> torch.Tensor: + """req_to_token reference: block_id * tokens_per_block + offset_in_block.""" + within = torch.arange(tokens_per_block, dtype=torch.int64) + grid = block_ids.to(torch.int64).unsqueeze(-1) * tokens_per_block + within + return grid.reshape(block_ids.shape[0], -1).to(torch.int32) + + +def test_build_kv_page_indices_matches_first_slot_of_each_page(): + """The host page table must equal the page ids each request's req_to_token + row holds at its page boundaries, since both use the manager's + tokens_per_block as the page size. Rows are ragged (0-padded block ids, + global and non-contiguous) and one request has no KV at all.""" + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import ( + build_kv_page_indices, + ) + + page_size = 8 + block_ids = torch.tensor( + [[11, 4, 7, 0], [5, 9, 0, 0], [3, 0, 0, 0], [21, 13, 6, 2]], + dtype=torch.int32, + ) + # 3 pages (partial last), 2 pages (exact), no pages, 4 pages. + kv_lens = torch.tensor([17, 16, 0, 32], dtype=torch.int32) + req_to_token = _expand_slot_rows(block_ids, page_size) + + reference = torch.cat( + [ + req_to_token[b, : int(kv_lens[b]) : page_size] // page_size + for b in range(block_ids.shape[0]) + ] + ) + page_indices = build_kv_page_indices(block_ids, kv_lens, page_size) + + assert page_indices.dtype == torch.int32 + assert page_indices.tolist() == [11, 4, 7, 5, 9, 21, 13, 6, 2] + torch.testing.assert_close(page_indices, reference, rtol=0, atol=0) + + +def test_build_paged_kv_slot_mapping_out_cache_loc_matches_slot_grid(): + """out_cache_loc must name the same slots as indexing req_to_token per new + token, for a mixed batch of one context request plus decode rows.""" + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import ( + build_paged_kv_slot_mapping, + ) + + tokens_per_block = 4 + block_ids = torch.tensor([[6, 2, 9], [4, 0, 0], [7, 1, 0]], dtype=torch.int32) + + class FakeCacheManager: + tokens_per_block = 4 + + def get_block_ids_per_seq(self, request_ids): + assert request_ids == [0, 1, 2] + return block_ids + + # Request 0 prefills 6 tokens over a 3-token prefix; 1 and 2 decode. + qo_lens_cpu = torch.tensor([6, 1, 1], dtype=torch.int32) + kv_lens_cpu = torch.tensor([9, 3, 5], dtype=torch.int32) + qo_offset_cpu = kv_lens_cpu - qo_lens_cpu + + mapping = build_paged_kv_slot_mapping( + kv_cache_manager=FakeCacheManager(), + request_ids=[0, 1, 2], + qo_lens_cpu=qo_lens_cpu, + qo_offset_cpu=qo_offset_cpu, + device=torch.device("cpu"), + ) + + req_to_token = _expand_slot_rows(block_ids, tokens_per_block) + reference = [ + int(req_to_token[b, int(qo_offset_cpu[b]) + offset]) + for b in range(3) + for offset in range(int(qo_lens_cpu[b])) + ] + + torch.testing.assert_close(mapping.req_to_token, req_to_token, rtol=0, atol=0) + assert mapping.slot_ids.tolist() == [0, 1, 2] + assert mapping.out_cache_loc.dtype == torch.int32 + assert mapping.out_cache_loc.tolist() == reference + assert mapping.block_ids_cpu.tolist() == block_ids.tolist() + + # A zero-length CUDA-graph padding row offsets to -1. Its slot is a + # placeholder no forward reads, so it only has to stay inside the row + # rather than index off the table. + padded = build_paged_kv_slot_mapping( + kv_cache_manager=FakeCacheManager(), + request_ids=[0, 1, 2], + qo_lens_cpu=torch.tensor([1, 1, 1], dtype=torch.int32), + qo_offset_cpu=torch.tensor([0, -1, -1], dtype=torch.int32), + device=torch.device("cpu"), + ) + for b, slot in enumerate(padded.out_cache_loc.tolist()): + assert slot in req_to_token[b].tolist()