[None][perf] Address inter-idle times and decode-first assumption in MSA

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
 
 

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,19 +196,27 @@ 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,
     request_ids,
     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,46 +225,66 @@ 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
     req_to_token = slot_grid.reshape(batch, max_kv_len).to(torch.int32)
     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",
 ]

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])

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(

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,