From 7fd49922ce891c2cc0684b2eb4a340accda58334 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Tue, 31 Mar 2026 08:05:54 -0700 Subject: [PATCH 1/4] [https://nvbugs/5983390][perf] Reduce host overhead in DSA MLA attention path Pass pre-computed num_contexts/num_ctx_tokens to thop::attention and trtllm_gen_attention to eliminate per-layer sum().item() calls that recompute batch structure from host_request_types/host_context_lengths. Move view/slice/reinterpret ops from Python _update_k_cache into the C++ indexer_k_cache_scatter_op kernel: accept original k_fp8 (FP8) and k_scale (float32) tensors directly with num_tokens, avoiding per-layer torch.empty, view, as_strided and slice overhead on the host. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/thop/bindings.cpp | 3 +- .../thop/IndexerKCacheScatterOp.cpp | 99 ++++++++----------- cpp/tensorrt_llm/thop/attentionOp.cpp | 26 +++-- cpp/tensorrt_llm/thop/attentionOp.h | 3 +- .../_torch/attention_backend/sparse/dsa.py | 32 ++---- .../_torch/attention_backend/trtllm.py | 13 ++- .../_torch/attention_backend/trtllm_gen.py | 6 +- .../attention/sparse/test_dsa_indexer.py | 9 +- 8 files changed, 93 insertions(+), 98 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index fc161ab4a6ca..66350f69d3fd 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -70,7 +70,8 @@ void initBindings(nb::module_& m) nb::arg("cu_kv_seqlens") = std::nullopt, nb::arg("fmha_scheduler_counter") = std::nullopt, nb::arg("mla_bmm1_scale") = std::nullopt, nb::arg("mla_bmm2_scale") = std::nullopt, nb::arg("quant_q_buffer") = std::nullopt, nb::arg("flash_mla_tile_scheduler_metadata") = std::nullopt, - nb::arg("flash_mla_num_splits") = std::nullopt, "Multi-head attention operation", + nb::arg("flash_mla_num_splits") = std::nullopt, nb::arg("opt_num_contexts") = std::nullopt, + nb::arg("opt_num_ctx_tokens") = std::nullopt, "Multi-head attention operation", nb::call_guard()); m.def( diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp index 940d59258ca8..2d7f7f79f88d 100644 --- a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp @@ -28,69 +28,56 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -void indexer_k_cache_scatter_op(th::Tensor const& k_fp8_bytes, th::Tensor const& k_scale_bytes, th::Tensor& k_cache, - th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale) +void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_scale, th::Tensor& k_cache, + th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t num_tokens) { - // Validate all tensors are CUDA tensors - TORCH_CHECK(k_fp8_bytes.is_cuda() && k_scale_bytes.is_cuda() && k_cache.is_cuda() && slot_mapping_fp8.is_cuda() + // k_fp8: [>=num_tokens, head_dim] in FP8 (1 byte/element) — reinterpreted as uint8 + // k_scale: [>=num_tokens, head_dim // quant_block_size] in float32 — reinterpreted as uint8 bytes + // slot_mapping_fp8, slot_mapping_scale: [>=num_tokens] int64 — only first num_tokens used + // k_cache: [num_blocks, block_size, 1, per_token_size] uint8 + + TORCH_CHECK(k_fp8.is_cuda() && k_scale.is_cuda() && k_cache.is_cuda() && slot_mapping_fp8.is_cuda() && slot_mapping_scale.is_cuda(), "All tensors must be CUDA tensors"); - // Validate tensor dimensions - TORCH_CHECK(k_fp8_bytes.dim() == 2, "k_fp8_bytes must be a 2D Tensor [num_tokens, head_dim]"); - TORCH_CHECK(k_scale_bytes.dim() == 2, "k_scale_bytes must be a 2D Tensor [num_tokens, scale_size]"); - TORCH_CHECK(slot_mapping_fp8.dim() == 1, "slot_mapping_fp8 must be a 1D Tensor [num_tokens]"); - TORCH_CHECK(slot_mapping_scale.dim() == 1, "slot_mapping_scale must be a 1D Tensor [num_tokens]"); - - // Enforce k_cache is 4D tensor - TORCH_CHECK(k_cache.dim() == 4, - "k_cache must be a 4D Tensor [num_blocks, block_size, 1, per_token_size], got %d dimensions", + TORCH_CHECK(k_fp8.dim() == 2, "k_fp8 must be 2D [num_tokens, head_dim]"); + TORCH_CHECK(k_scale.dim() == 2, "k_scale must be 2D [num_tokens, scale_elements]"); + TORCH_CHECK(k_cache.dim() == 4, "k_cache must be 4D [num_blocks, block_size, 1, per_token_size], got %d dims", static_cast(k_cache.dim())); - // Validate tensor dtypes - TORCH_CHECK(k_fp8_bytes.scalar_type() == torch::kUInt8, "k_fp8_bytes must be uint8"); - TORCH_CHECK(k_scale_bytes.scalar_type() == torch::kUInt8, "k_scale_bytes must be uint8"); - TORCH_CHECK(slot_mapping_fp8.scalar_type() == torch::kInt64, "slot_mapping_fp8 must be int64"); - TORCH_CHECK(slot_mapping_scale.scalar_type() == torch::kInt64, "slot_mapping_scale must be int64"); - - // Validate tensor shapes are consistent - auto num_tokens = static_cast(k_fp8_bytes.size(0)); - TORCH_CHECK( - k_scale_bytes.size(0) == num_tokens, "k_scale_bytes first dimension must equal k_fp8_bytes first dimension"); - TORCH_CHECK(slot_mapping_fp8.size(0) == num_tokens, "slot_mapping_fp8 length must equal num_tokens"); - TORCH_CHECK(slot_mapping_scale.size(0) == num_tokens, "slot_mapping_scale length must equal num_tokens"); - - // Validate tensors are contiguous (except k_cache which may be non-contiguous) - TORCH_CHECK(k_fp8_bytes.is_contiguous(), "k_fp8_bytes must be contiguous"); - TORCH_CHECK(k_scale_bytes.is_contiguous(), "k_scale_bytes must be contiguous"); - // k_cache can be non-contiguous - we handle this via strides + TORCH_CHECK(k_fp8.is_contiguous(), "k_fp8 must be contiguous"); + TORCH_CHECK(k_scale.is_contiguous(), "k_scale must be contiguous"); TORCH_CHECK(slot_mapping_fp8.is_contiguous(), "slot_mapping_fp8 must be contiguous"); TORCH_CHECK(slot_mapping_scale.is_contiguous(), "slot_mapping_scale must be contiguous"); - int32_t head_dim = static_cast(k_fp8_bytes.size(1)); // head_dim = quant_block_size = 128 - int32_t scale_size = static_cast(k_scale_bytes.size(1)); // scale_size = 4 bytes - - int32_t cache_dim_0 = static_cast(k_cache.size(0)); // num_blocks - int32_t cache_dim_1 = static_cast(k_cache.size(1)); // block_size - int32_t cache_dim_2 = static_cast(k_cache.size(2)); // num_kv_heads - int32_t cache_dim_3 = static_cast(k_cache.size(3)); // per_token_size - - // Validation for indexer k cache pool for DeepSeek-V3.2 constraints - TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1 for DeepSeek-V3.2, got %d", cache_dim_2); - TORCH_CHECK(head_dim == 128, "k_fp8_bytes head_dim must be 128 for DeepSeek-V3.2, got %d", head_dim); - TORCH_CHECK(scale_size == 4, "k_scale_bytes scale_size must be 4 bytes for DeepSeek-V3.2, got %d", scale_size); - - int64_t cache_stride_0 = static_cast(k_cache.stride(0)); - int64_t cache_stride_1 = static_cast(k_cache.stride(1)); - int64_t cache_stride_2 = static_cast(k_cache.stride(2)); - int64_t cache_stride_3 = static_cast(k_cache.stride(3)); - - auto stream = at::cuda::getCurrentCUDAStream(k_fp8_bytes.get_device()); - - tk::invokeIndexerKCacheScatter(k_fp8_bytes.data_ptr(), k_scale_bytes.data_ptr(), - k_cache.data_ptr(), slot_mapping_fp8.data_ptr(), slot_mapping_scale.data_ptr(), - num_tokens, head_dim, scale_size, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, - cache_stride_1, cache_stride_2, cache_stride_3, stream); + // FP8 is 1 byte per element, so head_dim in elements == head_dim in bytes. + int32_t const head_dim = static_cast(k_fp8.size(1)); + // float32 scale: each element is 4 bytes. + int32_t const scale_size = static_cast(k_scale.size(1)) * 4; + + int32_t const cache_dim_0 = static_cast(k_cache.size(0)); + int32_t const cache_dim_1 = static_cast(k_cache.size(1)); + int32_t const cache_dim_2 = static_cast(k_cache.size(2)); + int32_t const cache_dim_3 = static_cast(k_cache.size(3)); + + TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1, got %d", cache_dim_2); + TORCH_CHECK(head_dim == 128, "k_fp8 head_dim must be 128, got %d", head_dim); + TORCH_CHECK(scale_size == 4, "k_scale scale_size must be 4 bytes, got %d", scale_size); + + int64_t const cache_stride_0 = static_cast(k_cache.stride(0)); + int64_t const cache_stride_1 = static_cast(k_cache.stride(1)); + int64_t const cache_stride_2 = static_cast(k_cache.stride(2)); + int64_t const cache_stride_3 = static_cast(k_cache.stride(3)); + + auto stream = at::cuda::getCurrentCUDAStream(k_fp8.get_device()); + + // Reinterpret k_fp8 as uint8 bytes and k_scale as raw bytes via data_ptr. + // For slot mappings, use data_ptr directly — only the first num_tokens entries are read. + tk::invokeIndexerKCacheScatter(reinterpret_cast(k_fp8.data_ptr()), + reinterpret_cast(k_scale.data_ptr()), k_cache.data_ptr(), + slot_mapping_fp8.data_ptr(), slot_mapping_scale.data_ptr(), static_cast(num_tokens), + head_dim, scale_size, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, + cache_stride_2, cache_stride_3, stream); } } // namespace torch_ext @@ -100,8 +87,8 @@ TRTLLM_NAMESPACE_END TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( - "indexer_k_cache_scatter_op(Tensor k_fp8_bytes, Tensor k_scale_bytes, Tensor(a!) k_cache, " - "Tensor slot_mapping_fp8, Tensor slot_mapping_scale) -> ()"); + "indexer_k_cache_scatter_op(Tensor k_fp8, Tensor k_scale, Tensor(a!) k_cache, " + "Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 9a7af4da49f6..6f54d3288bd6 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -630,7 +630,8 @@ void attention(torch::Tensor q, std::optional k, std::optional cu_q_seqlens, std::optional cu_kv_seqlens, std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, std::optional mla_bmm2_scale, std::optional quant_q_buffer, - std::optional flash_mla_tile_scheduler_metadata, std::optional flash_mla_num_splits) + std::optional flash_mla_tile_scheduler_metadata, std::optional flash_mla_num_splits, + std::optional opt_num_contexts, std::optional opt_num_ctx_tokens) { TLLM_LOG_TRACE("Attention op starts at layer %d", layer_idx); // Use these tensors to infer if the attention is using KV cache @@ -833,19 +834,28 @@ void attention(torch::Tensor q, std::optional k, std::optional(opt_num_contexts.value()); + } + else + { + num_contexts = 0; + for (int32_t idx = 0; idx < num_seqs; idx++) { - break; + if (request_types[idx] != RequestType::kCONTEXT) + { + break; + } + ++num_contexts; } - ++num_contexts; } int32_t const num_generations = num_seqs - num_contexts; int32_t const num_tokens = qkv_or_q.size(0); - int32_t const num_ctx_tokens = host_context_lengths.slice(0, 0, num_contexts).sum().item(); + int32_t const num_ctx_tokens = opt_num_ctx_tokens.has_value() + ? static_cast(opt_num_ctx_tokens.value()) + : host_context_lengths.slice(0, 0, num_contexts).sum().item(); int32_t const num_gen_tokens = is_gen_only ? num_tokens : num_tokens - num_ctx_tokens; auto const ctx_total_kv_len = host_total_kv_lens.index({0}).item(); auto const gen_total_kv_len = host_total_kv_lens.index({1}).item(); diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index 0fc4788d6f0b..c98dec356d76 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -78,7 +78,8 @@ void attention(torch::Tensor q, std::optional k, std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, std::optional mla_bmm2_scale, std::optional quant_q_buffer, std::optional flash_mla_tile_scheduler_metadata = std::nullopt, - std::optional flash_mla_num_splits = std::nullopt); + std::optional flash_mla_num_splits = std::nullopt, + std::optional opt_num_contexts = std::nullopt, std::optional opt_num_ctx_tokens = std::nullopt); struct KvCachePoolPointers { diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index b3a5fbeb3136..566dbaf388f2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1392,34 +1392,18 @@ def _update_k_cache(self, k_fp8: torch.Tensor, k_scale: torch.Tensor, if metadata.kv_cache_manager is None or metadata.slot_mapping_fp8 is None: return - # [num_blocks, block_size, 1, per_token_size ] k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers( self.layer_idx) num_tokens = k_fp8.shape[0] - head_dim = k_fp8.shape[1] - scale_size = k_scale.shape[1] * 4 # Convert to bytes (float32 = 4 bytes) - - # Convert to bytes: flatten first, then view as uint8, then reshape - k_fp8_bytes = k_fp8.view(-1).view(torch.uint8).view( - num_tokens, head_dim) - - # k_scale: for single-element tensors, contiguous() may be no-op - # Fix stride(-1) for byte-level view - k_scale_flat = k_scale.view(-1) - if k_scale_flat.stride(-1) != 1: - k_scale_flat = torch.as_strided(k_scale_flat.contiguous(), - size=(k_scale_flat.numel(), ), - stride=(1, )) - k_scale_bytes = k_scale_flat.view(torch.uint8).view( - num_tokens, scale_size) - - # Use CUDA kernel to scatter FP8 and scale bytes into cache - flat_indices_fp8 = metadata.slot_mapping_fp8[:num_tokens] - flat_indices_scale = metadata.slot_mapping_scale[:num_tokens] - torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8_bytes, k_scale_bytes, - k_cache, flat_indices_fp8, - flat_indices_scale) + + # The C++ op reinterprets k_fp8 (FP8) and k_scale (float32) as raw + # bytes internally and only reads the first num_tokens entries from + # the slot mapping buffers, avoiding Python-side view/slice overhead. + torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache, + metadata.slot_mapping_fp8, + metadata.slot_mapping_scale, + num_tokens) def sparse_attn_indexer( self, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 271538190c8a..b4a76b0e6b79 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -407,6 +407,9 @@ def run( mla_bmm1_scale: Optional[torch.Tensor] = None, mla_bmm2_scale: Optional[torch.Tensor] = None, quant_q_buffer: Optional[torch.Tensor] = None, + num_contexts: int = 0, + num_generations: int = 0, + num_ctx_tokens: int = 0, ): """ Run the attention operation. @@ -639,6 +642,9 @@ def run( self.quant_config, self.kv_cache_manager, global_layer_idx=self.global_layer_idx, + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, ) else: thop.attention( @@ -722,6 +728,8 @@ def run( quant_q_buffer, self.flash_mla_tile_scheduler_metadata, self.flash_mla_num_splits, + num_contexts, + num_ctx_tokens, ) if self.print_skip_softmax_stat: @@ -2049,7 +2057,10 @@ def forward( fmha_scheduler_counter=fmha_scheduler_counter, mla_bmm1_scale=mla_bmm1_scale, mla_bmm2_scale=mla_bmm2_scale, - quant_q_buffer=quant_q_buffer) + quant_q_buffer=quant_q_buffer, + num_contexts=metadata.num_contexts, + num_generations=metadata.num_generations, + num_ctx_tokens=metadata.num_ctx_tokens) if output_sf is None: return output diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index a11d8f51fe19..40b2b430cd56 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -1562,6 +1562,9 @@ def trtllm_gen_attention( quant_config: Optional[QuantConfig], kv_cache_manager: Optional[KVCacheManager], global_layer_idx: Optional[int] = None, + num_contexts: int = 0, + num_generations: int = 0, + num_ctx_tokens: int = 0, ) -> None: """ TrtLLM-Gen attention using flashinfer backend. @@ -1691,9 +1694,6 @@ def trtllm_gen_attention( if attention_input_type is not None: attn_input_type = AttentionInputType(attention_input_type) - num_contexts, num_generations = _parse_request_types(host_request_types) - - num_ctx_tokens = int(host_context_lengths[:num_contexts].sum()) if num_contexts > 0 else 0 num_gen_tokens = num_tokens - num_ctx_tokens # Prepare Workspace diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py index 46bb102fcf7a..244c4d5dbc2a 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py @@ -703,7 +703,7 @@ def test_indexer_k_cache_scatter_custom_op(): dtype=torch.bfloat16) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_original) - # Prepare byte-level data + # Prepare byte-level data for the Python reference path scale_size = k_scale.shape[1] * 4 k_fp8_bytes = k_fp8.view(-1).view(torch.uint8).view(num_tokens, head_dim) k_scale_flat = k_scale.view(-1) @@ -742,9 +742,10 @@ def test_indexer_k_cache_scatter_custom_op(): # ========== Path 1: CUDA Kernel ========== print(f"\n=== Path 1: CUDA Kernel ===") - torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8_bytes, k_scale_bytes, - k_cache_cuda, flat_indices_fp8, - flat_indices_scale) + torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache_cuda, + metadata.slot_mapping_fp8, + metadata.slot_mapping_scale, + num_tokens) torch.cuda.synchronize() print(f"✓ CUDA kernel completed") From 4afa40d16a70a3348a94e4fb40dff45f587184ac Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Tue, 31 Mar 2026 19:34:19 -0700 Subject: [PATCH 2/4] [None][chore] Make trtllm_gen_attention batch params required Remove default values from num_contexts, num_generations, and num_ctx_tokens parameters in trtllm_gen_attention() so callers cannot silently omit them. Move them before the optional global_layer_idx parameter. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm.py | 6 +++--- tensorrt_llm/_torch/attention_backend/trtllm_gen.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index b4a76b0e6b79..e70bab6f1425 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -641,10 +641,10 @@ def run( quant_q_buffer, self.quant_config, self.kv_cache_manager, + num_contexts, + num_generations, + num_ctx_tokens, global_layer_idx=self.global_layer_idx, - num_contexts=num_contexts, - num_generations=num_generations, - num_ctx_tokens=num_ctx_tokens, ) else: thop.attention( diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index 40b2b430cd56..fc2831d900ff 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -1561,10 +1561,10 @@ def trtllm_gen_attention( quant_q_buffer: Optional[torch.Tensor], quant_config: Optional[QuantConfig], kv_cache_manager: Optional[KVCacheManager], + num_contexts: int, + num_generations: int, + num_ctx_tokens: int, global_layer_idx: Optional[int] = None, - num_contexts: int = 0, - num_generations: int = 0, - num_ctx_tokens: int = 0, ) -> None: """ TrtLLM-Gen attention using flashinfer backend. From ca6df439b14a441578a73075d907b61597314087 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Thu, 2 Apr 2026 05:17:34 -0700 Subject: [PATCH 3/4] [https://nvbugs/5983390][fix] Address PR review feedback - Remove std::optional from num_contexts/num_ctx_tokens in attention op (always passed by caller, never null) - Remove num_generations param from trtllm_gen_attention interface; compute it inside as host_request_types.size(0) - num_contexts - Remove unused _parse_request_types function - Restore tensor dimension/dtype validation in IndexerKCacheScatterOp with element_size checks to guard reinterpret_cast safety Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/thop/bindings.cpp | 5 ++-- .../thop/IndexerKCacheScatterOp.cpp | 14 ++++++++-- cpp/tensorrt_llm/thop/attentionOp.cpp | 26 +++---------------- cpp/tensorrt_llm/thop/attentionOp.h | 3 +-- .../_torch/attention_backend/trtllm.py | 3 --- .../_torch/attention_backend/trtllm_gen.py | 19 +------------- 6 files changed, 19 insertions(+), 51 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 66350f69d3fd..b71c39d40874 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -70,9 +70,8 @@ void initBindings(nb::module_& m) nb::arg("cu_kv_seqlens") = std::nullopt, nb::arg("fmha_scheduler_counter") = std::nullopt, nb::arg("mla_bmm1_scale") = std::nullopt, nb::arg("mla_bmm2_scale") = std::nullopt, nb::arg("quant_q_buffer") = std::nullopt, nb::arg("flash_mla_tile_scheduler_metadata") = std::nullopt, - nb::arg("flash_mla_num_splits") = std::nullopt, nb::arg("opt_num_contexts") = std::nullopt, - nb::arg("opt_num_ctx_tokens") = std::nullopt, "Multi-head attention operation", - nb::call_guard()); + nb::arg("flash_mla_num_splits") = std::nullopt, nb::arg("num_contexts") = 0, nb::arg("num_ctx_tokens") = 0, + "Multi-head attention operation", nb::call_guard()); m.def( "get_helix_workspace_size_per_rank", diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp index 2d7f7f79f88d..f5a1336ea3e1 100644 --- a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp @@ -40,11 +40,21 @@ void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_sca && slot_mapping_scale.is_cuda(), "All tensors must be CUDA tensors"); + // Validate tensor dimensions TORCH_CHECK(k_fp8.dim() == 2, "k_fp8 must be 2D [num_tokens, head_dim]"); TORCH_CHECK(k_scale.dim() == 2, "k_scale must be 2D [num_tokens, scale_elements]"); + TORCH_CHECK(slot_mapping_fp8.dim() == 1, "slot_mapping_fp8 must be 1D [num_tokens]"); + TORCH_CHECK(slot_mapping_scale.dim() == 1, "slot_mapping_scale must be 1D [num_tokens]"); TORCH_CHECK(k_cache.dim() == 4, "k_cache must be 4D [num_blocks, block_size, 1, per_token_size], got %d dims", static_cast(k_cache.dim())); + // Validate tensor dtypes — reinterpret_cast below assumes specific element sizes + TORCH_CHECK(k_fp8.element_size() == 1, "k_fp8 must have 1-byte elements (e.g. FP8), got %d", k_fp8.element_size()); + TORCH_CHECK(k_scale.element_size() == 4, "k_scale must have 4-byte elements (e.g. float32), got %d", + k_scale.element_size()); + TORCH_CHECK(slot_mapping_fp8.scalar_type() == torch::kInt64, "slot_mapping_fp8 must be int64"); + TORCH_CHECK(slot_mapping_scale.scalar_type() == torch::kInt64, "slot_mapping_scale must be int64"); + TORCH_CHECK(k_fp8.is_contiguous(), "k_fp8 must be contiguous"); TORCH_CHECK(k_scale.is_contiguous(), "k_scale must be contiguous"); TORCH_CHECK(slot_mapping_fp8.is_contiguous(), "slot_mapping_fp8 must be contiguous"); @@ -52,8 +62,8 @@ void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_sca // FP8 is 1 byte per element, so head_dim in elements == head_dim in bytes. int32_t const head_dim = static_cast(k_fp8.size(1)); - // float32 scale: each element is 4 bytes. - int32_t const scale_size = static_cast(k_scale.size(1)) * 4; + // Scale size in bytes: num_scale_elements * bytes_per_element. + int32_t const scale_size = static_cast(k_scale.size(1)) * static_cast(k_scale.element_size()); int32_t const cache_dim_0 = static_cast(k_cache.size(0)); int32_t const cache_dim_1 = static_cast(k_cache.size(1)); diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 6f54d3288bd6..b526310564e9 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -631,7 +631,7 @@ void attention(torch::Tensor q, std::optional k, std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, std::optional mla_bmm2_scale, std::optional quant_q_buffer, std::optional flash_mla_tile_scheduler_metadata, std::optional flash_mla_num_splits, - std::optional opt_num_contexts, std::optional opt_num_ctx_tokens) + int64_t num_contexts, int64_t num_ctx_tokens) { TLLM_LOG_TRACE("Attention op starts at layer %d", layer_idx); // Use these tensors to infer if the attention is using KV cache @@ -834,29 +834,9 @@ void attention(torch::Tensor q, std::optional k, std::optional(opt_num_contexts.value()); - } - else - { - num_contexts = 0; - for (int32_t idx = 0; idx < num_seqs; idx++) - { - if (request_types[idx] != RequestType::kCONTEXT) - { - break; - } - ++num_contexts; - } - } - int32_t const num_generations = num_seqs - num_contexts; + int32_t const num_generations = num_seqs - static_cast(num_contexts); int32_t const num_tokens = qkv_or_q.size(0); - int32_t const num_ctx_tokens = opt_num_ctx_tokens.has_value() - ? static_cast(opt_num_ctx_tokens.value()) - : host_context_lengths.slice(0, 0, num_contexts).sum().item(); - int32_t const num_gen_tokens = is_gen_only ? num_tokens : num_tokens - num_ctx_tokens; + int32_t const num_gen_tokens = is_gen_only ? num_tokens : num_tokens - static_cast(num_ctx_tokens); auto const ctx_total_kv_len = host_total_kv_lens.index({0}).item(); auto const gen_total_kv_len = host_total_kv_lens.index({1}).item(); diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index c98dec356d76..854de818292f 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -78,8 +78,7 @@ void attention(torch::Tensor q, std::optional k, std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, std::optional mla_bmm2_scale, std::optional quant_q_buffer, std::optional flash_mla_tile_scheduler_metadata = std::nullopt, - std::optional flash_mla_num_splits = std::nullopt, - std::optional opt_num_contexts = std::nullopt, std::optional opt_num_ctx_tokens = std::nullopt); + std::optional flash_mla_num_splits = std::nullopt, int64_t num_contexts, int64_t num_ctx_tokens); struct KvCachePoolPointers { diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index e70bab6f1425..d0b710cae5e2 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -408,7 +408,6 @@ def run( mla_bmm2_scale: Optional[torch.Tensor] = None, quant_q_buffer: Optional[torch.Tensor] = None, num_contexts: int = 0, - num_generations: int = 0, num_ctx_tokens: int = 0, ): """ @@ -642,7 +641,6 @@ def run( self.quant_config, self.kv_cache_manager, num_contexts, - num_generations, num_ctx_tokens, global_layer_idx=self.global_layer_idx, ) @@ -2059,7 +2057,6 @@ def forward( mla_bmm2_scale=mla_bmm2_scale, quant_q_buffer=quant_q_buffer, num_contexts=metadata.num_contexts, - num_generations=metadata.num_generations, num_ctx_tokens=metadata.num_ctx_tokens) if output_sf is None: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index fc2831d900ff..edb45df74059 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -1362,23 +1362,6 @@ def run_generation(self, params: EnqueueGenerationParams): ) -def _parse_request_types(host_request_types: torch.Tensor) -> Tuple[int, int]: - """ - Parse request types to count context and generation requests. - - Args: - host_request_types: Request types tensor (0=context, 1=generation). - num_seqs: Total number of sequences. - - Returns: - Tuple of (num_contexts, num_generations). - """ - - num_generations = host_request_types.sum().item() - num_contexts = host_request_types.size(0) - num_generations - return num_contexts, num_generations - - def is_supported( q: torch.Tensor, num_heads: int, @@ -1562,7 +1545,6 @@ def trtllm_gen_attention( quant_config: Optional[QuantConfig], kv_cache_manager: Optional[KVCacheManager], num_contexts: int, - num_generations: int, num_ctx_tokens: int, global_layer_idx: Optional[int] = None, ) -> None: @@ -1694,6 +1676,7 @@ def trtllm_gen_attention( if attention_input_type is not None: attn_input_type = AttentionInputType(attention_input_type) + num_generations = host_request_types.size(0) - num_contexts num_gen_tokens = num_tokens - num_ctx_tokens # Prepare Workspace From 756846f073dd8621030fda829ee81cc2c0098228 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Thu, 2 Apr 2026 20:40:34 -0700 Subject: [PATCH 4/4] [None][fix] Pass num_contexts and num_ctx_tokens in AutoDeploy thop.attention call Store num_contexts and num_ctx_tokens in _TrtllmPlanner during plan_host and pass them as keyword args to thop.attention, matching the updated C++ binding signature. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- .../auto_deploy/custom_ops/attention/trtllm_attention.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py index d0c8a1dd8c0f..d6d4b75da2af 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py @@ -82,6 +82,9 @@ def __init__(self): # keeping a separate copy here since we sometimes have to overwrite the original values self.host_past_kv_lengths: Optional[torch.Tensor] = None # [max_batch] int32 pinned self.host_context_lengths: Optional[torch.Tensor] = None # [max_batch] int32 pinned + # Batch counts for thop.attention (updated every forward in plan_host) + self.num_contexts: int = 0 + self.num_ctx_tokens: int = 0 # Persistent block_offsets buffer for CUDA graph compatibility. # Pre-allocated to max size so the tensor address is stable across replays. self.block_offsets: Optional[torch.Tensor] = None @@ -171,6 +174,10 @@ def plan_host( """ num_seq = num_prefill + num_decode + # Batch counts for thop.attention + self.num_contexts = num_prefill + self.num_ctx_tokens = int(seq_len_host[:num_prefill].sum()) if num_prefill > 0 else 0 + # host_request_types: 0 = prefill (context), 1 = decode (generation) self.host_request_types[:num_prefill].fill_(0) self.host_request_types[num_prefill:num_seq].fill_(1) @@ -500,6 +507,8 @@ def trtllm_mha_with_cache( None, # mla_bmm1_scale None, # mla_bmm2_scale None, # quant_q_buffer + num_contexts=_GlobalTrtllmPlanner.num_contexts, + num_ctx_tokens=_GlobalTrtllmPlanner.num_ctx_tokens, ) if out is not None: