[TensorRT] Add static-cache export support and numerical runtime coverage

[rui-ren]

Summary

• Add a standalone TensorRT target with heads-first static KV caches, axis-2 writes, rank-4 Attention, and explicit causal/valid-length masking.
• Derive workflow metadata layout and capacity axis from graph geometry, preserving flattened caches for existing targets.
• Add an opt-in, download-free TensorRT numerical regression covering prefill, single-token and chunked decode, KV parity, aliasing, slot preservation, and full-prefix equivalence.
• Include TensorRT diagnostic examples, export documentation, the corrected --features flag, and the existing README architecture documentation.

Validation

Validation performed on the identical pre-consolidation file tree:

• 78 focused cache/metadata tests passed with actual TensorRT execution enabled on Windows, TensorRT 11.3, and an RTX 4060 Laptop GPU.
• Default gated run: 77 passed, 1 skipped.
• Ruff checks passed for the metadata implementation and regression tests; git diff --check passed.
• Previously verified Qwen3-0.6B BF16 TensorRT generation: The capital of France is Paris.

The runtime regression is enabled with MOBIUS_TEST_TENSORRT=1 and TENSORRT_ROOT pointing to the SDK. These results do not claim a full repository test run.

Contribution History

This contribution was developed with GitHub Copilot assistance, including architecture documentation previously committed by the Copilot bot. That assistance is disclosed here and in the consolidated commit. The original PR is linked for provenance and review continuity. The replacement changes commit organization, not file contents; CLA acceptance remains subject to project policy.

[Copilot]

🟡 Changes recommended

Unresolved Gemma4 TensorRT compatibility and EpCapabilities positional-API issues remain, along with validation defects.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds standalone TensorRT static-cache export support, runtime coverage, diagnostics, and documentation.

Changes:

• Adds heads-first KV-cache layouts and rank-4 attention handling.
• Derives cache metadata from graph geometry.
• Adds opt-in numerical tests and TensorRT debugging tools.
• Updates architecture and execution-provider documentation.

File summaries

File	Summary
tests/static_cache_metadata_test.py	Adds TensorRT cache and runtime coverage.
src/mobius/tasks/_causal_lm.py	Selects provider-specific cache layouts.
src/mobius/models/base.py	Supports explicit static-cache bias.
src/mobius/integrations/onnx_genai/workflow_metadata.py	Derives cache metadata from geometry.
src/mobius/integrations/onnx_genai/workflow_metadata_test.py	Tests metadata geometry validation.
src/mobius/components/_attention.py	Adds rank-4 attention and cache support.
src/mobius/_execution_providers.py	Registers TensorRT capabilities.
README.md	Updates architecture documentation.
examples/tensorrt_static_cache_generation.py	Provides TensorRT generation and parity diagnostics.
examples/tensorrt_debug/runtime.py	Provides TensorRT runtime support.
examples/tensorrt_debug/inspect_cache.py	Inspects TensorRT caches.
examples/tensorrt_debug/inspect_attention.py	Inspects attention behavior.
examples/tensorrt_debug/full_prefix.py	Supports full-prefix diagnostics.
examples/tensorrt_debug/comparison.py	Supports runtime comparisons.
examples/tensorrt_debug/_runtime_support.py	Provides shared runtime helpers.
examples/tensorrt_debug/_attention_support.py	Provides shared attention helpers.
examples/tensorrt_debug/__init__.py	Defines the debugging package.
examples/tensorrt_attention_probe.py	Builds attention diagnostic engines.
docs/research/tensorrt-static-cache-debugging.md	Documents TensorRT debugging and validation.
docs/index.md	Links the debugging documentation.
docs/execution_providers.md	Documents TensorRT execution-provider behavior.

Review details

Suppressed comments (3)

docs/execution_providers.md:152

• The new standalone target is documented only in the registry example below, while the user-facing Supported Execution Providers table in this file and the Available Execution Providers table in docs/cli_reference.md still omit tensorrt. Users following the documented discovery tables will not find the --ep tensorrt route; add it to both tables.

EpCapabilities(name="tensorrt",
               static_cache_layout="heads_first",
               supports_attention_nonpad_kv_seqlen=False,
               gqa_dtypes=frozenset(), qkv_pack_dtypes=frozenset(),
               supports_skip_layer_norm=False, supports_matmul_nbits=False)

examples/tensorrt_static_cache_generation.py:69

• This validation rejects a valid diagnostic run with --compare-hf --max-new-tokens 1: the prefill is the only step, so --compare-steps 1 should be accepted. Please allow 1 here and update the error text; otherwise the smallest useful comparison cannot be run.

    if args.compare_hf and not 2 <= args.compare_steps <= args.max_new_tokens:
        parser.error("--compare-steps must be between 2 and --max-new-tokens")

examples/tensorrt_static_cache_generation.py:101

• The loop's final execution consumes prompt_length + max_new_tokens - 1 tokens; it predicts the last generated token without caching that token. This + max_new_tokens check therefore rejects the valid boundary case where the required prefix exactly fills capacity. Use the same - 1 bound as validate_full_prefix_budget.

    if current_ids.shape[1] + args.max_new_tokens > capacity:
        raise ValueError("Prompt and generation budget exceed cache capacity")

