# Prism ID Architecture - Options Analysis
**Last Updated: February 2026**

## Context
Pour TP/PP, Prism doit pouvoir:
1. Identifier de manière stable les ops/tenseurs à travers traces
2. Décider quels tenseurs slicer et sur quelle dimension
3. Placer les ops sur différents devices/stages
4. Insérer les ops de communication (all-reduce, all-gather)

**Note**: As of February 2026, NeuroBrix uses AllocationStrategy enum (SINGLE_GPU, PIPELINE_PARALLEL, PP_NVLINK, PP_PCIE, TP_INTENT, ZERO3_OFFLOAD, COMPONENT_AFFINITY) defined in `src/neurobrix/core/prism/structure.py`. This document describes architectural options for future TP/PP improvements.

## Option A: Weight-Anchored IDs (RECOMMANDÉ)

### Principe
Les poids (weights) ont déjà des noms stables provenant des safetensors:
- `transformer.blocks.0.attn.q_proj.weight`
- `transformer.blocks.0.attn.k_proj.weight`

Utiliser ces noms comme **ancres** pour identifier les ops.

### Implementation
```python
# In forge/tracer/capture.py, during tracing:
class StableOpIdentifier:
    def __init__(self):
        self.weight_to_op = {}  # weight_name -> first op that uses it

    def get_stable_id(self, op_type, input_tensors, weight_names):
        if weight_names:
            # Op consomme un poids → ID stable basé sur le poids
            primary_weight = weight_names[0]
            return f"{primary_weight}:{op_type}"
        else:
            # Op sans poids → ID basé sur son producteur
            producer_id = self.get_producer_stable_id(input_tensors[0])
            return f"{producer_id}:downstream:{op_type}:{local_idx}"
```

### Avantages
- **Stabilité maximale**: Les noms de poids ne changent JAMAIS entre traces
- **Sémantique claire**: `blocks.5.attn.q_proj:matmul` = projection Q du bloc 5
- **Prism-friendly**: Prism peut directement référencer les couches à sharder

### Inconvénients
- Ops sans poids (activations pures) ont des IDs dérivés
- Nécessite tracking du lineage des tenseurs

---

## Option B: Module Hierarchy IDs

### Principe
Utiliser la hiérarchie des modules PyTorch pendant le trace:
```python
# Pendant forward, on track le module courant
current_module = "transformer.blocks.5.attn"
op_uid = f"{current_module}:{op_type}:{local_counter}"
```

### Implementation
```python
class ModuleTracker:
    def __init__(self):
        self.module_stack = []
        self.local_counters = {}  # module_path -> {op_type -> count}

    def enter_module(self, name):
        self.module_stack.append(name)

    def exit_module(self):
        self.module_stack.pop()

    def get_op_id(self, op_type):
        path = ".".join(self.module_stack)
        key = f"{path}:{op_type}"
        self.local_counters[key] = self.local_counters.get(key, 0) + 1
        return f"{path}:{op_type}:{self.local_counters[key]}"
```

### Avantages
- Reflète la structure du modèle
- Human-readable

### Inconvénients
- Dépend du forward order (peut varier avec control flow)
- Le local_counter peut diverger entre stimuli

---

## Option C: Content-Hash IDs

### Principe
Hash basé sur le contenu structurel de l'op:
```python
def content_hash(op_type, input_shapes, output_shapes, weight_names):
    content = f"{op_type}|{input_shapes}|{output_shapes}|{sorted(weight_names)}"
    return hashlib.md5(content.encode()).hexdigest()[:12]
```

### Avantages
- Complètement déterministe
- Indépendant de l'ordre d'exécution

### Inconvénients
- Collisions possibles pour ops identiques (ex: 28 blocs transformer identiques)
- Non human-readable

---

## Option D: Hybrid (RECOMMANDATION FINALE)

### Architecture en Couches

```
┌─────────────────────────────────────────────────────────────────┐
│  Layer 3: PRISM DIRECTIVES                                      │
│  ├── shard_tensor("blocks.5.attn.q_proj", dim=0, devices=[0,1]) │
│  ├── place_stage("blocks.0-7", device=0)                        │
│  └── insert_comm("all_reduce", after="blocks.7.output")         │
├─────────────────────────────────────────────────────────────────┤
│  Layer 2: SEMANTIC ANNOTATIONS (metadata on ops/tensors)        │
│  ├── layer_type: "attention" | "ffn" | "norm" | "embed"         │
│  ├── layer_index: 0-27                                          │
│  ├── shardable_dims: [0, 2]  # batch, heads                     │
│  └── memory_bytes: 1234567                                      │
├─────────────────────────────────────────────────────────────────┤
│  Layer 1: STABLE STRUCTURAL IDS                                 │
│  ├── op_uid: "blocks.5.attn.q_proj.weight:aten::linear"         │
│  ├── tensor_id: "blocks.5.attn.q_proj.weight:aten::linear:o0"   │
│  └── Based on weight anchors + derivation chain                 │
├─────────────────────────────────────────────────────────────────┤
│  Layer 0: RAW GRAPH (current)                                   │
│  ├── ops: sequential execution record                           │
│  └── tensors: data flow                                         │
└─────────────────────────────────────────────────────────────────┘
```

### Implementation Strategy

#### Phase 1: Stable IDs (Layer 1)
```python
# In forge/tracer/capture.py, during tracing:

class StableIDGenerator:
    def __init__(self, weight_names: Set[str]):
        self.weight_names = weight_names
        self.tensor_lineage = {}  # tensor_ptr -> stable_id

    def register_weight(self, tensor_ptr, weight_name):
        """Poids = ancre stable"""
        self.tensor_lineage[tensor_ptr] = f"W:{weight_name}"

    def get_tensor_id(self, tensor_ptr, producer_op_id, output_idx):
        """Tensor ID = producer_op:output_index"""
        return f"{producer_op_id}:o{output_idx}"

    def get_op_id(self, op_type, input_stable_ids, weight_consumed=None):
        """Op ID basé sur poids consommé ou inputs"""
        if weight_consumed:
            return f"{weight_consumed}:{op_type}"
        else:
            # Derive from first input's lineage
            primary_input = input_stable_ids[0]
            # Extract the weight ancestor
            weight_ancestor = self._find_weight_ancestor(primary_input)
            return f"{weight_ancestor}:derived:{op_type}:{hash(input_stable_ids)[:8]}"
```

#### Phase 2: Semantic Annotations (Layer 2)
```python
# Post-trace analysis adds semantic metadata:

def annotate_graph(graph, model_config):
    """Add semantic annotations for Prism"""
    for op_uid, op in graph['ops'].items():
        # Detect layer type from weight name pattern
        if 'q_proj' in op_uid or 'k_proj' in op_uid or 'v_proj' in op_uid:
            op['semantic'] = {
                'layer_type': 'attention',
                'component': 'qkv_projection',
                'shardable': True,
                'shard_dim': 0,  # output features
            }
        elif 'mlp' in op_uid or 'ffn' in op_uid:
            op['semantic'] = {
                'layer_type': 'ffn',
                'shardable': True,
                'shard_dim': 0,
            }
        # ... etc
```

