Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 56 additions & 14 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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<br/>Transformers · Diffusers · GGUF · NeMo"]
Config["ArchitectureConfig<br/>Normalizes source configuration"]
Registry["Registry<br/>Selects the model class and task"]
Components["Reusable components<br/>Attention · MLP · Norm · RoPE · MoE · Vision · Audio"]
Models["Model modules<br/>Compose components into architectures"]
Tasks["Tasks<br/>Define ONNX inputs, outputs, caches, and model splits"]
Graph["ONNX graph construction<br/>onnxscript.nn + onnx_ir"]
Optimize["EP-aware optimization<br/>Cleanup · Fusion · Lowering · Folding"]
Weights["Weight pipeline<br/>Download · Rename · Transform · Cast · Apply"]
Package["ModelPackage<br/>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:
Expand All @@ -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<br/>__init__.py · _builder.py · _model_package.py"]
Root --> Configs["_configs/<br/>Normalized architecture configuration"]
Root --> Components["components/<br/>Reusable ONNX building blocks"]
Root --> Models["models/<br/>Architecture implementations"]
Root --> Tasks["tasks/<br/>Graph I/O contracts"]
Root --> Registry["_registry.py<br/>Model and task lookup"]
Root --> Optimizations["_optimizations.py · rewrite_rules/ · _passes/<br/>Graph optimization"]
Root --> Integrations["integrations/<br/>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
Expand Down
17 changes: 17 additions & 0 deletions docs/execution_providers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions docs/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,4 +52,5 @@ ai-model-support-strategy
:caption: Research

research/testing-strategy-analysis
research/tensorrt-static-cache-debugging
```
Loading
Loading