diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 7caf1c6ecc28..7806cb204007 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/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) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index 0784d7256264..3ba703006935 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/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: