[None][perf] Address inter-iter idle times

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:

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
 
 

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,19 +191,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
@@ -210,44 +220,65 @@ 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
     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]
+        )
+        # 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",
 ]