diff --git a/README.md b/README.md
index d2aa51543..c4a82ab0a 100644
--- a/README.md
+++ b/README.md
@@ -145,20 +145,31 @@ See the [CLI Reference](https://onnxruntime.github.io/mobius/cli_reference.html)
## Architecture
-```
-HuggingFace Hub
- │
- ▼
- ArchitectureConfig ◄── from_transformers() / from_diffusers()
- │
- ▼
- Model Module ◄── Reusable Components (Attention, MLP, RMSNorm, RoPE, …)
- │
- ▼
- Task ◄── CausalLMTask, VisionLanguageTask, VAETask, DenoisingTask, …
- │
- ▼
- ONNX Model ◄── preprocess_weights() + apply_weights()
+```mermaid
+flowchart TD
+ Sources["Model sources
Transformers · Diffusers · GGUF · NeMo"]
+ Config["ArchitectureConfig
Normalizes source configuration"]
+ Registry["Registry
Selects the model class and task"]
+ Components["Reusable components
Attention · MLP · Norm · RoPE · MoE · Vision · Audio"]
+ Models["Model modules
Compose components into architectures"]
+ Tasks["Tasks
Define ONNX inputs, outputs, caches, and model splits"]
+ Graph["ONNX graph construction
onnxscript.nn + onnx_ir"]
+ Optimize["EP-aware optimization
Cleanup · Fusion · Lowering · Folding"]
+ Weights["Weight pipeline
Download · Rename · Transform · Cast · Apply"]
+ Package["ModelPackage
One or more deployable ONNX models"]
+ Runtime["ONNX Runtime / ONNX Runtime GenAI"]
+
+ Sources --> Config
+ Config --> Registry
+ Registry --> Models
+ Registry --> Tasks
+ Components --> Models
+ Models --> Tasks
+ Tasks --> Graph
+ Graph --> Optimize
+ Optimize --> Weights
+ Weights --> Package
+ Package --> Runtime
```
The package is organised into four layers:
@@ -171,6 +182,37 @@ The package is organised into four layers:
See the [design document](https://onnxruntime.github.io/mobius/design.html) for details.
+### Repository organization
+
+```mermaid
+flowchart LR
+ Root["src/mobius/"]
+ Root --> API["Public API and build orchestration
__init__.py · _builder.py · _model_package.py"]
+ Root --> Configs["_configs/
Normalized architecture configuration"]
+ Root --> Components["components/
Reusable ONNX building blocks"]
+ Root --> Models["models/
Architecture implementations"]
+ Root --> Tasks["tasks/
Graph I/O contracts"]
+ Root --> Registry["_registry.py
Model and task lookup"]
+ Root --> Optimizations["_optimizations.py · rewrite_rules/ · _passes/
Graph optimization"]
+ Root --> Integrations["integrations/
Transformers · Diffusers · GGUF · NeMo · ORT GenAI"]
+
+ Configs --> Models
+ Components --> Models
+ Registry --> Models
+ Registry --> Tasks
+ Models --> Tasks
+ Tasks --> API
+ Integrations --> API
+ API --> Optimizations
+```
+
+Supporting directories:
+
+- `tests/` contains graph-construction, integration, parity, generation, and runtime tests.
+- `src/mobius/**/*_test.py` contains unit tests co-located with their implementation.
+- `examples/` demonstrates text, multimodal, speech, and diffusion workflows.
+- `docs/` contains user guides, design documentation, and API reference material.
+
## Development
```bash
diff --git a/docs/execution_providers.md b/docs/execution_providers.md
index f22c52b36..3036e9e7b 100644
--- a/docs/execution_providers.md
+++ b/docs/execution_providers.md
@@ -145,6 +145,11 @@ EpCapabilities(name="webgpu", gqa_dtypes={FLOAT, FLOAT16},
EpCapabilities(name="trt-rtx", gqa_dtypes={FLOAT16, BFLOAT16},
supports_skip_layer_norm=False, enable_graph_capture=True,
provider_options={"enable_cuda_graph": "1"})
+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)
EpCapabilities(name="onnx-standard",
gqa_dtypes=frozenset(), qkv_pack_dtypes=frozenset(),
supports_fused_rope=False, # not used by default for this EP, since it does not enable GQA fusion
@@ -153,6 +158,18 @@ EpCapabilities(name="onnx-standard",
supports_packed_multi_head_attention=False)
```
+Standalone `tensorrt` is separate from ORT's `trt-rtx` provider. Static-cache
+exports use `[B, kv_heads, capacity, head_dim]` caches and an explicit additive
+attention bias, with `is_causal=0`. Query cache slots are
+`write_indices + arange(query_length)`; allowed keys must also be below
+`nonpad_kv_seqlen`. The length remains a graph input used by the bias, but is
+omitted from native Attention input #6: the tested TensorRT 11.3 path ignores
+that input and otherwise applies incorrect top-left causality during decode.
+Default/ORT masking behavior is unchanged. No extra `static-cache-bias` flag
+is required for standalone TensorRT. Other model backbones must supply a full
+static-cache bias or export fails explicitly instead of emitting maskless
+attention on this provider.
+
Out-of-tree EPs can register at runtime via `register_ep()`:
```python
diff --git a/docs/index.md b/docs/index.md
index 103db3353..7ab0a645f 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -52,4 +52,5 @@ ai-model-support-strategy
:caption: Research
research/testing-strategy-analysis
+research/tensorrt-static-cache-debugging
```
diff --git a/docs/research/tensorrt-static-cache-debugging.md b/docs/research/tensorrt-static-cache-debugging.md
new file mode 100644
index 000000000..f87bcde70
--- /dev/null
+++ b/docs/research/tensorrt-static-cache-debugging.md
@@ -0,0 +1,535 @@
+# Debugging TensorRT Static-Cache Decode: From Symptom to Fix
+
+This is a worked debugging case, recorded on 2026-09-15, for a Mobius-exported
+Qwen3-0.6B model. It follows the evidence in the order we collected it: an engine
+that built successfully, incorrect generated text, competing hypotheses,
+progressively smaller comparisons, and a verified exporter correction.
+
+The central lesson is **find the first observable divergence, then design a
+check that can distinguish its possible causes**. Do not change the exporter
+just because a symptom sounds like a familiar bug.
+
+## Reading the Runner
+
+Start with `generate()` in
+[the generation script](../../examples/tensorrt_static_cache_generation.py).
+Its loop keeps the order visible: prepare inputs, snapshot the cache, execute
+and wait, snapshot again, run the HF reference, compare, then select a token.
+`parse_args()` validates options; `run()` prepares the tokenizer and inspectors.
+
+Read only the helper for the question you are investigating:
+
+| Module | Responsibility |
+|---|---|
+| [full_prefix.py](../../examples/tensorrt_debug/full_prefix.py) | Build prompt + generated token IDs and check the input-length budget. The loop visibly resets both runtimes' caches and position to zero. |
+| [comparison.py](../../examples/tensorrt_debug/comparison.py) | Run the independent HF reference cache and report whole-vector logits or tensor errors. |
+| [inspect_cache.py](../../examples/tensorrt_debug/inspect_cache.py) | Validate paired cache layouts, take read-only snapshots, report changed slots, and compare valid K/V against HF. |
+| [inspect_attention.py](../../examples/tensorrt_debug/inspect_attention.py) | Keep TRT Q/K/V fixed, reconstruct attention with each mask hypothesis, and compare with the actual output. |
+| [_attention_support.py](../../examples/tensorrt_debug/_attention_support.py) | Layer-0 HF hooks, RoPE alignment, and reference comparisons. Skip this setup while reading the mask experiment. |
+| [runtime.py](../../examples/tensorrt_debug/runtime.py) | Follow just four per-step operations: prepare inputs, execute, read tensors, and reset caches. |
+| [_runtime_support.py](../../examples/tensorrt_debug/_runtime_support.py) | Engine loading, allocation, aliased cache addresses, shape/layout validation, binding, and cleanup. Skip this plumbing while following the exporter investigation. |
+
+The command-line options and commands below are unchanged. Inspectors do not
+select tokens or modify the GPU caches. `read_tensor()` waits for the runner's
+stream and returns an independent CPU array with the original dtype, so BF16
+byte comparisons remain meaningful. The separate
+[attention probe builder](../../examples/tensorrt_attention_probe.py) still builds
+the diagnostic engine; `--inspect-attention` only reads its exposed outputs.
+
+## 1. The Result in One Paragraph
+
+The original engine generated `TheQuestionQuestion...`. Prefill was close to
+HuggingFace, but cached decode was wrong. Cache snapshots showed correct write
+locations and preserved history. A separate diagnostic engine revealed that
+layer-0 decode attention exactly matched a top-left causal mask: the single
+query attended only to cache slot 0. The diagnostic build also reported
+`nonpad_kv_seqlen` as unused. Explicit position-aware causal/valid-length masking
+corrected the computation. The integrated exporter then generated
+`The capital of France is Paris.`, matching HuggingFace's top token on every
+step through EOS in this test.
+
+This identifies the observed execution behavior and a working correction. It
+does not determine whether the original behavior is an unsupported TensorRT
+parser feature or an implementation defect, nor prove behavior on all versions.
+
+## 2. Environment and Controlled Variables
+
+| Item | Tested configuration |
+|---|---|
+| Model | `Qwen/Qwen3-0.6B` |
+| HF reference revision | `c1899de289a04d12100db370d81485cdf75e47ca` |
+| Device | NVIDIA GeForce RTX 4060 Laptop GPU, SM 8.9, 8 GiB |
+| TensorRT | 11.3.0.99, standalone runtime, not ORT TRT-RTX |
+| Export/runtime dtype | BF16 |
+| Reference | HuggingFace Transformers, FP32, eager attention, CPU PyTorch |
+| Static cache | 28 layers, 8 KV heads, head dimension 128, capacity 4096 |
+| Batch/profile | Batch 1; token input min/opt/max lengths 1/32/128 |
+| Engine build | `--decomposableAttentions=*` |
+| Prompt | `What is the capital of France? Answer briefly.` |
+| Tokenization | Qwen chat template, generation prompt, `enable_thinking=False` |
+
+The rendered prompt has 22 tokens in this configuration. Reuse the actual token
+IDs rather than assuming a string always tokenizes to the same sequence.
+
+The runner's `--revision` option can pin tokenizer/reference downloads. The
+original engine does not independently prove which checkpoint revision produced
+its weights. Record export provenance for future investigations.
+
+Preserve separate artifacts for the failing baseline, instrumented baseline,
+experimental correction, and fresh exporter-generated correction. Do not
+overwrite the only reproducible failing engine.
+
+## 3. Understand What Crosses the Prefill/Decode Boundary
+
+Each layer owns key and value tensors shaped `[B, Hkv, capacity, D]`. The cache
+contains projected values and normalized, RoPE-transformed keys. It does not
+contain token IDs. The query is computed for the current input chunk and is not
+stored in this cache.
+
+For the first two calls in this case:
+
+| Input/state | Prefill | First cached decode |
+|---|---|---|
+| `input_ids` | 22 prompt tokens | `[[785]]`, the token `The` |
+| `position_ids` | `[[0, ..., 21]]` | `[[22]]` |
+| `write_indices` | `[0]` | `[22]` |
+| `nonpad_kv_seqlen` | `[22]` | `[23]` |
+| Slots written | 0 through 21 | 22 |
+| Last query should attend | Slots 0 through 21 | Slots 0 through 22 |
+
+Prefill predicts `The`; the decode call consumes `The` and predicts the token
+after it. This distinction prevents a common off-by-one mistake in comparisons.
+
+`TensorScatter` writes the new K/V into the cache. `Attention` reads that cache.
+Correct writes do not guarantee correct reads or masking.
+
+For an unpadded chunk of width S starting at slot W, valid length is W + S.
+Capacity is fixed at 4096; valid length is not capacity. With padding, the
+relationship involves the valid token count, not necessarily the padded width.
+This walkthrough's real-model run is unpadded, batch one.
+
+## 4. The Investigation Roadmap
+
+| Stage | Question | Evidence | Next decision |
+|---|---|---|---|
+| Build/load | Can TensorRT accept the artifact? | Engine built and deserialized | Run meaningful inputs |
+| Generation | Does it produce sensible output? | `TheQuestionQuestion...` | Compare numerical outputs |
+| Identical-token reference | When does it first diverge? | Prefill close, first decode wrong | Isolate incremental execution |
+| Full-prefix control | Does the same token history work without reuse? | Full prefix predicts ` capital` | Inspect persisted state |
+| Cache snapshots | Are slots missing or overwritten? | Correct slots, history preserved | Inspect computation using the cache |
+| Neighboring layer | Where does the error become large? | Layer-0 K/V close; layer-1 new K/V diverge | Probe layer-0 attention |
+| Intermediate outputs | Are query or attention outputs wrong? | Query close; attention wrong | Test explicit mask hypotheses |
+| Own-Q/K/V reconstruction | Which mask explains the result? | Top-left mask matches exactly | Try position-aware masking |
+| Corrective experiment | Does changing masking remove the symptom? | Correct attention and generation | Integrate and regression-test |
+
+### Stage A: Build Success Is Not Correctness
+
+Deserialization establishes that the runtime can load the engine. Finite logits
+establish that the checked numbers are not NaN or infinity. Neither establishes
+that the graph computes the intended function.
+
+The Python TensorRT binding executes the native engine through
+`context.execute_async_v3(...)`. Python handles tokenization, buffer allocation,
+input preparation, next-token selection, and decoding text. It is not running
+the model through PyTorch or ONNX Runtime.
+
+The first attempt used separate input/output cache buffers. TensorRT rejected
+execution because this engine required those pairs to alias. We then bound each
+cache output to its input address. That was an explicit runtime requirement,
+not evidence that aliasing was numerically correct; the snapshot tests checked
+the resulting writes later.
+
+**Earlier clue we should have prioritized:** an input that controls correctness
+being reported as unused deserves immediate investigation, even if build passes.
+
+### Stage B: Compare Identical Tokens, Not Independent Stories
+
+Let both implementations consume the same prompt. Save their last-token logits.
+Choose one next token, here TensorRT's `785`, and feed that exact token to both.
+Each runtime maintains its own cache. Compare the next logit vectors.
+
+If the implementations independently select different tokens, later differences
+can be explained by different histories. Teacher forcing removes that ambiguity.
+
+| Baseline measurement | Prefill | First decode |
+|---|---:|---:|
+| Mean absolute logit error | 0.051921 | 4.260486 |
+| Maximum absolute error | 0.268508 | 28.258783 |
+| Correlation | 0.999899 | 0.073147 |
+| TRT top token | `The` | `Question` |
+| HF top token | `The` | ` capital` |
+
+This says the first tested cached call is wrong. It does **not** prove that the
+cache is corrupted: write logic, read logic, positions, masking, shape-dependent
+kernels, or other decode computation could explain it.
+
+### Stage C: Recompute the Same History Without Cache Reuse
+
+Compare two paths:
+
+1. Process the prompt, preserve caches, then process only `[[785]]`.
+2. Clear caches and process the prompt concatenated with `[[785]]` in one call.
+
+The second path computes the same next-token prediction, but does not depend on
+state persisted across calls. Both baseline runs selected `785` first, so their
+second-call token histories were identical.
+
+| Second-call result | Cached | Full prefix |
+|---|---:|---:|
+| TRT top token | `Question` | ` capital` |
+| Correlation with HF | 0.073147 | 0.999931 |
+| Mean absolute logit error | 4.260486 | 0.052296 |
+
+This narrows the failure to differences in incremental execution. It still
+does not distinguish storage from attention computation. Full-prefix mode also
+changes query length and may select different runtime implementations.
+
+For longer cross-mode comparisons, explicitly share continuation tokens. Two
+independent runs are comparable only while their token histories remain equal.
+
+### Stage D: Read the Cache, Do Not Guess About It
+
+Before interpreting memory, check both cache bindings' location, dtype, format,
+shape, and runtime strides. The probe verified device-resident LINEAR BF16 data
+with shape `(1, 8, 4096, 128)` and element strides
+`(4194304, 524288, 128, 1)`.
+
+Synchronize GPU execution before copying. Keep independent CPU snapshots so
+in-place updates cannot overwrite the evidence. Preserve raw BF16 bytes for
+unchanged-region tests, then convert to FP32 for numerical comparisons.
+
+For prompt length P = 22:
+
+| Region | Expected after prefill | Expected after decode |
+|---|---|---|
+| `[:, :, :P, :]` | Prompt K/V | Bit-for-bit unchanged |
+| `[:, :, P:P+1, :]` | Zero | New token K/V |
+| `[:, :, P+1:, :]` | Zero | Still zero and unchanged |
+
+The HF Qwen3 cache stores `[B, Hkv, S, D]`, with keys after Q/K normalization and
+RoPE, matching the intended Mobius representation. Compare equivalent stages;
+comparing pre-RoPE keys with post-RoPE keys would manufacture a false mismatch.
+
+Observed for both inspected layers: prefill changed only slots 0-21, decode
+changed only slot 22, old slots were unchanged, and unused slots stayed zero.
+
+| New decode slot vs HF | Layer 0 mean abs | Layer 1 mean abs |
+|---|---:|---:|
+| Keys | 0.007317 | 1.325746 |
+| Values | 0.000327 | 0.122456 |
+
+Layer-0 prefill key maximum error was 2.673859, but keys reached magnitude 520
+and mean error was 0.009226. Do not interpret a maximum absolute error without
+the tensor's scale and dtype. Small averages are not a formal parity guarantee,
+either: a few important entries can affect attention disproportionately.
+
+The strong conclusion here is about storage: the inspected slots were written
+and preserved correctly. The next large observed numerical mismatch was in
+layer-1 K/V, which depend on computations after layer-0 K/V formation.
+
+### Stage E: Expose the First Suspect Computation
+
+A serialized engine does not automatically expose every intermediate tensor.
+We built a separate diagnostic engine by loading the ONNX graph with `onnx_ir`
+and adding graph outputs for:
+
+- Layer-0 post-RoPE query: `[B, Hq, S, D]`.
+- Layer-0 attention output immediately before `o_proj`: `[B, S, Hq * D]`.
+
+The HF probe captures `o_proj` input with a forward pre-hook. It also captures
+normalized Q and applies the same HF RoPE function/positions to compare queries.
+
+Exposing outputs can change optimization and fusion. We therefore verified that
+the instrumented engine still reproduced `TheQuestion`. Its logits differed
+slightly from the original engine, so keep instrumented measurements separate.
+
+| Instrumented layer-0 comparison vs HF | Prefill mean abs | Decode mean abs |
+|---|---:|---:|
+| Post-RoPE query | 0.004103 | 0.004588 |
+| Pre-projection attention output | 0.000375 | 0.124763 |
+
+The first directly observed large mismatch is now inside attention, before its
+output projection. If attention had matched, the next probes would have followed
+the output projection, residual addition, normalization, and MLP instead.
+
+### Stage F: Test Two Masks Using the Engine's Own Q/K/V
+
+An HF output comparison still mixes input differences with computation
+differences. To isolate the mask, reconstruct attention on CPU from the
+instrumented engine's actual Q and cache K/V, converted to FP32.
+
+For grouped-query attention, repeat each KV head to match its associated query
+heads. This model has 16 query heads and 8 KV heads, so each KV head serves two
+query heads. The reconstruction computes:
+
+$$
+A = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{D}} + M\right)V
+$$
+
+Then transpose/reshape A to the same pre-projection layout. Test two hypotheses:
+
+| Hypothesis | Allowed keys for local query offset t |
+|---|---|
+| Correct cached causality | `key_slot <= write_indices[b] + t`, with valid-prefix bound |
+| Top-left causality | `key_slot <= t` |
+
+At prefill, write index is zero, so these hypotheses make the same prediction.
+At single-token decode, t = 0: top-left masking allows only slot 0, while the
+correct mask allows slots 0 through 22. This is why decode is discriminating.
+
+| Original instrumented decode attention vs own-Q/K/V reconstruction | Mean abs |
+|---|---:|
+| Correct cached mask | 0.124732 |
+| Top-left mask | **0.000000** |
+
+The top-left reconstruction matched exactly, not merely in its top token. With
+only one key allowed, softmax has a single probability of one, so the output is
+the corresponding slot-0 value. This is particularly strong evidence.
+
+The build log independently reported `Unused Input: nonpad_kv_seqlen`.
+Together, these observations explain the failure: the tested path does not use
+the valid-length information to obtain the intended cached causal behavior.
+
+The reconstruction in the current probe assumes unpadded inputs, so the correct
+causal frontier also bounds valid length. Extend it with an explicit validity
+condition before using it for padded experiments.
+
+## 5. The Correct Mask, With Shapes
+
+Let B be batch size, S the current query width, and C the cache capacity.
+
+```text
+write_indices: [B]
+query_offsets: [S] = arange(S)
+query_slots: [B, S] = write_indices[:, None] + query_offsets[None, :]
+key_slots: [C] = arange(C)
+nonpad_kv_seqlen: [B]
+
+causal: [B, S, C] = key_slots <= query_slots[:, :, None]
+valid: [B, 1, C] = key_slots < nonpad_kv_seqlen[:, None, None]
+allowed: [B, S, C] = causal AND valid
+additive bias: [B, 1, S, C]
+```
+
+Use zero bias for allowed positions and a sufficiently negative value for
+blocked positions. Broadcasting the head dimension shares the mask across heads.
+
+For a three-token chunk starting at 22, with valid length 25:
+
+```text
+query slot 22 -> keys 0..22
+query slot 23 -> keys 0..23
+query slot 24 -> keys 0..24
+all queries -> block slots 25..4095
+```
+
+Setting `is_causal=0` alone is not a fix: without a full explicit mask, queries
+could attend future tokens and unused capacity. Likewise, retaining the faulty
+implicit causal mask on top of a correct bias would reapply the wrong restriction.
+
+## 6. Corrective Experiment, Then Integration
+
+### Separate graph experiment
+
+We applied an explicit BF16 bias to all 28 attention layers, set `is_causal=0`,
+and removed native Attention input #6. The graph-level `nonpad_kv_seqlen` input
+remained because the explicit mask consumes it. Cache scatter behavior was not
+changed.
+
+The prototype used RoPE `position_ids` as query positions and `-inf` as blocked
+bias. That is sufficient for this specific prompt, where RoPE positions equal
+cache slots. It is not a general reason to equate those concepts.
+
+After the experiment, first-decode attention mean error vs HF fell from
+0.124763 to 0.000312. Logit correlation became 0.999909 and the next token became
+` capital`. Generation matched all eight HF top-token predictions through EOS.
+
+### Exporter implementation
+
+The permanent code path reuses existing abstractions rather than copying the
+prototype graph rewrite:
+
+| Location | Responsibility |
+|---|---|
+| [EP capabilities](../../src/mobius/_execution_providers.py) | `supports_attention_nonpad_kv_seqlen` defaults true; false for standalone `tensorrt` |
+| [TextModel](../../src/mobius/models/base.py) | `_maybe_static_cache_bias()` enables explicit bias when the EP needs it, even without a sliding window or extra feature flag |
+| [Mask helper](../../src/mobius/components/_common.py) | `create_static_cache_attention_bias()` uses write indices, query offsets, and valid length |
+| [Attention emission](../../src/mobius/components/_attention.py) | Requires a supplied bias on unsupported EPs, selects `is_causal=0`, omits native nonpad input |
+| [Regression tests](../../tests/static_cache_metadata_test.py) | Provider/dtype contracts, mask geometry, and missing-bias guard |
+
+Important boundaries:
+
+- Cache layout and masking support are separate capabilities. A rank-4 cache is
+ not itself a reason to disable native masking.
+- Default/ORT providers keep their existing native nonpad behavior.
+- The dynamic-cache branch remains unchanged.
+- The reusable bias uses cache-slot positions, not an assumption about RoPE IDs.
+- Other backbones that fail to provide a required bias are rejected explicitly;
+ this change does not claim universal model support.
+- The existing helper uses `dtype.min`, not the prototype's `-inf`. Both worked
+ for the tested unpadded inputs. Fully masked rows require separate analysis;
+ the two conventions can behave differently there.
+
+### Validation of the integrated exporter
+
+We exported the real model using the updated Mobius CLI, checked all 28 Attention
+nodes, and built an uninstrumented engine directly from that export. No
+experimental ONNX mask rewrite was applied to this final artifact.
+
+```text
+Response: The capital of France is Paris.
+Top-1 agreement: all 8 executed steps, including EOS
+Logit correlation: approximately 0.999799 to 0.999931
+```
+
+The neighboring test file passed 40 tests, including 17 new cases. New coverage
+includes default/TensorRT, FLOAT/FLOAT16/BFLOAT16 graph contracts, static/dynamic
+paths, prefill/decode/chunked mask geometry, valid-prefix clamping, and a
+missing-bias guard. Numerical mask tests run on CPU ORT; they do not substitute
+for TensorRT engine execution. Ruff and editor checks also passed.
+
+## 7. Reproduce the Investigation
+
+Run commands from the repository root in PowerShell. These examples refer to
+local artifacts created during this case; engines and weights are not portable
+source files or guaranteed to exist in a fresh checkout.
+
+Prerequisites include the TensorRT SDK/runtime DLLs, the matching TensorRT Python
+wheel, CUDA runtime bindings, NumPy with BF16 support via `ml_dtypes`, Transformers,
+and PyTorch for the CPU reference. The runner's default SDK path is
+`C:/TensorRT-11.3.0.99`; use `--sdk` when it differs. The build helper accepts
+`--trtexec` for an alternative executable path.
+
+The [generation runner](../../examples/tensorrt_static_cache_generation.py)
+contains the comparisons. The [diagnostic builder](../../examples/tensorrt_attention_probe.py)
+adds intermediate outputs and optionally applies the experimental mask.
+
+### Baseline, full-prefix, and cache probes
+
+```powershell
+$python = '.\.venv\Scripts\python.exe'
+$runner = 'examples\tensorrt_static_cache_generation.py'
+$baseline = 'qwen3-06B\models\tensorrt-4k\model.engine'
+
+& $python -X utf8 $runner --engine $baseline --compare-hf --max-new-tokens 2
+& $python -X utf8 $runner --engine $baseline --compare-hf --max-new-tokens 2 --full-prefix
+& $python -X utf8 $runner --engine $baseline --compare-hf --max-new-tokens 2 --inspect-cache
+& $python -X utf8 $runner --engine $baseline --compare-hf --max-new-tokens 2 --inspect-cache --cache-layer 1
+```
+
+Confirm both two-step paths select the same first token. Cache inspection is
+restricted to cached mode, requires `--compare-hf`, and snapshots the first two
+steps. Full-prefix mode resets both caches each step and must fit the engine's
+maximum query length, not merely its larger cache capacity.
+
+### Attention probe and experimental engine
+
+```powershell
+& $python -X utf8 $runner --engine qwen3-06B\models\tensorrt-4k-attention-probe\model.engine --compare-hf --max-new-tokens 2 --inspect-attention
+& $python -X utf8 $runner --engine qwen3-06B\models\tensorrt-4k-explicit-mask\model.engine --compare-hf --max-new-tokens 2 --inspect-attention
+```
+
+To rebuild diagnostic copies from the retained failing ONNX graph, choose new
+output directories. The builder deliberately refuses an existing output directory:
+
+```powershell
+& $python -X utf8 examples\tensorrt_attention_probe.py --onnx qwen3-06B\models\tensorrt-4k\model.onnx --output qwen3-06B\models\attention-probe-new
+& $python -X utf8 examples\tensorrt_attention_probe.py --onnx qwen3-06B\models\tensorrt-4k\model.onnx --output qwen3-06B\models\explicit-mask-new --explicit-mask
+```
+
+Do not apply `--explicit-mask` to an already corrected export: the prototype
+expects the original maskless BF16 graph. Without that option, the builder only
+adds diagnostic outputs; the corrected exporter does not require this tool.
+
+### Integrated export and final validation
+
+```powershell
+.\.venv\Scripts\mobius.exe build --model Qwen/Qwen3-0.6B --dtype bf16 --ep tensorrt --features static-cache --max-seq-len 4096 qwen3-06B/models/tensorrt-4k-exporter-new
+
+& $python -X utf8 $runner --engine qwen3-06B\models\tensorrt-4k-exporter-fixed\model.engine --compare-hf --compare-steps 16 --max-new-tokens 16
+& $python -m pytest tests/static_cache_metadata_test.py -q --tb=short
+```
+
+The first command creates ONNX, not a TensorRT engine. The second runs the
+already-built final engine from this session. To compile the new ONNX yourself,
+use `trtexec` with `--decomposableAttentions=*` and min/opt/max input profiles
+1/32/128 for tokens/positions, batch-one fixed cache shapes, and `[1]` length/index
+inputs, as done by the diagnostic builder. Use a new engine filename. Successful
+compilation must still be followed by the runtime comparison.
+
+`--compare-steps` defaults to two. Increasing only `--max-new-tokens` does not
+increase numerical comparison coverage. `Compared 8/16` in this example means
+EOS arrived at step 8; it is normal completion, not a runtime failure. It also
+means the test did not exercise 16 execution steps.
+
+## 8. How to Read the Metrics Without Overclaiming
+
+- **Mean absolute error** measures average numerical distance, but can hide a
+ few important wrong entries.
+- **Maximum absolute error** exposes outliers, but must be interpreted relative
+ to tensor magnitude and precision.
+- **Relative RMS error** compares error magnitude with reference magnitude; it
+ is helpful for intermediate tensors with different scales.
+- **Correlation** measures similarity of variation, not equality. A constant
+ offset or scaling can retain high correlation. Inspect other metrics too.
+- **Top-1 agreement** checks the selected next token. It is not sufficient to
+ prove that the rest of the distribution or intermediate computation is right.
+- **Bitwise equality** is appropriate for cache regions that must not change.
+ It is not appropriate for BF16-versus-FP32 computed tensor comparisons.
+
+BF16 and FP32 outputs need not match exactly. There is no universal tolerance
+that makes every model/runtime pair correct. The large decode error, exact
+wrong-mask reconstruction, and targeted correction together are much stronger
+evidence than any individual threshold.
+
+## 9. What Is Proven, and What Remains Open
+
+**Supported by this case:** the failing single-token attention matched top-left
+causality; inspected cache writes preserved state; explicit masking corrected the
+tested intermediate and final outputs; the integrated exporter reproduced that
+correction; default/dynamic graph contracts passed focused regressions.
+
+**Not established:** universal TensorRT behavior, the upstream parser/kernel
+classification, all architectures, long-context generation, real multi-batch or
+padded execution, all precision modes on GPU, fully masked-row semantics,
+performance impact, or the entire repository's test suite.
+
+Before broadly deploying the fix, expand tests according to those risks. Keep
+correctness and performance investigations separate: an explicit mask can alter
+kernel selection even when it corrects the results.
+
+## 10. A Reusable Debugging Checklist
+
+1. Preserve the failing artifact and record model revision, runtime, dtype, and profile.
+2. Read warnings about ignored inputs or unsupported operators before blaming model quality.
+3. Run meaningful inputs; build/load success is only the first gate.
+4. Compare identical token IDs against a trusted reference.
+5. Find the earliest failing call: prefill, decode, or a particular shape transition.
+6. Remove state reuse with a full-prefix control, keeping the token history fixed.
+7. Validate memory layout before interpreting raw buffers.
+8. Snapshot state and verify write regions independently of numerical computation.
+9. Compare equivalent intermediates and move to the first large mismatch.
+10. Test competing explanations using the same actual inputs when possible.
+11. Change one semantic factor in a separate experiment; verify the baseline failure still reproduces under instrumentation.
+12. Integrate through the owning abstraction, preserve unrelated runtimes, and add regression tests.
+13. Rebuild from the real exporter and verify again; a patched diagnostic graph is not the final product.
+14. Report the exact evidence and remaining limits, not just "it works."
+
+### Check Your Understanding
+
+**Why did prefill work with the wrong mask?** Query offsets start at zero in
+prefill, so top-left and intended cache-slot causal frontiers coincide.
+
+**Why did correct cache snapshots not rule out an attention bug?** Storage can
+hold the right values while the attention mask selects the wrong subset.
+
+**Why was the own-Q/K/V test decisive?** It removed most projection/reference
+differences and asked which masking rule reproduces the observed output.
+
+**Why not change every provider to explicit masking?** The observed limitation
+is provider-specific; other providers depend on working native behavior and may
+lose performance or change edge-case semantics under a global rewrite.
+
+**Why rebuild from the exporter after the experiment?** The experiment proves
+the proposed correction can work. A fresh export proves the implementation
+actually emits that correction through the normal user workflow.
\ No newline at end of file
diff --git a/examples/tensorrt_attention_probe.py b/examples/tensorrt_attention_probe.py
new file mode 100644
index 000000000..9e164b8b3
--- /dev/null
+++ b/examples/tensorrt_attention_probe.py
@@ -0,0 +1,152 @@
+"""Build a separate TensorRT engine exposing Qwen3 layer-zero attention tensors."""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+from pathlib import Path
+
+import ml_dtypes
+import numpy as np
+import onnx_ir as ir
+
+
+def add_explicit_mask(model):
+ attentions = [node for node in model.graph if node.op_type == "Attention"]
+ inputs = {value.name: value for value in model.graph.inputs}
+ capacity = inputs["key_cache.0"].shape[2]
+ if not isinstance(capacity, int):
+ raise TypeError("Explicit-mask experiment requires a fixed cache capacity")
+ if any(
+ node.inputs[1].shape[2] != capacity
+ or node.inputs[3] is not None
+ or node.inputs[0].dtype != ir.DataType.BFLOAT16
+ for node in attentions
+ ):
+ raise ValueError("Expected uniform BF16 static-cache attention without existing masks")
+ nodes = []
+
+ def constant(name, array):
+ value = ir.Value(name=f"experiment.{name}", const_value=ir.tensor(array))
+ model.graph.initializers[value.name] = value
+ return value
+
+ def operation(name, op_type, arguments, dtype, shape):
+ output = ir.Value(
+ name=f"experiment.{name}", type=ir.TensorType(dtype), shape=ir.Shape(shape)
+ )
+ nodes.append(ir.Node("", op_type, inputs=arguments, outputs=[output]))
+ return output
+
+ key_positions = constant(
+ "key_positions", np.arange(capacity, dtype=np.int64)[None, None, None, :]
+ )
+ query_axes = constant("query_axes", np.asarray([1, 3], dtype=np.int64))
+ length_axes = constant("length_axes", np.asarray([1, 2, 3], dtype=np.int64))
+ batch, sequence = inputs["position_ids"].shape
+ query_positions = operation(
+ "query_positions",
+ "Unsqueeze",
+ [inputs["position_ids"], query_axes],
+ ir.DataType.INT64,
+ [batch, 1, sequence, 1],
+ )
+ valid_length = operation(
+ "valid_length",
+ "Unsqueeze",
+ [inputs["nonpad_kv_seqlen"], length_axes],
+ ir.DataType.INT64,
+ [batch, 1, 1, 1],
+ )
+ causal = operation(
+ "causal",
+ "LessOrEqual",
+ [key_positions, query_positions],
+ ir.DataType.BOOL,
+ [batch, 1, sequence, capacity],
+ )
+ valid = operation(
+ "valid",
+ "Less",
+ [key_positions, valid_length],
+ ir.DataType.BOOL,
+ [batch, 1, 1, capacity],
+ )
+ allowed = operation(
+ "allowed", "And", [causal, valid], ir.DataType.BOOL, [batch, 1, sequence, capacity]
+ )
+ zero = constant("zero", np.asarray(0, dtype=ml_dtypes.bfloat16))
+ negative_inf = constant("negative_inf", np.asarray(-np.inf, dtype=ml_dtypes.bfloat16))
+ mask = operation(
+ "mask",
+ "Where",
+ [allowed, zero, negative_inf],
+ ir.DataType.BFLOAT16,
+ [batch, 1, sequence, capacity],
+ )
+ model.graph.insert_before(attentions[0], nodes)
+ for node in attentions:
+ node.replace_input_with(3, mask)
+ node.replace_input_with(6, None)
+ node.attributes["is_causal"] = ir.AttrInt64("is_causal", 0)
+ print(f"Applied explicit position/valid-length mask to {len(attentions)} layers.")
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--onnx", type=Path, required=True)
+ parser.add_argument("--output", type=Path, required=True)
+ parser.add_argument(
+ "--explicit-mask",
+ action="store_true",
+ help="Experiment with explicit causal/valid-length bias in every layer",
+ )
+ parser.add_argument(
+ "--trtexec", type=Path, default=Path("C:/TensorRT-11.3.0.99/bin/trtexec.exe")
+ )
+ args = parser.parse_args()
+ args.output.mkdir(parents=True, exist_ok=False)
+ model = ir.load(args.onnx)
+ if args.explicit_mask:
+ add_explicit_mask(model)
+ attention = next(node for node in model.graph if node.op_type == "Attention")
+ projection = next(
+ node
+ for node in model.graph
+ if node.op_type == "MatMul" and "layers.0.self_attn.o_proj" in node.outputs[0].name
+ )
+ for label, value in (("query", attention.inputs[0]), ("attention", projection.inputs[0])):
+ output = ir.Value(name=f"probe.{label}", shape=value.shape, type=value.type)
+ model.graph.append(ir.Node("", "Identity", inputs=[value], outputs=[output]))
+ model.graph.outputs.append(output)
+ print("Attention attributes:", attention.attributes, flush=True)
+ destination = args.output / "model.onnx"
+ ir.save(model, destination, external_data="model.onnx.data")
+ command = [
+ str(args.trtexec),
+ f"--onnx={destination}",
+ f"--saveEngine={args.output / 'model.engine'}",
+ "--skipInference",
+ "--decomposableAttentions=*",
+ "--profilingVerbosity=detailed",
+ ]
+ for option, length in (("minShapes", 1), ("optShapes", 32), ("maxShapes", 128)):
+ shapes = []
+ for value in model.graph.inputs:
+ shape = [dim if isinstance(dim, int) else 1 for dim in value.shape]
+ if value.name in ("input_ids", "position_ids"):
+ shape = [1, length]
+ shapes.append(f"{value.name}:{'x'.join(map(str, shape))}")
+ command.append(f"--{option}={','.join(shapes)}")
+ print("Building diagnostic engine; original engine is unchanged.", flush=True)
+ with (args.output / "build.log").open("w", encoding="utf-8") as log:
+ result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=False)
+ print(
+ "\n".join((args.output / "build.log").read_text(encoding="utf-8").splitlines()[-20:])
+ )
+ if result.returncode:
+ raise RuntimeError(f"TensorRT build failed; see {args.output / 'build.log'}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/tensorrt_debug/__init__.py b/examples/tensorrt_debug/__init__.py
new file mode 100644
index 000000000..51000eb02
--- /dev/null
+++ b/examples/tensorrt_debug/__init__.py
@@ -0,0 +1,3 @@
+"""Learning-oriented helpers for the TensorRT static-cache generation example."""
+
+from __future__ import annotations
diff --git a/examples/tensorrt_debug/_attention_support.py b/examples/tensorrt_debug/_attention_support.py
new file mode 100644
index 000000000..3c14ee4fa
--- /dev/null
+++ b/examples/tensorrt_debug/_attention_support.py
@@ -0,0 +1,53 @@
+"""HF hook setup and RoPE alignment for the layer-0 Qwen3 attention probe."""
+
+from __future__ import annotations
+
+import torch
+
+from tensorrt_debug.comparison import compare_tensor
+
+
+class AttentionReference:
+ """Capture HF intermediates; leave the mask experiment in inspect_attention.py."""
+
+ def __init__(self, runner, reference_model):
+ if not {"probe.query", "probe.attention"}.issubset(runner.names):
+ raise ValueError("Use an engine built by tensorrt_attention_probe.py")
+ self.runner = runner
+ self.model = reference_model
+ self.attention = reference_model.model.layers[0].self_attn
+ self.intermediates = {}
+ self.hooks = [
+ self.attention.o_proj.register_forward_pre_hook(self._capture_attention),
+ self.attention.q_norm.register_forward_hook(self._capture_query),
+ ]
+
+ def _capture_attention(self, module, inputs):
+ self.intermediates["attention"] = inputs[0].detach().float().cpu().numpy().copy()
+
+ def _capture_query(self, module, inputs, output):
+ self.intermediates["query"] = output.detach().clone()
+
+ def compare_with_hf(self, feeds, actual_query, actual_attention):
+ from transformers.models.qwen3.modeling_qwen3 import apply_rotary_pos_emb
+
+ with torch.inference_mode():
+ # HF captures Q before RoPE; the diagnostic engine exposes it after RoPE.
+ query = self.intermediates["query"].transpose(1, 2)
+ cos, sin = self.model.model.rotary_emb(
+ query, torch.from_numpy(feeds["position_ids"].copy())
+ )
+ query, _ = apply_rotary_pos_emb(query, query, cos, sin)
+ compare_tensor(
+ "Layer-0 post-RoPE query vs HF", actual_query, query.float().cpu().numpy()
+ )
+ compare_tensor(
+ "Layer-0 pre-projection attention vs HF",
+ actual_attention,
+ self.intermediates["attention"],
+ )
+
+ def close(self):
+ for handle in self.hooks:
+ handle.remove()
+ self.hooks.clear()
diff --git a/examples/tensorrt_debug/_runtime_support.py b/examples/tensorrt_debug/_runtime_support.py
new file mode 100644
index 000000000..cb08341f2
--- /dev/null
+++ b/examples/tensorrt_debug/_runtime_support.py
@@ -0,0 +1,181 @@
+"""TensorRT setup, buffer allocation, validation, and cleanup.
+
+These mechanics are separate from the per-step debugging flow in runtime.py.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import ml_dtypes
+import numpy as np
+
+
+class RuntimeResources:
+ """Own the engine and GPU resources used by TensorRTRunner."""
+
+ def __init__(self, engine_path: Path, sdk: Path):
+ self.dll_handles = []
+ self.allocations = []
+ self.stream = None
+ self.context = self.engine = self.runtime = None
+ directories = [sdk / name for name in ("bin", "lib") if (sdk / name).exists()]
+ os.environ["PATH"] = (
+ os.pathsep.join(map(str, directories)) + os.pathsep + os.environ["PATH"]
+ )
+ try:
+ if os.name == "nt":
+ for directory in directories:
+ self.dll_handles.append(os.add_dll_directory(str(directory)))
+ self._load_engine(engine_path)
+ self.buffers = {}
+ self.stream = self.checked(self.cuda.cudaStreamCreate())
+ self._allocate_buffers()
+ except BaseException:
+ self.close()
+ raise
+
+ def _load_engine(self, engine_path):
+ # Load TensorRT only after the SDK DLL directories are available.
+ import tensorrt as trt
+ from cuda.bindings import runtime as cuda
+
+ self.trt, self.cuda = trt, cuda
+ self.logger = trt.Logger(trt.Logger.WARNING)
+ trt.init_libnvinfer_plugins(self.logger, "")
+ self.runtime = trt.Runtime(self.logger)
+ self.engine = self.runtime.deserialize_cuda_engine(engine_path.read_bytes())
+ if self.engine is None:
+ raise RuntimeError("Engine deserialization failed")
+ self.context = self.engine.create_execution_context()
+ if self.context is None:
+ raise RuntimeError("Execution context creation failed")
+ self.names = [
+ self.engine.get_tensor_name(index) for index in range(self.engine.num_io_tensors)
+ ]
+ self.inputs = [
+ name
+ for name in self.names
+ if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT
+ ]
+ self.caches = [
+ name for name in self.inputs if name.startswith(("key_cache.", "value_cache."))
+ ]
+ if not self.caches:
+ raise ValueError("Engine has no static KV cache inputs")
+ self.capacity = int(self.engine.get_tensor_shape(self.caches[0])[2])
+ self.max_chunk = int(self.engine.get_tensor_profile_shape("input_ids", 0)[2][1])
+
+ def checked(self, result):
+ status, *values = result
+ if status != self.cuda.cudaError_t.cudaSuccess:
+ raise RuntimeError(f"CUDA call failed: {status}")
+ return values[0] if len(values) == 1 else values
+
+ def numpy_dtype(self, name):
+ dtype = self.engine.get_tensor_dtype(name)
+ if dtype == self.trt.bfloat16:
+ return np.dtype(ml_dtypes.bfloat16)
+ return np.dtype(self.trt.nptype(dtype))
+
+ def validate_tensor(self, name, show_layout=False):
+ """Reject layouts that cannot be read as a contiguous NumPy array."""
+ shape = tuple(self.context.get_tensor_shape(name))
+ dtype = self.numpy_dtype(name)
+ strides = tuple(self.context.get_tensor_strides(name))
+ location = self.engine.get_tensor_location(name)
+ tensor_format = self.engine.get_tensor_format(name)
+ expected_strides = tuple(int(np.prod(shape[axis + 1 :])) for axis in range(len(shape)))
+ if show_layout:
+ print(
+ f"Cache layout {name}: shape={shape}; dtype={self.engine.get_tensor_dtype(name)}; "
+ f"location={location}; format={tensor_format}; strides={strides}",
+ flush=True,
+ )
+ if (
+ any(dim <= 0 for dim in shape)
+ or location != self.trt.TensorLocation.DEVICE
+ or tensor_format != self.trt.TensorFormat.LINEAR
+ or strides != expected_strides
+ ):
+ raise ValueError(f"{name}: runner requires resolved contiguous device I/O")
+ return shape, dtype
+
+ def _allocate_buffers(self):
+ """Allocate once for the profile's maximum chunk length."""
+ total_bytes = 0
+ for name in self.names:
+ if name.startswith(("updated_key_cache.", "updated_value_cache.")):
+ continue
+ shape = tuple(self.engine.get_tensor_shape(name))
+ if name in self.inputs:
+ shape = tuple(self.engine.get_tensor_profile_shape(name, 0)[2])
+ else:
+ shape = tuple(
+ 1 if dim < 0 and axis == 0 else self.max_chunk if dim < 0 else dim
+ for axis, dim in enumerate(shape)
+ )
+ if any(dim < 1 for dim in shape):
+ raise ValueError(f"Unresolved allocation shape for {name}: {shape}")
+ size = int(np.prod(shape)) * self.numpy_dtype(name).itemsize
+ pointer = self.checked(self.cuda.cudaMalloc(size))
+ self.allocations.append(pointer)
+ self.buffers[name] = (pointer, size)
+ self.checked(self.cuda.cudaMemset(pointer, 0, size))
+ total_bytes += size
+ # The tested TensorRT engines require each updated cache to alias its input.
+ for name in self.caches:
+ self.buffers["updated_" + name] = self.buffers[name]
+ print(f"Allocated I/O buffers: {total_bytes / 1024**2:.1f} MiB", flush=True)
+
+ def set_input_shapes(self, feeds):
+ for name in self.inputs:
+ shape = (
+ feeds[name].shape
+ if name in feeds
+ else tuple(1 if dim < 0 else dim for dim in self.engine.get_tensor_shape(name))
+ )
+ minimum, _, maximum = self.engine.get_tensor_profile_shape(name, 0)
+ if len(shape) != len(minimum) or any(
+ actual < low or actual > high
+ for actual, low, high in zip(shape, minimum, maximum)
+ ):
+ raise ValueError(
+ f"Input {name} shape {shape} outside profile {minimum}..{maximum}"
+ )
+ if not self.context.set_input_shape(name, shape):
+ raise RuntimeError(f"Cannot set input shape for {name}")
+
+ def bind_buffers(self):
+ for name in self.names:
+ shape, dtype = self.validate_tensor(name)
+ if int(np.prod(shape)) * dtype.itemsize > self.buffers[name][1]:
+ raise ValueError(f"Buffer too small for {name}")
+ if not self.context.set_tensor_address(name, int(self.buffers[name][0])):
+ raise RuntimeError(f"Cannot bind {name}")
+
+ def synchronize(self):
+ self.checked(self.cuda.cudaStreamSynchronize(self.stream))
+
+ def close(self):
+ if self.stream is not None:
+ self.cuda.cudaStreamSynchronize(self.stream)
+ self.context = None
+ self.engine = None
+ self.runtime = None
+ for pointer in self.allocations:
+ self.cuda.cudaFree(pointer)
+ self.allocations.clear()
+ if self.stream is not None:
+ self.cuda.cudaStreamDestroy(self.stream)
+ self.stream = None
+ for handle in self.dll_handles:
+ handle.close()
+ self.dll_handles.clear()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ self.close()
diff --git a/examples/tensorrt_debug/comparison.py b/examples/tensorrt_debug/comparison.py
new file mode 100644
index 000000000..2aa4b8b7d
--- /dev/null
+++ b/examples/tensorrt_debug/comparison.py
@@ -0,0 +1,96 @@
+"""Numerical reports for logits and intermediate tensors.
+
+Top-five output is only a display limit: error metrics compare the entire
+last-token logit vector. These reports are diagnostics, not pass/fail gates.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def compare_logits(label, actual, expected, tokenizer):
+ actual = np.asarray(actual, dtype=np.float32)
+ expected = np.asarray(expected, dtype=np.float32)
+ if actual.shape != expected.shape:
+ raise ValueError(f"{label}: logit shapes differ: {actual.shape} vs {expected.shape}")
+ if not np.isfinite(actual).all() or not np.isfinite(expected).all():
+ raise ValueError(f"{label}: non-finite logits")
+ difference = np.abs(actual - expected)
+ correlation = (
+ float(np.corrcoef(actual, expected)[0, 1])
+ if actual.std() > 0 and expected.std() > 0
+ else float("nan")
+ )
+ print(
+ f"\n{label}: mean_abs={difference.mean():.6f}; "
+ f"max_abs={difference.max():.6f}; correlation={correlation:.6f}"
+ )
+ print(f" Top-1 match: {actual.argmax() == expected.argmax()}")
+ for source, logits in (("TRT", actual), ("HF", expected)):
+ top = np.argsort(logits)[-5:][::-1]
+ predictions = ", ".join(
+ f"{int(token)} {tokenizer.decode([int(token)])!r} ({logits[token]:.4f})"
+ for token in top
+ )
+ print(f" {source} top-5: {predictions}", flush=True)
+
+
+def compare_tensor(label, actual, expected):
+ actual = np.asarray(actual, dtype=np.float32)
+ expected = np.asarray(expected, dtype=np.float32)
+ if actual.shape != expected.shape:
+ raise ValueError(f"{label}: shapes differ: {actual.shape} vs {expected.shape}")
+ if not np.isfinite(actual).all() or not np.isfinite(expected).all():
+ raise ValueError(f"{label}: non-finite values")
+ error = actual - expected
+ relative_rms = np.linalg.norm(error) / max(float(np.linalg.norm(expected)), 1e-12)
+ print(
+ f" {label}: mean_abs={np.abs(error).mean():.6f}; "
+ f"max_abs={np.abs(error).max():.6f}; relative_rms={relative_rms:.6f}",
+ flush=True,
+ )
+
+
+class HuggingFaceReference:
+ """Run identical inputs on CPU while keeping an independent HF cache."""
+
+ def __init__(self, model_id, revision=None):
+ import torch
+ from transformers import AutoModelForCausalLM
+
+ self.torch = torch
+ self.cache = None
+ print("Loading HuggingFace FP32/eager reference on CPU...", flush=True)
+ self.model = (
+ AutoModelForCausalLM.from_pretrained(
+ model_id,
+ revision=revision,
+ dtype=torch.float32,
+ attn_implementation="eager",
+ )
+ .cpu()
+ .eval()
+ )
+ print(f"Reference revision: {getattr(self.model.config, '_commit_hash', None)}")
+ print("Ensure this checkpoint matches the export; engine provenance is not verified.")
+ print(
+ "Both runtimes receive TRT-selected tokens; metrics are diagnostic, not a parity gate."
+ )
+
+ def reset_cache(self):
+ self.cache = None
+
+ def forward(self, feeds):
+ torch = self.torch
+ valid_length = int(feeds["nonpad_kv_seqlen"][0])
+ with torch.inference_mode():
+ output = self.model(
+ input_ids=torch.from_numpy(feeds["input_ids"].copy()),
+ position_ids=torch.from_numpy(feeds["position_ids"].copy()),
+ attention_mask=torch.ones((1, valid_length), dtype=torch.long),
+ past_key_values=self.cache,
+ use_cache=True,
+ )
+ self.cache = output.past_key_values
+ return output.logits[0, -1].float().cpu().numpy()
diff --git a/examples/tensorrt_debug/full_prefix.py b/examples/tensorrt_debug/full_prefix.py
new file mode 100644
index 000000000..447fd7562
--- /dev/null
+++ b/examples/tensorrt_debug/full_prefix.py
@@ -0,0 +1,26 @@
+"""Full-prefix control: resend all tokens instead of relying on cached history.
+
+The caller resets the TRT and HF caches and uses position zero for this input.
+This module only constructs and validates the complete token history.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def build_full_prefix(prompt_ids: np.ndarray, generated: list[int]) -> np.ndarray:
+ """Join [1, prompt_length] and [1, generated_length] along the token axis."""
+ generated_ids = np.asarray([generated], dtype=np.int64)
+ return np.concatenate([prompt_ids, generated_ids], axis=1)
+
+
+def validate_full_prefix_budget(
+ prompt_length: int, max_new_tokens: int, max_chunk: int
+) -> None:
+ """The final prediction consumes all preceding tokens, not itself."""
+ if prompt_length + max_new_tokens - 1 > max_chunk:
+ raise ValueError(
+ f"Full-prefix generation exceeds the engine's maximum input length "
+ f"of {max_chunk}; reduce --max-new-tokens or shorten the prompt"
+ )
diff --git a/examples/tensorrt_debug/inspect_attention.py b/examples/tensorrt_debug/inspect_attention.py
new file mode 100644
index 000000000..42e7380e1
--- /dev/null
+++ b/examples/tensorrt_debug/inspect_attention.py
@@ -0,0 +1,52 @@
+"""Keep TRT's Q/K/V fixed; find which mask reproduces its attention output.
+
+For the unpadded, batch-one Qwen3 diagnostic engine only.
+HF hooks and RoPE alignment live in _attention_support.py.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from tensorrt_debug._attention_support import AttentionReference
+from tensorrt_debug.comparison import compare_tensor
+from tensorrt_debug.inspect_cache import snapshot_cache
+
+
+class AttentionInspector(AttentionReference):
+ """Compare after HF forward; hook setup and cleanup are inherited."""
+
+ def compare(self, feeds, position):
+ actual_query = self.runner.read_tensor("probe.query").astype(np.float32)
+ actual_attention = self.runner.read_tensor("probe.attention").astype(np.float32)
+ self.compare_with_hf(feeds, actual_query, actual_attention)
+ self._compare_masks(actual_query, actual_attention, position)
+
+ def _compare_masks(self, actual_query, actual_attention, position):
+ key = torch.from_numpy(snapshot_cache(self.runner, "key_cache.0").astype(np.float32))
+ value = torch.from_numpy(
+ snapshot_cache(self.runner, "value_cache.0").astype(np.float32)
+ )
+ # GQA: repeat each KV head for its group of query heads.
+ groups = actual_query.shape[1] // key.shape[1]
+ key = key.repeat_interleave(groups, dim=1)
+ value = value.repeat_interleave(groups, dim=1)
+ scores = torch.from_numpy(actual_query) @ key.transpose(-1, -2)
+ scores *= self.attention.scaling
+ length = actual_query.shape[2]
+ key_positions = torch.arange(key.shape[2])[None, :]
+ query_positions = torch.arange(length)[:, None]
+ # At cached decode position 22, these hypotheses allow slots 0..22 vs only 0.
+ for hypothesis, offset in (
+ ("correct cached causal mask", position),
+ ("top-left causal mask", 0),
+ ):
+ mask = key_positions <= query_positions + offset
+ masked_scores = scores.masked_fill(~mask, -torch.inf)
+ probabilities = torch.softmax(masked_scores, dim=-1)
+ reconstructed = probabilities @ value
+ expected = reconstructed.transpose(1, 2).reshape(1, length, -1).numpy()
+ compare_tensor(
+ f"TRT attention vs own QKV + {hypothesis}", actual_attention, expected
+ )
diff --git a/examples/tensorrt_debug/inspect_cache.py b/examples/tensorrt_debug/inspect_cache.py
new file mode 100644
index 000000000..635fadb35
--- /dev/null
+++ b/examples/tensorrt_debug/inspect_cache.py
@@ -0,0 +1,64 @@
+"""Read-only cache probe: check writes, preserved history, and agreement with HF."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def snapshot_cache(runner, name, show_layout=False):
+ """Validate paired cache I/O and return an independent CPU copy."""
+ shape, dtype = runner.validate_tensor(name, show_layout)
+ updated_shape, updated_dtype = runner.validate_tensor("updated_" + name, show_layout)
+ if len(shape) != 4 or shape != updated_shape or dtype != updated_dtype:
+ raise ValueError(f"{name}: cache probe requires matching [B,H,S,D] device I/O")
+ return runner.read_tensor(name)
+
+
+def report_cache(name, before, after, expected, position, length):
+ """Report changed storage slots separately from numerical agreement with HF."""
+ end = position + length
+ valid = after[:, :, :end, :].astype(np.float32)
+ if valid.shape != expected.shape:
+ raise ValueError(f"{name}: TRT/HF cache shapes differ: {valid.shape}/{expected.shape}")
+ # [B,H,S,D,bytes_per_element]: reduce every axis except the sequence slots.
+ raw_before = before.view(np.uint8).reshape(*before.shape, before.dtype.itemsize)
+ raw_after = after.view(np.uint8).reshape(*after.shape, after.dtype.itemsize)
+ changed = np.any(raw_before != raw_after, axis=(0, 1, 3, 4))
+ slots = np.flatnonzero(changed)
+ print(
+ f"\nCache {name}: write=[{position}:{end}], "
+ f"changed_slots={slots[:32].tolist()} (total={slots.size})"
+ )
+ print(
+ f" Old prefix bitwise unchanged: {not changed[:position].any()}; "
+ f"unused tail bitwise unchanged: {not changed[end:].any()}"
+ )
+ for region, actual, reference_values in (
+ ("valid prefix", valid, expected),
+ ("new slots", valid[:, :, position:end, :], expected[:, :, position:end, :]),
+ ):
+ if not np.isfinite(actual).all() or not np.isfinite(reference_values).all():
+ print(f" {region} vs HF: non-finite values; numerical comparison invalid")
+ continue
+ error = np.abs(actual - reference_values)
+ print(f" {region} vs HF: mean_abs={error.mean():.6f}; max_abs={error.max():.6f}")
+
+
+class CacheInspector:
+ """Inspect one layer's K and V; the generation loop controls snapshot timing."""
+
+ def __init__(self, runner, layer):
+ self.layer = layer
+ self.names = [f"{kind}_cache.{layer}" for kind in ("key", "value")]
+ if any(name not in runner.caches for name in self.names):
+ raise ValueError(f"Engine does not contain cache layer {layer}")
+ self.runner = runner
+
+ def snapshot(self, show_layout=False):
+ return {name: snapshot_cache(self.runner, name, show_layout) for name in self.names}
+
+ def compare(self, before, after, reference_cache, position, length):
+ reference_layer = reference_cache.layers[self.layer]
+ for name, tensor in zip(self.names, (reference_layer.keys, reference_layer.values)):
+ expected = tensor.detach().float().cpu().numpy().copy()
+ report_cache(name, before[name], after[name], expected, position, length)
diff --git a/examples/tensorrt_debug/runtime.py b/examples/tensorrt_debug/runtime.py
new file mode 100644
index 000000000..5b50f95f0
--- /dev/null
+++ b/examples/tensorrt_debug/runtime.py
@@ -0,0 +1,71 @@
+"""Per-step GPU flow: prepare inputs -> execute -> read outputs.
+
+Setup, allocation, safety checks, and cleanup live in _runtime_support.py.
+Cache input/output buffers share GPU addresses: execution updates them in place.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from tensorrt_debug._runtime_support import RuntimeResources
+
+
+class TensorRTRunner(RuntimeResources):
+ """The inference steps; resource management is inherited from RuntimeResources."""
+
+ def prepare_inputs(self, feeds):
+ """Set shapes, bind GPU addresses, then upload feeds. Leave caches untouched."""
+ self.set_input_shapes(feeds)
+ self.bind_buffers()
+ for name, array in feeds.items():
+ array = np.ascontiguousarray(array, dtype=self.numpy_dtype(name))
+ self.checked(
+ self.cuda.cudaMemcpy(
+ self.buffers[name][0],
+ array.ctypes.data,
+ array.nbytes,
+ self.cuda.cudaMemcpyKind.cudaMemcpyHostToDevice,
+ )
+ )
+
+ def execute(self):
+ """Compute next-token scores. The caller selects a token with argmax()."""
+ # Run the compiled model on the GPU: update KV caches and compute logits.
+ inference_started = self.context.execute_async_v3(int(self.stream))
+ if not inference_started:
+ raise RuntimeError("TensorRT execution failed")
+
+ # The launch is asynchronous; wait before reading the GPU results.
+ self.synchronize()
+ logits_on_cpu = self.read_tensor("logits") # [batch, input_length, vocabulary]
+
+ # Batch 0, last input position: one score per possible next-token ID.
+ next_token_scores = logits_on_cpu[0, -1, :].astype(np.float32)
+ if not np.isfinite(next_token_scores).all():
+ raise RuntimeError("Non-finite logits")
+ return next_token_scores
+
+ def read_tensor(self, name):
+ """Return an independent CPU copy, preserving the original dtype/bytes."""
+ self.synchronize()
+ shape, dtype = self.validate_tensor(name)
+ snapshot = np.empty(shape, dtype=dtype)
+ if snapshot.nbytes > self.buffers[name][1]:
+ raise ValueError(f"{name}: snapshot exceeds allocated buffer")
+ self.checked(
+ self.cuda.cudaMemcpy(
+ snapshot.ctypes.data,
+ self.buffers[name][0],
+ snapshot.nbytes,
+ self.cuda.cudaMemcpyKind.cudaMemcpyDeviceToHost,
+ )
+ )
+ return snapshot
+
+ def reset_caches(self):
+ """Full-prefix control: erase cached history before recomputing all tokens."""
+ for name in self.caches:
+ pointer, size = self.buffers[name]
+ self.checked(self.cuda.cudaMemsetAsync(pointer, 0, size, self.stream))
+ self.synchronize()
diff --git a/examples/tensorrt_static_cache_generation.py b/examples/tensorrt_static_cache_generation.py
new file mode 100644
index 000000000..3d90b1e4c
--- /dev/null
+++ b/examples/tensorrt_static_cache_generation.py
@@ -0,0 +1,205 @@
+"""Run batch-one greedy generation with a TensorRT heads-first static-cache engine.
+
+Read generate() first for the token loop. Independent diagnostic modules live
+in tensorrt_debug/: full_prefix, inspect_cache, inspect_attention, comparison.
+CUDA and TensorRT details live in tensorrt_debug/runtime.py.
+"""
+
+from __future__ import annotations
+
+import argparse
+from contextlib import ExitStack
+from pathlib import Path
+
+import numpy as np
+from tensorrt_debug.comparison import HuggingFaceReference, compare_logits
+from tensorrt_debug.full_prefix import build_full_prefix, validate_full_prefix_budget
+from tensorrt_debug.inspect_cache import CacheInspector
+from tensorrt_debug.runtime import TensorRTRunner
+from transformers import AutoTokenizer
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--engine", type=Path, required=True)
+ parser.add_argument("--sdk", type=Path, default=Path("C:/TensorRT-11.3.0.99"))
+ parser.add_argument("--model", default="Qwen/Qwen3-0.6B")
+ parser.add_argument("--prompt", default="What is the capital of France? Answer briefly.")
+ parser.add_argument("--max-new-tokens", type=int, default=32)
+ parser.add_argument(
+ "--compare-hf",
+ action="store_true",
+ help="Compare identical-token steps against HuggingFace on CPU",
+ )
+ parser.add_argument(
+ "--compare-steps",
+ type=int,
+ default=2,
+ help="Number of steps to compare, including prefill (default: 2)",
+ )
+ parser.add_argument(
+ "--revision",
+ default=None,
+ help="HF checkpoint revision; must match the engine's source weights",
+ )
+ parser.add_argument(
+ "--full-prefix",
+ action="store_true",
+ help="Recompute the entire prefix with zeroed caches at every step",
+ )
+
+ parser.add_argument(
+ "--inspect-cache",
+ action="store_true",
+ help="Snapshot cache before/after prefill and first decode; requires --compare-hf",
+ )
+ parser.add_argument(
+ "--cache-layer", type=int, default=0, help="Layer to inspect (default: 0)"
+ )
+ parser.add_argument(
+ "--inspect-attention",
+ action="store_true",
+ help="Compare layer-0 intermediates from a diagnostic engine with HF",
+ )
+
+ args = parser.parse_args()
+ if args.max_new_tokens < 1:
+ parser.error("--max-new-tokens must be positive")
+ 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")
+ if args.inspect_cache and (not args.compare_hf or args.full_prefix):
+ parser.error(
+ "--inspect-cache requires --compare-hf and cached mode (no --full-prefix)"
+ )
+ if args.cache_layer < 0:
+ parser.error("--cache-layer must be nonnegative")
+ if args.inspect_attention and not args.compare_hf:
+ parser.error("--inspect-attention requires --compare-hf")
+
+ return args
+
+
+def run(args, runner):
+ tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision)
+ rendered = tokenizer.apply_chat_template(
+ [{"role": "user", "content": args.prompt}],
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=False,
+ )
+ current_ids = np.asarray(
+ [tokenizer.encode(rendered, add_special_tokens=False)], dtype=np.int64
+ )
+ prompt_ids = current_ids.copy()
+ if args.inspect_attention and not {"probe.query", "probe.attention"}.issubset(
+ runner.names
+ ):
+ raise ValueError("Use an engine built by tensorrt_attention_probe.py")
+ cache_inspector = CacheInspector(runner, args.cache_layer) if args.inspect_cache else None
+ capacity = runner.capacity
+ if current_ids.shape[1] + args.max_new_tokens > capacity:
+ raise ValueError("Prompt and generation budget exceed cache capacity")
+
+ if args.full_prefix:
+ validate_full_prefix_budget(prompt_ids.shape[1], args.max_new_tokens, runner.max_chunk)
+ reference = HuggingFaceReference(args.model, args.revision) if args.compare_hf else None
+ if args.inspect_cache or args.inspect_attention:
+ if reference.model.config.model_type != "qwen3":
+ raise ValueError("Cache/RoPE comparison currently verified only for Qwen3")
+
+ with ExitStack() as resources:
+ attention_inspector = None
+ if args.inspect_attention:
+ from tensorrt_debug.inspect_attention import AttentionInspector
+
+ attention_inspector = AttentionInspector(runner, reference.model)
+ resources.callback(attention_inspector.close)
+ print(f"Prompt tokens: {prompt_ids.shape[1]}; cache capacity: {capacity}", flush=True)
+ generate(
+ args,
+ runner,
+ tokenizer,
+ prompt_ids,
+ reference,
+ cache_inspector,
+ attention_inspector,
+ )
+
+
+def generate(
+ args, runner, tokenizer, prompt_ids, reference, cache_inspector, attention_inspector
+):
+ """Prepare -> snapshot -> execute -> snapshot -> compare -> choose next token."""
+ current_ids = prompt_ids.copy()
+ generated = []
+ position = 0
+ compared_steps = 0
+
+ for step in range(args.max_new_tokens):
+ if args.full_prefix:
+ current_ids = build_full_prefix(prompt_ids, generated)
+ position = 0
+ runner.reset_caches()
+ if reference is not None:
+ reference.reset_cache()
+
+ phase = "full-prefix" if args.full_prefix else "prefill" if step == 0 else "decode"
+ length = current_ids.shape[1]
+ feeds = {
+ "input_ids": current_ids,
+ "position_ids": np.arange(position, position + length, dtype=np.int64)[None, :],
+ "write_indices": np.asarray([position], dtype=np.int64),
+ "nonpad_kv_seqlen": np.asarray([position + length], dtype=np.int64),
+ }
+ runner.prepare_inputs(feeds)
+
+ inspect_this_step = cache_inspector is not None and step < 2
+ if inspect_this_step:
+ before = cache_inspector.snapshot(show_layout=step == 0)
+
+ last_logits = runner.execute()
+
+ if inspect_this_step:
+ after = cache_inspector.snapshot()
+
+ if reference is not None and step < args.compare_steps:
+ reference_logits = reference.forward(feeds)
+ label = f"{phase} step {step} (position {position}, input length {length})"
+ compare_logits(label, last_logits, reference_logits, tokenizer)
+ compared_steps += 1
+ if attention_inspector is not None:
+ attention_inspector.compare(feeds, position)
+ if inspect_this_step:
+ cache_inspector.compare(before, after, reference.cache, position, length)
+
+ token = int(last_logits.argmax())
+ generated.append(token)
+ print(
+ f"Step {step}: {phase}; token={token}; text={tokenizer.decode([token])!r}",
+ flush=True,
+ )
+ position += length
+ if token == tokenizer.eos_token_id:
+ break
+ current_ids = np.asarray([[token]], dtype=np.int64)
+ print("\nPrompt:", args.prompt)
+ print("Response:", tokenizer.decode(generated, skip_special_tokens=True))
+ if args.compare_hf:
+ print(
+ f"Compared {compared_steps}/{args.compare_steps} identical-token steps "
+ "against FP32 HF. BF16 engine outputs need not match exactly."
+ )
+ if compared_steps < args.compare_steps:
+ print("Comparison truncated by EOS; requested decode coverage is incomplete.")
+ else:
+ print("Real inference completed; numerical reference parity has not been checked.")
+
+
+def main() -> None:
+ args = parse_args()
+ with TensorRTRunner(args.engine, args.sdk) as runner:
+ run(args, runner)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/mobius/_execution_providers.py b/src/mobius/_execution_providers.py
index a878e8547..ea8648c42 100644
--- a/src/mobius/_execution_providers.py
+++ b/src/mobius/_execution_providers.py
@@ -31,6 +31,7 @@
import dataclasses
import logging
+from typing import Literal
import onnx_ir as ir
@@ -77,6 +78,10 @@ class EpCapabilities:
DecomposeAttention. ``True`` leaves the fused op unchanged. Set
``False`` only for runtimes without an ``Attention`` kernel (QNN
HTP), where the fused op would otherwise be forced onto CPU.
+ supports_attention_nonpad_kv_seqlen: Whether native Attention consumes
+ valid static-cache lengths. When ``False``, static-cache exports
+ require an explicit causal/valid-length bias and omit Attention's
+ nonpad input. Standalone TensorRT 11.3 ignores that native input.
supports_rotary_embedding: ``False`` decomposes the opset-24
``RotaryEmbedding`` op into rotate-half primitives (Reshape/Slice/
Mul/Sub/Add/Concat) via DecomposeRotaryEmbedding. ``True`` leaves
@@ -147,6 +152,7 @@ class EpCapabilities:
"""
name: str
+ static_cache_layout: Literal["flattened", "heads_first"] = "flattened"
gqa_dtypes: frozenset[ir.DataType] = dataclasses.field(default_factory=frozenset)
qkv_pack_dtypes: frozenset[ir.DataType] = dataclasses.field(default_factory=frozenset)
supports_fused_rope: bool = True
@@ -160,6 +166,7 @@ class EpCapabilities:
supports_tensor_scatter: bool = True
supports_range: bool = True
supports_fp8_kv_cache: bool = False
+ supports_attention_nonpad_kv_seqlen: bool = True
default_int4_accuracy_level: int = 0
provider_options: dict[str, str] = dataclasses.field(default_factory=dict)
enable_graph_capture: bool = False
@@ -366,6 +373,15 @@ def _register_builtins() -> None:
enable_graph_capture=True,
supports_past_present_share_buffer=True,
),
+ 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,
+ ),
# Qualcomm Hexagon NPU via the QNN EP (onnxruntime-qnn QAIRT plugin),
# HTP backend. The HTP runs a static-shaped, QDQ-quantized QNN context
# binary with no kernels for ORT contrib fused ops, so everything is
diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py
index 2fdd1ee40..ed40a2176 100644
--- a/src/mobius/components/_attention.py
+++ b/src/mobius/components/_attention.py
@@ -9,6 +9,7 @@
import onnx_ir as ir
from onnxscript import OpBuilder, nn
+from mobius._build_context import ep_capabilities
from mobius._configs import ArchitectureConfig
from mobius.components._common import Linear
from mobius.components._rms_norm import OffsetRMSNorm, RMSNorm
@@ -73,6 +74,18 @@ class StaticCacheState(NamedTuple):
write_indices: ir.Value
nonpad_kv_seqlen: ir.Value
+ @property
+ def sequence_axis(self) -> int:
+ shape = self.key_cache.shape
+ if shape is None:
+ raise ValueError("Key cache shape is not defined")
+ if len(shape) == 3:
+ return 1
+ elif len(shape) == 4:
+ return 2
+ else:
+ raise ValueError(f"Static cache must have rank 3 or 4, got rank {len(shape)}.")
+
def _apply_attention(
op: OpBuilder,
@@ -143,17 +156,31 @@ def _apply_attention(
# for all t in range(seq_len). This handles both prefill
# (write_indices=0, seq_len=N) and decode (write_indices=N,
# seq_len=1) with the same graph.
+
+ sequence_axis = static_cache.sequence_axis
+ heads_first = sequence_axis == 2
+
+ if heads_first:
+ query = op.Reshape(query, [0, 0, num_attention_heads, -1])
+ query = op.Transpose(query, perm=[0, 2, 1, 3])
+
+ key = op.Reshape(key, [0, 0, num_key_value_heads, -1])
+ key = op.Transpose(key, perm=[0, 2, 1, 3])
+
+ value = op.Reshape(value, [0, 0, num_key_value_heads, -1])
+ value = op.Transpose(value, perm=[0, 2, 1, 3])
+
updated_k = op.TensorScatter(
static_cache.key_cache,
key,
static_cache.write_indices,
- axis=1,
+ axis=sequence_axis,
) # [B, max_seq, kv_hidden]
updated_v = op.TensorScatter(
static_cache.value_cache,
value,
static_cache.write_indices,
- axis=1,
+ axis=sequence_axis,
) # [B, max_seq, kv_hidden]
# External-cache masking. Two modes, selected by whether the caller
@@ -170,8 +197,8 @@ def _apply_attention(
# bidirectional unmasking encoded in the bias. This routes ORT to
# the MEA external-cache path (Flash is precluded by any bias).
#
- # nonpad_kv_seqlen stays as input #6 in BOTH modes: it bounds the valid
- # KV prefix and, on the CUDA Flash path, drives the fully-masked-row
+ # On supporting EPs, nonpad_kv_seqlen stays as input #6 in BOTH modes.
+ # It bounds the valid KV prefix and, on CUDA Flash, drives the fully-masked-row
# zero guard (LaunchZeroFullyMaskedRows). In bias mode the additive
# bias already encodes the same ``slot < nonpad`` validity. The
# cross-repo invariant is ``nonpad == write_indices + valid_token_count``
@@ -190,21 +217,41 @@ def _apply_attention(
mask_arg, causal = attn_mask, 0
else:
mask_arg, causal = None, 1
- attn_output, _, _ = op.Attention(
+
+ supports_nonpad = ep_capabilities().supports_attention_nonpad_kv_seqlen
+ if not supports_nonpad and attn_mask is None:
+ raise ValueError(
+ "This execution provider requires an explicit static-cache attention bias"
+ )
+ # Unsupported EPs use the bias's slot-validity test instead of input #6.
+ nonpad_input = static_cache.nonpad_kv_seqlen if supports_nonpad else None
+
+ head_attrs: dict[str, int] = {}
+ if not heads_first:
+ head_attrs = {
+ "q_num_heads": num_attention_heads,
+ "kv_num_heads": num_key_value_heads,
+ }
+
+ attn_output = op.Attention(
query,
updated_k,
updated_v,
mask_arg,
- None, # no past_key (full cache is already provided)
- None, # no past_value
- static_cache.nonpad_kv_seqlen,
- q_num_heads=num_attention_heads,
- kv_num_heads=num_key_value_heads,
+ None,
+ None,
+ nonpad_input,
scale=scale,
softcap=softcap,
is_causal=causal,
- _outputs=3,
+ _outputs=1,
+ **head_attrs,
)
+
+ if heads_first:
+ attn_output = op.Transpose(attn_output, perm=[0, 2, 1, 3])
+ attn_output = op.Reshape(attn_output, [0, 0, -1])
+
return attn_output, updated_k, updated_v
# Dynamic cache mode: standard Attention with past KV concatenation.
diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py
index 989cfd698..3b6520abf 100644
--- a/src/mobius/integrations/onnx_genai/workflow_metadata.py
+++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py
@@ -19,8 +19,6 @@
from mobius._constants import (
STATIC_CACHE_KV_SEQUENCE_LENGTH,
- STATIC_CACHE_LAYOUT,
- STATIC_CACHE_SEQUENCE_AXIS,
STATIC_CACHE_WRITE_INDICES,
)
from mobius.generation import (
@@ -2377,33 +2375,46 @@ def _static_cache_ports(model: ir.Model) -> dict[str, Any] | None:
"no paired cache buffer to scatter into; regenerate the package with "
"updated_ outputs for every static cache input"
)
- axes = {
- node.attributes.get_int("axis", 0)
+ scatters = [
+ node
for node in ir.traversal.RecursiveGraphIterator(model.graph)
if node.op_type == "TensorScatter"
- }
- if axes - {STATIC_CACHE_SEQUENCE_AXIS}:
- raise ValueError(
- f"static cache buffers are addressed on axes {sorted(axes)}, but the mobius "
- f"static-cache ABI scatters along axis {STATIC_CACHE_SEQUENCE_AXIS}; the "
- "declared capacity axis and the graph disagree"
- )
+ ]
capacities = set()
+ layouts = set()
for buffer in buffers.values():
shape = list(buffer.shape or [])
- if len(shape) <= STATIC_CACHE_SEQUENCE_AXIS:
+ geometry = {3: (1, "bsh"), 4: (2, "bnsh")}.get(len(shape))
+ if geometry is None:
+ raise ValueError(
+ f"static cache buffer {buffer.name!r} has unsupported rank {len(shape)}; "
+ "expected flattened [B, S, H*D] or heads-first [B, H, S, D]"
+ )
+ sequence_axis, layout = geometry
+ layouts.add(geometry)
+ axes = {
+ node.attributes.get_int("axis", 0)
+ for node in scatters
+ if node.inputs and node.inputs[0] is buffer
+ }
+ if axes - {sequence_axis, sequence_axis - len(shape)}:
raise ValueError(
- f"static cache buffer {buffer.name!r} has rank {len(shape)}, which cannot "
- f"carry a capacity on axis {STATIC_CACHE_SEQUENCE_AXIS}"
+ f"static cache buffer {buffer.name!r} is addressed on axes {sorted(axes)}, "
+ f"but its declared capacity axis is {sequence_axis} for layout {layout}"
)
- capacity = _constant_extent(shape[STATIC_CACHE_SEQUENCE_AXIS])
+ capacity = _constant_extent(shape[sequence_axis])
if capacity is None:
raise ValueError(
f"static cache buffer {buffer.name!r} declares a symbolic extent "
- f"{shape[STATIC_CACHE_SEQUENCE_AXIS]!r} on its capacity axis; an "
+ f"{shape[sequence_axis]!r} on its capacity axis; an "
"indexed scatter is only meaningful against one constant capacity"
)
capacities.add(capacity)
+ if len(layouts) != 1:
+ raise ValueError(
+ "static cache buffers declare conflicting layouts; "
+ "one indexed-scatter group requires a consistent sequence axis and layout"
+ )
if len(capacities) != 1:
raise ValueError(
f"static cache buffers declare conflicting capacities {sorted(capacities)}; "
@@ -2414,6 +2425,8 @@ def _static_cache_ports(model: ir.Model) -> dict[str, Any] | None:
"kv_sequence_length": STATIC_CACHE_KV_SEQUENCE_LENGTH,
"buffers": buffers,
"capacity": capacities.pop(),
+ "sequence_axis": sequence_axis,
+ "layout": layout,
}
@@ -2737,8 +2750,10 @@ def _state_service_groups(
name = names[(kind, update)]
group: dict[str, Any] = {
"kind": kind,
- "sequence_axis": (STATIC_CACHE_SEQUENCE_AXIS if is_scattered else sequence_axis),
- "layout": STATIC_CACHE_LAYOUT if is_scattered else "bnsh",
+ "sequence_axis": (
+ indexed_scatter["sequence_axis"] if is_scattered else sequence_axis
+ ),
+ "layout": indexed_scatter["layout"] if is_scattered else "bnsh",
}
group_lengths = indexed_scatter["logical_lengths"] if is_scattered else logical_lengths
if group_lengths:
@@ -6941,6 +6956,8 @@ def build_vlm_workflow_metadata(
indexed_scatter=(
{
"buffers": static_cache["buffers"],
+ "sequence_axis": static_cache["sequence_axis"],
+ "layout": static_cache["layout"],
"capacity": "package.cache_capacity",
# The write cursor and the logical length are one quantity: a
# row's next write lands exactly where its valid prefix ends.
@@ -9151,6 +9168,8 @@ def _build_autoregressive_workflow_metadata(
indexed_scatter=(
{
"buffers": static_cache["buffers"],
+ "sequence_axis": static_cache["sequence_axis"],
+ "layout": static_cache["layout"],
"capacity": "package.cache_capacity",
# The write cursor and the logical length are the same quantity:
# a row's next write lands exactly where its valid prefix ends.
diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py
index 08fe71233..676093ef9 100644
--- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py
+++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py
@@ -1037,6 +1037,7 @@ def _static_cache_model(
scatter_axis: int | None = None,
paired: bool = True,
control_ports: bool = True,
+ heads_first: bool = False,
) -> ir.Model:
"""A minimal graph shaped like a mobius static-cache decoder export."""
capacities = capacities or [32, 32]
@@ -1055,9 +1056,14 @@ def _static_cache_model(
(f"updated_key_cache.{layer}", ir.DataType.FLOAT, ["batch", capacity, 16])
)
model = _model("decoder", inputs, outputs)
+ if heads_first:
+ for value in (*model.graph.inputs, *model.graph.outputs):
+ if "key_cache" in value.name:
+ value.shape = ir.Shape(["batch", 2, value.shape[1], 8])
if scatter_axis is not None:
cache = model.graph.inputs[1]
scattered = _value("scattered", ir.DataType.FLOAT, ["batch", capacities[0], 16])
+ scattered.shape = cache.shape
model.graph.append(
ir.Node(
"",
@@ -1109,3 +1115,32 @@ def test_rejects_a_scatter_that_disagrees_with_the_declared_axis(self):
def test_accepts_a_scatter_on_the_declared_axis(self):
assert _static_cache_ports(_static_cache_model(scatter_axis=1))["capacity"] == 32
+
+ @pytest.mark.parametrize("axis", [2, -2])
+ def test_discovers_heads_first_geometry(self, axis):
+ ports = _static_cache_ports(_static_cache_model(heads_first=True, scatter_axis=axis))
+ assert ports["capacity"] == 32
+ assert ports["sequence_axis"] == 2
+ assert ports["layout"] == "bnsh"
+
+ def test_rejects_heads_first_scatter_on_head_axis(self):
+ with pytest.raises(ValueError, match="declared capacity axis"):
+ _static_cache_ports(_static_cache_model(heads_first=True, scatter_axis=1))
+
+ def test_rejects_heads_first_symbolic_capacity(self):
+ model = _static_cache_model(heads_first=True, scatter_axis=2)
+ model.graph.inputs[1].shape = ir.Shape(["batch", 2, "capacity", 8])
+ with pytest.raises(ValueError, match="symbolic extent"):
+ _static_cache_ports(model)
+
+ def test_rejects_mixed_static_layouts(self):
+ model = _static_cache_model()
+ model.graph.inputs[1].shape = ir.Shape(["batch", 2, 32, 8])
+ with pytest.raises(ValueError, match="conflicting layouts"):
+ _static_cache_ports(model)
+
+ def test_rejects_unsupported_cache_rank(self):
+ model = _static_cache_model()
+ model.graph.inputs[1].shape = ir.Shape(["batch", 32])
+ with pytest.raises(ValueError, match="unsupported rank"):
+ _static_cache_ports(model)
diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py
index 86873c2d7..cfd9daae9 100644
--- a/src/mobius/models/base.py
+++ b/src/mobius/models/base.py
@@ -121,8 +121,8 @@ def __init__(self, config: ArchitectureConfig, mlp_class: type | None = None):
# Sliding-window models declare a local-attention span; it drives the
# optional static-cache float bias (flags.static_cache_bias). Standard
- # full-attention models leave this None, so the bias path is a no-op
- # for them even when the flag is set.
+ # full-attention models leave this None, but still need a causal bias
+ # when the EP cannot consume native static-cache valid lengths.
self._sliding_window: int | None = getattr(config, "sliding_window", None)
def _maybe_static_cache_bias(
@@ -133,10 +133,9 @@ def _maybe_static_cache_bias(
) -> ir.Value | None:
"""Optionally build the static-cache float additive attention bias.
- Returns ``None`` (maskless ``is_causal=1`` default) unless ALL hold:
- * ``flags.static_cache_bias`` is set, AND
- * the model declares a bias need (``self._sliding_window`` is set), AND
- * the cache is the opset-24 external cache (``StaticCacheState``).
+ Requires an external ``StaticCacheState`` and either an EP without
+ native Attention valid-length support, or ``flags.static_cache_bias``
+ together with a declared sliding window. Otherwise returns ``None``.
When emitted, the bias is a ``(B, 1, S_q, max_seq_len)`` additive mask
keyed on absolute query positions with KV validity
@@ -151,7 +150,9 @@ def _maybe_static_cache_bias(
this instead of ``input_ids`` keeps the bias enabled for
``inputs_embeds``-driven forwards (where ``input_ids`` is None).
"""
- if not flags.static_cache_bias or self._sliding_window is None:
+ requires_explicit_bias = not ep_capabilities().supports_attention_nonpad_kv_seqlen
+ requested_sliding_bias = flags.static_cache_bias and self._sliding_window is not None
+ if not (requires_explicit_bias or requested_sliding_bias):
return None
if not past_key_values:
return None
@@ -162,12 +163,15 @@ def _maybe_static_cache_bias(
# Static cache KV axis width is a concrete int: [B, max_seq_len, kv_hidden].
# Guard against a symbolic dim, which would otherwise raise an opaque
# TypeError downstream. Static-cache always allocates a fixed width today.
- max_seq_len = first.key_cache.shape[1]
+ sequence_axis = first.sequence_axis
+ cache_shape = first.key_cache.shape
+ assert cache_shape is not None
+ max_seq_len = cache_shape[sequence_axis]
if not isinstance(max_seq_len, int):
raise TypeError(
"static-cache bias requires a concrete key_cache KV dimension "
- f"(axis 1), but got symbolic dim {max_seq_len!r}. The static "
- "cache must be allocated with a fixed max_seq_len."
+ f"(axis {sequence_axis}), but got symbolic dim {max_seq_len!r}. "
+ "The static cache must be allocated with a fixed max_seq_len."
)
# S_q lives at dim 1 of both input_ids ([B, S_q]) and hidden_states
# ([B, S_q, hidden]), so the bias works for either forward entry point.
diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py
index 68f0d0836..d9b7b1464 100644
--- a/src/mobius/tasks/_causal_lm.py
+++ b/src/mobius/tasks/_causal_lm.py
@@ -8,7 +8,7 @@
import onnx_ir as ir
from onnxscript import GraphBuilder, nn
-from mobius._build_context import prefill_prefix_pruning
+from mobius._build_context import prefill_prefix_pruning, ep_capabilities
from mobius._configs import ArchitectureConfig
from mobius._constants import (
STATIC_CACHE_KV_SEQUENCE_LENGTH,
@@ -503,18 +503,25 @@ def _make_static_cache_inputs(
if cache_specs is None:
cache_specs = [(num_key_value_heads, head_dim)] * num_layers
+ layout = ep_capabilities().static_cache_layout
+
cache_pairs: list[tuple[ir.Value, ir.Value]] = []
for i, (kv_heads, layer_head_dim) in enumerate(cache_specs):
- kv_hidden = kv_heads * layer_head_dim
+
+ if layout == "heads_first":
+ cache_shape = [batch, kv_heads, max_seq_len, layer_head_dim]
+ else:
+ cache_shape = [batch, max_seq_len, kv_heads * layer_head_dim]
+
key_cache = builder.input(
f"key_cache.{i}",
dtype=dtype,
- shape=[batch, max_seq_len, kv_hidden],
+ shape=cache_shape,
)
value_cache = builder.input(
f"value_cache.{i}",
dtype=dtype,
- shape=[batch, max_seq_len, kv_hidden],
+ shape=cache_shape,
)
cache_pairs.append((key_cache, value_cache))
diff --git a/tests/static_cache_metadata_test.py b/tests/static_cache_metadata_test.py
index eef86097f..bfd269284 100644
--- a/tests/static_cache_metadata_test.py
+++ b/tests/static_cache_metadata_test.py
@@ -14,24 +14,44 @@
These tests pin that contract against real exported packages, not synthetic
graphs, so a change to the exporter's port names or scatter axis fails here
rather than at runtime.
+
+Opt in to real TensorRT execution with ``MOBIUS_TEST_TENSORRT=1`` and set
+``TENSORRT_ROOT`` to a TensorRT 11.3+ SDK containing ``bin/trtexec``. The
+``tensorrt_static_cache_runtime`` test requires CUDA, TensorRT Python bindings,
+cuda-python and Transformers. It builds a tiny seeded model without downloads;
+once enabled, missing prerequisites or engine failures are test failures.
"""
from __future__ import annotations
+import os
+import subprocess
+from pathlib import Path
from typing import Any
+import numpy as np
import onnx_ir as ir
import pytest
+import yaml
+from onnxscript import GraphBuilder
from mobius import registry
+from mobius._build_context import build_context
from mobius._configs import ArchitectureConfig
from mobius._constants import (
+ OPSET_VERSION,
STATIC_CACHE_KV_SEQUENCE_LENGTH,
STATIC_CACHE_SEQUENCE_AXIS,
STATIC_CACHE_WRITE_INDICES,
)
+from mobius._execution_providers import get_ep
+from mobius._flags import override_flags
+from mobius._testing.ort_inference import OnnxModelSession
+from mobius.components import create_static_cache_attention_bias
+from mobius.components._attention import StaticCacheState, _apply_attention
from mobius.integrations.onnx_genai.workflow_metadata import (
build_decoder_workflow_metadata,
+ write_decoder_workflow_metadata,
)
from mobius.tasks import CausalLMTask
@@ -55,11 +75,283 @@ def _text_config(**overrides) -> ArchitectureConfig:
return ArchitectureConfig(**params)
-def _static_package(**overrides):
+def _static_package(*, ep_name="default", **overrides):
config = _text_config(**overrides)
- module = registry.get("qwen2")(config)
- task = CausalLMTask(static_cache=True, max_seq_len=CAPACITY)
- return task.build(module, config), config
+ with build_context(get_ep(ep_name), dtype=config.dtype):
+ module = registry.get("qwen2")(config)
+ task = CausalLMTask(static_cache=True, max_seq_len=CAPACITY)
+ return task.build(module, config), config
+
+
+@pytest.mark.parametrize(
+ "ep_name,axis,layout", [("default", 1, "bsh"), ("tensorrt", 2, "bnsh")]
+)
+def test_static_cache_metadata_layout(ep_name, axis, layout, tmp_path):
+ package, config = _static_package(ep_name=ep_name)
+ metadata = build_decoder_workflow_metadata(package, config)
+ _, group = _scatter_group(metadata)
+ assert group["sequence_axis"] == axis
+ assert group["layout"] == layout
+ assert _static_cache_abi(metadata)["capacity"] == CAPACITY
+ path = write_decoder_workflow_metadata(package, str(tmp_path), config)
+ saved = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
+ _, saved_group = _scatter_group(saved)
+ assert saved_group["sequence_axis"] == axis
+ assert saved_group["layout"] == layout
+ abi = _static_cache_abi(saved)
+ assert abi == _static_cache_abi(metadata)
+ ports = _graph_ports(package)
+ for input_name, output_name in zip(abi["cache_inputs"], abi["cache_outputs"], strict=True):
+ assert ports[input_name].shape[axis] == CAPACITY
+ assert ports[output_name].shape == ports[input_name].shape
+
+
+@pytest.mark.parametrize("ep_name", ["default", "tensorrt"])
+@pytest.mark.parametrize(
+ "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16]
+)
+@pytest.mark.parametrize("static_cache", [False, True])
+def test_attention_cache_mask_ep_contract(ep_name, dtype, static_cache):
+ config = _text_config(dtype=dtype)
+ with build_context(get_ep(ep_name), dtype=dtype), override_flags(static_cache_bias=False):
+ module = registry.get("qwen2")(config)
+ package = CausalLMTask(static_cache=static_cache, max_seq_len=CAPACITY).build(
+ module, config
+ )
+ graph = package["model"].graph
+ attentions = [node for node in graph if node.op_type == "Attention"]
+ assert len(attentions) == 2
+ explicit_bias = static_cache and ep_name == "tensorrt"
+ for node in attentions:
+ if static_cache:
+ assert (node.inputs[3] is not None) == explicit_bias
+ assert node.attributes["is_causal"].as_int() == (0 if explicit_bias else 1)
+ native_nonpad = node.inputs[6] if len(node.inputs) > 6 else None
+ assert (native_nonpad is not None) == (static_cache and not explicit_bias)
+ if explicit_bias:
+ assert node.inputs[3].dtype == dtype
+ assert len(node.outputs) == 1
+ assert "q_num_heads" not in node.attributes
+ assert "kv_num_heads" not in node.attributes
+ if static_cache:
+ assert node.inputs[4] is None and node.inputs[5] is None
+ cache_input = next(value for value in graph.inputs if value.name == "key_cache.0")
+ assert len(cache_input.shape) == (4 if ep_name == "tensorrt" else 3)
+ else:
+ assert node.inputs[4] is not None and node.inputs[5] is not None
+ assert len(node.outputs) == 3
+ if explicit_bias:
+ assert attentions[0].inputs[3] is attentions[1].inputs[3]
+ graph_inputs = {value.name: value for value in graph.inputs}
+ assert graph_inputs["nonpad_kv_seqlen"].uses()
+ assert graph_inputs["write_indices"].uses()
+
+
+@pytest.mark.skipif(
+ os.environ.get("MOBIUS_TEST_TENSORRT") != "1",
+ reason="Set MOBIUS_TEST_TENSORRT=1 to build and execute a TensorRT engine",
+)
+def test_tensorrt_static_cache_runtime(tmp_path, monkeypatch):
+ import torch
+ from transformers import Qwen2Config, Qwen2ForCausalLM
+
+ monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[1] / "examples"))
+ from tensorrt_debug.runtime import TensorRTRunner
+
+ assert torch.cuda.is_available(), "The enabled TensorRT runtime test requires CUDA"
+ sdk = Path(os.environ["TENSORRT_ROOT"])
+ trtexec = sdk / "bin" / ("trtexec.exe" if os.name == "nt" else "trtexec")
+ assert trtexec.is_file(), f"TensorRT builder not found: {trtexec}"
+ monkeypatch.setenv(
+ "PATH", os.pathsep.join([str(sdk / "bin"), str(sdk / "lib"), os.environ["PATH"]])
+ )
+ hf_config = Qwen2Config(
+ num_hidden_layers=2,
+ hidden_size=64,
+ intermediate_size=128,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ vocab_size=256,
+ max_position_embeddings=128,
+ tie_word_embeddings=False,
+ )
+ hf_config._attn_implementation = "eager"
+ with torch.random.fork_rng(devices=[]):
+ torch.manual_seed(42)
+ reference = Qwen2ForCausalLM(hf_config).float().eval()
+ config = ArchitectureConfig.from_transformers(hf_config)
+ config.dtype = ir.DataType.FLOAT
+ capacity = 16
+ with build_context(get_ep("tensorrt"), dtype=config.dtype):
+ module = registry.get("qwen2")(config)
+ package = CausalLMTask(static_cache=True, max_seq_len=capacity).build(module, config)
+ package.apply_weights(module.preprocess_weights(reference.state_dict()))
+ model = package["model"]
+ model_path = tmp_path / "model.onnx"
+ engine_path = tmp_path / "model.engine"
+ ir.save(model, model_path, external_data="model.onnx.data")
+ command = [
+ str(trtexec),
+ f"--onnx={model_path}",
+ f"--saveEngine={engine_path}",
+ "--skipInference",
+ "--noTF32",
+ "--decomposableAttentions=*",
+ ]
+ for option, length in (("minShapes", 1), ("optShapes", 4), ("maxShapes", 8)):
+ shapes = []
+ for value in model.graph.inputs:
+ shape = [
+ dimension if isinstance(dimension, int) else 1 for dimension in value.shape
+ ]
+ if value.name in ("input_ids", "position_ids"):
+ shape = [1, length]
+ shapes.append(f"{value.name}:{'x'.join(map(str, shape))}")
+ command.append(f"--{option}={','.join(shapes)}")
+ build_log = tmp_path / "build.log"
+ with build_log.open("w", encoding="utf-8") as log:
+ result = subprocess.run(
+ command, stdout=log, stderr=subprocess.STDOUT, timeout=300, check=False
+ )
+ assert result.returncode == 0, build_log.read_text(encoding="utf-8")[-12000:]
+
+ history = []
+ with TensorRTRunner(engine_path, sdk) as runner, torch.inference_mode():
+ assert runner.capacity == capacity
+ for tokens in ([5, 17, 23, 42], [71], [19, 11]):
+ position = len(history)
+ history.extend(tokens)
+ length = len(history)
+ feeds = {
+ "input_ids": np.asarray([tokens], dtype=np.int64),
+ "position_ids": np.arange(position, length, dtype=np.int64)[None, :],
+ "write_indices": np.asarray([position], dtype=np.int64),
+ "nonpad_kv_seqlen": np.asarray([length], dtype=np.int64),
+ }
+ runner.prepare_inputs(feeds)
+ before = {name: runner.read_tensor(name) for name in runner.caches}
+ for name in runner.caches:
+ assert runner.context.get_tensor_address(
+ name
+ ) == runner.context.get_tensor_address("updated_" + name)
+ runner.execute()
+ actual_logits = runner.read_tensor("logits")
+ expected = reference(torch.tensor([history]), use_cache=True)
+ np.testing.assert_allclose(
+ actual_logits, expected.logits[:, position:].numpy(), atol=1e-3, rtol=1e-3
+ )
+ for layer in range(config.num_hidden_layers):
+ cache_layer = expected.past_key_values.layers[layer]
+ key, value = cache_layer.keys, cache_layer.values
+ for role, target in (("key", key), ("value", value)):
+ name = f"{role}_cache.{layer}"
+ actual = runner.read_tensor(name)
+ assert actual.shape == (
+ 1,
+ config.num_key_value_heads,
+ capacity,
+ config.head_dim,
+ )
+ np.testing.assert_allclose(
+ actual[:, :, :length], target.numpy(), atol=1e-3, rtol=1e-3
+ )
+ np.testing.assert_array_equal(
+ actual[:, :, :position], before[name][:, :, :position]
+ )
+ np.testing.assert_array_equal(
+ actual[:, :, length:], before[name][:, :, length:]
+ )
+ cached_logits = actual_logits.copy()
+ runner.reset_caches()
+ runner.prepare_inputs(
+ {
+ "input_ids": np.asarray([history], dtype=np.int64),
+ "position_ids": np.arange(len(history), dtype=np.int64)[None, :],
+ "write_indices": np.asarray([0], dtype=np.int64),
+ "nonpad_kv_seqlen": np.asarray([len(history)], dtype=np.int64),
+ }
+ )
+ runner.execute()
+ np.testing.assert_allclose(
+ cached_logits,
+ runner.read_tensor("logits")[:, -len(tokens) :],
+ atol=1e-3,
+ rtol=1e-3,
+ )
+
+
+@pytest.mark.parametrize("write,length,nonpad", [(0, 4, 4), (4, 1, 5), (4, 3, 7), (4, 3, 6)])
+def test_tensorrt_static_bias_slot_geometry(write, length, nonpad):
+ def input_value(name):
+ return ir.Value(name=name, shape=ir.Shape([1]), type=ir.TensorType(ir.DataType.INT64))
+
+ cursor = input_value("write_indices")
+ valid_length = input_value("nonpad_kv_seqlen")
+ graph = ir.Graph(
+ inputs=[cursor, valid_length],
+ outputs=[],
+ nodes=[],
+ opset_imports={"": OPSET_VERSION},
+ name="static_bias_geometry",
+ )
+ op = GraphBuilder(graph).op
+ with build_context(get_ep("tensorrt")):
+ bias = create_static_cache_attention_bias(
+ op,
+ write_indices=cursor,
+ seq_len=op.Constant(value_ints=[length]),
+ nonpad_kv_seqlen=valid_length,
+ max_seq_len=8,
+ )
+ bias.name = "bias"
+ graph.outputs.append(bias)
+ model = ir.Model(graph, ir_version=10)
+ actual = OnnxModelSession(model).run(
+ {
+ "write_indices": np.asarray([write], dtype=np.int64),
+ "nonpad_kv_seqlen": np.asarray([nonpad], dtype=np.int64),
+ }
+ )["bias"]
+ key_slots = np.arange(8)[None, :]
+ query_slots = write + np.arange(length)[:, None]
+ allowed = (key_slots <= query_slots) & (key_slots < nonpad)
+ expected = np.where(allowed, 0, np.finfo(np.float32).min).astype(np.float32)[None, None]
+ np.testing.assert_array_equal(actual, expected)
+
+
+def test_tensorrt_static_attention_rejects_missing_bias():
+ def value(name, shape):
+ return ir.Value(
+ name=name, shape=ir.Shape(shape), type=ir.TensorType(ir.DataType.FLOAT)
+ )
+
+ query = value("query", [1, 1, 64])
+ key = value("key", [1, 1, 32])
+ cache = value("cache", [1, 2, 8, 16])
+ index = ir.Value(name="index", shape=ir.Shape([1]), type=ir.TensorType(ir.DataType.INT64))
+ graph = ir.Graph(
+ inputs=[query, key, cache, index],
+ outputs=[],
+ nodes=[],
+ opset_imports={"": OPSET_VERSION},
+ )
+ with (
+ build_context(get_ep("tensorrt")),
+ pytest.raises(ValueError, match="explicit static-cache"),
+ ):
+ _apply_attention(
+ GraphBuilder(graph).op,
+ query,
+ key,
+ key,
+ None,
+ None,
+ None,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ scale=0.25,
+ static_cache=StaticCacheState(cache, cache, index, index),
+ )
def _cache_cells(workflow) -> list[str]: