[None][perf] Avoid implicit device-scalar syncs in DeepSeek-V4 ctx sparse metadata

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py

@@ -274,6 +274,7 @@ def __post_init__(self):
         super().__post_init__()
         self.num_total_compressed_tokens = {}
         self.max_ctx_compressed_tokens = {}
+        self._ctx_output_sizes: Optional[Dict[int, int]] = None
         sparse_metadata_params = self.sparse_metadata_params
         if not isinstance(sparse_metadata_params, DeepSeekV4MetadataParams):
             raise ValueError("DeepSeek-V4 sparse attention metadata params are not set")
@@ -736,12 +737,18 @@ def prepare(self):
         kv_lens_slice = kv_lens[:num_requests]
         cached_slice = cached_token_lens[:num_requests]
 
+        # Host-side per-ratio ctx compressed-token counts (Python ints), so
+        # _compute_ctx_compressed_position_ids never reads a device scalar
+        # (implicit D2H + stream sync) for its arange size / slice bound.
+        ctx_output_sizes: Optional[Dict[int, int]] = None
         if num_contexts > 0:
             # Prefill path: need per-request tensor ops for ctx scalar metadata.
+            ctx_output_sizes = {}
             for compress_ratio in self.compress_ratio_set:
                 new_comp_kv_lens = kv_lens_slice // compress_ratio - cached_slice // compress_ratio
                 cu_new = new_comp_kv_lens.cumsum(0)
                 num_ctx_compressed_tokens = cu_new[num_contexts - 1].item()
+                ctx_output_sizes[compress_ratio] = num_ctx_compressed_tokens
                 num_gen_compressed_tokens = num_generations * (
                     (num_gen_tokens_per_seq + compress_ratio - 1) // compress_ratio
                 )
@@ -760,12 +767,15 @@ def prepare(self):
                 )
                 self.max_ctx_compressed_tokens[compress_ratio] = 0
 
+        # Cached for on_update_kv_lens(); see the reuse gate there.
+        self._ctx_output_sizes = ctx_output_sizes
+
         # 2) CUDA-side: fill *_cuda buffers on device.
         kv_lens_cuda = (
             self.cached_token_lens_cuda[:num_requests] + self._seq_lens_cuda[:num_requests]
         )
         cached_tokens_cuda = self.cached_token_lens_cuda[:num_requests]
-        self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda)
+        self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda, ctx_output_sizes)
 
         self._compute_compressed_mask(
             self.new_comp_kv_lens_cuda,
@@ -780,6 +790,7 @@ def prepare_compressed_kv_metadata(
         self,
         kv_lens: torch.Tensor,
         cached_tokens: torch.Tensor,
+        ctx_output_sizes: Optional[Dict[int, int]] = None,
     ):
         """Compute per-ratio compressed KV lens and position IDs on device.
 
@@ -788,6 +799,12 @@ def prepare_compressed_kv_metadata(
         Args:
             kv_lens: Total KV lengths per request (device tensor, [batch_size]).
             cached_tokens: Cached token counts per request (device tensor, [batch_size]).
+            ctx_output_sizes: Optional per-ratio host-computed ctx
+                compressed-token counts (Python ints); avoids implicit
+                device-scalar reads (D2H + stream sync) in the ctx position-id
+                computation. prepare() always passes it; on_update_kv_lens()
+                reuses the cached copy unless the extend_ctx path may have
+                mutated ctx-row kv_lens on device.
         """
         batch_size = kv_lens.shape[0]
         num_contexts = self.num_contexts
@@ -811,6 +828,7 @@ def prepare_compressed_kv_metadata(
                 self.compressed_position_ids_cuda,
                 num_contexts,
                 self._compress_ratios_sorted,
+                ctx_output_sizes,
             )
 
         if self.num_gen_tokens_per_seq > 0 and num_generations > 0:
@@ -847,7 +865,13 @@ def on_update_kv_lens(self):
             num_gen_tokens // self.num_generations if self.num_generations > 0 else 0
         )
 