#### Phase 3: Prism Directives (Layer 3)
```python
# Prism solver outputs placement decisions:
# Location: src/neurobrix/core/prism/solver.py

from neurobrix.core.prism.structure import AllocationStrategy

class PrismTPPPSolver:
    def solve(self, graph, device_topology):
        directives = []

        # TP: Shard large linear layers
        for op_uid, op in graph['ops'].items():
            if op.get('semantic', {}).get('shardable'):
                if self._should_shard(op, device_topology):
                    directives.append({
                        'type': 'tensor_parallel',
                        'strategy': AllocationStrategy.TP_INTENT,
                        'op': op_uid,
                        'shard_dim': op['semantic']['shard_dim'],
                        'devices': [0, 1, 2, 3],
                    })

        # PP: Assign layers to stages
        stages = self._partition_into_stages(graph, num_stages=4)
        for stage_idx, stage_ops in enumerate(stages):
            directives.append({
                'type': 'pipeline_stage',
                'strategy': AllocationStrategy.PP_NVLINK,  # or PP_PCIE
                'ops': stage_ops,
                'device': stage_idx,
            })

        # Communication: Insert all-reduce after TP ops
        for d in directives:
            if d['type'] == 'tensor_parallel':
                directives.append({
                    'type': 'insert_allreduce',
                    'after': d['op'],
                    'devices': d['devices'],
                })

        return directives
```

---

## Recommendation

### Court Terme (Fix Merger)
1. Implémenter **Weight-Anchored IDs** (Option A) dans `capture.py`
2. Les ops consommant des poids ont des IDs stables automatiquement
3. Les ops dérivés héritent du lineage

### Moyen Terme (Prism TP/PP)
1. Ajouter **Semantic Annotations** post-trace
2. Prism utilise les annotations pour décisions de sharding
3. Runtime interprète les directives Prism

### Long Terme (Full Automation)
1. Prism analyse automatiquement le graph pour détecter patterns shardables
2. Cost model basé sur profiling réel
3. Auto-tuning des configurations TP/PP

---

## Impact sur Fichiers Existants (February 2026)

| Fichier | Changement | Status |
|---------|------------|--------|
| `forge/tracer/capture.py` | Nouvelle classe `StableIDGenerator` | Future |
| `~/.neurobrix/cache/{model}/components/{name}/graph.json` | Nouveau champ `stable_id` sur ops et tensors | Future |
| `src/neurobrix/core/prism/solver.py` | Lecture des semantic annotations | Current file exists |
| `src/neurobrix/core/prism/structure.py` | AllocationStrategy enum (SINGLE_GPU, PP_NVLINK, PP_PCIE, TP_INTENT, ZERO3_OFFLOAD, COMPONENT_AFFINITY) | Implemented |
| Nouveau: `src/neurobrix/core/prism/tp_solver.py` | Logique tensor parallelism | Future |
| Nouveau: `src/neurobrix/core/prism/pp_solver.py` | Logique pipeline parallelism | Future |

**Current Prism Implementation (Feb 2026):**
- `src/neurobrix/core/prism/solver.py` — PrismSolver with strategy selection
- `src/neurobrix/core/prism/structure.py` — DeviceSpec, InterconnectLink, AllocationStrategy
- `src/neurobrix/core/prism/profiler.py` — Hardware profiling
- `src/neurobrix/core/prism/memory_estimator.py` — Memory estimation
- `src/neurobrix/core/prism/loader.py` — Profile loading
- `src/neurobrix/core/prism/common/dtype_resolver.py` — Dtype resolution

---

## Example: PixArt Transformer Block

```json
{
  "ops": {
    "blocks.5.attn.q_proj.weight:aten::linear": {
      "stable_id": "blocks.5.attn.q_proj.weight:aten::linear",
      "op_type": "aten::linear",
      "semantic": {
        "layer_type": "attention",
        "layer_index": 5,
        "component": "q_projection",
        "shardable": true,
        "shard_dim": 0,
        "memory_bytes": 4718592
      }
    }
  },
  "tensors": {
    "blocks.5.attn.q_proj.weight:aten::linear:o0": {
      "stable_id": "blocks.5.attn.q_proj.weight:aten::linear:o0",
      "shape": [1152, 4096],
      "semantic": {
        "is_activation": true,
        "sharded_across": [0, 1, 2, 3]
      }
    }
  }
}
```

Cette architecture permet à Prism de:
1. **Identifier** les couches par nom stable
2. **Décider** quoi sharder via semantic annotations
3. **Émettre** des directives que le runtime exécute
4. **Rester découplé** de la logique de graph

---

## Current System Status (February 2026)

### Package Structure
NeuroBrix is now a pip-installable package:
```
src/neurobrix/
├── cli.py                    # Entry point: neurobrix run/import/list/remove/info/inspect/validate
├── core/
│   ├── prism/
│   │   ├── solver.py        # PrismSolver.solve_from_container()
│   │   ├── structure.py     # AllocationStrategy enum, DeviceSpec, InterconnectLink
│   │   ├── profiler.py      # Hardware profiling
│   │   ├── memory_estimator.py
│   │   ├── loader.py        # load_profile(hardware_id)
│   │   └── common/dtype_resolver.py
│   ├── runtime/
│   │   ├── executor.py      # RuntimeExecutor (orchestrates from execution.json)
│   │   ├── factory.py       # ExecutorFactory (creates component executors with Prism dtype)
│   │   └── graph_executor.py
│   └── strategies/          # Strategy implementations (single_gpu, zero3, pipeline, etc.)
├── kernels/                 # Triton kernels
├── nbx/                     # NBX container format
└── config/                  # Hardware/vendor YAML configs

forge/                       # PRIVATE — Requires pip install neurobrix
├── forge.py                 # CLI: snap, build, trace, publish
├── importer/                # NBX builder (from forge/importer/)
├── tracer/                  # Graph tracing (from trace/ → forge/tracer/)
│   ├── capture.py           # Main trace orchestration
│   └── ...
├── vendors/                 # Vendored diffusers
└── config/families/         # Family YAML templates
```

### AllocationStrategy Enum (Current)
Located in `src/neurobrix/core/prism/structure.py`:
```python
class AllocationStrategy(Enum):
    SINGLE_GPU = "single_gpu"
    PIPELINE_PARALLEL = "pipeline_parallel"
    PP_NVLINK = "pp_nvlink"
    PP_PCIE = "pp_pcie"
    TP_INTENT = "tp_intent"
    ZERO3_OFFLOAD = "zero3_offload"
    COMPONENT_AFFINITY = "component_affinity"
```

### Cache Structure
Models are extracted to `~/.neurobrix/cache/` (not `.cache/models/`):
```
~/.neurobrix/
├── store/                   # Downloaded .nbx files (9+ GB each)
└── cache/                   # Extracted models (runtime reads from here)
    └── {model_name}/
        ├── manifest.json
        ├── topology.json
        ├── runtime/
        │   ├── defaults.json
        │   └── variables.json
        ├── components/
        │   └── {name}/
        │       ├── graph.json     # TensorDAG (from trace)
        │       ├── runtime.json   # Component attributes
        │       └── weights/
        └── modules/               # Scheduler, tokenizer
```

### Workflow
```
TRACE:   forge/forge.py trace   → Captures TensorDAG to graphs/
BUILD:   forge/forge.py build   → Creates .nbx from snapshot + graphs
PUBLISH: forge/forge.py publish → Uploads .nbx to MinIO registry
IMPORT:  neurobrix import org/model → Downloads .nbx → extracts to ~/.neurobrix/cache/
RUN:     neurobrix run --model ... → PrismSolver → RuntimeExecutor
```

### Key Imports (February 2026)
```python
from neurobrix.core.prism.solver import PrismSolver
from neurobrix.core.prism.structure import AllocationStrategy, DeviceSpec
from neurobrix.core.prism.loader import load_profile
from neurobrix.core.runtime.executor import RuntimeExecutor
from neurobrix.core.runtime.factory import ExecutorFactory
from neurobrix.nbx.container import NBXContainer
```

### Notes
- **NO ONNX**: ONNX references are obsolete
- **NO C++**: NeuroBrix runtime is pure Python/PyTorch
- Graph merge stability issues described in this document remain OPEN (future work)
- TP/PP solver architecture described here is FUTURE WORK (current system uses strategy fallback with score ranking)
