# uwm-se/stageml architecture

> Architecture diagram of the GitHub repository uwm-se/stageml, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-26). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/uwm-se/stageml
- Repository: https://github.com/uwm-se/stageml

## Overview

StageML specializes inference by separating static model or adapter work from request-time computation, then representing the remaining work as residuals for execution. The main user-facing path stages and traces a model, propagates binding-time annotations, rewrites eligible computation, and evaluates or runs the residual. LoRA/PEFT and MoE serving are important branches, with MLIR and GPU backends providing lowering options. The README documents these capabilities; sampled excerpts confirm selected rewrite, adapter, lowering, kernel-generation, and routing behavior, but not all subsystem wiring.

Main entry points: [`stageml/vllm_stage_router.py`](https://github.com/uwm-se/stageml/blob/main/stageml/vllm_stage_router.py) (Tenant stage router), [`stageml/tracer.py`](https://github.com/uwm-se/stageml/blob/main/stageml/tracer.py) (FX tracing), [`stageml/peft_bridge.py`](https://github.com/uwm-se/stageml/blob/main/stageml/peft_bridge.py) (PEFT bridge), [`stageml/__init__.py`](https://github.com/uwm-se/stageml/blob/main/stageml/__init__.py) (Staging API)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_staging["Staging and Residuals"]
  node_public["Staging API<br/>[__init__.py]"]
  node_annotations["Stage annotations<br/>[annotations.py]"]
  node_tracer["FX tracing<br/>[tracer.py]"]
  node_rewrite["Graph rewrites<br/>[rewrite.py]"]
  node_planner["Residual planning"]
  node_validation["Plan validation"]
  node_policy["Cost policy"]
  node_evaluator["Staged evaluator<br/>[evaluator.py]"]
  node_runtime["Residual runtime<br/>[runtime.py]"]
end

subgraph group_adapters["Adapters and MoE"]
  node_peft["PEFT bridge<br/>[peft_bridge.py]"]
  node_adapterbank["Adapter bank<br/>[adapter_bank.py]"]
  node_adaptercache[("Adapter cache<br/>[adapter_cache.py]")]
  node_moelayers["MoE LoRA layers<br/>[moe_lora_layers.py]"]
  node_moestages["MoE staging<br/>[moe_stages.py]"]
  node_moeir["MoE residual plans<br/>[moe_ir.py]"]
  node_moeexport["MoE plan export<br/>[moe_plan_export.py]"]
end

subgraph group_backends["Lowering and Backends"]
  node_mlirlower["MLIR lowering<br/>[mlir_lower.py]"]
  node_canonical["Canonical MLIR"]
  node_moemlir["MoE MLIR<br/>[moe_mlir_lower.py]"]
  node_torchmlir["Torch MLIR backend"]
  node_iree["IREE residual"]
  node_triton["Triton generator"]
  node_native["SM90 backend"]
end

subgraph group_serving["Serving"]
  node_router["Tenant stage router"]
  node_customop["Residual custom op"]
  node_nativeop["Native op skeleton"]
end

node_caller(("Model developer"))
node_vllm["vLLM fused experts"]
node_tenant(("Tenant request"))

node_caller -->|"uses staging API"| node_public
node_public -->|"traces models"| node_tracer
node_public -->|"declares stages"| node_annotations
node_tracer -->|"propagates stages"| node_annotations
node_rewrite -->|"recomputes stages"| node_tracer
node_peft -->|"traces layer"| node_tracer
node_peft -->|"optimizes graph"| node_rewrite
node_peft -->|"wraps residual"| node_runtime
node_planner -->|"validates plan"| node_validation
node_planner -->|"selects plan"| node_policy
node_moestages -->|"forms plans"| node_moeir
node_moeir -->|"lowers plans"| node_moemlir
node_moeir -->|"exports plans"| node_moeexport
node_caller -->|"serves requests"| node_router
node_tenant -->|"provides tenant IDs"| node_router
node_router -->|"dispatches fallback tokens"| node_vllm
node_router -->|"dispatches staged tokens"| node_customop
node_customop -.->|"registers native op"| node_nativeop
node_router -.->|"resolves fused experts"| node_vllm

click node_public "https://github.com/uwm-se/stageml/blob/main/stageml/__init__.py"
click node_annotations "https://github.com/uwm-se/stageml/blob/main/stageml/annotations.py"
click node_tracer "https://github.com/uwm-se/stageml/blob/main/stageml/tracer.py"
click node_rewrite "https://github.com/uwm-se/stageml/blob/main/stageml/rewrite.py"
click node_planner "https://github.com/uwm-se/stageml/blob/main/stageml/residual_planner.py"
click node_validation "https://github.com/uwm-se/stageml/blob/main/stageml/planner_validation.py"
click node_policy "https://github.com/uwm-se/stageml/blob/main/stageml/policy_optimizer.py"
click node_evaluator "https://github.com/uwm-se/stageml/blob/main/stageml/evaluator.py"
click node_runtime "https://github.com/uwm-se/stageml/blob/main/stageml/runtime.py"
click node_peft "https://github.com/uwm-se/stageml/blob/main/stageml/peft_bridge.py"
click node_adapterbank "https://github.com/uwm-se/stageml/blob/main/stageml/adapter_bank.py"
click node_adaptercache "https://github.com/uwm-se/stageml/blob/main/stageml/adapter_cache.py"
click node_moelayers "https://github.com/uwm-se/stageml/blob/main/stageml/moe_lora_layers.py"
click node_moestages "https://github.com/uwm-se/stageml/blob/main/stageml/moe_stages.py"
click node_moeir "https://github.com/uwm-se/stageml/blob/main/stageml/moe_ir.py"
click node_moeexport "https://github.com/uwm-se/stageml/blob/main/stageml/moe_plan_export.py"
click node_mlirlower "https://github.com/uwm-se/stageml/blob/main/stageml/mlir_lower.py"
click node_canonical "https://github.com/uwm-se/stageml/blob/main/stageml/canonical_mlir_lower.py"
click node_moemlir "https://github.com/uwm-se/stageml/blob/main/stageml/moe_mlir_lower.py"
click node_torchmlir "https://github.com/uwm-se/stageml/blob/main/stageml/torch_mlir_backend.py"
click node_iree "https://github.com/uwm-se/stageml/blob/main/stageml/iree_residual_mlir.py"
click node_triton "https://github.com/uwm-se/stageml/blob/main/stageml/triton_generator.py"
click node_native "https://github.com/uwm-se/stageml/blob/main/stageml/sm90_native_backend.py"
click node_router "https://github.com/uwm-se/stageml/blob/main/stageml/vllm_stage_router.py"
click node_customop "https://github.com/uwm-se/stageml/blob/main/stageml/custom_op_registry.py"
click node_nativeop "https://github.com/uwm-se/stageml/blob/main/csrc/stageml_residual_op.cpp"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_public,node_annotations,node_tracer,node_rewrite,node_planner,node_validation,node_policy,node_evaluator,node_runtime toneBlue
class node_peft,node_adapterbank,node_adaptercache,node_moelayers,node_moestages,node_moeir,node_moeexport toneAmber
class node_mlirlower,node_canonical,node_moemlir,node_torchmlir,node_iree,node_triton,node_native,node_tenant toneMint
class node_router,node_customop,node_nativeop toneRose
class node_caller,node_vllm toneIndigo
```

## Components

### Staging and Residuals

- [Staging API](https://github.com/uwm-se/stageml/blob/main/stageml/__init__.py): `stageml/__init__.py`
- [Stage annotations](https://github.com/uwm-se/stageml/blob/main/stageml/annotations.py): `stageml/annotations.py`
- [FX tracing](https://github.com/uwm-se/stageml/blob/main/stageml/tracer.py): `stageml/tracer.py`
- [Graph rewrites](https://github.com/uwm-se/stageml/blob/main/stageml/rewrite.py): `stageml/rewrite.py`
- [Residual planning](https://github.com/uwm-se/stageml/blob/main/stageml/residual_planner.py): `stageml/residual_planner.py`
- [Plan validation](https://github.com/uwm-se/stageml/blob/main/stageml/planner_validation.py): `stageml/planner_validation.py`
- [Cost policy](https://github.com/uwm-se/stageml/blob/main/stageml/policy_optimizer.py): `stageml/policy_optimizer.py`
- [Staged evaluator](https://github.com/uwm-se/stageml/blob/main/stageml/evaluator.py): `stageml/evaluator.py`
- [Residual runtime](https://github.com/uwm-se/stageml/blob/main/stageml/runtime.py): `stageml/runtime.py`

### Adapters and MoE

- [PEFT bridge](https://github.com/uwm-se/stageml/blob/main/stageml/peft_bridge.py): `stageml/peft_bridge.py`
- [Adapter bank](https://github.com/uwm-se/stageml/blob/main/stageml/adapter_bank.py): `stageml/adapter_bank.py`
- [Adapter cache](https://github.com/uwm-se/stageml/blob/main/stageml/adapter_cache.py): `stageml/adapter_cache.py`
- [MoE LoRA layers](https://github.com/uwm-se/stageml/blob/main/stageml/moe_lora_layers.py): `stageml/moe_lora_layers.py`
- [MoE staging](https://github.com/uwm-se/stageml/blob/main/stageml/moe_stages.py): `stageml/moe_stages.py`
- [MoE residual plans](https://github.com/uwm-se/stageml/blob/main/stageml/moe_ir.py): `stageml/moe_ir.py`
- [MoE plan export](https://github.com/uwm-se/stageml/blob/main/stageml/moe_plan_export.py): `stageml/moe_plan_export.py`

### Lowering and Backends

- [MLIR lowering](https://github.com/uwm-se/stageml/blob/main/stageml/mlir_lower.py): `stageml/mlir_lower.py`
- [Canonical MLIR](https://github.com/uwm-se/stageml/blob/main/stageml/canonical_mlir_lower.py): `stageml/canonical_mlir_lower.py`
- [MoE MLIR](https://github.com/uwm-se/stageml/blob/main/stageml/moe_mlir_lower.py): `stageml/moe_mlir_lower.py`
- [Torch MLIR backend](https://github.com/uwm-se/stageml/blob/main/stageml/torch_mlir_backend.py): `stageml/torch_mlir_backend.py`
- [IREE residual](https://github.com/uwm-se/stageml/blob/main/stageml/iree_residual_mlir.py): `stageml/iree_residual_mlir.py`
- [Triton generator](https://github.com/uwm-se/stageml/blob/main/stageml/triton_generator.py): `stageml/triton_generator.py`
- [SM90 backend](https://github.com/uwm-se/stageml/blob/main/stageml/sm90_native_backend.py): `stageml/sm90_native_backend.py`

### Serving

- [Tenant stage router](https://github.com/uwm-se/stageml/blob/main/stageml/vllm_stage_router.py): `stageml/vllm_stage_router.py`
- [Residual custom op](https://github.com/uwm-se/stageml/blob/main/stageml/custom_op_registry.py): `stageml/custom_op_registry.py`
- [Native op skeleton](https://github.com/uwm-se/stageml/blob/main/csrc/stageml_residual_op.cpp): `csrc/stageml_residual_op.cpp`

### Other components

- **Model developer**
- **vLLM fused experts**
- **Tenant request**

## Connections

- Model developer → Staging API: uses staging API
- Staging API → FX tracing: traces models
- Staging API → Stage annotations: declares stages
- FX tracing → Stage annotations: propagates stages
- Graph rewrites → FX tracing: recomputes stages
- PEFT bridge → FX tracing: traces layer
- PEFT bridge → Graph rewrites: optimizes graph
- PEFT bridge → Residual runtime: wraps residual
- Residual planning → Plan validation: validates plan
- Residual planning → Cost policy: selects plan
- MoE staging → MoE residual plans: forms plans
- MoE residual plans → MoE MLIR: lowers plans
- MoE residual plans → MoE plan export: exports plans
- Model developer → Tenant stage router: serves requests
- Tenant request → Tenant stage router: provides tenant IDs
- Tenant stage router → vLLM fused experts: dispatches fallback tokens
- Tenant stage router → Residual custom op: dispatches staged tokens
- Residual custom op → Native op skeleton: registers native op
- Tenant stage router → vLLM fused experts: resolves fused experts