How uwm-se/stageml works
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: vllm_stage_router.py Tenant stage router · tracer.py FX tracing · peft_bridge.py PEFT bridge · __init__.py Staging API
Components 29
Staging and Residuals
- Staging API
stageml/__init__.py - Stage annotations
stageml/annotations.py - FX tracing
stageml/tracer.py - Graph rewrites
stageml/rewrite.py - Residual planning
stageml/residual_planner.py - Plan validation
stageml/planner_validation.py - Cost policy
stageml/policy_optimizer.py - Staged evaluator
stageml/evaluator.py - Residual runtime
stageml/runtime.py
Adapters and MoE
- PEFT bridge
stageml/peft_bridge.py - Adapter bank
stageml/adapter_bank.py - Adapter cache
stageml/adapter_cache.py - MoE LoRA layers
stageml/moe_lora_layers.py - MoE staging
stageml/moe_stages.py - MoE residual plans
stageml/moe_ir.py - MoE plan export
stageml/moe_plan_export.py
Lowering and Backends
- MLIR lowering
stageml/mlir_lower.py - Canonical MLIR
stageml/canonical_mlir_lower.py - MoE MLIR
stageml/moe_mlir_lower.py - Torch MLIR backend
stageml/torch_mlir_backend.py - IREE residual
stageml/iree_residual_mlir.py - Triton generator
stageml/triton_generator.py - SM90 backend
stageml/sm90_native_backend.py
Serving
- Tenant stage router
stageml/vllm_stage_router.py - Residual custom op
stageml/custom_op_registry.py - Native op skeleton
csrc/stageml_residual_op.cpp
Other components
- Model developer
- vLLM fused experts
- Tenant request
Connections 19
- 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