# PRISM - Hardware Supervisor Logic

**Version:** 0.2.0
**Last Updated:** 2026-02-08
**Status:** Production

---

## Table of Contents

1. [Overview](#overview)
2. [Architecture](#architecture)
3. [Core Components](#core-components)
4. [Runtime Allocation (PrismSolver)](#runtime-allocation)
5. [Data Structures](#data-structures)
6. [Decision Flowcharts](#decision-flowcharts)
7. [Rules & Constraints](#rules--constraints)
8. [Archived - Import Planning](#archived---import-planning)

---

## Overview

**Prism** is NeuroBrix's hardware supervisor. It serves as the single brain for all hardware-related decisions during runtime execution.

**Runtime Phase:** Distributes model components across available GPUs based on hardware profiles and component memory requirements.

### Design Principles

1. **Hardware Agnostic:** Runtime dtype decided based on actual hardware capabilities
2. **Zero Hardcoding:** Decisions based on tensor shapes and hardware capabilities
3. **Safety First:** Pre-flight memory checks before allocation
4. **Strategy Intelligence:** Automatic fallback from optimal to feasible strategies

---

## Architecture

```
┌─────────────────────────────────────────────────────────────────┐
│                         PRISM                                    │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │                    PrismSolver                               │ │
│  │                                                              │ │
│  │  - Strategy Selection (SINGLE_GPU → ZERO3)                  │ │
│  │  - Component Distribution                                   │ │
│  │  - Multi-GPU Allocation                                     │ │
│  │  - Dtype Resolution                                         │ │
│  │  - KV Cache Estimation                                      │ │
│  │  - Lifecycle Classification (Persistent/Transient)          │ │
│  └─────────────────────────────────────────────────────────────┘ │
│                                                                  │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │                    Hardware Profile                          │ │
│  │                                                              │ │
│  │  - DeviceSpec[]    (per-GPU specs)                          │ │
│  │  - InterconnectTopology (NVLink groups)                     │ │
│  │  - Dtype Support   (bf16, fp16, fp32)                       │ │
│  └─────────────────────────────────────────────────────────────┘ │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘
```

---

## Core Components

### File: `src/neurobrix/core/prism/solver.py`

| Component | Description |
|-----------|-------------|
| `ComponentAllocation` | Allocation plan for a single component |
| `PipelineExecutionPlan` | Complete execution plan for all components |
| `PrismSolver` | Runtime allocation solver with strategy intelligence |

### File: `src/neurobrix/core/prism/structure.py`

| Component | Description |
|-----------|-------------|
| `DeviceSpec` | Single GPU device specification |
| `InterconnectLink` | NVLink/PCIe connection between GPUs |
| `InterconnectTopology` | Complete interconnect mapping |
| `AllocationStrategy` | Enum: SINGLE_GPU, PIPELINE_PARALLEL, etc. |

### File: `src/neurobrix/core/prism/loader.py`

| Component | Description |
|-----------|-------------|
| `load_profile()` | Load hardware profile from YAML |

### File: `src/neurobrix/core/prism/profiler.py`

| Component | Description |
|-----------|-------------|
| Profiling logic | Symbolic activation profiling |

### File: `src/neurobrix/core/prism/memory_estimator.py`

| Component | Description |
|-----------|-------------|
| Memory estimation | Component memory requirements |

### File: `src/neurobrix/core/prism/common/dtype_resolver.py`

| Component | Description |
|-----------|-------------|
| Dtype resolution | Hardware-aware dtype selection |

---

## [ARCHIVED] Import Planning

**Note:** As of February 2026, NeuroBrix no longer uses ONNX export. Model import is handled by the `forge/` subsystem, which directly traces PyTorch graphs to NBX format. The PrismImportPlanner class has been removed. This section is preserved for historical reference.

### [ARCHIVED] Class: `PrismImportPlanner`

**Purpose:** [ARCHIVED] Planned hardware allocation for ONNX export during model import.

This functionality was part of the legacy ONNX-based workflow and has been replaced by the trace-based workflow in `forge/tracer/`. The forge system handles model graph capture independently.

---

## Runtime Allocation

### Class: `PrismSolver`

**Purpose:** Distribute model components across GPUs for execution with intelligent strategy selection.

**Location:** `src/neurobrix/core/prism/solver.py`

#### Method: `solve_from_container(container, profile)`

**Main entry point for runtime allocation planning.**

**Algorithm:**

1. **Load Component Metadata**
   ```
   FOR each component in container:
       Read manifest.json
       Read runtime.json (attributes)
       Read defaults.json (max_tokens for KV cache)
       Extract weight sizes
   ```

2. **KV Cache Estimation (LLM Models)**
   ```
   IF model has language_model component:
       cache_len = max_tokens + prompt_margin (128)
       per_token = num_layers × 2 × num_kv_heads × head_dim × dtype_bytes
       kv_cache_bytes = cache_len × per_token
   ```

3. **Lifecycle Classification (Autoregressive Models)**
   ```
   IF generation_type == "autoregressive_image":
       Persistent: language_model, gen_head, gen_embed, gen_aligner
       Transient: vision_model, gen_vision_model, aligner
   ```

4. **Strategy Selection (Scored, Best First)**
   ```
   Strategies (descending score):
     SINGLE_GPU (1000)
     PIPELINE_PARALLEL variants (795-800)
     TENSOR_PARALLEL (600)
     ZERO3_OFFLOAD (1)

   FOR each strategy (by score):
       IF strategy.validate(components, profile, kv_cache):
           RETURN strategy allocation

   Fallback: ZERO3_OFFLOAD (always fits - activations only)
   ```

5. **Dtype Resolution**
   ```
   requested_dtype = profile.preferred_dtype or "float16"

   IF bf16 requested AND all GPUs support bf16:
       use torch.bfloat16
   ELSE IF fp16 AND all GPUs support fp16:
       use torch.float16
   ELSE:
       use torch.float32

   bf16 → fp16: ALLOWED after _scan_bf16_fp16_safety() verifies all values ≤ 65504
   ```

#### Strategy Descriptions

| Strategy | Score | Condition | Description |
|----------|-------|-----------|-------------|
| `SINGLE_GPU` | 1000 | All components fit on largest GPU | Fastest, no inter-GPU transfers |
| `PP_NVLINK` | 800 | Pipeline stages fit across NVLink GPUs | High-bandwidth interconnect |
| `PP_PCIE` | 795 | Pipeline stages fit across PCIe GPUs | Standard interconnect |
| `TENSOR_PARALLEL` | 600 | Components support sharding | Distributes single component |
| `ZERO3_OFFLOAD` | 1 | Always valid | Activations only, weights on CPU |

#### KV Cache Validation (NEW - January 2026)

**Problem Solved:** Previous solver used `max_position_embeddings` (16384) instead of actual generation length (`max_tokens`=576 for Janus), causing 28× overestimate and false OOM predictions.

**Solution:**
```python
def _estimate_kv_cache_bytes(component_name, manifest, defaults):
    """Data-driven from defaults.json, NOT architectural limits."""
    max_tokens = defaults.get("max_tokens", 256)  # Generation tokens
    prompt_margin = 128  # Headroom for prefill
    cache_len = max_tokens + prompt_margin

    lm_config = defaults.get("lm_config", {})
    num_layers = lm_config["num_layers"]
    num_kv_heads = lm_config["num_kv_heads"]
    head_dim = lm_config["head_dim"]

    per_token = num_layers * 2 * num_kv_heads * head_dim
    return cache_len * per_token * dtype_bytes
```

**Strategy Validation:**
```python
# Single GPU must fit: weights + activations + KV cache
required = weights_mb + activation_mb + kv_cache_mb
IF required > largest_gpu_vram:
    Try next strategy (PP_NVLINK)
```

### Runtime Allocation Diagram

```
┌───────────────────────────────────────────────────┐
│              NBX Model Container                  │
│                                                   │
│  Components:                                      │
│    transformer (12 GB weights + 2 GB activations) │
│    text_encoder (500 MB weights + 100 MB acts)    │
│    vae (160 MB weights + 50 MB activations)       │
│                                                   │
│  LLM: language_model (6 GB) + KV cache (660 MB)   │
└───────────────────────────────────────────────────┘
                        │
                        ▼
┌───────────────────────────────────────────────────┐
│         PrismSolver.solve_from_container()        │
│                                                   │
│  1. Estimate KV cache (LLM models)                │
│  2. Classify lifecycle (autoregressive)           │
│  3. Try strategies by score:                      │
│     - SINGLE_GPU (all on largest GPU)             │
│     - PP_NVLINK (pipeline across NVLink)          │
│     - ZERO3_OFFLOAD (activations only)            │
│  4. Validate: weights + acts + KV ≤ VRAM          │
└───────────────────────────────────────────────────┘
                        │
                        ▼
┌───────────────────────────────────────────────────┐
│         Strategy: SINGLE_GPU (V100 32GB)          │
│                                                   │
│  GPU 0:                                           │
│    transformer: 12 GB weights + 2 GB acts         │
│    text_encoder: 500 MB weights + 100 MB acts     │
│    vae: 160 MB weights + 50 MB acts               │
│  ────────────────────────────────────────────     │
│  Total: 14.8 GB (fits in 32 GB)                   │
└───────────────────────────────────────────────────┘
                        │
                        ▼
┌───────────────────────────────────────────────────┐
│    Strategy: single_gpu_lifecycle (Janus)         │
│                                                   │
│  GPU 0:                                           │
│    Persistent: language_model (6 GB) + KV (660 MB)│
│    Persistent: gen_head (200 MB)                  │
│    Transient: gen_vision_model (4 GB, unloads)    │
│  ────────────────────────────────────────────     │
│  Peak: 10.86 GB, Decode: 6.86 GB                  │
└───────────────────────────────────────────────────┘
```

---

## Data Structures

### ComponentAllocation

```python
@dataclass
class ComponentAllocation:
    """Allocation plan for a single component."""
    name: str
    device: str              # Primary device: "cuda:0", "cpu"
    dtype: torch.dtype       # Resolved from hardware profile
    weight_memory_mb: float  # Weight size
    activation_memory_mb: float  # Peak activation size
    kv_cache_mb: float       # KV cache size (LLM only)
```

### PipelineExecutionPlan

```python
@dataclass
class PipelineExecutionPlan:
    """Complete execution plan for all components."""
    components: Dict[str, ComponentAllocation]
    strategy: AllocationStrategy  # Enum: SINGLE_GPU, PIPELINE_PARALLEL, etc.
    target_dtype: torch.dtype
    total_weight_mb: float
    total_activation_mb: float
    total_kv_cache_mb: float
```

### DeviceSpec

```python
@dataclass
class DeviceSpec:
    """Single GPU device specification."""
    index: int
    name: str               # "Tesla V100-PCIE-32GB"
    memory_mb: int          # Total VRAM
    compute_capability: str # "7.0" (Volta), "8.0" (Ampere)
    supports_dtypes: List[str]  # ["float32", "float16", "bfloat16"]
    architecture: str       # "volta", "ampere", "hopper"
    brand: str              # "nvidia", "amd", "intel"
```

### AllocationStrategy (Enum)

```python
class AllocationStrategy(Enum):
    """Available allocation strategies, scored by preference."""
    SINGLE_GPU = 1000           # All components on one GPU
    PIPELINE_PARALLEL = 800     # Components pipelined across GPUs
    PP_NVLINK = 800             # Pipeline with NVLink interconnect
    PP_PCIE = 795               # Pipeline with PCIe interconnect
    TENSOR_PARALLEL = 600       # Component sharding
    TP_INTENT = 600             # Tensor parallel intent
    COMPONENT_AFFINITY = 500    # Component-to-GPU affinity
    ZERO3_OFFLOAD = 1           # ZeRO-3 style offload (last resort)
```

### InterconnectLink

```python
@dataclass
class InterconnectLink:
    """Connection between two GPUs."""
    device_a: int
    device_b: int
    link_type: str  # "nvlink", "pcie"
    bandwidth_gbps: float
```

### Hardware Profile (YAML)

**Location:** `src/neurobrix/config/hardware/{profile_id}.yml`

```yaml
id: v100-32g
vendor: nvidia
devices:
  - index: 0
    name: "Tesla V100-PCIE-32GB"
    memory_mb: 32768
    compute_capability: "7.0"
    supports_dtypes: ["float32", "float16"]
    architecture: volta
    brand: nvidia

topology:
  links: []  # Single GPU, no interconnect
```

---

## Decision Flowcharts

### Dtype Resolution (Hardware-Driven)

```
                    ┌─────────────────────────┐
                    │  Hardware Profile YAML  │
                    │  supports_dtypes: [...]  │
                    └────────┬────────────────┘
                             │
                             ▼
                    ┌─────────────────────────┐
                    │ Check requested dtype   │
                    │ (from CLI or defaults)  │
                    └────────┬────────────────┘
                             │
               ┌─────────────┴─────────────┐
               │                           │
               ▼                           ▼
        ┌──────────────┐           ┌──────────────┐
        │ bf16 in list?│           │ bf16 NOT in  │
        │   Use bf16   │           │  list        │
        └──────────────┘           └──────┬───────┘
                                           │
                                           ▼
                                   ┌──────────────┐
                                   │ fp16 in list?│
                                   │   Use fp16   │
                                   └──────┬───────┘
                                           │
                                           ▼
                                   ┌──────────────┐
                                   │ Fallback to  │
                                   │   fp32       │
                                   └──────────────┘

CRITICAL: bf16 → fp16 conversion ALLOWED after range verification.
          PrismSolver._scan_bf16_fp16_safety() checks all bf16 weights ≤ 65504.
          If safe → convert to fp16 (saves memory, V100 compatible).
          If unsafe → keep bf16 or upcast to fp32.
          Always safe without scan: bf16 → fp32 (same exponent range).
```

### Memory Safety Margins

| Phase | Margin | Reason |
|-------|--------|--------|
| Runtime Allocation | 0.85x (15%) | Fragmentation + safety buffer |
| KV Cache | prompt_margin: 128 tokens | Headroom for prefill phase |

---

## Rules & Constraints

### Runtime Rules (February 2026)

| # | Rule | Reason |
|---|------|--------|
| 1 | bf16 not supported → fp32 | NEVER fall back to fp16 (exponent mismatch) |
| 2 | Strategy scoring: SINGLE_GPU → ZERO3 | Try optimal first, fallback to feasible |
| 3 | KV cache from `max_tokens`, NOT `max_position_embeddings` | Actual generation length, not architectural limit |
| 4 | Lifecycle classification for autoregressive | Persistent vs transient components |
| 5 | KV cache counted once (not double) | Already in activation_bytes |
| 6 | Validate KV cache fit per strategy | Prevents false OOM predictions |
| 7 | ZERO3_OFFLOAD always valid | Activations only, last resort |

### Dtype Hierarchy

```
Preferred               Fallback
    ▼                      ▼
┌──────────┐          ┌──────────┐
│ bfloat16 │ ───────► │ float32  │
└──────────┘          └──────────┘
                           │
                           │ NEVER
                           ▼
                      ┌──────────┐
                      │ float16  │ ✗
                      └──────────┘
```

---

## Usage Examples

### Runtime Phase

```python
from neurobrix.core.prism.solver import PrismSolver
from neurobrix.core.prism.loader import load_profile
from neurobrix.nbx.container import NBXContainer

# Load NBX container
container = NBXContainer.from_path("~/.neurobrix/cache/Janus-Pro-7B")

# Load hardware profile
profile = load_profile("v100-32g")  # From src/neurobrix/config/hardware/

# Create solver
solver = PrismSolver()

# Solve allocation from container
plan = solver.solve_from_container(container, profile)

print(f"Strategy: {plan.strategy.name}")  # SINGLE_GPU, PP_NVLINK, etc.
print(f"Dtype: {plan.target_dtype}")      # torch.float16
print(f"Total weights: {plan.total_weight_mb/1024:.1f}GB")
print(f"Total activations: {plan.total_activation_mb/1024:.1f}GB")
print(f"Total KV cache: {plan.total_kv_cache_mb/1024:.1f}GB")

# Component allocations
for comp_name, alloc in plan.components.items():
    print(f"{comp_name}:")
    print(f"  Device: {alloc.device}")
    print(f"  Weights: {alloc.weight_memory_mb:.1f} MB")
    print(f"  Activations: {alloc.activation_memory_mb:.1f} MB")
    if alloc.kv_cache_mb > 0:
        print(f"  KV Cache: {alloc.kv_cache_mb:.1f} MB")
```

### CLI Usage

```bash
# Run model with automatic Prism allocation
neurobrix run \
  --model Janus-Pro-7B \
  --hardware v100-32g \
  --prompt "A sunset over mountains"

# Prism will:
# 1. Load container from ~/.neurobrix/cache/Janus-Pro-7B/
# 2. Load profile from src/neurobrix/config/hardware/v100-32g.yml
# 3. Estimate KV cache (660 MB for Janus, max_tokens=576)
# 4. Try strategies: SINGLE_GPU → PP_NVLINK → ZERO3_OFFLOAD
# 5. Select best valid strategy
# 6. Allocate components to GPU(s)
```

---

## Dtype Flow (CRITICAL - 2025-12-30 Fix)

### Complete Dtype Flow

```
┌────────────────────────────────────────────────────────────────────────┐
│                          DTYPE FLOW (ZERO HARDCODE)                     │
├────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│  1. Hardware Profile YAML                                               │
│     src/neurobrix/config/hardware/v100-32g.yml                          │
│     ┌─────────────────────────────────┐                                 │
│     │ supports_dtypes:                │                                 │
│     │   - float32                     │                                 │
│     │   - float16                     │   ◄── V100 supports fp16/fp32   │
│     │ # bfloat16 NOT listed           │                                 │
│     └─────────────────────────────────┘                                 │
│                         │                                               │
│                         ▼                                               │
│  2. Prism Dtype Resolution                                              │
│     src/neurobrix/core/prism/common/dtype_resolver.py                   │
│     ┌─────────────────────────────────┐                                 │
│     │ def resolve_dtype():            │                                 │
│     │   if "float16" in supports:     │                                 │
│     │     return torch.float16        │   ◄── Best supported dtype      │
│     └─────────────────────────────────┘                                 │
│                         │                                               │
│                         ▼                                               │
│  3. ComponentAllocation                                                 │
│     src/neurobrix/core/prism/solver.py                                  │
│     ┌─────────────────────────────────┐                                 │
│     │ allocation.dtype = torch.float16│   ◄── Embedded in allocation    │
│     └─────────────────────────────────┘                                 │
│                         │                                               │
│                         ▼                                               │
│  4. ExecutorFactory                                                     │
│     src/neurobrix/core/runtime/factory.py                               │
│     ┌─────────────────────────────────┐                                 │
│     │ dtype = allocation.dtype        │                                 │
│     │ executor = GraphExecutor(       │                                 │
│     │   graph, weights_path,          │                                 │
│     │   dtype=dtype  # PASSED!        │   ◄── Factory extracts dtype    │
│     │ )                               │                                 │
│     └─────────────────────────────────┘                                 │
│                         │                                               │
│                         ▼                                               │
│  5. DtypeEngine (AMP Rules)                                             │
│     src/neurobrix/core/dtype/engine.py                                  │
│     ┌─────────────────────────────────┐                                 │
│     │ compute_dtype = torch.float16   │   ◄── From allocation            │
│     │ Apply PyTorch AMP rules:        │                                 │
│     │   pow/rsqrt → fp32              │   ◄── Stability                 │
│     │   mm/conv → compute_dtype       │   ◄── Performance               │
│     └─────────────────────────────────┘                                 │
│                         │                                               │
│                         ▼                                               │
│  6. WeightLoader                                                        │
│     src/neurobrix/core/io/weight_loader.py                              │
│     ┌─────────────────────────────────┐                                 │
│     │ weights = load_safetensors()    │   ◄── Source is fp32/bf16       │
│     │ weights = convert(weights,      │                                 │
│     │                   self.dtype)   │   ◄── Convert to fp16           │
│     └─────────────────────────────────┘                                 │
│                         │                                               │
│                         ▼                                               │
│  7. GPU Memory: ~12GB (not 20GB)                                        │
│                                                                         │
└────────────────────────────────────────────────────────────────────────┘
```

### Key Principle

**NEVER hardcode dtype.** It must flow:
1. Hardware profile defines supported dtypes
2. Prism resolves best dtype
3. Allocation carries dtype
4. Factory passes to executor
5. WeightLoader converts

### Hardware Profile Examples

**Location:** `src/neurobrix/config/hardware/`

```yaml
# v100-32g.yml (Volta)
id: v100-32g
vendor: nvidia
devices:
  - index: 0
    name: "Tesla V100-PCIE-32GB"
    memory_mb: 32768
    compute_capability: "7.0"
    supports_dtypes:
      - float32
      - float16
    # NO bfloat16 - Volta doesn't support it
    architecture: volta
    brand: nvidia

# a100-80g.yml (Ampere)
id: a100-80g
vendor: nvidia
devices:
  - index: 0
    name: "A100-SXM4-80GB"
    memory_mb: 81920
    compute_capability: "8.0"
    supports_dtypes:
      - float32
      - float16
      - bfloat16  # Ampere supports bf16
    architecture: ampere
    brand: nvidia
```

---

## Changelog

| Date | Version | Changes |
|------|---------|---------|
| 2026-02-08 | 0.2.0 | Package Restructure + KV Cache Fixes |
|  |  | - Updated paths: `core/` → `src/neurobrix/core/` |
|  |  | - Updated imports: `from core.` → `from neurobrix.core.` |
|  |  | - Removed ONNX/PrismImportPlanner (archived) |
|  |  | - Added KV cache estimation (max_tokens, not max_position_embeddings) |
|  |  | - Added lifecycle classification (persistent/transient) |
|  |  | - Added strategy fallback loop |
|  |  | - Updated execution modes: compiled, native, triton |
|  |  | - Cache paths: `.cache/models/` → `~/.neurobrix/cache/` |
| 2026-01-12 | 0.1.0 | Prism SmartSolver |
|  |  | - SmartSolver (Best-Fit-Decreasing) |
|  |  | - Interconnect-aware block placement |
|  |  | - FGP_NVLINK, FGP_PCIE strategies |
|  |  | - InterconnectTopology graph |
| 2025-12-30 | 3.0 | Dtype Flow Documentation |
|  |  | - Complete dtype flow diagram |
|  |  | - Hardware → Prism → Factory → WeightLoader chain |
|  |  | - Memory fix verification |
| 2025-12-03 | 2.0 | Universal Detection Integration |
|  |  | - Added `solve_from_container()` method |
|  |  | - Memory from detection.json input_spec |
|  |  | - dtype from detection.json resolution_sources |
| 2024-12-01 | 1.0 | [ARCHIVED] Initial documentation |
|  |  | - Added PrismImportPlanner (removed in 0.2.0) |
|  |  | - fp32 universal graph rule (ONNX-based) |
|  |  | - Pre-flight VRAM verification |

---

## Prism SmartSolver

### Overview

SmartSolver is the allocation engine with interconnect-aware placement.

### Key Features

| Feature | Description |
|---------|-------------|
| Best-Fit-Decreasing | Allocate largest blocks first |
| Interconnect-aware spillover | Prefer same NVLink group when GPU full |
| Symbolic activation profiling | Accurate memory estimation |
| Multi-strategy selection | SINGLE_GPU → FGP → ZERO3_OFFLOAD |

### Strategies

```
SINGLE_GPU        → Model fits on 1 GPU
COMPONENT_AFFINITY → 1 component = 1 GPU
PP_NVLINK         → Pipeline parallel (NVLink)
PP_PCIE           → Pipeline parallel (PCIe)
FGP_NVLINK        → Fine-grained pipeline (NVLink)
FGP_PCIE          → Fine-grained pipeline (PCIe)
TP_INTENT         → Tensor parallel (placeholder)
ZERO3_OFFLOAD     → CPU offload
```

### Interconnect-Aware Spillover

```python
def _get_spillover_device(self, current_dev, available_devices, profile):
    """When GPU full, prefer same NVLink group over larger capacity."""
    current_idx = int(current_dev.device_string.split(':')[1])

    def spillover_score(dev):
        dev_idx = int(dev.device_string.split(':')[1])
        same_group = profile.topology.devices_have_fast_interconnect(
            [current_idx, dev_idx]
        )
        return (same_group, dev.free_mb)

    return max(candidates, key=spillover_score)
```

### File Structure (February 2026)

| File | Purpose |
|------|---------|
| `src/neurobrix/core/prism/solver.py` | PrismSolver with strategy intelligence |
| `src/neurobrix/core/prism/structure.py` | DeviceSpec, InterconnectTopology, AllocationStrategy |
| `src/neurobrix/core/prism/loader.py` | YAML profile loading |
| `src/neurobrix/core/prism/profiler.py` | Symbolic activation profiler |
| `src/neurobrix/core/prism/memory_estimator.py` | Component memory estimation |
| `src/neurobrix/core/prism/common/dtype_resolver.py` | Hardware-aware dtype resolution |
| `src/neurobrix/config/hardware/` | Hardware profile YAML files |

---

## Execution Modes & Integration (February 2026)

### Runtime Modes

NeuroBrix supports three execution modes, controlled by CLI flags:

| Mode | Flag | Dispatcher | Description |
|------|------|------------|-------------|
| **Compiled** | Default | `CompiledSequenceV2` | Zero-overhead arena-based execution with DtypeEngine |
| **Native** | `--seq_aten` | `NativeATenDispatcher` | Dict-based PyTorch ATen ops (debugging) |
| **Triton** | `--triton` | `TritonAdapter` | Custom Triton kernels for compute ops |

**Execution Flow:**

```
CLI: neurobrix run --model ... --hardware ...
  ↓
Load NBX container (~/.neurobrix/cache/model_name/)
  ↓
Load hardware profile (src/neurobrix/config/hardware/{id}.yml)
  ↓
PrismSolver.solve_from_container()
  ├─ Estimate KV cache (LLM models)
  ├─ Classify lifecycle (autoregressive)
  ├─ Try strategies (SINGLE_GPU → ZERO3)
  └─ Return PipelineExecutionPlan
  ↓
RuntimeExecutor(pkg, plan, mode="compiled")
  ├─ ExecutorFactory creates GraphExecutor per component
  ├─ GraphExecutor(dtype=plan.target_dtype)
  └─ WeightLoader converts weights to dtype
  ↓
GraphExecutor.load_weights()
  ├─ Load safetensors (fp32/bf16 source)
  └─ Convert to allocation.dtype (fp16)
  ↓
GraphExecutor.run(inputs)
  ├─ CompiledSequenceV2 (compiled mode)
  │   ├─ TensorArena (integer slot addressing)
  │   ├─ DtypeEngine (PyTorch AMP rules)
  │   └─ Zero dict lookups
  ├─ NativeATenDispatcher (native mode)
  │   └─ Direct PyTorch ATen ops
  └─ TritonAdapter (triton mode)
      ├─ TRITON ops → Custom kernels
      └─ METADATA ops → PyTorch native
```

### DtypeEngine Integration (Compiled Mode)

**Location:** `src/neurobrix/core/dtype/engine.py`

The DtypeEngine implements PyTorch AMP autocast rules for stability:

```python
# AMP categories (source: PyTorch aten/src/ATen/autocast_mode.h)
AMP_FP32_OPS = {"pow", "rsqrt", "softmax", "sum", "norm", ...}  # 41 ops
AMP_FP16_OPS = {"mm", "bmm", "conv2d", "linear", ...}           # 67 ops
AMP_PROMOTE_OPS = {"add", "mul", "cat", ...}                    # 18 ops

# Example: RMSNorm stability
# Without AMP: pow(fp16) → overflow → rsqrt(Inf) = 0 → ALL ZEROS
# With AMP: pow → upcast to fp32 → mean → rsqrt → mm brings back to fp16
```

**Critical Rule:** Complex tensors (RoPE freq via `polar`, `view_as_complex`) MUST NOT be cast to real dtype. Guard in `_to_copy`:

```python
if inp.is_complex() and not target_dtype.is_complex:
    return inp  # Preserve complex tensor
```

### Package Structure (pip install neurobrix)

```
src/neurobrix/
├── cli.py                     # Entry point: neurobrix run/import/list
├── core/                      # Runtime engine
│   ├── prism/                 # Prism solver
│   ├── runtime/               # Execution
│   │   ├── executor.py        # RuntimeExecutor
│   │   ├── factory.py         # ExecutorFactory
│   │   ├── graph_executor.py  # GraphExecutor
│   │   └── graph/
│   │       ├── compiled_sequence.py  # CompiledSequenceV2
│   │       └── ...
│   ├── dtype/                 # DtypeEngine
│   │   ├── engine.py          # AMP rules
│   │   └── converter.py
│   ├── io/
│   │   └── weight_loader.py   # WeightLoader
│   └── ...
├── kernels/                   # Triton kernels
│   ├── adapter.py
│   ├── classification.py
│   └── ops/
├── nbx/                       # NBX format
│   ├── container.py
│   └── cache.py               # ~/.neurobrix/cache/
└── config/                    # Bundled YAML
    ├── hardware/              # v100-32g.yml, a100-80g.yml
    └── vendors/
```

**Import Convention:**
```python
# Public API
from neurobrix.core.prism.solver import PrismSolver
from neurobrix.core.runtime.executor import RuntimeExecutor
from neurobrix.nbx.container import NBXContainer

# Private (forge/)
from neurobrix.nbx.neurotax import Neurotax  # Used by forge/importer/
```

---

## KV Cache & Lifecycle Intelligence (January 2026)

### Problem: KV Cache 28× Overestimate

**Old Behavior:**
```python
# WRONG: Used architectural limit
cache_len = lm_config["max_position_embeddings"]  # 16384 for Janus
kv_cache_bytes = cache_len × per_token × dtype_bytes
# Result: 18 GB estimate for Janus (28× too large!)
```

**Root Cause:** `max_position_embeddings` is the architectural maximum the model CAN handle, not what it WILL generate. Janus generates 384px images = (384/16)² = 576 tokens.

**Fix:**
```python
# CORRECT: Use actual generation length
max_tokens = defaults.get("max_tokens", 256)  # 576 for Janus
prompt_margin = 128  # Headroom for prefill
cache_len = max_tokens + prompt_margin  # 704 tokens
kv_cache_bytes = cache_len × per_token × dtype_bytes
# Result: 660 MB for Janus (realistic)
```

### Lifecycle Classification (Autoregressive Models)

**Problem:** Single GPU lifecycle strategy treated all components equally, causing OOM on decode loop.

**Solution:** Classify components based on `topology.json → execution.generation`:

```python
def _classify_lifecycle(container):
    """Classify components as persistent or transient."""
    execution = container.get_execution()
    gen_flow = execution.get("generation", {})

    if not gen_flow:
        return {"persistent": [], "transient": []}

    # Components in decode loop = persistent
    persistent = gen_flow.get("iterates_over", [])

    # Vision models = transient (prefill only)
    transient = [c for c in container.components
                 if "vision" in c.lower() or "aligner" in c.lower()]

    return {"persistent": persistent, "transient": transient}
```

**Example: Janus-Pro-7B**
- **Persistent:** language_model (6 GB), gen_head (200 MB), KV cache (660 MB)
  - Loaded simultaneously during decode loop
  - Peak: 6.86 GB
- **Transient:** gen_vision_model (4 GB), aligner (500 MB)
  - Loaded during prefill, unloaded before decode
  - Not counted in decode memory

**Strategy Validation:**
```python
# single_gpu_lifecycle strategy
persistent_mem = sum(persistent_weights + activations) + kv_cache
transient_mem = max(transient_weights + activations)
peak_mem = persistent_mem + transient_mem  # Prefill peak
decode_mem = persistent_mem                # Decode steady-state

if peak_mem ≤ gpu_vram:
    return VALID
```

### Strategy Fallback Loop

**Problem:** Old solver picked best strategy WITHOUT validating KV cache fit.

**Solution:** Try strategies in score order, validate each:

```python
strategies = [
    (SINGLE_GPU, 1000),
    (PP_NVLINK, 800),
    (PP_PCIE, 795),
    (TENSOR_PARALLEL, 600),
    (ZERO3_OFFLOAD, 1),
]

for strategy, score in sorted(strategies, key=lambda x: -x[1]):
    allocation = strategy.allocate(components, profile)

    # Validate KV cache fits
    if strategy == SINGLE_GPU:
        required = weights + activations + kv_cache
        if required > largest_gpu_vram:
            continue  # Try next strategy

    return allocation  # First valid strategy wins

# ZERO3_OFFLOAD always valid (activations only)
return zero3_allocate(components, profile)
```

**Key Points:**
- KV cache already counted in `activation_bytes` (LM component)
- Remaining VRAM calculation subtracts KV cache to avoid double-counting
- Each strategy validates total memory (weights + acts + KV) before committing
- Fallback ensures ALWAYS a valid allocation

---

## References

- `src/neurobrix/core/prism/solver.py` - Main implementation
- `src/neurobrix/core/prism/structure.py` - Data structures
- `src/neurobrix/core/prism/loader.py` - Profile loading
- `src/neurobrix/core/prism/common/dtype_resolver.py` - Dtype resolution
- `src/neurobrix/core/runtime/factory.py` - ExecutorFactory integration
- `src/neurobrix/core/dtype/engine.py` - DtypeEngine (AMP rules)
- `forge/importer/` - Import subsystem (separate from runtime)
- `forge/tracer/` - Graph tracing (separate from runtime)
