The NeuroBrix runtime engine executes AI models efficiently across diverse hardware configurations through a choice of execution modes and a hardware solver.
The production execution mode:
--sequential)Sequential ATen dispatcher for debugging and development. Useful for tracing execution issues.
--triton, --triton-sequential)Triton kernels. Optional: Triton is not installed with NeuroBrix (neurobrix doctor reports whether it is present).
The Prism solver automatically selects the optimal execution strategy based on your hardware:
| Strategy | Description |
|---|---|
single_gpu |
Everything fits on one GPU |
lifecycle |
Component lifecycle management |
pipeline_parallel |
Pipeline across GPUs |
component_placement |
Components assigned to specific GPUs |
block_scatter |
Transformer blocks distributed across GPUs |
weight_sharding |
Weights sharded across GPUs |
component_placement_lazy |
Lazy component loading |
lazy_sequential |
Sequential with lazy loading |
zero3 |
ZeRO Stage 3 with CPU offload |
Memory estimation is done from safetensors headers without loading tensors, making strategy selection fast and memory-efficient.
Loop-based execution for models like Sana, PixArt, Flex.1. The topology defines loop count (diffusion steps) and data routing between denoising iterations.
Token-by-token generation with KV cache management. Session-based caching for multi-turn conversations.
One-pass execution for models like Whisper (speech-to-text).
Automatic mixed precision with operation-level control:
| Category | Operations | Dtype |
|---|---|---|
| FP32 ops | pow, rsqrt, softmax, layer_norm | Always FP32 |
| FP16 ops | mm, conv2d, linear, bmm | Hardware-preferred |
| Promote ops | add, mul, cat | Match inputs |
| Special | bmm (FP32 for T5-XXL stability), polar/view_as_complex (FP32 for RoPE) | Case-specific |
The serve mode runs as a background daemon:
--timeout)neurobrix serve --model <name> # Start daemon
neurobrix run --prompt "..." # Send inference request
neurobrix chat # Interactive LLM chat
neurobrix stop # Shutdown daemon
Next: ZERO Principles