diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index 6f32f6b5a470..beed6c6b3bdb 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -1,3 +1,19 @@ +diff --git a/python/fmha_sm100/api.py b/python/fmha_sm100/api.py +index b02f747..b69cd8b 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, diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py index d72b17a..e4ad52e 100644 --- a/python/fmha_sm100/cute/interface.py @@ -183,10 +199,65 @@ index 21c777e..c22beef 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..81bf1ff 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: @@ -285,7 +356,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: 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 2baefd5cc6c1..6fb73c055888 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 @@ -469,20 +474,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 c508563bb720..b1f244313809 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: @@ -189,6 +191,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, @@ -196,12 +207,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 @@ -210,21 +220,26 @@ 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. - 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 @@ -232,22 +247,38 @@ 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] + ) + # Optimistic (full-acceptance) offsets can overhang the last allocated + # slot and a zero-length padding row lands at -1. Those slots are + # placeholders that on_update_kv_lens re-derives before any forward + # reads them, so keep them in bounds 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 501a1b925697..4c1b7410b008 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 .common import ( MiniMaxM3SparseConfig, @@ -60,31 +61,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: - """Keep the sparse-prefill plan's kv_segment_lens on the host. - - The fmha_sm100 sparse-prefill plan (MM-SA-Nv) rebuilds its KV page table - every sparse layer in sparse_fmha._build_page_table, which reads - plan["kv_segment_lens"] on the host via .tolist(). The planner stores it on - the device, so that read is a blocking D2H sync per sparse layer. Store a - host copy of the per-request KV lengths (already available as kv_lens_cpu) - instead; the page table still builds on the device from kv_indices, so no - copy is added. Mirrors the in-place plan-dict edits in on_update_kv_lens. - - Only the sparse-prefill sub-plan is MM-SA-Nv. Decode and dense plans keep - their device kv_segment_lens for the kernel. - """ - has_mixed, split = plan[0], plan[1] - # Non-mixed batches are one sparse plan (plan[3]). Mixed batches place the - # sparse prefill rows in the second sub-plan (plan[4]), after 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, *, @@ -125,6 +101,9 @@ def _worst_case_proxy_max_k_tiles( "qo_segment_lens", "kv_segment_lens", "qo_offset", + # Sparse plans mask with seqused_k rather than kv_segment_lens, so it is a + # length mirror on_update_kv_lens patches and must be equally graph-stable. + "seqused_k", ) _MSA_PLAN_INT64_KEYS = ("packed_work_range", "packed_work_info") # fmha_sm100 sizes packed_work_info at 131072 * max(num_kv_splits, 1); forcing @@ -138,6 +117,22 @@ def _worst_case_proxy_max_k_tiles( "workspace_o", "workspace_lse", ) +# The per-row lengths on_update_kv_lens patches, by plan flavour. Dense plans +# mask with kv_segment_lens/qo_offset; sparse plans mask with seqused_k and keep +# kv_segment_lens on the host for the page-table builder, which consumed it at +# plan time. cu_seqlens_k is excluded: it describes the page-table layout staged +# by prepare(), which stays valid because corrections only shrink lengths. +_MSA_DENSE_LENGTH_KEYS = ("kv_segment_lens", "qo_offset") +_MSA_SPARSE_LENGTH_KEYS = ("seqused_k",) + + +def _msa_plan_length_keys(sub_plan: dict) -> tuple: + """The length mirrors to patch in one fmha_sm100 sub-plan. + + fmha_sm100 tags sparse plans with "MM-SA-Nv"; the two flavours expose the + attended length under different keys. + """ + return _MSA_SPARSE_LENGTH_KEYS if sub_plan.get("MM-SA-Nv") else _MSA_DENSE_LENGTH_KEYS class _MsaGraphSafePlan: @@ -254,6 +249,9 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): msa_kv_indices: Optional[torch.Tensor] = None msa_max_score: Optional[torch.Tensor] = None msa_n_valid_blocks: Optional[torch.Tensor] = None + # Per-request kv_lens as staged by prepare(), before the overlap scheduler + # corrects them. on_update_kv_lens clamps against this; see there. + msa_kv_lens_staged: Optional[torch.Tensor] = None # Layer whose K/V/index-K caches were already written this step by the # fused scatter (msa_write_layer_caches); run_msa_paged_gqa consumes and # clears it so the legacy per-cache writes are skipped exactly once. @@ -280,12 +278,10 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # Eager (prefill/mixed) per-token valid-block count. It is layer-invariant # (a function of qo/kv lengths and page size), so it is computed on the host # and staged to the device once per step via a non-blocking copy_, then - # reused by every sparse layer's indexer. _msa_eager_all_blocks_empty caches - # the host-side empty-selection result so the indexer can short-circuit without - # a device read. _msa_eager_n_valid_buf is the persistent backing store for the view. + # reused by every sparse layer's indexer. _msa_eager_n_valid_buf is the + # persistent backing store for the view. _msa_eager_n_valid_buf: Optional[torch.Tensor] = None _msa_eager_n_valid_blocks: Optional[torch.Tensor] = None - _msa_eager_all_blocks_empty: bool = False def __post_init__(self) -> None: super().__post_init__() @@ -295,17 +291,24 @@ def __post_init__(self) -> None: self._msa_live_batch = 0 self._msa_live_total_q = 0 self._msa_page_size = 0 - self._msa_corrected_kv_lens_cpu: Optional[torch.Tensor] = None + self._msa_q_token_starts = (0,) self._create_msa_buffers() @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]: @@ -316,8 +319,6 @@ def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: slots and page counts need the true attended length, so it is excluded here. """ - if self._msa_corrected_kv_lens_cpu is not None: - return self._msa_corrected_kv_lens_cpu kv_lens = getattr(self, "kv_lens", None) if self.seq_lens is None or kv_lens is None: return None @@ -326,7 +327,9 @@ def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]: extra = params.num_extra_kv_tokens if params is not None else 0 if extra: out = out - extra - 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]: @@ -335,7 +338,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]: @@ -376,11 +379,6 @@ def msa_eager_n_valid_blocks(self) -> Optional[torch.Tensor]: step was prepared (a decode step or a structural test).""" return self._msa_eager_n_valid_blocks - @property - def msa_eager_all_blocks_empty(self) -> bool: - """Whether the eager step has no valid KV blocks for any query token.""" - return self._msa_eager_all_blocks_empty - def _msa_main_kv_is_fp8(self) -> bool: """Whether the main paged K/V cache is stored as FP8 E4M3. @@ -467,6 +465,13 @@ def _create_msa_buffers(self) -> None: dtype=torch.int32, capture_graph=capture_graph, ) + self.msa_kv_lens_staged = self.get_empty( + buffers, + (max_num_sequences,), + cache_name="msa_kv_lens_staged", + dtype=torch.int32, + capture_graph=capture_graph, + ) # The proxy scratch needs the fmha_sm100 plan geometry. This metadata # exists only for the MSA backend, whose selection already required the # kernels, so a failed import here is a hard error rather than a reason @@ -593,19 +598,43 @@ def _ensure_eager_n_valid_buffer(self, total_q: int, device: torch.device) -> to def prepare(self) -> None: super().prepare() - self._msa_corrected_kv_lens_cpu = None self._build_msa_fields() self._build_step_plans() + def _msa_live_plans(self) -> tuple: + """The fmha_sm100 plans in play for this step. + + Decode steps populate the graph-safe owners, prefill and mixed steps the + plain eager tuples. _build_step_plans clears whichever set does not + apply, so only one set is ever live. + """ + plans = [ + owner.plan + for owner in (self._msa_proxy_plan, self._msa_gqa_plan, self._msa_dense_plan) + if owner is not None and owner.plan is not None + ] + plans.extend( + plan + for plan in ( + self._msa_eager_proxy_plan, + self._msa_eager_gqa_plan, + self._msa_eager_dense_plan, + ) + if plan is not None + ) + return tuple(plans) + def on_update_kv_lens(self) -> None: """Re-derive length-dependent MSA state from the corrected kv_lens_cuda. The overlap scheduler corrects kv_lens_cuda on device after prepare() - staged optimistic (full-acceptance) lens. Decode steps are patched - with pure device ops (capture-safe); the correction only shrinks - lengths, so the host-baked plan worklists and the page-table layout - stay valid. Eager mixed batches install corrected host lens and - rebuild (small D2H sync). Idempotent. + staged optimistic (full-acceptance) lens. The correction only shrinks + lengths, so the host-baked plan worklists and the page-table layout stay + valid and only the per-row lengths need patching. The clamp against + msa_kv_lens_staged enforces that bound rather than trusting it: a longer + length would drive the kernels past the extent prepare() planned for. + Decode, prefill and mixed steps share one device-side patch, so this is + capture-safe, sync-free and idempotent. """ super().on_update_kv_lens() if not self._msa_fields_ready: @@ -614,18 +643,7 @@ def on_update_kv_lens(self) -> None: total_q = self._msa_live_total_q if batch <= 0 or total_q <= 0: return - if self.msa_decode_proxy_plan is None: - # Prefill/mixed step: the eager plans and the page-table layout - # were built from the optimistic host mirrors, so install the - # corrected lens and rebuild both. Mixed steps are never captured. - if torch.cuda.is_current_stream_capturing(): - return - self._msa_corrected_kv_lens_cpu = self.kv_lens_cuda[:batch].to("cpu", torch.int32) - self._build_msa_fields() - self._build_step_plans() - return - - kv_true = self.kv_lens_cuda[:batch] + kv_true = torch.minimum(self.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch]) qbr = self.msa_q_batch_row[:total_q].to(torch.long) qo_dev = self.msa_qo_lens_dev[:batch] kv_true_tok = kv_true[qbr] @@ -642,31 +660,61 @@ def on_update_kv_lens(self) -> None: # block, which would NaN the fully-masked GQA row. page = self._msa_page_size n_valid = torch.div((pos + 1).clamp_min(1) + (page - 1), page, rounding_mode="floor") - self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32)) - - # Plan mirrors: proxy and dense keep one row per request; the sparse - # GQA plan is row-expanded per token with the ladder in the per-row - # offset. qo_offset must stay non-negative: negative values hit the - # kernel's packed-length sentinel fallback. - off_req = (kv_true - qo_dev).clamp_min(0) - for owner, expanded in ( - (self._msa_proxy_plan, False), - (self._msa_gqa_plan, True), - (self._msa_dense_plan, False), - ): - decode = owner.plan[3] - seg_lens = decode.get("kv_segment_lens") - qo_off = decode.get("qo_offset") - if expanded: - if seg_lens is not None: - seg_lens[:total_q].copy_(kv_true_tok.to(seg_lens.dtype)) - if qo_off is not None: - qo_off[:total_q].copy_(pos.clamp_min(0).to(qo_off.dtype)) - else: - if seg_lens is not None: - seg_lens[:batch].copy_(kv_true.to(seg_lens.dtype)) - if qo_off is not None: - qo_off[:batch].copy_(off_req.to(qo_off.dtype)) + n_valid_buf = ( + self.msa_n_valid_blocks + if self.msa_decode_proxy_plan is not None + else self._msa_eager_n_valid_blocks + ) + if n_valid_buf is not None: + n_valid_buf[:total_q].copy_(n_valid.to(torch.int32)) + + # Plan length mirrors. A plan is (has_mixed, split, batch, decode_sub, + # prefill_sub); a mixed batch is split at the first long request, so + # each sub-plan mirrors only its own request range. Within a range a + # sub-plan holds either one row per request or one per query token, + # since the planner row-expands dense plans over query tokens, so the row + # count selects the source. qo_offset must stay non-negative: negative + # values hit the kernel's packed-length sentinel fallback. + per_request = { + "kv_segment_lens": kv_true, + "qo_offset": (kv_true - qo_dev).clamp_min(0), + "seqused_k": kv_true, + } + per_token = { + "kv_segment_lens": kv_true_tok, + "qo_offset": pos.clamp_min(0), + "seqused_k": (pos + 1).clamp_min(0), + } + starts = self._msa_q_token_starts + for has_mixed, split, _, decode_sub, prefill_sub in self._msa_live_plans(): + subs = ( + ((decode_sub, 0, split), (prefill_sub, split, batch)) + if has_mixed + else ((decode_sub, 0, batch),) + ) + for sub, first, last in subs: + if sub is None: + continue + tok_first, tok_last = starts[first], starts[last] + for key in _msa_plan_length_keys(sub): + dst = sub.get(key) + if dst is None: + raise RuntimeError( + f"MSA plan has no length mirror {key!r}, so the corrected " + "kv_lens cannot reach the kernel." + ) + rows = int(dst.shape[0]) + if rows == last - first: + src = per_request[key][first:last] + elif rows == tok_last - tok_first: + src = per_token[key][tok_first:tok_last] + else: + raise RuntimeError( + f"MSA plan {key!r} has {rows} rows for requests " + f"[{first}, {last}); expected {last - first} (one per " + f"request) or {tok_last - tok_first} (one per query token)." + ) + dst.copy_(src) def _build_step_plans(self) -> None: """Build the three layer-invariant fmha_sm100 plans once per step. @@ -693,7 +741,6 @@ def _build_step_plans(self) -> None: self._msa_eager_gqa_plan = None self._msa_eager_dense_plan = None self._msa_eager_n_valid_blocks = None - self._msa_eager_all_blocks_empty = False if not self._msa_fields_ready: return # Geometry is captured in __post_init__; skip when it is unavailable. @@ -768,19 +815,19 @@ 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 - # Avoid the per-layer page-table D2H sync in sparse prefill. - _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). The empty-selection check runs - # here on the host, so run_indexer can short-circuit cheaply. + # (see _msa_eager_n_valid_blocks). clamp_min(1) matches the decode + # path and on_update_kv_lens: a zero-valid row would -inf-mask every + # block and NaN the GQA row that consumes the selection. n_valid_host = per_token_valid_blocks( qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size ) total_q = int(n_valid_host.shape[0]) - self._msa_eager_all_blocks_empty = total_q == 0 or int(n_valid_host.max().item()) <= 0 if total_q > 0: dev_buf = self._ensure_eager_n_valid_buffer(total_q, device) - dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True) + dev_buf[:total_q].copy_( + n_valid_host.clamp_min(1).to(torch.int32), non_blocking=True + ) self._msa_eager_n_valid_blocks = dev_buf[:total_q] return @@ -866,14 +913,18 @@ 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) + req_to_token = mapping.req_to_token + 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]) @@ -903,9 +954,22 @@ def _build_msa_fields(self) -> None: torch.arange(total_new_tokens, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) ).to(torch.int32) - self.msa_q_batch_row[:total_new_tokens].copy_(batch_row_cpu) - self.msa_q_intra[:total_new_tokens].copy_(intra_cpu) - self.msa_qo_lens_dev[:batch_size].copy_(qo_lens_cpu) + # Pinned and non-blocking: a plain copy_ from pageable host memory ends + # in cudaStreamSynchronize, which drains the whole queue. + self.msa_q_batch_row[:total_new_tokens].copy_( + maybe_pin_memory(batch_row_cpu), non_blocking=True + ) + self.msa_q_intra[:total_new_tokens].copy_(maybe_pin_memory(intra_cpu), non_blocking=True) + self.msa_qo_lens_dev[:batch_size].copy_(maybe_pin_memory(qo_lens_cpu), non_blocking=True) + # Snapshot the staged lens as the upper bound on_update_kv_lens clamps + # to. Device-to-device, so it stays sync-free. + kv_lens_cuda = getattr(self, "kv_lens_cuda", None) + if kv_lens_cuda is not None: + self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cuda[:batch_size], non_blocking=True) + # Token offset of each request, plus total_new_tokens as the tail. Host + # side, so on_update_kv_lens can slice a sub-plan's token range without + # a device read. + self._msa_q_token_starts = (0, *torch.cumsum(qo_long, 0).tolist()) self._msa_live_batch = batch_size self._msa_live_total_q = total_new_tokens self._msa_page_size = page_size @@ -1132,20 +1196,9 @@ def run_indexer( else: proxy_plan = metadata.msa_eager_proxy_plan max_score = None - # The empty case was resolved on the host, so short-circuit here. - if metadata.msa_eager_all_blocks_empty: - output_shape = ( - (config.num_kv_heads, num_tokens, MSA_REQUIRED_TOPK) - if head_major_output - else (num_tokens, config.num_kv_heads, MSA_REQUIRED_TOPK) - ) - output = torch.full( - output_shape, - -1, - dtype=torch.int32, - device=idx_q.device, - ) - return output.permute(1, 0, 2) if head_major_output else output + # No host-side empty check: the staged counts are clamped to at + # least one block, and the kernel masks each query to its own + # valid-block extent. n_valid_blocks = metadata.msa_eager_n_valid_blocks if n_valid_blocks is not None: n_valid_blocks = n_valid_blocks[:num_tokens] 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 ee3e1dc9ca9a..fa9b4c8e4f5c 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 @@ -145,36 +147,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 a57cf5eac2e1..6c8e84a3b04b 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 @@ -332,7 +332,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, @@ -359,7 +359,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 f3c3ee9f35ae..cea8ea419b70 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 @@ -216,7 +216,10 @@ def select_blocks(self, idx_q, idx_k, **kwargs): class FakeMetadata: msa_decode_proxy_plan = None msa_eager_proxy_plan = (False, 0, 2, {}, None) - msa_eager_all_blocks_empty = False + # Two decode requests. run_indexer reads both counts to route the + # selection output head-major or token-major. + num_contexts = 0 + num_generations = 2 msa_eager_n_valid_blocks = torch.ones(2, dtype=torch.int32, device="cuda") msa_kv_indices = torch.arange(2, dtype=torch.int32, device="cuda") msa_qo_lens_cpu = torch.ones(2, dtype=torch.int32) @@ -273,7 +276,6 @@ def select_blocks(self, *args, **kwargs): class FakeMetadata: msa_decode_proxy_plan = None msa_eager_proxy_plan = ("eager",) - msa_eager_all_blocks_empty = False msa_eager_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32) msa_kv_indices = torch.arange(num_tokens, dtype=torch.int32) msa_qo_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32) @@ -283,11 +285,14 @@ class FakeMetadata: def __init__(self): self.num_contexts = num_contexts self.num_generations = num_generations - self.idx_k_cache = None + # run_indexer reads the index-K cache before it writes this layer's + # index-K, so the fake has to hold a tensor from the start, as the + # persistent cache does in production. + self.idx_k_cache = torch.zeros(num_tokens, 1, sparse_index_dim) def msa_write_idx_k(self, layer_idx, idx_k): del layer_idx - self.idx_k_cache = idx_k + self.idx_k_cache.copy_(idx_k) def msa_idx_k_cache(self, layer_idx): del layer_idx @@ -423,6 +428,101 @@ def test_per_token_valid_blocks_multi_token_decode(): assert n_valid.tolist() == [3, 1, 2, 2] +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 that on_update_kv_lens re-derives, 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() + + def _reference_scatter_write(k_cache, v_cache, idx_cache, slots, k, v, idx_k): num_tokens = int(slots.shape[0]) num_heads, head_dim = int(k_cache.shape[1]), int(k_cache.shape[3])