[#10244][feat] AutoDeploy: separate prefill/decode in flashinfer #10252
Conversation
|
/bot run |
📝 WalkthroughWalkthroughThis pull request refactors the AutoDeploy system to standardize metadata parameter naming by introducing Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Pre-merge checks and finishing touches❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/auto_deploy/models/patches/bamba.py (1)
1-1: Add missing NVIDIA copyright header.Per coding guidelines, all TensorRT-LLM source files must include an NVIDIA copyright header with the year of latest meaningful modification (2025 in this case).
🔎 Suggested copyright header
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """A patch for the Bamba model to make it compatible with torch.export."""As per coding guidelines.
🧹 Nitpick comments (8)
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (1)
292-292: LGTM: Appropriate logging level adjustment.The change from
errortoinfois reasonable since this represents a graceful fallback when CUDA graph padding cannot be applied, rather than a critical failure. The system continues to function correctly without the optimization.Optionally, consider using
warninglevel if the performance impact of missing CUDA graph padding is significant enough to warrant operator attention.tensorrt_llm/_torch/auto_deploy/custom_ops/mla.py (1)
58-58: Prefix unused variable with underscore.The unpacked variable
num_prefill_tokensis extracted frombatch_info_hostbut never used in the function body. According to Python conventions and the static analysis hint, prefix it with an underscore to indicate it's intentionally unused.🔎 Proposed fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, num_decode = batch_info_host.tolist()tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py (1)
213-213: Prefix unused variable with underscore.The unpacked variable
num_prefill_tokensis extracted frombatch_info_hostbut never used in the function. Following Python conventions, prefix it with an underscore.🔎 Proposed fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, num_decode = batch_info_host.tolist()tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.py (1)
177-177: Prefix unused variable with underscore.The variable
num_prefill_tokensis unpacked but never used. Prefix it with an underscore to indicate intentional non-use.🔎 Proposed fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, num_decode = batch_info_host.tolist()tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.py (1)
148-148: Prefix unused variable with underscore.The variable
num_prefill_tokensis unpacked frombatch_info_hostbut never used in the function. Prefix it with an underscore.🔎 Proposed fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, num_decode = batch_info_host.tolist()tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py (1)
281-281: Prefix unused variable with underscore.The
num_prefill_tokensvariable is unpacked but never used in this function. Per Python convention, prefix it with an underscore to indicate it's intentionally unused.🔎 Proposed fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, num_decode = batch_info_host.tolist()tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py (1)
56-56: Prefix unused variables with underscores.
num_prefill_tokensandnum_decodeare unpacked but never used in_triton_ssm_prepare_metadata. Prefix them with underscores to indicate they're intentionally unused.🔎 Proposed fix
- num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_prefill, _num_prefill_tokens, _num_decode = batch_info_host.tolist()tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py (1)
98-98: Consider removing unusedseq_len_hostparameter.The
seq_len_hostparameter is passed toplan_prefillbut never used. If it was intended for future use (e.g., the commented-outmax_token_per_sequenceline), consider either removing it or documenting why it's kept.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (24)
tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.pytensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.pytensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.pytensorrt_llm/_torch/auto_deploy/models/patches/bamba.pytensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.pytests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/unittest/_torch/auto_deploy/_utils_test/torch_attention_reference.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_cuda_causal_conv_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_causal_conv_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gather_logits_before_lm_head.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces. Do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used
Python files should use snake_case naming:some_file.py
Python classes should use PascalCase naming:class SomeClass
Python functions and methods should use snake_case naming:def my_awesome_function():
Python local variables should use snake_case naming:my_variable = ...
Python variable names that start with a number should be prefixed with 'k':k_99th_percentile = ...
Python global variables should use upper snake_case with prefix 'G':G_MY_GLOBAL = ...
Python constants should use upper snake_case naming:MY_CONSTANT = ...
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings in Python for classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except to the smallest set of errors possible
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible, using the else block for logic
Files:
tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.pytensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.pytensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.pytests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gather_logits_before_lm_head.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_attention_op.pytensorrt_llm/_torch/auto_deploy/models/patches/bamba.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_cuda_causal_conv_cached_op.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla.pytensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_causal_conv_cached_op.pytests/unittest/_torch/auto_deploy/_utils_test/torch_attention_reference.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py
**/*.{cpp,h,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the year of its latest meaningful modification
Files:
tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.pytensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.pytensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.pytests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gather_logits_before_lm_head.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_attention_op.pytensorrt_llm/_torch/auto_deploy/models/patches/bamba.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_cuda_causal_conv_cached_op.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla.pytensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_attention_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_causal_conv_cached_op.pytests/unittest/_torch/auto_deploy/_utils_test/torch_attention_reference.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py
🧠 Learnings (9)
📚 Learning: 2025-08-08T04:10:19.038Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6728
File: cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp:966-966
Timestamp: 2025-08-08T04:10:19.038Z
Learning: TensorRT plugins currently don't support padding functionality, and TensorRT is not getting new features (in maintenance mode). This means that duplicating parameters like mExpertHiddenSize in function calls, even with TODO comments, can be acceptable as pragmatic solutions within these constraints.
Applied to files:
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
📚 Learning: 2025-08-19T12:45:11.997Z
Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.
Applied to files:
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
📚 Learning: 2025-08-14T21:04:50.248Z
Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Applied to files:
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py
📚 Learning: 2025-12-19T06:31:54.973Z
Learnt from: nvyocox
Repo: NVIDIA/TensorRT-LLM PR: 10117
File: tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rope_attention.py:336-339
Timestamp: 2025-12-19T06:31:54.973Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rope_attention.py, the cast to torch.float16 for qkv_node before creating the AttentionPlugin is intentional and required because DriveOS LLM expects float16 dtype specifically. This should not be changed to preserve original dtype or made configurable for bfloat16 models in the DriveOS LLM ONNX export path.
Applied to files:
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py
📚 Learning: 2025-10-20T16:54:09.824Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py:6-6
Timestamp: 2025-10-20T16:54:09.824Z
Learning: In tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py, the import `from ...modules.mamba.layernorm_gated import _layer_norm_fwd` is correct and should not be changed to modules.fla.layernorm_gated. The _layer_norm_fwd function exists in both modules/mamba/layernorm_gated.py and modules/fla/layernorm_gated.py, but the mamba version is the intended implementation for this use case.
Applied to files:
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.pytensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.pytensorrt_llm/_torch/auto_deploy/models/patches/bamba.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla.py
📚 Learning: 2025-08-15T06:46:53.813Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:53.813Z
Learning: In the TensorRT-LLM KV cache manager, SWA (Sliding Window Attention) combined with beam search is currently in a broken/non-functional state and is planned for future rework. During preparatory refactoring phases, code related to SWA+beam search may intentionally remain in a non-working state until the broader rework is completed.
Applied to files:
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py
📚 Learning: 2025-11-14T11:22:03.729Z
Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.
Applied to files:
tensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.pytensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.pytensorrt_llm/_torch/auto_deploy/models/patches/bamba.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_cuda_causal_conv_cached_op.pytensorrt_llm/_torch/auto_deploy/custom_ops/mla.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_causal_conv_cached_op.py
📚 Learning: 2025-10-20T17:07:18.745Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py:98-116
Timestamp: 2025-10-20T17:07:18.745Z
Learning: In NemotronH models (tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py), the gate (self.gate) returns topk_indices and topk_weights that are already in the correct shape to be passed directly to torch_ops.auto_deploy.torch_moe without needing to reshape them when hidden_states is flattened.
Applied to files:
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gather_logits_before_lm_head.py
📚 Learning: 2025-10-20T17:09:21.560Z
Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py:180-182
Timestamp: 2025-10-20T17:09:21.560Z
Learning: In tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py, the _gated_rmsnorm_replacement function does not need to cast the output of torch.ops.auto_deploy.torch_rmsnorm_gated back to the input dtype, even though the custom op returns fp32. The dtype handling is managed elsewhere or the fp32 output is acceptable for downstream consumers.
Applied to files:
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.py
🧬 Code graph analysis (3)
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gather_logits_before_lm_head.py (1)
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.py (1)
gather_logits_before_lm_head(5-30)
tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py (3)
tensorrt_llm/_torch/attention_backend/flashinfer.py (2)
page_size(200-204)kv_indptr(112-113)tensorrt_llm/models/modeling_utils.py (1)
kv_dtype(470-480)tensorrt_llm/_torch/flashinfer_utils.py (1)
get_env_enable_pdl(10-15)
tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (1)
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.py (1)
gather_logits_before_lm_head(5-30)
🪛 Ruff (0.14.10)
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py
281-281: Unpacked variable num_prefill_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
355-355: Unused function argument: batch_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.py
37-37: Unused function argument: logits_gather_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.py
123-123: Unused function argument: batch_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.py
148-148: Unpacked variable num_prefill_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
249-249: Unused function argument: batch_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.py
177-177: Unpacked variable num_prefill_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
250-250: Unused function argument: batch_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py
213-213: Unpacked variable num_prefill_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
293-293: Unused function argument: batch_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/mla.py
58-58: Unpacked variable num_prefill_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
169-169: Unused function argument: batch_info_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py
98-98: Unused method argument: seq_len_host
(ARG002)
215-215: Unused function argument: batch_info_host
(ARG001)
379-379: Unused function argument: batch_info_host
(ARG001)
380-380: Unused function argument: cu_seqlen_host
(ARG001)
381-381: Unused function argument: cu_num_pages
(ARG001)
382-382: Unused function argument: cu_num_pages_host
(ARG001)
383-383: Unused function argument: cache_loc
(ARG001)
384-384: Unused function argument: last_page_len
(ARG001)
385-385: Unused function argument: last_page_len_host
(ARG001)
386-386: Unused function argument: seq_len_with_cache_host
(ARG001)
387-387: Unused function argument: seq_len_host
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py
56-56: Unpacked variable num_prefill_tokens is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
56-56: Unpacked variable num_decode is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
77-77: Unused function argument: batch_info_host
(ARG001)
248-248: Unused function argument: batch_info_host
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (34)
tensorrt_llm/_torch/auto_deploy/models/patches/bamba.py (1)
52-62: Variable renaming improves clarity and is compatible with downstream custom ops.The
batch_info_host_tvariable properly indicates host-side metadata. Bothtorch_cached_causal_conv1dandtorch_cached_ssmcorrectly define thebatch_info_hostparameter, and the calls pass it consistently. The conditional logic correctly distinguishes decode (seq_len == 1) from prefill paths.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_triton_mamba_cached_op.py (1)
137-171: LGTM! Consistent naming convention update.The rename from
batch_info_tensortobatch_info_hostaligns with the broader PR objective to standardize host-side metadata naming across the codebase. The change is applied consistently to both the Torch reference and Triton test paths.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_cuda_causal_conv_cached_op.py (1)
62-75: LGTM! Consistent host metadata naming.The rename from
batch_infotobatch_info_hostis correctly applied in the generate-only test path, maintaining consistency with the broader naming convention update across the test suite.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_causal_conv_cached_op.py (2)
62-72: LGTM! Consistent naming in generate-only path.The rename to
batch_info_hostis correctly applied with updated comments and call sites.
127-140: LGTM! Consistent naming in context/prefill path.The rename to
batch_info_hostis correctly applied across the context phase test, maintaining consistency with the generate-only path.tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py (3)
50-83: LGTM! Public API updated with host-prefixed metadata.The parameter rename from
batch_infotobatch_info_hostin the function signature and its usage is correct. This aligns with the broader convention to explicitly denote host-side metadata tensors.
134-157: LGTM! Fake implementation signature updated.The fake function signature correctly mirrors the main implementation's parameter rename to
batch_info_host.
190-192: LGTM! Metadata args list updated.The
get_standard_metadata_argsreturn value correctly reflects the newbatch_info_hostnaming, ensuring consistency across the attention interface.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_attention_op.py (1)
128-153: LGTM! Consistent naming in flattened GQA test.The rename from
batch_infotobatch_info_hostis correctly applied to both the test operation and reference implementation calls, maintaining consistency across the validation path.tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py (1)
70-77: LGTM! Logits gather metadata renamed to host variant.The rename from
logits_gather_infotologits_gather_info_hostis correctly applied in both the input retrieval and the subsequent function call, maintaining consistency with the broader host-prefixed naming convention.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_attention_op.py (1)
249-291: LGTM! Test data flow updated with host-prefixed naming.The rename from
batch_infotobatch_info_hostis correctly propagated through the test data creation (_create_test_data), storage (data dict), and usage (_run_attention), ensuring consistency across all test cases that use these helper methods.tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mamba_cached_op.py (2)
69-86: LGTM! Generate-only path updated with host metadata naming.The rename to
batch_info_hostis correctly applied in the generate-only test with appropriate comment updates.
144-161: LGTM! Context/prefill path updated with host metadata naming.The rename to
batch_info_hostis correctly applied in the context phase test, maintaining consistency with the generate-only path and the broader naming convention.tensorrt_llm/_torch/auto_deploy/custom_ops/mla.py (1)
34-34: LGTM: Consistent host-suffix metadata naming.The parameter rename from
batch_infotobatch_info_hostaligns with the broader PR objective to standardize on host-side metadata tensors across attention backends. The update toget_standard_metadata_argsensures the public API contract reflects this change.Also applies to: 169-169, 215-215
tests/unittest/_torch/auto_deploy/_utils_test/torch_attention_reference.py (1)
43-51: LGTM: Test utilities properly updated.The test reference implementation correctly adopts the
batch_info_hostnaming convention across all method signatures and usages, maintaining consistency with the production code changes.Also applies to: 63-63, 87-87, 104-104, 147-148, 155-155, 173-173, 192-192
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gather_logits_before_lm_head.py (1)
59-59: LGTM: Test properly adopts host-side metadata.The test correctly updates to use
logits_gather_info_hostand creates these tensors on the CPU (device="cpu"), which aligns with the host-metadata naming convention introduced across the codebase. The changes are consistent with the production op signature intorch_gather_logits.py.Also applies to: 62-62, 85-85, 88-88, 108-108, 116-116, 135-135, 143-143, 220-220, 226-226, 281-281, 287-287
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_attention_op.py (2)
98-101: Verify device placement for batch_info_host.The test creates
batch_info_hosttensors on the CUDA device, while other host metadata tensors (qo_indptr_host,paged_kv_indptr_host, etc.) are explicitly created on CPU via.cpu(). This inconsistency may be intentional, but it's worth verifying thatbatch_info_hostshould remain on device rather than CPU like the other host-suffixed metadata.Example from line 99-101:
batch_info_host = torch.tensor( [BATCH_SIZE, BATCH_SIZE * SEQ_LEN, 0], dtype=torch.int32, device=device )Compare to lines 68-72 where other host tensors are on CPU:
qo_indptr_host = qo_indptr.cpu() paged_kv_indptr_host = paged_kv_indptr.cpu()Should
batch_info_hostalso be created on CPU for consistency?Also applies to: 245-247, 379-382, 473-474, 598-600, 761-764, 874-875, 972-973
67-72: LGTM: Host metadata tensors properly created.The test correctly creates host-side copies of metadata tensors (
qo_indptr_host,paged_kv_indptr_host,paged_kv_last_page_len_host, etc.) using.cpu()and passes them to the FlashInfer attention op, aligning with the new host-metadata interface.Also applies to: 182-187, 348-353, 451-456, 566-571, 700-705, 856-861, 955-960
tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py (1)
191-191: LGTM: Consistent host-suffix metadata naming.The parameter rename and metadata API update align with the PR's standardization of host-side metadata tensors across attention backends.
Also applies to: 293-293, 340-340
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.py (1)
150-150: LGTM: Consistent host-suffix metadata naming.The changes align with the broader refactoring to standardize on host-side metadata tensors for the cached causal convolution backend.
Also applies to: 250-250, 299-299
tensorrt_llm/_torch/auto_deploy/custom_ops/fla/fla_backend_delta.py (2)
61-61: LGTM: num_prefill_tokens is correctly used.Unlike some other files in this PR, this implementation correctly uses
num_prefill_tokensfor slicing operations (lines 79, 89, 98, 107). The static analysis hint suggesting it's unused appears to be incorrect.Also applies to: 78-79, 89-89, 98-98, 107-107
38-38: LGTM: Consistent host-suffix metadata naming.The parameter rename and metadata API update follow the established pattern across the codebase.
Also applies to: 123-123, 163-163
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_mamba.py (1)
124-124: LGTM: Consistent host-suffix metadata naming.The changes properly standardize on host-side metadata tensors for the Mamba SSM backend, consistent with the broader refactoring effort.
Also applies to: 249-249, 296-296
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py (1)
256-256: Consistent rename tobatch_info_hostacross the interface.The parameter rename from
batch_infotobatch_info_hostis applied consistently across the main function, its fake variant, andget_standard_metadata_args. This aligns with the broader PR goal of standardizing host-prefixed metadata naming.Also applies to: 355-355, 403-403
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_gather_logits.py (1)
8-8: Consistent rename tologits_gather_info_host.The parameter rename from
logits_gather_infotologits_gather_info_hostis applied consistently across the function signature, docstring, implementation, and fake variant. This aligns with the host-prefixed metadata naming convention.Also applies to: 15-15, 24-24, 37-37
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py (1)
45-45: Consistent rename tobatch_info_hostacross Triton SSM backend.The parameter rename is applied consistently across
_triton_ssm_prepare_metadata,_triton_cached_ssm, their fake variants, andget_standard_metadata_args. This aligns with the host-prefixed metadata naming convention used across the codebase.Also applies to: 113-113, 143-143, 248-248, 297-297
tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (4)
499-501: Host-variant argument expansion is well-designed.The
_available_argsset now includes both device and host variants for all buffer tensors (e.g.,batch_infoandbatch_info_host). This enables the attention interface to expose host-side tensors when needed for planning operations that require CPU-side data.
547-551: Host-aware argument retrieval correctly strips suffix.The
_get_argmethod properly handles host variants by stripping the_hostsuffix and fetching fromget_host_viewinstead of the device view. The shaping logic is also correctly applied based on whether the argument is in_shapeable_args.
795-796: Storage activation check includes host variants.The
_store_argmethod now checks if the host variant (f"{name}_host") is in_active_args, ensuring data is stored to the buffer when either the device or host variant is activated. This is necessary for the host-prefixed metadata to work correctly.
1016-1022: Host-side logits gather info used correctly.The
maybe_gather_and_squeeze_logitsmethod now useslogits_gather_info_hostto retrieve the host-side tensor and call.tolist()on it. This is correct since the gather info needs to be read on the CPU to determine the gather parameters.tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py (4)
91-123: Newplan_prefillmethod cleanly separates prefill planning.The prefill planning is now isolated in its own method with host-side tensor parameters. The method correctly:
- Checks for re-planning by comparing against
plan_params_prefill- Passes host tensors to the FlashInfer wrapper's plan method
- Returns the prefill wrapper for execution
This separation enables the FlashInfer backend to handle prefill and decode independently.
290-368: Split prefill/decode execution correctly assembles outputs.The execution flow properly handles three cases:
- Mixed batch (prefill + decode): Pre-allocates output, fills from both paths
- Prefill-only: Returns prefill output directly
- Decode-only: Returns decode output directly
The token indexing (
num_prefill_tokens,num_total_tokens) correctly partitions the query tensor.
436-446: Expanded standard metadata args include host variants.The
get_standard_metadata_argsnow returns both device and host variants for tensors that need both (e.g.,cu_num_pagesandcu_num_pages_host). This enables the attention interface to wire up both variants to the FlashInfer kernel.
350-355: The indexing pattern is correct and intentionally designed this way. Thecu_num_pagesarray contains absolute cumulative indices intocache_loc, which means slicingcu_num_pageswhile keepingcache_locfull produces the correct behavior. When a mixed prefill+decode batch exists, the fullcache_loccontains pages for all sequences in order, and the slicedcu_num_pages[num_prefill : num_seq + 1]still provides valid absolute indices into that full array. This pattern correctly handles FlashInfer's paged KV cache indexing semantics.
|
PR_Github #29687 [ run ] triggered by Bot. Commit: |
|
No perf implications for 1k/1k for nano 3 fp8 tp1: before (concurrency 1 and 384) vs after Also tested llama3-8b and also didn't notice any difference. Nsys trace looks reasonable: 
There are some (expected) bubbles at the beginning of forward from calling the plan i/f in flashinfer |
tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py
Outdated
Show resolved
Hide resolved
|
PR_Github #29687 [ run ] completed with state |
lucaslie
force-pushed
the
ll/fi_decode_prefill
branch
from
746b71f to
b85376f
Compare
|
/bot run |
|
PR_Github #30196 [ run ] triggered by Bot. Commit: |
|
PR_Github #30196 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #30205 [ run ] triggered by Bot. Commit: |
|
PR_Github #30205 [ run ] completed with state
|
Signed-off-by: Lucas Liebenwein <[email protected]>
Signed-off-by: Lucas Liebenwein <[email protected]>
Signed-off-by: Lucas Liebenwein <[email protected]>
Signed-off-by: Lucas Liebenwein <[email protected]>
lucaslie
force-pushed
the
ll/fi_decode_prefill
branch
from
d628399 to
0dc78ea
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #30298 [ run ] triggered by Bot. Commit: |
|
PR_Github #30298 [ run ] completed with state |
3 participants
Summary by CodeRabbit
New Features
Refactor
✏️ Tip: You can customize this high-level summary in your review settings.
Description
fixes #10244
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.