[TensorRT] Add TensorRT target and export qwen3-06b model

[rui-ren]

Register a standalone TensorRT target with a heads-first static cache layout while preserving flattened caches for existing targets. Convert Q/K/V around external-cache attention, scatter along the selected sequence axis, emit one Attention output, and restore the output projection layout.

Update static-cache capacity lookup for both layouts and add a TensorRT documentation entry. Draft follow-up work includes numerical prefill/decode parity, permanent regression tests, and documentation/formatting cleanup.

Validation

1. Build .engine

 "C:\TensorRT-11.3.0.99\bin\trtexec.exe" `
>>     --onnx=C:\Users\**\mobius\qwen3-06B\models\tensorrt\model.onnx `
>>     --verbose

2. Test with CUDA

 $python -X utf8 $runner --engine qwen3-06B\models\tensorrt-4k-exporter-fixed\model.engine                                                     
Allocated I/O buffers: 485.1 MiB  
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
Prompt tokens: 22; cache capacity: 4096
Step 0: prefill; token=785; text='The'
Step 1: decode; token=6722; text=' capital'
Step 2: decode; token=315; text=' of'
Step 3: decode; token=9625; text=' France'
Step 4: decode; token=374; text=' is'
Step 5: decode; token=12095; text=' Paris'
Step 6: decode; token=13; text='.'
Step 7: decode; token=151645; text='<|im_end|>'

Prompt: What is the capital of France? Answer briefly.
Response: The capital of France is Paris.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
0 out of 3 committers have signed the CLA.

❌ ruiren@microsoft.com
❌ rui-ren
❌ Copilot

---

ruiren@microsoft.com seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Copilot]

🟡 Changes recommended

Critical workflow metadata handling still rejects or mispublishes heads-first TensorRT caches, with additional validation, coverage, and documentation fixes unresolved.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR adds standalone TensorRT support for Qwen3 using heads-first static KV caches and explicit attention handling.

Changes:

• Registers TensorRT capabilities and adapts cache allocation, attention conversion, scattering, and output projection.
• Adds cache contract tests and TensorRT generation/debugging utilities.
• Updates architecture, execution-provider, and research documentation.

File summaries

File	Summary
tests/static_cache_metadata_test.py	Adds TensorRT cache and bias contract tests; lacks TensorRT-gated numerical prefill/decode coverage.
src/mobius/tasks/_causal_lm.py	Adds provider-specific cache layouts; workflow metadata still assumes flattened caches and can fail TensorRT exports.
src/mobius/models/base.py	Adds explicit static-cache bias selection; documentation should scope behavior to providers with native length support.
src/mobius/components/_attention.py	Adds heads-first cache scatter, attention handling, and output conversion.
src/mobius/_execution_providers.py	Registers TensorRT; metadata generation does not derive the heads-first axis/layout correctly.
README.md	Updates architecture diagrams.
examples/tensorrt_static_cache_generation.py	Adds generation and comparison tooling; two validation guards reject valid configurations.
examples/tensorrt_debug/runtime.py	Provides TensorRT runtime helpers.
examples/tensorrt_debug/inspect_cache.py	Provides cache inspection utilities.
examples/tensorrt_debug/inspect_attention.py	Provides attention inspection utilities.
examples/tensorrt_debug/full_prefix.py	Provides full-prefix debugging support.
examples/tensorrt_debug/comparison.py	Provides comparison helpers.
examples/tensorrt_debug/_runtime_support.py	Provides shared runtime support.
examples/tensorrt_debug/_attention_support.py	Provides shared attention support.
examples/tensorrt_debug/__init__.py	Initializes TensorRT debugging utilities.
examples/tensorrt_attention_probe.py	Builds diagnostic TensorRT attention engines.
docs/research/tensorrt-static-cache-debugging.md	Documents debugging workflows; one export command uses invalid --feature syntax.
docs/index.md	Links the TensorRT research documentation.
docs/execution_providers.md	Documents TensorRT capabilities; public provider and CLI discovery tables omit the standalone tensorrt target.

Review details

Suppressed comments (4)

docs/execution_providers.md:152

• The standalone tensorrt target is shown only in the registry example below, while the Supported Execution Providers table and docs/cli_reference.md still list only trt-rtx. Users following the public EP/CLI tables will not discover the new --ep tensorrt route; add it to those discovery tables as well.

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

• With --compare-hf --max-new-tokens 1, the only meaningful run is the prefill step, but this guard rejects every possible --compare-steps value because it requires a minimum of 2. Since the option counts steps including prefill, allow a minimum of 1.

    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:100

• This capacity guard is one token too strict. The final loop iteration consumes prompt_length + max_new_tokens - 1 input tokens before producing the last token, so a prompt plus generation budget exactly equal to capacity is valid but rejected here; validate_full_prefix_budget already uses the correct - 1 bound.

    if current_ids.shape[1] + args.max_new_tokens > capacity:

src/mobius/_execution_providers.py:379

• Setting supports_attention_nonpad_kv_seqlen=False makes TextModel._maybe_static_cache_bias() emit an explicit bias even when flags.static_cache_bias is false. The static_cache_bias docstring still says the default leaves every shipped model graph unchanged unless that flag is set, which is no longer true for this new provider; update the wording to scope that statement to providers with native length support.

            supports_attention_nonpad_kv_seqlen=False,

• Files reviewed: 19/19 changed files
• Comments generated: 4
• Review effort level: Lite

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.