[docs] Add VLA full hidden_states export example - #17420
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughThis change adds VLA hidden-state export support through backend patches, engine build instructions, Python inference, validation tests, and documentation. ChangesVLA hidden-state export
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant trtllm_build
participant TensorRT_Engine
participant HiddenStatesEngine
participant ValidationTests
CLI->>trtllm_build: build engine with patched model
trtllm_build->>TensorRT_Engine: write rank0.engine
HiddenStatesEngine->>TensorRT_Engine: load and execute input_ids
TensorRT_Engine-->>HiddenStatesEngine: return full_hidden_states
HiddenStatesEngine-->>CLI: return hidden states and waypoint feature
ValidationTests->>TensorRT_Engine: verify output and shape
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: 10
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/vla_hidden_states_export/export_hidden_states.py`:
- Line 35: All functions in the VLA hidden-states export example lack required
type annotations. Add complete parameter and return annotations to
build_engine_v0x and main in
examples/vla_hidden_states_export/export_hidden_states.py (lines 35-35 and
66-66), HiddenStatesEngine.__init__, infer_and_extract_hidden_states, and main
in examples/vla_hidden_states_export/inference_python.py (lines 29-29, 57-57,
and 124-124), and test_engine_has_full_hidden_states, test_hidden_states_is_3d,
test_token_extraction, and main in
examples/vla_hidden_states_export/tests/test_hidden_states.py (lines 19-19,
41-41, 58-58, and 72-72).
In `@examples/vla_hidden_states_export/inference_python.py`:
- Around line 129-138: Validate args.waypoint_idx against the dummy input
sequence length before infer_and_extract_hidden_states indexes full_hs,
rejecting values outside the valid range instead of using the default 32000 with
599 tokens. In examples/vla_hidden_states_export/tests/test_hidden_states.py
lines 60-69, replace modulo-based indexing with an explicit out-of-range test
that verifies invalid waypoint indices are rejected.
- Around line 80-103: Update the hidden-state and logits buffer handling around
full_hidden_states to derive each tensor’s dtype from engine.get_tensor_dtype(),
convert it to the corresponding NumPy dtype, and allocate using that dtype’s
item size. Use the same dtype when creating host arrays and reshaping copied
data, preserving correct fp16 and fp32 behavior for all outputs.
- Around line 72-98: Update the inference setup around the execution flow to set
runtime shapes for every tensor in self.input_names before querying any output
shapes, and bind each required input tensor rather than only input_ids. Remove
the set_tensor_shape call for full_hidden_states, validate that TensorRT shape
inference succeeds and that execute_async_v3 returns true, then synchronize only
after successful execution.
In `@examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch`:
- Line 1: Add the repository-standard NVIDIA copyright header, using the year of
the latest meaningful modification, at the beginning of each affected file:
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch:1-1 and
modeling_utils_v1x.patch:1-1 before the patch descriptions;
export_hidden_states.py:1-1, inference_python.py:1-1, and
tests/test_hidden_states.py:1-1 before their module docstrings.
- Around line 25-33: Align the hidden-state output contract across all affected
sites: in examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch
lines 25-33, either export a padded 3D tensor or explicitly rename and document
the packed 2D output; in export_hidden_states.py lines 41-44, disable
remove_input_padding when 3D output is required; in README.md lines 26-50,
document packed output if retained; in inference_python.py lines 80-111, map
packed-token offsets per request; and in tests/test_hidden_states.py lines
41-55, assert the rank corresponding to the selected build mode.
In `@examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch`:
- Around line 26-32: Define a v1.x additional-output bridge that consumes
DecoderModelForCausalLM.forward() tensors, converts context outputs to [tokens,
hidden_dim] and generation outputs to [tokens, beam_width, hidden_dim], and
publishes them through the output dictionary consumed by HandleAdditionalOutputs
instead of assigning unused _additional_outputs. In
examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch lines 26-32,
implement this bridge for hidden_states and preserve the SamplingParams
additional-output contract; update examples/vla_hidden_states_export/README.md
lines 57-76 to document both supported shapes and add batched context and
generation runtime coverage.
In `@examples/vla_hidden_states_export/README.md`:
- Around line 7-10: Add the Markdown language identifier text to the fenced
diagram code block in the README, preserving the existing diagram content
unchanged.
- Around line 10-12: The documented invocation must match the parser: update
examples/vla_hidden_states_export/README.md lines 10-12 to use the supported
--engine_path argument and remove unsupported --engine_dir, --input_ids, and
--hidden_dim flags. In examples/vla_hidden_states_export/inference_python.py
lines 124-138, add corresponding parser arguments only if configurable engine
input is intended; otherwise leave this site unchanged because the README should
reflect the existing dummy-input behavior.
In `@examples/vla_hidden_states_export/tests/test_hidden_states.py`:
- Around line 19-25: Update the engine-dependent tests around
test_engine_has_full_hidden_states to obtain engine_path through a defined
pytest fixture or command-line option, and apply the same change to the related
tests. When no engine path is configured, skip these tests cleanly instead of
allowing fixture setup to fail.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 195c7d8c-f621-43e4-ae85-ddcf642e8ce3
📒 Files selected for processing (6)
examples/vla_hidden_states_export/README.mdexamples/vla_hidden_states_export/export_hidden_states.pyexamples/vla_hidden_states_export/inference_python.pyexamples/vla_hidden_states_export/patches/modeling_utils_v0x.patchexamples/vla_hidden_states_export/patches/modeling_utils_v1x.patchexamples/vla_hidden_states_export/tests/test_hidden_states.py
| import sys | ||
|
|
||
|
|
||
| def build_engine_v0x(model_dir, output_dir, dtype="fp16"): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add type annotations to every function.
The Python coding guidelines require annotations for every function, including constructors, test functions, and main functions.
examples/vla_hidden_states_export/export_hidden_states.py#L35-L35: annotatebuild_engine_v0x.examples/vla_hidden_states_export/export_hidden_states.py#L66-L66: annotatemain.examples/vla_hidden_states_export/inference_python.py#L29-L29: annotateHiddenStatesEngine.__init__.examples/vla_hidden_states_export/inference_python.py#L57-L57: annotateinfer_and_extract_hidden_states.examples/vla_hidden_states_export/inference_python.py#L124-L124: annotatemain.examples/vla_hidden_states_export/tests/test_hidden_states.py#L19-L19: annotatetest_engine_has_full_hidden_states.examples/vla_hidden_states_export/tests/test_hidden_states.py#L41-L41: annotatetest_hidden_states_is_3d.examples/vla_hidden_states_export/tests/test_hidden_states.py#L58-L58: annotatetest_token_extraction.examples/vla_hidden_states_export/tests/test_hidden_states.py#L72-L72: annotatemain.
📍 Affects 3 files
examples/vla_hidden_states_export/export_hidden_states.py#L35-L35(this comment)examples/vla_hidden_states_export/export_hidden_states.py#L66-L66examples/vla_hidden_states_export/inference_python.py#L29-L29examples/vla_hidden_states_export/inference_python.py#L57-L57examples/vla_hidden_states_export/inference_python.py#L124-L124examples/vla_hidden_states_export/tests/test_hidden_states.py#L19-L19examples/vla_hidden_states_export/tests/test_hidden_states.py#L41-L41examples/vla_hidden_states_export/tests/test_hidden_states.py#L58-L58examples/vla_hidden_states_export/tests/test_hidden_states.py#L72-L72
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/export_hidden_states.py` at line 35, All
functions in the VLA hidden-states export example lack required type
annotations. Add complete parameter and return annotations to build_engine_v0x
and main in examples/vla_hidden_states_export/export_hidden_states.py (lines
35-35 and 66-66), HiddenStatesEngine.__init__, infer_and_extract_hidden_states,
and main in examples/vla_hidden_states_export/inference_python.py (lines 29-29,
57-57, and 124-124), and test_engine_has_full_hidden_states,
test_hidden_states_is_3d, test_token_extraction, and main in
examples/vla_hidden_states_export/tests/test_hidden_states.py (lines 19-19,
41-41, 58-58, and 72-72).
Source: Coding guidelines
| # Input: input_ids | ||
| input_ids_np = np.ascontiguousarray(input_ids.astype(np.int32)) | ||
| d_input = cuda.mem_alloc(input_ids_np.nbytes) | ||
| cuda.memcpy_htod(d_input, input_ids_np) | ||
| self.context.set_tensor_address("input_ids", int(d_input)) | ||
| buffers["input_ids"] = d_input | ||
|
|
||
| # Output: full_hidden_states | ||
| hs_shape = tuple(self.context.get_tensor_shape("full_hidden_states")) | ||
| hs_size = int(np.prod([max(1, s) for s in hs_shape])) | ||
| hs_dtype = np.float16 # fp16 engine | ||
| d_hs = cuda.mem_alloc(hs_size * 2) # fp16 = 2 bytes | ||
| self.context.set_tensor_shape("full_hidden_states", hs_shape) | ||
| self.context.set_tensor_address("full_hidden_states", int(d_hs)) | ||
| buffers["full_hidden_states"] = d_hs | ||
|
|
||
| # Output: logits (standard) | ||
| if "logits" in self.output_names: | ||
| logits_shape = tuple(self.context.get_tensor_shape("logits")) | ||
| logits_size = int(np.prod([max(1, s) for s in logits_shape])) | ||
| d_logits = cuda.mem_alloc(logits_size * 4) | ||
| self.context.set_tensor_address("logits", int(d_logits)) | ||
| buffers["logits"] = d_logits | ||
|
|
||
| # --- Execute --- | ||
| self.context.execute_async_v3(self.stream.handle) | ||
| self.stream.synchronize() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python - <<'PY'
import tensorrt as trt
context = trt.IExecutionContext
print("TensorRT:", trt.__version__)
print("set_input_shape:", hasattr(context, "set_input_shape"))
print("set_tensor_shape:", hasattr(context, "set_tensor_shape"))
print("infer_shapes:", hasattr(context, "infer_shapes"))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 278
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="examples/vla_hidden_states_export/inference_python.py"
printf '%s\n' "== target file outline =="
ast-grep outline "$file" || true
printf '%s\n' "== target lines =="
sed -n '1,150p' "$file"
printf '%s\n' "== related tensor setup and engine construction =="
rg -n -C 4 'input_names|output_names|set_input_shape|set_tensor_shape|get_tensor_shape|execute_async_v3|full_hidden_states|logits' examples/vla_hidden_states_export
printf '%s\n' "== all references to inference class and input names =="
rg -n -C 3 'class |input_names|input_ids|pixel|image|vision|inference_python' examples/vla_hidden_states_exportRepository: NVIDIA/TensorRT-LLM
Length of output: 44188
🌐 Web query:
TensorRT Python IExecutionContext set_input_shape set_tensor_shape get_tensor_shape execute_async_v3 output shape inference documentation
💡 Result:
In TensorRT's Python API, managing dynamic shapes and executing inference with execute_async_v3 involves a standardized workflow within the IExecutionContext interface [1][2]. Key methods for dynamic shape management: - context.set_input_shape(name, shape): Sets the input dimensions for a specific tensor before inference [3][2]. This must be performed after selecting an appropriate optimization profile [3][4]. - context.get_tensor_shape(name): Returns the runtime dimensions of a tensor [3]. After setting input shapes, this method is used to query the computed output shapes for the given input [3][4]. Inference execution: - context.set_tensor_address(name, ptr): Assigns GPU memory buffers (as integer pointers) to input and output tensors [1][2]. - context.execute_async_v3(stream_handle): Enqueues the inference task to the specified CUDA stream [1][2]. Unlike deprecated V2 methods that used bindings, this V3 method relies on the tensor names previously configured with set_tensor_address [3][5]. Workflow Summary: 1. Create an execution context: context = engine.create_execution_context [1][2]. 2. Set the desired optimization profile (if multiple exist) [3][4]. 3. Set input dimensions: context.set_input_shape(input_name, dimensions) [3][2]. 4. Query/calculate output shape: output_shape = context.get_tensor_shape(output_name) [2][6]. 5. Allocate device memory for inputs and outputs [2][6]. 6. Map memory to tensors: context.set_tensor_address(name, ptr) [1][2]. 7. Execute: context.execute_async_v3(stream_handle) [1][2]. 8. Synchronize the stream to ensure inference completion [1][2][7]. For complex networks or when output shapes are strictly data-dependent, context.infer_shapes() may be used to validate or compute shapes [3][8]. Note that set_tensor_shape is not the standard API for setting input dimensions; set_input_shape is the correct method for defining input runtime dimensions [3][5].
Citations:
- 1: https://docs.nvidia.com/deeplearning/tensorrt/10.x.x/inference-library/python-api-docs.html
- 2: https://docs.nvidia.com/deeplearning/tensorrt/latest/getting-started/quick-start-runtime-tutorial.html
- 3: https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/dynamic-shapes-basics.html
- 4: https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/work-dynamic-shapes.html
- 5: https://docs.nvidia.com/deeplearning/tensorrt/10.16.0/api/migration-guide.html
- 6: https://docs.nvidia.com/deeplearning/tensorrt/11.2.1/getting-started/quick-start-runtime-tutorial.html
- 7: https://github.com/NVIDIA/TensorRT/blob/HEAD/quickstart/IntroNotebooks/onnx_helper.py
- 8: https://docs.nvidia.com/deeplearning/tensorrt/10.x.x/_static/python-api/infer/Core/pyCore.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' "== repository guidance and file headers =="
fd -i 'CODING_GUIDELINES.md' . -t f -x sh -c 'echo "--- $1"; sed -n "1,220p" "$1"' sh {}
printf '%s\n' "== engine build command and input-related options =="
sed -n '1,110p' examples/vla_hidden_states_export/export_hidden_states.py
printf '%s\n' "== inference documentation around invocation =="
sed -n '1,120p' examples/vla_hidden_states_export/README.md
printf '%s\n' "== related engine-input conventions in examples =="
rg -n -C 5 'trtllm-build|remove_input_padding|set_input_shape|set_tensor_address|input_names|position_ids|attention_mask|last_token_ids|past_key_value' examples -g '*.py' -g '*.md' | head -n 500Repository: NVIDIA/TensorRT-LLM
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("examples/vla_hidden_states_export/inference_python.py")
tree = ast.parse(path.read_text())
method = next(
node for node in ast.walk(tree)
if isinstance(node, ast.FunctionDef)
and node.name == "infer_and_extract_hidden_states"
)
calls = []
for node in ast.walk(method):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
owner = node.func.value
if isinstance(owner, ast.Attribute) and owner.attr == "context":
calls.append((node.lineno, node.func.attr))
print("context calls in AST order:")
for item in sorted(calls):
print(item)
ordered_names = [name for _, name in sorted(calls)]
print("set_input_shape present:", "set_input_shape" in ordered_names)
print("set_tensor_shape present:", "set_tensor_shape" in ordered_names)
print("first get_tensor_shape index:", next(
(i for i, name in enumerate(ordered_names) if name == "get_tensor_shape"), None
))
print("first execute_async_v3 index:", next(
(i for i, name in enumerate(ordered_names) if name == "execute_async_v3"), None
))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 492
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("examples/vla_hidden_states_export/inference_python.py")
tree = ast.parse(path.read_text())
method = next(
node for node in ast.walk(tree)
if isinstance(node, ast.FunctionDef)
and node.name == "infer_and_extract_hidden_states"
)
calls = []
for node in ast.walk(method):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
owner = node.func.value
if isinstance(owner, ast.Attribute) and owner.attr == "context":
calls.append((node.lineno, node.func.attr))
print("context calls in AST order:")
for item in sorted(calls):
print(item)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 366
Set input shapes and bind all required tensors before execution.
Call set_input_shape() for each runtime input before querying output shapes. Remove set_tensor_shape() for full_hidden_states; TensorRT derives output shapes from input shapes. Bind every tensor in self.input_names, and check shape inference and the boolean result of execute_async_v3() before synchronizing.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/inference_python.py` around lines 72 - 98,
Update the inference setup around the execution flow to set runtime shapes for
every tensor in self.input_names before querying any output shapes, and bind
each required input tensor rather than only input_ids. Remove the
set_tensor_shape call for full_hidden_states, validate that TensorRT shape
inference succeeds and that execute_async_v3 returns true, then synchronize only
after successful execution.
| hs_shape = tuple(self.context.get_tensor_shape("full_hidden_states")) | ||
| hs_size = int(np.prod([max(1, s) for s in hs_shape])) | ||
| hs_dtype = np.float16 # fp16 engine | ||
| d_hs = cuda.mem_alloc(hs_size * 2) # fp16 = 2 bytes | ||
| self.context.set_tensor_shape("full_hidden_states", hs_shape) | ||
| self.context.set_tensor_address("full_hidden_states", int(d_hs)) | ||
| buffers["full_hidden_states"] = d_hs | ||
|
|
||
| # Output: logits (standard) | ||
| if "logits" in self.output_names: | ||
| logits_shape = tuple(self.context.get_tensor_shape("logits")) | ||
| logits_size = int(np.prod([max(1, s) for s in logits_shape])) | ||
| d_logits = cuda.mem_alloc(logits_size * 4) | ||
| self.context.set_tensor_address("logits", int(d_logits)) | ||
| buffers["logits"] = d_logits | ||
|
|
||
| # --- Execute --- | ||
| self.context.execute_async_v3(self.stream.handle) | ||
| self.stream.synchronize() | ||
|
|
||
| # --- Read full_hidden_states --- | ||
| full_hs = np.zeros(hs_size, dtype=np.float16) | ||
| cuda.memcpy_dtoh(full_hs, d_hs) | ||
| full_hs = full_hs.reshape([max(1, s) for s in hs_shape]) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Allocate outputs from the engine tensor dtype.
export_hidden_states.py accepts fp32, but this code always allocates two bytes per hidden-state element and reads it as np.float16. An fp32 engine can write past d_hs and the host copy uses the wrong dtype. Get each output dtype from engine.get_tensor_dtype() and allocate with its NumPy item size.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/inference_python.py` around lines 80 - 103,
Update the hidden-state and logits buffer handling around full_hidden_states to
derive each tensor’s dtype from engine.get_tensor_dtype(), convert it to the
corresponding NumPy dtype, and allocate using that dtype’s item size. Use the
same dtype when creating host arrays and reshaping copied data, preserving
correct fp16 and fp32 behavior for all outputs.
| parser.add_argument("--waypoint_idx", type=int, default=32000, | ||
| help="Token position to extract hidden state from (default: 32000)") | ||
| args = parser.parse_args() | ||
|
|
||
| engine = HiddenStatesEngine(args.engine_path) | ||
|
|
||
| # Dummy input for demonstration | ||
| input_ids = np.ones((1, 599), dtype=np.int32) # seq_len=599 | ||
| full_hs, ego_feature = engine.infer_and_extract_hidden_states( | ||
| input_ids, args.waypoint_idx) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not use token ID 32000 as a sequence index.
The default waypoint_idx is 32000, but the dummy input has 599 tokens. The default invocation indexes outside full_hs. The test hides this failure with 32000 % 599. Validate the index against the sequence length, or locate the waypoint token during tokenization.
examples/vla_hidden_states_export/inference_python.py#L129-L138: reject an out-of-range waypoint index before indexing.examples/vla_hidden_states_export/tests/test_hidden_states.py#L60-L69: test the out-of-range case instead of reducing it with modulo arithmetic.
📍 Affects 2 files
examples/vla_hidden_states_export/inference_python.py#L129-L138(this comment)examples/vla_hidden_states_export/tests/test_hidden_states.py#L60-L69
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/inference_python.py` around lines 129 -
138, Validate args.waypoint_idx against the dummy input sequence length before
infer_and_extract_hidden_states indexes full_hs, rejecting values outside the
valid range instead of using the default 32000 with 599 tokens. In
examples/vla_hidden_states_export/tests/test_hidden_states.py lines 60-69,
replace modulo-based indexing with an explicit out-of-range test that verifies
invalid waypoint indices are rejected.
| @@ -0,0 +1,34 @@ | |||
| # Patch: Add full_hidden_states output to TRT-LLM v0.7–v0.21 (all TRT-backend versions) | |||
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the NVIDIA copyright header to each new non-Markdown file.
The repository guidelines require the NVIDIA copyright header with the year of the latest meaningful modification.
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch#L1-L1: add the header before the patch description.examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch#L1-L1: add the header before the patch description.examples/vla_hidden_states_export/export_hidden_states.py#L1-L1: add the header before the module docstring.examples/vla_hidden_states_export/inference_python.py#L1-L1: add the header before the module docstring.examples/vla_hidden_states_export/tests/test_hidden_states.py#L1-L1: add the header before the module docstring.
📍 Affects 5 files
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch#L1-L1(this comment)examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch#L1-L1examples/vla_hidden_states_export/export_hidden_states.py#L1-L1examples/vla_hidden_states_export/inference_python.py#L1-L1examples/vla_hidden_states_export/tests/test_hidden_states.py#L1-L1
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch` at line
1, Add the repository-standard NVIDIA copyright header, using the year of the
latest meaningful modification, at the beginning of each affected file:
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch:1-1 and
modeling_utils_v1x.patch:1-1 before the patch descriptions;
export_hidden_states.py:1-1, inference_python.py:1-1, and
tests/test_hidden_states.py:1-1 before their module docstrings.
Source: Coding guidelines
| + # [VLA] Attach full 3D hidden_states before LogitsProcessor compresses it. | ||
| + # This enables VLA models to access hidden_states at any token position | ||
| + # for downstream task heads (e.g., planning, action prediction). | ||
| + # Usage: SamplingParams(additional_model_outputs=["hidden_states"]) | ||
| + # See: https://github.com/NVIDIA/TensorRT-LLM/issues/4414 | ||
| + if hasattr(self, '_additional_output_names') and 'hidden_states' in self._additional_output_names: | ||
| + self._additional_outputs = {'hidden_states': output} |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/_torch/models/modeling_utils.py \
--match DecoderModelForCausalLM --view expanded
rg -n -C 4 --glob '*.py' \
'\b_additional_outputs\b|\b_additional_output_names\b|HandleAdditionalOutputs' \
tensorrt_llm/_torchRepository: NVIDIA/TensorRT-LLM
Length of output: 3867
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- modeling_utils.py ---'
sed -n '740,790p' tensorrt_llm/_torch/models/modeling_utils.py
printf '%s\n' '--- handle_additional_outputs.py ---'
wc -l tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py
sed -n '1,220p' tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py
printf '%s\n' '--- py_executor call site ---'
sed -n '6535,6595p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- patch and README ---'
sed -n '1,100p' examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch
sed -n '1,100p' examples/vla_hidden_states_export/README.md
printf '%s\n' '--- related symbols and request/output plumbing ---'
rg -n -C 5 --glob '*.py' \
'additional_model_outputs|additional_generation_outputs|_additional_outputs|_additional_output_names|batch_outputs' \
tensorrt_llm/_torchRepository: NVIDIA/TensorRT-LLM
Length of output: 45503
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- all additional-output symbols ---'
rg -n -C 4 --glob '*.py' \
'additional_model_outputs|_additional_output_names|_additional_outputs|additional_outputs' .
printf '%s\n' '--- forward-step implementation ---'
rg -n -C 8 'def _forward_step|DecoderModelForCausalLM\(' \
tensorrt_llm/_torch/pyexecutor tensorrt_llm/_torch/models
printf '%s\n' '--- model output consumers ---'
rg -n -C 6 --glob '*.py' \
'model\(|model_forward|forward\(' tensorrt_llm/_torch/pyexecutor/py_executor.py \
| head -n 240
printf '%s\n' '--- additional-output tests and API definitions ---'
rg -n -C 6 --glob '*.py' --glob '*.cpp' --glob '*.h' --glob '*.md' \
'additional_generation_outputs|additional_context_outputs|additional_model_outputs' \
. | head -n 320Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- ModelEngine._forward_step ---'
sed -n '6588,6630p' tensorrt_llm/_torch/pyexecutor/model_engine.py
printf '%s\n' '--- PyExecutor._forward_step ---'
sed -n '6424,6505p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- existing additional-output integration test ---'
sed -n '160,285p' tests/unittest/llmapi/test_additional_model_outputs.py
printf '%s\n' '--- output container shapes and append methods ---'
sed -n '430,575p' tensorrt_llm/_torch/pyexecutor/llm_request.py
printf '%s\n' '--- exact definitions of unpublished attributes ---'
rg -n --glob '*.py' \
'(^|[^A-Za-z0-9_])_additional_output_names([^A-Za-z0-9_]|$)|(^|[^A-Za-z0-9_])_additional_outputs([^A-Za-z0-9_]|$)' .Repository: NVIDIA/TensorRT-LLM
Length of output: 17470
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from dataclasses import dataclass
`@dataclass`
class Shape:
dims: tuple[int, ...]
`@property`
def shape0(self):
return self.dims[0]
def slice0(self, begin, end):
return Shape((end - begin, *self.dims[1:]))
def reshape(self, *dims):
known = 1
unknown = None
for i, dim in enumerate(dims):
if dim == -1:
unknown = i
else:
known *= dim
total = 1
for dim in self.dims:
total *= dim
result = list(dims)
if unknown is not None:
result[unknown] = total // known
assert result[unknown] * known == total if unknown is not None else known == total
return Shape(tuple(result))
def tile(self, reps):
assert len(reps) == len(self.dims)
return Shape(tuple(dim * rep for dim, rep in zip(self.dims, reps)))
# Mirror HandleAdditionalOutputs for two context requests and one generation request.
batch, seq_len, hidden, beam = 2, 5, 4, 1
num_context_tokens = 2 * seq_len
output_length_with_context = num_context_tokens + beam
raw = Shape((batch, seq_len, hidden))
gather_context = raw.shape0 == output_length_with_context
print("raw output:", raw.dims)
print("expected token-major leading length:", output_length_with_context)
print("gather_context:", gather_context)
# Generation path used when gather_context is false.
generation_slice = raw.slice0(0, beam).reshape(1, beam, -1)
print("generation output after handler:", generation_slice.dims)
assert generation_slice.dims == (1, beam, seq_len * hidden)
assert generation_slice.dims != (seq_len, beam, hidden)
PY
printf '%s\n' '--- copyright headers in affected files ---'
for file in examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch \
examples/vla_hidden_states_export/README.md; do
printf '%s: ' "$file"
head -n 5 "$file" | tr '\n' ' '
printf '\n'
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 1057
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from dataclasses import dataclass
`@dataclass`
class Shape:
dims: tuple[int, ...]
`@property`
def shape0(self):
return self.dims[0]
def slice0(self, begin, end):
return Shape((end - begin, *self.dims[1:]))
def reshape(self, *dims):
total = 1
for dim in self.dims:
total *= dim
known = 1
unknown = None
result = list(dims)
for i, dim in enumerate(result):
if dim == -1:
unknown = i
else:
known *= dim
if unknown is not None:
result[unknown] = total // known
assert known * (result[unknown] if unknown is not None else 1) == total
return Shape(tuple(result))
batch, seq_len, hidden, beam = 2, 5, 4, 1
num_context_tokens = 2 * seq_len
expected_length = num_context_tokens + beam
raw = Shape((batch, seq_len, hidden))
gather_context = raw.shape0 == expected_length
print("raw output:", raw.dims)
print("expected token-major leading length:", expected_length)
print("gather_context:", gather_context)
generation_slice = raw.slice0(0, beam).reshape(1, beam, -1)
print("generation output after handler:", generation_slice.dims)
assert generation_slice.dims == (1, beam, seq_len * hidden)
assert generation_slice.dims != (seq_len, beam, hidden)
PY
printf '%s\n' '--- copyright headers in affected files ---'
for file in examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch \
examples/vla_hidden_states_export/README.md; do
printf '%s: ' "$file"
head -n 5 "$file" | tr '\n' ' '
printf '\n'
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 1057
Define and test one v1.x additional-output contract.
DecoderModelForCausalLM.forward()returns a tensor._additional_output_namesis not defined, and_additional_outputshas no consumer. The assignment does not addhidden_statesto the output dictionary.HandleAdditionalOutputsexpects token-major data. A raw[batch, seq_len, hidden_dim]tensor is misclassified and reshaped to[1, beam_width, seq_len * hidden_dim].- Add the bridge, publish the supported shapes (
[tokens, hidden_dim]for context and[tokens, beam_width, hidden_dim]for generation), update the README, and add a batched context/generation runtime test.
examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch#L26-L32
examples/vla_hidden_states_export/README.md#L57-L76
📍 Affects 2 files
examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch#L26-L32(this comment)examples/vla_hidden_states_export/README.md#L57-L76
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch` around
lines 26 - 32, Define a v1.x additional-output bridge that consumes
DecoderModelForCausalLM.forward() tensors, converts context outputs to [tokens,
hidden_dim] and generation outputs to [tokens, beam_width, hidden_dim], and
publishes them through the output dictionary consumed by HandleAdditionalOutputs
instead of assigning unused _additional_outputs. In
examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch lines 26-32,
implement this bridge for hidden_states and preserve the SamplingParams
additional-output contract; update examples/vla_hidden_states_export/README.md
lines 57-76 to document both supported shapes and add batched context and
generation runtime coverage.
| ``` | ||
| Standard LLM: input → LLM → logits → token sampling → text | ||
| VLA model: images+text → LLM → hidden_states → planning head → trajectory | ||
| ``` |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add a language identifier to this code block.
Use text for the diagram block. This satisfies MD040.
🧰 Tools
🪛 markdownlint-cli2 (0.23.2)
[warning] 7-7: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/README.md` around lines 7 - 10, Add the
Markdown language identifier text to the fenced diagram code block in the
README, preserving the existing diagram content unchanged.
Source: Linters/SAST tools
| def test_engine_has_full_hidden_states(engine_path): | ||
| """Verify the engine exports full_hidden_states tensor.""" | ||
| import tensorrt as trt | ||
|
|
||
| logger = trt.Logger(trt.Logger.WARNING) | ||
| with open(engine_path, "rb") as f: | ||
| engine = trt.Runtime(logger).deserialize_cuda_engine(f.read()) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Provide engine_path to pytest.
Pytest resolves engine_path as a fixture. This repository does not define that fixture in the supplied test, so python -m pytest fails during setup. Add a fixture or pytest option that supplies the engine path, or skip these engine-dependent tests when no path is configured.
Also applies to: 41-47
🧰 Tools
🪛 ast-grep (0.45.0)
[warning] 23-23: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(engine_path, "rb")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/tests/test_hidden_states.py` around lines
19 - 25, Update the engine-dependent tests around
test_engine_has_full_hidden_states to obtain engine_path through a defined
pytest fixture or command-line option, and apply the same change to the related
tests. When no engine path is configured, skip these tests cleanly instead of
allowing fixture setup to fail.
ecf4f76 to
30f7336
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Two things need sorting before this can be looked at on the merits.
1. The branch reverts merged work. The single commit here contains, besides the new example, a full revert of #12733 ("Unify sparse attention framework with clean backend interfaces") and its follow-up #17416 — 80+ files across cpp/tensorrt_llm/, tensorrt_llm/_torch/attention_backend/sparse/, tensorrt_llm/_torch/modules/mla.py, and the sparse-attention tests/docs. It re-collapses sparse/dsa/, sparse/rocket/, sparse/skip_softmax/ back into single files, restores sparse_attn_kv_lens / aux_kv_cache_pool_ptr naming in the thop signature, and deletes sparse/hooks.py and registry.py. This is almost certainly a bad rebase/merge rather than intent. Please rebase onto current main so the PR contains only examples/vla_hidden_states_export/.
2. Solution A does not apply to main. The TRT network-build backend is gone: there is no trtllm-build command under tensorrt_llm/commands/, gather_last_token_logits no longer exists anywhere in tensorrt_llm/, and tensorrt_llm/models/modeling_utils.py has no DecoderModelForCausalLM. So export_hidden_states.py, inference_python.py, tests/test_hidden_states.py, and patches/modeling_utils_v0x.patch describe a workflow that cannot be run against the branch they're being added to. An examples/ directory whose entry point is trtllm-build will be a support burden. My suggestion: drop Solution A to a short "historical note (v0.x)" paragraph in the README, and make the example PyTorch-backend-only, with runnable code rather than patch files.
On Solution B, the underlying idea is right and the framework already supports it: HandleAdditionalOutputs (tensorrt_llm/_torch/pyexecutor/py_executor.py:6572) picks up any key named in SamplingParams.additional_model_outputs from the dict a model's forward returns (model_engine.py:6611). So a model that returns {"logits": ..., "hidden_states": ...} already works end-to-end today — no core patch needed. That would make a genuinely useful example: a small modeling_*.py subclass that returns the extra key, plus a llm.generate script reading outputs[0].additional_generation_outputs["hidden_states"]. That's testable in CI and doesn't ask users to patch installed source.
Also: new files need the NVIDIA copyright header (see AGENTS.md), and a new examples/ subdir should be listed in the table in examples/README.md.
| + # for downstream task heads (e.g., planning, action prediction). | ||
| + # Usage: SamplingParams(additional_model_outputs=["hidden_states"]) | ||
| + # See: https://github.com/NVIDIA/TensorRT-LLM/issues/4414 | ||
| + if hasattr(self, '_additional_output_names') and 'hidden_states' in self._additional_output_names: |
There was a problem hiding this comment.
This patch is a no-op as written: nothing in the codebase ever sets _additional_output_names, and nothing ever reads self._additional_outputs, so additional_model_outputs=["hidden_states"] would still come back empty.
The framework's actual contract is dict-return from forward. model_engine.py:6611 checks isinstance(outputs, dict), and HandleAdditionalOutputs (py_executor.py:6572) looks up each name in SamplingParams.additional_model_outputs as a key of that dict. So the working version is what your README snippet already shows:
return {
"logits": self.logits_processor.forward(output, self.lm_head, attn_metadata, return_context_logits),
"hidden_states": output,
}Since that path already exists, a model subclass overriding forward this way needs no core patch at all — which would make a much better example than a patch file.
|
|
||
| --- a/tensorrt_llm/models/modeling_utils.py | ||
| +++ b/tensorrt_llm/models/modeling_utils.py | ||
| @@ -XXX,7 +XXX,11 @@ class DecoderModelForCausalLM(DecoderModel): |
There was a problem hiding this comment.
@@ -XXX,7 +XXX,11 @@ isn't a valid hunk header — git apply and patch will both reject this file. If a patch file is kept at all, generate it with git diff against a real checkout of the target version so the line numbers are concrete.
| ```python | ||
| # tensorrt_llm/_torch/models/modeling_utils.py — DecoderModelForCausalLM.forward() | ||
| def forward(self, ..., additional_model_outputs=None): | ||
| output = self.model(...) # Full 3D [batch, seq_len, hidden_dim] |
There was a problem hiding this comment.
# Full 3D [batch, seq_len, hidden_dim] is wrong for the PyTorch backend. With remove-input-padding (always on in the PyTorch path) the model returns packed tokens, shape [num_tokens, hidden_dim], where num_tokens is the flattened sum over the batch. HandleAdditionalOutputs then slices per request using the context/generation token counts, so what a caller gets back is [seq_len, hidden_dim] for that request — no batch dimension. Please correct the shapes throughout the README; the "complete 3D tensor" framing in the intro sets up the wrong expectation.
| hs_size = int(np.prod([max(1, s) for s in hs_shape])) | ||
| hs_dtype = np.float16 # fp16 engine | ||
| d_hs = cuda.mem_alloc(hs_size * 2) # fp16 = 2 bytes | ||
| self.context.set_tensor_shape("full_hidden_states", hs_shape) |
There was a problem hiding this comment.
Three problems make this unrunnable against a real engine:
set_tensor_shapeis only valid for input tensors; calling it onfull_hidden_states(an output) is an error.get_tensor_shapeis queried before any input shape is bound, so dynamic dims come back as-1.max(1, s)then turns those into 1, sohs_sizeis far smaller than the real output andcuda.mem_alloc(hs_size * 2)under-allocates — the engine writes past the buffer.- Only
input_idsis bound; a TRT-LLM engine needs position ids, sequence lengths, KV-cache pointers etc., soexecute_async_v3will fail on unbound tensors.
Also pycuda is not a TensorRT-LLM dependency. Rather than repairing this, the PyTorch-backend llm.generate + additional_generation_outputs path gives a working ~20-line example with no manual buffer management.
| import numpy as np | ||
|
|
||
|
|
||
| def test_engine_has_full_hidden_states(engine_path): |
There was a problem hiding this comment.
The module docstring documents python -m pytest tests/test_hidden_states.py -v, but test_engine_has_full_hidden_states(engine_path) and test_hidden_states_is_3d(engine_path) take a parameter that is not a fixture — pytest will error both with fixture 'engine_path' not found. Either add a --engine-path option plus an engine_path fixture in a conftest.py, or rename these to non-test_ helpers and drop the pytest instruction.
| ego_feature = full_hs[0, waypoint_idx, :] | ||
|
|
||
| assert ego_feature.shape == (4096,), f"Expected (4096,), got {ego_feature.shape}" | ||
| assert np.allclose(ego_feature, full_hs[0, waypoint_idx, :]), "Extraction mismatch" |
There was a problem hiding this comment.
assert np.allclose(ego_feature, full_hs[0, waypoint_idx, :]) compares the slice to itself — it can never fail, so this test verifies nothing beyond the shape assert on the line above. Relatedly, waypoint_idx = 32000 % 599 on line 62 is arbitrary; if the intent is "a caller-supplied token index", just pick one and say so.
| def build_engine_v0x(model_dir, output_dir, dtype="fp16"): | ||
| """Build engine for v0.x TRT backend (requires mark_output patch).""" | ||
| cmd = [ | ||
| "trtllm-build", |
There was a problem hiding this comment.
trtllm-build no longer exists on main — tensorrt_llm/commands/ has only bench, eval, and serve. This script can't run against the branch it's being added to. Please either drop it or move the whole v0.x flow into a clearly-labelled historical note in the README.
| @@ -0,0 +1,83 @@ | |||
| """ | |||
There was a problem hiding this comment.
New files need the NVIDIA copyright header (AGENTS.md, "CRITICAL" section). Applies to inference_python.py and tests/test_hidden_states.py too.
30f7336 to
329ffd2
Compare
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/vla_hidden_states_export/export_hidden_states.py`:
- Line 1: Update the SPDX copyright header at the top of export_hidden_states.py
so the NVIDIA copyright year range ends in 2026 instead of 2025.
- Around line 6-9: Update the version/CLI note in the script to pin the legacy
TRT backend to TRT-LLM v0.13.0, keep trtllm-build scoped to that legacy path,
and describe v1.x as loading Hugging Face checkpoints directly without an
engine-build step. Use the correct v1.x commands trtllm-serve for serving and
trtllm-bench for benchmarking, while preserving the reference to the README’s
Solution B.
In `@examples/vla_hidden_states_export/inference_python.py`:
- Around line 1-2: Update the NVIDIA copyright header’s latest year to 2026 in
examples/vla_hidden_states_export/inference_python.py lines 1-2 and
examples/vla_hidden_states_export/tests/test_hidden_states.py lines 1-2; no
other changes are needed.
- Line 118: Both main functions lack explicit return annotations. Update
examples/vla_hidden_states_export/inference_python.py lines 118-118 and
examples/vla_hidden_states_export/tests/test_hidden_states.py lines 68-68 so
each main signature declares a None return type.
- Around line 40-43: Replace the assertion checking "full_hidden_states" in
self.output_names with an explicit ValueError when the output is missing,
preserving the existing diagnostic message and preventing execution from
continuing without the required engine output.
In `@examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch`:
- Around line 6-22: Align full_hidden_states export with remove_input_padding:
update examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch at
lines 6-22 to either disable packed input for the 3D workflow or explicitly
support/document the packed 2D shape; update
examples/vla_hidden_states_export/export_hidden_states.py at lines 11-14 to
remove the unconditional [batch, seq_len, hidden_dim] claim and document the
required build setting. Preserve the mark_output call in
DecoderModelForCausalLM.forward().
- Around line 1-22: Replace the commit-message prose in
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch:1-22 with a
valid version-specific unified diff targeting DecoderModelForCausalLM.forward(),
adding mark_output for full_hidden_states immediately before
gather_last_token_logits. Also replace the prose in
examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch:6-24 with an
applicable unified diff implementing the required dictionary-return behavior,
including test context that verifies requested additional outputs and generation
output shapes.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dfbebe57-3dec-4002-bf52-94c61a371492
📒 Files selected for processing (6)
examples/vla_hidden_states_export/README.mdexamples/vla_hidden_states_export/export_hidden_states.pyexamples/vla_hidden_states_export/inference_python.pyexamples/vla_hidden_states_export/patches/modeling_utils_v0x.patchexamples/vla_hidden_states_export/patches/modeling_utils_v1x.patchexamples/vla_hidden_states_export/tests/test_hidden_states.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/vla_hidden_states_export/README.md
| @@ -0,0 +1,28 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the copyright range to include 2026.
Line 1 ends at 2025, but the patch metadata is dated August 7, 2026. Update the ending year to 2026.
As per coding guidelines, source files must contain the NVIDIA copyright header with the year of the latest meaningful modification.
Proposed fix
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/export_hidden_states.py` at line 1, Update
the SPDX copyright header at the top of export_hidden_states.py so the NVIDIA
copyright year range ends in 2026 instead of 2025.
Source: Coding guidelines
| NOTE: This script applies to the legacy TRT backend (v0.7-v0.21) which uses | ||
| `trtllm-build`. On v1.x (main branch), the build command has changed to | ||
| `trtllm serve` / `trtllm bench`. For v1.x, refer to Solution B in the README | ||
| (modify model forward to return hidden_states dict). |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- target file ---'
cat -n examples/vla_hidden_states_export/export_hidden_states.py | sed -n '1,30p'
printf '%s\n' '--- repository references ---'
rg -n --glob '!node_modules' --glob '!build' --glob '!dist' 'trtllm-build|trtllm-serve|trtllm-bench|trtllm serve|trtllm bench|v0\.7|v0\.21|v1\.x' .
printf '%s\n' '--- referenced upstream README ---'
curl -L --fail --silent --show-error --max-time 20 \
'https://raw.githubusercontent.com/NVIDIA/TensorRT-LLM/main/examples/models/core/qwen/README.md' |
rg -n -C 3 'trtllm-build|trtllm-serve|trtllm-bench|trtllm serve|trtllm bench|version|release' || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- example directory ---'
git ls-files examples/vla_hidden_states_export
printf '%s\n' '--- example README and related references ---'
if test -f examples/vla_hidden_states_export/README.md; then
cat -n examples/vla_hidden_states_export/README.md | sed -n '1,240p'
fi
rg -n -C 4 'Solution B|hidden_states|v0\.7|v0\.21|v1\.x|trtllm-build|trtllm-serve|trtllm-bench' \
examples/vla_hidden_states_export README.md docs examples/models/core/qwen/README.md
printf '%s\n' '--- focused local Qwen README ---'
cat -n examples/models/core/qwen/README.md | sed -n '1,190p'
printf '%s\n' '--- upstream Qwen README focused extract ---'
curl -L --fail --silent --show-error --max-time 20 \
'https://raw.githubusercontent.com/NVIDIA/TensorRT-LLM/main/examples/models/core/qwen/README.md' |
sed -n '1,190p'
printf '%s\n' '--- package/version and release references ---'
rg -n --glob '!.git/**' '__version__|version\s*=|v0\.[0-9]+|v1\.[0-9]+' setup.py pyproject.toml tensorrt_llm 2>/dev/null |
head -200 || true
printf '%s\n' '--- recent tags/releases visible locally ---'
git tag --sort=-version:refname | head -30Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- tracked files in example ---'
git ls-files examples/vla_hidden_states_export
printf '%s\n' '--- example README ---'
if test -f examples/vla_hidden_states_export/README.md; then
cat -n examples/vla_hidden_states_export/README.md | sed -n '1,220p'
else
printf '%s\n' 'README.md not present'
fi
printf '%s\n' '--- exact local Qwen command references ---'
rg -n -C 5 'trtllm-build|trtllm-serve|trtllm-bench|Solution B|hidden_states' \
examples/models/core/qwen/README.md examples/vla_hidden_states_export/export_hidden_states.py
printf '%s\n' '--- exact legacy benchmark build documentation ---'
rg -n -C 4 'trtllm-bench.*build|trtllm-build' \
docs/source/legacy/performance/perf-benchmarking.md \
docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md |
head -120
printf '%s\n' '--- release/version metadata ---'
git tag --sort=-version:refname | head -40
printf '%s\n' '--- repository version declarations ---'
rg -n --glob 'setup.py' --glob 'pyproject.toml' --glob 'version.py' \
'version\s*=|__version__' . | head -80 || true
printf '%s\n' '--- upstream Qwen README command lines ---'
curl -L --fail --silent --show-error --max-time 20 \
'https://raw.githubusercontent.com/NVIDIA/TensorRT-LLM/main/examples/models/core/qwen/README.md' |
rg -n -C 5 'trtllm-build|trtllm-serve|trtllm-bench|Solution B|hidden_states'Repository: NVIDIA/TensorRT-LLM
Length of output: 39018
Correct the v1.x CLI description and pin the legacy release.
For v1.x, the PyTorch backend loads Hugging Face checkpoints directly and has no engine-build step. Use trtllm-serve for serving and trtllm-bench for benchmarking. Keep trtllm-build scoped to the legacy TRT backend and pin this script to TRT-LLM v0.13.0, the version tested in the README.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/export_hidden_states.py` around lines 6 -
9, Update the version/CLI note in the script to pin the legacy TRT backend to
TRT-LLM v0.13.0, keep trtllm-build scoped to that legacy path, and describe v1.x
as loading Hugging Face checkpoints directly without an engine-build step. Use
the correct v1.x commands trtllm-serve for serving and trtllm-bench for
benchmarking, while preserving the reference to the README’s Solution B.
| From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 | ||
| From: Hou Song <hous.ailab@gmail.com> | ||
| Date: Thu, 7 Aug 2026 22:00:00 +0800 | ||
| Subject: [PATCH] Export full 3D hidden_states via mark_output | ||
|
|
||
| Insert mark_output('full_hidden_states') before gather_last_token_logits | ||
| in DecoderModelForCausalLM.forward(), so that the complete [batch, seq_len, | ||
| hidden_dim] tensor is exported from the engine. | ||
|
|
||
| This addresses the need of VLA models (Orion, OpenVLA) whose downstream | ||
| planning heads require hidden_states at a specific token position, not | ||
| just the last token's logits. | ||
|
|
||
| Related issue: https://github.com/NVIDIA/TensorRT-LLM/issues/4414 | ||
|
|
||
| Apply to: tensorrt_llm/models/modeling_utils.py | ||
| Compatible with: v0.7.0 through v0.21.0 (verified) | ||
|
|
||
| NOTE: Line numbers vary by version. Search for `gather_last_token_logits` | ||
| in DecoderModelForCausalLM.forward() and add the mark_output line above it: | ||
|
|
||
| hidden_states.mark_output('full_hidden_states', self.config.dtype) |
There was a problem hiding this comment.
Provide real patch hunks for both backend patches.
Both files contain commit-message prose, not source diffs. Applying them cannot implement the described backend changes. The v1.x executor preserves extra dictionary keys, and HandleAdditionalOutputs reads requested names from that dictionary, so the v1.x contract requires actual code and runtime coverage. (raw.githubusercontent.com)
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch#L1-L22: add a concrete version-specific diff containing themark_outputchange.examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch#L6-L24: add the real dictionary-return implementation and test context and generation output shapes.
📍 Affects 2 files
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch#L1-L22(this comment)examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch#L6-L24
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch` around
lines 1 - 22, Replace the commit-message prose in
examples/vla_hidden_states_export/patches/modeling_utils_v0x.patch:1-22 with a
valid version-specific unified diff targeting DecoderModelForCausalLM.forward(),
adding mark_output for full_hidden_states immediately before
gather_last_token_logits. Also replace the prose in
examples/vla_hidden_states_export/patches/modeling_utils_v1x.patch:6-24 with an
applicable unified diff implementing the required dictionary-return behavior,
including test context that verifies requested additional outputs and generation
output shapes.
329ffd2 to
94421d4
Compare
329ffd2 to
7ecfc46
Compare
Fixes NVIDIA#4414 Signed-off-by: Hou Song <hous.ailab@gmail.com>
7ecfc46 to
d5deab9
Compare
2 participants
Summary
Adds
examples/vla_hidden_states_export/- a guide for exporting full 3D hidden_states from TensorRT-LLM for VLA (Vision-Language-Action) models.Fixes #4414
Background
VLA models (Orion, OpenVLA, RT-2) use the LLM's hidden states as input to downstream task heads, rather than for token generation:
The downstream head needs the hidden state at a specific token position (e.g. a "waypoint" token), which requires the complete
[batch, seq_len, hidden_dim]tensor - not just the last token.Why existing APIs fall short
DecoderModelForCausalLM.forwardonly returns logits; hidden_states is consumed by LogitsProcessor before it can be exposedSolutions
Solution A: TRT backend mark_output (v0.7-v0.21, verified)
Insert
mark_outputbeforegather_last_token_logits:Compatibility verified across v0.7.0 to v0.21.0 (15 versions).
Reading at inference:
Solution B: PyTorch backend forward modification (v1.x)
Attach the raw 3D tensor before LogitsProcessor compresses it, enabling
additional_model_outputs=["hidden_states"].Verification (Solution A)
Tested on Orion VLA model (ICCV'25) with TRT-LLM v0.13.0:
[1, 599, 4096](full 3D)Notes
Dev Engineer Review
mark_outputbeforegather_last_token_logits.additional_model_outputs.full_hidden_states, and extracts token features.HiddenStatesEngine.QA Engineer Review
test_engine_has_full_hidden_states.test_hidden_states_is_3d.test_token_extraction.tests/integration/test_lists/test-db or QA files.