• Files reviewed: 21/21 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

The generation example rejects a valid exact-capacity configuration, and the static-cache contract documentation is outdated.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 22/22 changed files
• Comments generated: 2
• Review effort level: Balanced

[titaiwangms]

I completed a full static review of the standalone TensorRT static-cache
export path at cc7cf75f.

The core graph transformation looks sound. In particular:

• rank-4 heads-first caches use [B, kv_heads, capacity, head_dim] and
  TensorScatter(axis=2) consistently;
• omitting q_num_heads / kv_num_heads for rank-4 Q/K/V follows the ONNX
  Attention schema;
• the explicit bias correctly derives query slots from
  write_indices + arange(S_q) and valid keys from
  slot < nonpad_kv_seqlen;
• the complete additive bias is correctly paired with is_causal=0;
• reducing static-cache Attention to one output removes a schema-invalid
  combination of nonpad_kv_seqlen with present-cache outputs;
• incompatible Gemma4 static-cache paths fail closed.

I did not find an unconditional core numerical blocker, but the following
issues should be addressed.

Conditional major: padded-batch contract

src/mobius/components/_common.py:311-318,355-420

The explicit bias can represent only a cache whose valid KV entries form a
dense prefix. It cannot represent ordinary left-padded prompts or another
non-prefix-valid physical cache layout.

For example, with [PAD, PAD, A, B], write_indices=0, and
nonpad_kv_seqlen=2, the graph scatters padding into slots 0-1 and real tokens
into slots 2-3, while the bias declares only slots 0-1 valid. This can produce
finite but silently incorrect output.

If padded batching is intended to be supported, this is a blocker and the
graph needs per-token validity/compacted scatter indices. If the intended
contract is compact, unpadded chunks only, please enforce that at the serving
boundary and state it explicitly in the public TensorRT documentation. The
currently documented valid-length invariant alone cannot validate or recover a
left-padded layout.

Confirmed minor issues

1. An invalid SDK path prepends the current directory to PATH

   examples/tensorrt_debug/_runtime_support.py:23-26

   If neither <sdk>/bin nor <sdk>/lib exists, directories is empty and
   the new value begins with os.pathsep. On POSIX, an empty PATH component
   means the current working directory. Modify PATH only when at least one
   directory exists, and restore the prior value when resources are closed.

2. Static-cache metadata validation passes when no matching scatter is found

   src/mobius/integrations/onnx_genai/workflow_metadata.py,
   _static_cache_ports

   The new per-buffer filtering is correct, but if a Cast, Identity, Q/DQ, or
   another node is inserted between the graph input and TensorScatter,
   axes becomes empty and the axis check passes vacuously. Fail explicitly
   when a declared static-cache buffer has no matching scatter consumer.

3. The sliding-window warning contradicts the generated TensorRT graph

   src/mobius/models/base.py:272-279

   TensorRT static-cache export now applies the sliding window through the
   explicit bias, but the non-GQA warning still says the graph uses full causal
   attention and diverges from Hugging Face. Skip or reword the warning when
   the explicit static-cache bias is emitted.

4. The public static-cache task documentation describes only the old ABI

   src/mobius/tasks/_causal_lm.py:58-69

   The docstring still specifies rank-3 flattened caches and is_causal=1.
   Document both contracts:

   • flattened: [B, capacity, kv_heads * head_dim], axis 1;
   • heads-first: [B, kv_heads, capacity, head_dim], axis 2, explicit bias,
     is_causal=0.

5. The new capability and flag behavior are under-documented

   src/mobius/_execution_providers.py
src/mobius/_flags.py

   static_cache_layout is a public graph-ABI capability but has no Args:
   documentation. The static_cache_bias flag documentation also still says
   shipped graphs do not change unless the flag is enabled, which is no longer
   true for an EP with supports_attention_nonpad_kv_seqlen=False.

Runtime coverage still needed

These are not confirmed static findings, but they are the main residual risks:

• nonpad_kv_seqlen=0, or padding combined with a sliding window, can create
  an all-dtype.min bias row. Without Attention input #6, this may aggregate
  unwritten or stale cache values rather than produce a zero row.
• The new tensorrt EP can also be selected for dynamic-cache export, but the
  PR's investigation validates only static cache. TensorRT's bottom-right
  causal offset with dynamic past_key / past_value should be tested or the
  supported scope should be limited.
• There is no TensorRT runtime coverage for batch greater than one, unequal
  write_indices, padded inputs, or cache-capacity overflow.
• --ep tensorrt --runtime ort-genai is not explicitly rejected, although
  standalone TensorRT is distinct from ORT's TensorRT/TRT-RTX providers.

I recommend defining this path explicitly as:

standalone TensorRT + static cache + compact/unpadded input

unless the additional batching and dynamic-cache cases are validated and
supported.

This was a static review. I did not execute the TensorRT engine tests.