-        self.prepare_compressed_kv_metadata(kv_lens, cached_tokens)
+        # Reuse prepare()'s host-computed ctx sizes unless the extend_ctx path
+        # (num_chunked_ctx_requests > 0) may have mutated ctx-row kv_lens on
+        # device; every other path only changes gen rows.
+        ctx_output_sizes = (
+            self._ctx_output_sizes if getattr(self, "num_chunked_ctx_requests", 0) == 0 else None
+        )
+        self.prepare_compressed_kv_metadata(kv_lens, cached_tokens, ctx_output_sizes)
 
         self._compute_compressed_mask(
             self.new_comp_kv_lens_cuda,
@@ -1003,14 +1027,24 @@ def _compute_ctx_compressed_position_ids(
         compressed_position_ids_bufs: Dict[int, torch.Tensor],
         num_contexts: int,
         compress_ratios: list,
+        ctx_output_sizes: Optional[Dict[int, int]] = None,
     ):
-        """Context-only compressed position IDs (eager, data-dependent shapes)."""
+        """Context-only compressed position IDs (eager, data-dependent shapes).
+
+        ctx_output_sizes (host ints) keeps the arange size and slice bound off
+        the device; the 0-dim-CUDA fallback costs two implicit D2H syncs per
+        ratio.
+        """
         device = past_kv_lens_bufs[compress_ratios[0]].device
         for compress_ratio in compress_ratios:
             past_kv = past_kv_lens_bufs[compress_ratio]
             cu_new_comp = cu_new_comp_kv_bufs[compress_ratio]
 
-            total_ctx_comp = cu_new_comp[num_contexts]
+            total_ctx_comp = (
+                ctx_output_sizes[compress_ratio]
+                if ctx_output_sizes is not None
+                else cu_new_comp[num_contexts]
+            )
             ctx_idx = torch.arange(total_ctx_comp, dtype=torch.int32, device=device)
             ctx_cu = cu_new_comp[: num_contexts + 1].to(torch.int32)
             ctx_req = torch.searchsorted(ctx_cu[1:], ctx_idx, right=True)

tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py

@@ -449,6 +449,50 @@ def test_mixed_context_generation_position_ids_follow_compact_output():
     assert actual_position_ids == [0, 4, 4, 4]
 
 
+@pytest.mark.parametrize(
+    "compress_ratio,cached_tokens,kv_lens",
+    [
+        pytest.param(1, [0], [4096], id="cr1_single"),
+        pytest.param(4, [0, 75, 4000], [4096, 4171, 8096], id="cr4_multi_boundary"),
+        pytest.param(128, [0], [64], id="cr128_empty_output"),
+        pytest.param(128, [973632, 8064], [990016, 8192], id="cr128_chunked_long_ctx"),
+    ],
+)
+def test_ctx_position_ids_host_sizes_match_device_scalar_fallback(
+    compress_ratio, cached_tokens, kv_lens
+):
+    """Host-int ctx_output_sizes must reproduce the device-scalar fallback exactly.
+
+    prepare() threads host-computed ctx compressed-token counts into
+    _compute_ctx_compressed_position_ids so the arange size and slice bound
+    are Python ints (no implicit D2H + stream sync). Both paths must produce
+    identical position IDs, including the untouched padding tail.
+    """
+    num_contexts = len(kv_lens)
+    cached = torch.tensor(cached_tokens, dtype=torch.int32, device=DEVICE)
+    kv = torch.tensor(kv_lens, dtype=torch.int32, device=DEVICE)
+    past = (cached // compress_ratio).to(torch.int32)
+    new_comp = (kv // compress_ratio).to(torch.int32) - past
+    cu = F.pad(torch.cumsum(new_comp, dim=0), (1, 0)).to(torch.int32)
+    total = int(cu[num_contexts].item())
+
+    def _run(ctx_output_sizes):
+        out = torch.full((total + 8,), -1, dtype=torch.int32, device=DEVICE)
+        DeepseekV4TrtllmAttentionMetadata._compute_ctx_compressed_position_ids(
+            {compress_ratio: past},
+            {compress_ratio: cu},
+            {compress_ratio: out},
+            num_contexts,
+            [compress_ratio],
+            ctx_output_sizes,
+        )
+        return out
+
+    golden = _run(None)
+    fast = _run({compress_ratio: total})
+    assert torch.equal(golden, fast)
+
+
 def precompute_freqs_cis(
     dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow
 ) -> torch.Tensor: