[#11529][perf] AD host time attention MD optimization for large context

[MrGeva]

Summary by CodeRabbit

Release Notes

• New Features

  • Added helper function for efficient list-to-tensor conversion with numpy-backed optimization.

• Refactor

  • Updated InputBuffer.store() method signature to accept torch.Tensor instead of lists.
  • Converted internal data structures from Python lists to tensors.
  • Optimized data operations using vectorized torch functions.

Description

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 the stage-list parameter 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.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip 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-pipeline

Reuse 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR refactors the attention interface module to use PyTorch tensors internally instead of Python lists. It changes InputBuffer.store to accept tensors directly, introduces a _list_to_tensor helper for efficient conversions, replaces Python loops with vectorized torch operations for index computations, and optimizes memory transfers using numpy memcpy.

Changes

Cohort / File(s)

Summary

Attention Interface Refactoring tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

Modified InputBuffer.store signature to accept torch.Tensor instead of List[Number]. Added _list_to_tensor helper function and _TORCH_TO_NUMPY_DTYPE mapping for efficient conversions. Changed internal storage in _args_list from lists to tensors. Replaced Python loops with vectorized torch operations (e.g., torch.arange, torch.repeat_interleave) for page index computations. Optimized memory transfers using numpy memcpy. Updated properties to return Python lists via tensor conversions.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description is essentially empty with only the template structure included. Required sections like Description, Test Coverage, and substantive checklist responses were not filled out.	Add a description explaining the issue and solution, list relevant test cases that validate the tensor conversion changes, and confirm checklist items (coding guidelines compliance, tests provided, dependencies reviewed, documentation updated).

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title specifically references 'AD host time attention MD optimization for large context', which directly aligns with the summary describing host-to-device transfer optimizations and tensor-based performance improvements in the attention interface module.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (3)

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (3)

44-65: Consider documenting the bfloat16 limitation.

The _TORCH_TO_NUMPY_DTYPE mapping excludes torch.bfloat16 since numpy doesn't support it natively. The fallback at line 65 handles this correctly, but a brief comment would clarify this intentional omission for future maintainers.

📝 Suggested documentation

 # Torch dtype → numpy dtype for fast list-to-tensor conversion.
 # numpy's list→array conversion is ~2-3x faster than torch.tensor(list) for large lists.
+# Note: bfloat16 is intentionally excluded (numpy lacks native support); falls back to torch.tensor().
 _TORCH_TO_NUMPY_DTYPE: Dict[torch.dtype, np.dtype] = {

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py` around
lines 44 - 65, Add a short comment above the _TORCH_TO_NUMPY_DTYPE mapping (or
near the _list_to_tensor function) noting that torch.bfloat16 is intentionally
omitted because NumPy lacks native bfloat16 support, and that the fallback path
in _list_to_tensor (torch.tensor(...)) will handle bfloat16 tensors; reference
the _TORCH_TO_NUMPY_DTYPE constant and the _list_to_tensor function so
maintainers know why bfloat16 isn't mapped to a NumPy dtype.

---

606-619: Properties trigger repeated tensor-to-list conversions.

The properties seq_len, input_pos, cache_loc, and pages_per_seq now call .tolist() on each access. In nest_sequences, these are accessed multiple times (e.g., self.input_pos at lines 990 and 1044, self.seq_len at lines 990 and 1044). Each access triggers a new conversion.

For large batches, consider caching the list locally within nest_sequences to avoid repeated conversions.

Based on learnings: "ensure it does not access torch.Tensor objects (CPU or GPU) inside the loop. Instead, the code should use .tolist() to convert batched data tensors to Python lists beforehand."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py` around
lines 606 - 619, The property accessors seq_len, input_pos, cache_loc, and
pages_per_seq currently call .tolist() on every access causing repeated
tensor-to-list conversions; modify nest_sequences to call these properties once
at the start (e.g., seq_len_list = self.seq_len, input_pos_list =
self.input_pos, cache_loc_list = self.cache_loc, pages_per_seq_list =
self.pages_per_seq) and use those local lists throughout the method (instead of
repeated self.<prop> accesses) so no tensor .tolist() is performed inside loops;
alternatively, compute and store the .tolist() results in local variables
immediately after entering nest_sequences and reference those locals everywhere
in the method.

---

214-235: Docstring should specify CPU tensor requirement.

The implementation at lines 233-234 calls .numpy() on the input tensor, which requires the tensor to be on CPU. While internal callers handle this correctly, the method's docstring should clarify this constraint to prevent misuse.

📝 Suggested docstring update

         """Store a tensor into the pinned host buffer.
 
         Args:
             name: Name of the tensor to store to.
-            data: 1-D torch.Tensor to store.
+            data: 1-D torch.Tensor on CPU to store (GPU tensors not supported).
             fill_value: Optional value to fill the entire buffer with before storing.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py` around
lines 214 - 235, Update the docstring for the "Store a tensor into the pinned
host buffer." method to state that the data argument must be a CPU torch.Tensor
(1-D) because the implementation calls .numpy() on the input (which requires the
tensor to be on CPU); explicitly document any expectations about dtype
conversion (the method will convert to the buffer dtype if needed) and mention
that callers should provide data on CPU to avoid runtime errors when invoking
host_view[:length].numpy() and data.numpy().

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py`:
- Around line 44-65: Add a short comment above the _TORCH_TO_NUMPY_DTYPE mapping
(or near the _list_to_tensor function) noting that torch.bfloat16 is
intentionally omitted because NumPy lacks native bfloat16 support, and that the
fallback path in _list_to_tensor (torch.tensor(...)) will handle bfloat16
tensors; reference the _TORCH_TO_NUMPY_DTYPE constant and the _list_to_tensor
function so maintainers know why bfloat16 isn't mapped to a NumPy dtype.
- Around line 606-619: The property accessors seq_len, input_pos, cache_loc, and
pages_per_seq currently call .tolist() on every access causing repeated
tensor-to-list conversions; modify nest_sequences to call these properties once
at the start (e.g., seq_len_list = self.seq_len, input_pos_list =
self.input_pos, cache_loc_list = self.cache_loc, pages_per_seq_list =
self.pages_per_seq) and use those local lists throughout the method (instead of
repeated self.<prop> accesses) so no tensor .tolist() is performed inside loops;
alternatively, compute and store the .tolist() results in local variables
immediately after entering nest_sequences and reference those locals everywhere
in the method.
- Around line 214-235: Update the docstring for the "Store a tensor into the
pinned host buffer." method to state that the data argument must be a CPU
torch.Tensor (1-D) because the implementation calls .numpy() on the input (which
requires the tensor to be on CPU); explicitly document any expectations about
dtype conversion (the method will convert to the buffer dtype if needed) and
mention that callers should provide data on CPU to avoid runtime errors when
invoking host_view[:length].numpy() and data.numpy().

[MrGeva]

/bot run

[tensorrt-cicd]

PR_Github #36413 [ run ] triggered by Bot. Commit: 7953ba4 Link to invocation

[MrGeva]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #36428 [ run ] triggered by Bot. Commit: 7953ba4 Link to invocation

[tensorrt-cicd]

PR_Github #36428 [ run ] completed with state SUCCESS. Commit: 7953ba4
/LLM/main/L0_MergeRequest_PR pipeline #28187 completed with status: 'SUCCESS'

Link to invocation

[MrGeva]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #36654 [ run ] triggered by Bot. Commit: e9edfcc Link to invocation

[MrGeva]

@lucaslie please approve the PR

[tensorrt-cicd]

PR_Github #36654 [ run ] completed with state SUCCESS. Commit: e9edfcc
/LLM/main/L0_MergeRequest_PR pipeline #28375 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[MrGeva]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #36768 [ run ] triggered by Bot. Commit: cd84888 Link to invocation

[MrGeva]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #36783 [ run ] triggered by Bot. Commit: 433645b Link to invocation

[tensorrt-cicd]

PR_Github #36783 [ run ] completed with state SUCCESS. Commit: 433645b
/LLM/main/L0_MergeRequest_PR pipeline #28486 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[MrGeva]

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast

[tensorrt-cicd]

PR_Github #36814 [ run ] triggered by Bot. Commit: 433645b Link to invocation

[MrGeva]

@lucaslie fixed the issue, can you please review?

[tensorrt-cicd]

PR_Github #36814 [ run ] completed with state SUCCESS. Commit: 433645b
/LLM/main/L0_MergeRequest_PR pipeline #28506 completed with status: 'SUCCESS'

Link to invocation