# KERNELS OPERATOR SYSTEM - Architecture & Specification

## Version: 0.2.0
## Date: 2026-02-08
## Status: APPROVED

---

## 1. PHILOSOPHY

```
╔═══════════════════════════════════════════════════════════════════════════════════════════╗
║                                                                                           ║
║   "Un kernel n'est qu'une fonction optimisée pour une architecture spécifique"           ║
║                                                                                           ║
║   Le Kernels Operator System charge dynamiquement les kernels selon:                      ║
║   - FAMILY: image, llm, video, audio                                                      ║
║   - VENDOR: nvidia, amd, metal                                                            ║
║   - TIER: common, volta, ampere, hopper (NVIDIA) | common, cdna2, cdna3 (AMD)             ║
║                                                                                           ║
║   Le nom du tier correspond à l'architecture MINIMUM requise.                            ║
║   Une cascade de résolution assure la compatibilité ascendante.                          ║
║                                                                                           ║
╚═══════════════════════════════════════════════════════════════════════════════════════════╝
```

---

## 2. PRINCIPLES

| Principe | Description | Exemple |
|----------|-------------|---------|
| **ZERO HARDCODE** | Aucune valeur en dur dans le code | `if vendor == config["vendor"]` ✅ |
| **VARIABLE EMBARQUÉE** | Config hardware = source de vérité | `hardware.yml → vendor, architecture` |
| **1 OP = 1 FICHIER** | Chaque opération dans son propre fichier | `softmax.py`, `matmul.py` |
| **SCALABILITÉ** | Ajouter une op = ajouter un fichier | Pas de modification centrale |
| **TIERS COMPATIBILITÉ** | Kernels classés par requis minimum | `ampere` = nécessite Ampere ou mieux |

---

## 3. ARCHITECTURE

### 3.1 Structure des Répertoires

```
src/neurobrix/kernels/          # Kernels package inside neurobrix
│
├── __init__.py                 # Entry point
├── adapter.py                  # ATen → Kernel translation
├── resolver.py                 # Op → Kernel resolution
├── registry.py                 # @register_kernel decorator
├── exceptions.py               # KernelNotFoundError (ZERO FALLBACK)
├── classification.py           # METADATA vs TRITON classification
├── mapping.py                  # ATen → kernel name mapping
├── metadata_ops.py             # PyTorch metadata ops (--triton mode)
├── spec.py                     # Kernel specification
│
├── ops/
│   ├── image/
│   │   ├── nvidia/
│   │   │   ├── common/         # Compatible Volta, Ampere, Hopper
│   │   │   │   ├── add.py
│   │   │   │   └── softmax.py
│   │   │   │
│   │   │   ├── volta/          # Requis: Volta minimum (CC 7.0)
│   │   │   │
│   │   │   ├── ampere/         # Requis: Ampere minimum (CC 8.0, BF16)
│   │   │   │   └── flash_attention.py
│   │   │   │
│   │   │   └── hopper/         # Requis: Hopper minimum (CC 9.0, FP8, TMA)
│   │   │       └── fp8_matmul.py
│   │   │
│   │   └── amd/
│   │       ├── common/
│   │       ├── cdna2/          # Requis: CDNA2 minimum (MI200)
│   │       └── cdna3/          # Requis: CDNA3 minimum (MI300)
│   │
│   ├── llm/
│   │   └── nvidia/
│   │       ├── common/
│   │       └── ampere/
│   │
│   └── ... (video, audio)
```

### 3.2 Tiers de Compatibilité (NVIDIA)

| Tier | Architecture Minimum | Compatible Avec | Exemple Kernels |
|------|---------------------|-----------------|-----------------|
| `common` | Toutes | volta, ampere, hopper | add, mul, softmax |
| `volta` | Volta (CC 7.0) | volta, ampere, hopper | Tensor Cores v1 |
| `ampere` | Ampere (CC 8.0) | ampere, hopper | FlashAttention, BF16 natif |
| `hopper` | Hopper (CC 9.0) | hopper uniquement | FP8, TMA, WGMMA |

### 3.3 Cascade de Résolution

Quand un kernel est demandé, le système cherche du **plus optimisé au plus général**.

**Ordre de recherche par architecture:**
*   **Hopper:** `hopper` → `ampere` → `volta` → `common`
*   **Ampere:** `ampere` → `volta` → `common`
*   **Volta:** `volta` → `common`

**Exemple : Demande `matmul` sur GPU Hopper (sm90)**
1. `ops/image/nvidia/hopper/matmul.py` ? (Non)
2. `ops/image/nvidia/ampere/matmul.py` ? (Non)
3. `ops/image/nvidia/volta/matmul.py` ? (Non)
4. `ops/image/nvidia/common/matmul.py` ? (Oui) ✅

**Exemple : Demande `flash_attention` sur GPU Volta (sm70)**
1. `ops/image/nvidia/volta/flash_attention.py` ? (Non)
2. `ops/image/nvidia/common/flash_attention.py` ? (Non)
3. **ZERO FALLBACK** ❌ (Pas de kernel compatible)

---

## 4. API SPECIFICATION

### 4.1 Registre & Décorateur

```python
# neurobrix/kernels/registry.py

def register_kernel(
    family: str,
    vendor: str,
    tier: str,
    op_name: str,
    **constraints
):
    """Enregistre un kernel dans le registre global."""
    ...
```

### 4.2 Utilisation dans un Kernel

```python
# neurobrix/kernels/ops/image/nvidia/common/add.py

from neurobrix.kernels.registry import register_kernel

@register_kernel(
    family="image",
    vendor="nvidia",
    tier="common",
    op_name="add"
)
def add_kernel(x, y):
    ...
```

### 4.3 Résolution (get_kernel)

```python
# neurobrix/kernels/__init__.py

def get_kernel(op_name: str, hw, family: str) -> Callable:
    """
    Résout le kernel en appliquant la cascade de tiers.

    1. Détermine les tiers éligibles pour l'architecture hw.
    2. Cherche dans l'ordre de priorité (plus optimisé -> plus général).
    3. Si aucun trouvé -> KernelNotFoundError (ZERO FALLBACK).
    """
    ...
```

---

## 5. HARDWARE MAPPING

### 5.1 Architecture -> Tiers (NVIDIA)

```python
NVIDIA_ARCH_TO_TIERS = {
    "hopper": ["hopper", "ampere", "volta", "common"],
    "ada":    ["ada", "ampere", "volta", "common"],
    "ampere": ["ampere", "volta", "common"],
    "turing": ["turing", "volta", "common"],
    "volta":  ["volta", "common"],
}
```

### 5.2 Architecture -> Tiers (AMD)

```python
AMD_ARCH_TO_TIERS = {
    "cdna3": ["cdna3", "cdna2", "cdna", "common"],
    "cdna2": ["cdna2", "cdna", "common"],
    "rdna3": ["rdna3", "common"],
}
```

---

## 6. TYPES DE KERNELS

### 6.1 Définition "Zero Compute PyTorch"

NeuroBrix applique **Zero Compute PyTorch**:
- ❌ `torch.matmul()`, `torch.softmax()` → INTERDIT (compute PyTorch)
- ✅ `@triton.jit` kernels → AUTORISÉ (compute Triton)
- ✅ `torch.empty()` → AUTORISÉ (allocation mémoire)
- ✅ `torch.Tensor` → AUTORISÉ (structure de données)
- ✅ `x.squeeze()`, `x.reshape()` → AUTORISÉ (metadata, ZERO compute)

### 6.2 Interface Standard

**TOUS les kernels utilisent la même interface:**
```python
@register_kernel(family="image", vendor="nvidia", tier="common", op_name="nom_op")
def kernel_name(x: torch.Tensor, ...) -> torch.Tensor:
    """
    Input:  torch.Tensor (sur GPU)
    Output: torch.Tensor (sur GPU)
    """
    return output
```

### 6.3 Catégories de Kernels

| Type | Description | Compute | Exemple |
|------|-------------|---------|---------|
| **Triton Compute** | Calcul GPU via `@triton.jit` | Triton | matmul, softmax, gelu |
| **Triton Memory** | Réorganisation mémoire | Triton | transpose, concat, gather |
| **Triton Creation** | Création tenseurs | Triton | constant, range, cast |
| **Metadata Ops** | Manipulation view | ZERO | squeeze, unsqueeze, expand |

### 6.4 Exemples par Type

**Triton Compute (calcul GPU):**
```python
@register_kernel(family="image", vendor="nvidia", tier="common", op_name="add")
def add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    # Allocation output
    out = torch.empty_like(x)

    # Kernel Triton
    _add_kernel[grid](x, y, out, n_elements, BLOCK_SIZE=1024)

    return out
```

**Triton Memory (réorganisation):**
```python
@register_kernel(family="image", vendor="nvidia", tier="common", op_name="transpose")
def transpose_kernel(x: torch.Tensor, perm: List[int]) -> torch.Tensor:
    # Transpose = TOUJOURS copie pour garantir contiguous
    out = torch.empty(new_shape, dtype=x.dtype, device=x.device)

    # Kernel Triton qui copie dans le nouvel ordre
    _transpose_kernel[grid](x, out, ...)

    return out  # Contiguous garanti
```

**Metadata Ops (ZERO compute):**
```python
@register_kernel(family="image", vendor="nvidia", tier="common", op_name="squeeze")
def squeeze_kernel(x: torch.Tensor, dim: int = None) -> torch.Tensor:
    # Utilise torch.squeeze() - c'est juste une view, ZERO compute GPU
    return x.squeeze(dim)
```

**Note:** Les metadata ops comme `squeeze`, `reshape`, `flatten` utilisent les méthodes PyTorch correspondantes car elles ne font AUCUN calcul GPU - elles manipulent uniquement les métadonnées du tensor (shape, strides).

---

## 7. CLASSIFICATION DES OPS ATen (January 2026)

**NOTE IMPORTANTE**: La classification actuelle optimise pour stabilité numérique en fp16.

### 7.1 Classification Overview

| Category | Count | Description |
|----------|-------|-------------|
| **TRITON** | 36 | Custom Triton kernels working correctly |
| **METADATA (precision)** | 26 | PyTorch native - Triton has fp16 precision drift |
| **METADATA (optimized)** | ~50 | PyTorch/cuDNN highly optimized |
| **METADATA (no compute)** | ~50 | View ops, tensor creation |

### 7.2 TRITON Ops (36 total)

| Catégorie | Ops | Status |
|-----------|-----|--------|
| **Activations simples** | relu, sigmoid, tanh, leaky_relu, elu, mish, softplus, hardsigmoid, hardswish | ✅ |
| **Normalisation** | rms_norm, batch_norm, instance_norm | ✅ |
| **Pooling** | max_pool2d, avg_pool2d, adaptive_avg_pool2d, adaptive_max_pool2d | ✅ |
| **Loss** | cross_entropy_loss, nll_loss, mse_loss, l1_loss, binary_cross_entropy, kl_div | ✅ |
| **Scatter/Gather** | scatter, scatter_add, gather, index_select, index_add | ✅ |
| **Sorting** | sort, topk, argsort | ✅ |
| **Other** | triu, tril, masked_fill, one_hot, log_softmax | ✅ |

### 7.3 METADATA Ops - Precision Workaround

Ces ops ont des kernels Triton mais causent une dérive de précision fp16 qui s'accumule sur 1000+ ops (400x divergence dans le text encoder).

| Catégorie | Ops | Raison |
|-----------|-----|--------|
| **Matrice** | mm, bmm, addmm, linear, matmul | Triton precision drift |
| **Binaire** | add, sub, rsub, mul, div, pow, maximum, minimum | Triton precision drift |
| **Unaire** | abs, neg, exp, log, sqrt, rsqrt, sin, cos, erf, reciprocal | Triton precision drift |

**TODO:** Fix Triton kernels for fp16 precision, then switch to TRITON.

### 7.4 METADATA Ops - PyTorch Optimized

| Catégorie | Ops | Raison |
|-----------|-----|--------|
| **Normalisation** | layer_norm, group_norm | PyTorch highly optimized |
| **Activations** | gelu, silu, softmax | cuDNN optimized |
| **Attention** | scaled_dot_product_attention | Flash Attention optimal |
| **Convolution** | conv2d, conv1d, conv_transpose2d | cuDNN optimal |
| **Réduction** | sum, mean, max, min, argmax | PyTorch optimized |
| **Comparaison** | eq, ne, lt, gt, le, ge, where, clamp | PyTorch native |

### 7.5 METADATA Ops - No Compute (View Ops)

| Catégorie | Ops | Raison |
|-----------|-----|--------|
| **View** | view, reshape, flatten, squeeze, unsqueeze, permute, transpose | O(1) pointer math |
| **Creation** | zeros, ones, empty, full, arange, linspace | PyTorch native |
| **Type** | to, _to_copy, float, half, bfloat16 | PyTorch native |
| **Clone** | clone, copy, detach, alias | Memory copy |

### 7.6 Classification File

See `kernels/classification.py` for complete list (~165 ops classified).

```python
# Example classification
ATEN_CLASSIFICATION = {
    # TRITON - working custom kernels
    "aten::relu": OpExecution.TRITON,
    "aten::sigmoid": OpExecution.TRITON,

    # METADATA - precision workaround
    "aten::matmul": OpExecution.METADATA,
    "aten::add": OpExecution.METADATA,

    # METADATA - no compute
    "aten::view": OpExecution.METADATA,
}
```

### 7.7 inf/nan Guard

The inf/nan guard only checks TRITON ops:

```python
# METADATA ops may produce transient inf values (e.g., T5 GELU pow(x,3))
if exec_type != OpExecution.METADATA and op_uid in self._op_outputs:
    # ... check for nan/inf ...
```

---

## 8. EXECUTION MODES (Default: Compiled, Debug: Native, R&D: Triton)

### Purpose

NeuroBrix supports three execution modes optimized for different purposes:

| Mode | Flag | Purpose | Performance |
|------|------|---------|-------------|
| **Compiled** | (default) | Production - zero Python overhead | Fastest (95% GPU util) |
| **Native** | `--seq_aten` | Debugging - PyTorch ATen fallback | Fast (80% GPU util) |
| **Triton** | `--triton` | R&D - custom kernel testing | Variable (testing) |

### Architecture

```
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│                                   EXECUTION MODES                                        │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│                                                                                          │
│  ┌─────────────────────┐    ┌─────────────────────┐    ┌──────────────────────────┐   │
│  │ COMPILED (default)  │    │ NATIVE (--seq_aten) │    │  TRITON (--triton)       │   │
│  │                     │    │                     │    │                          │   │
│  │  Graph Load         │    │  ATen Op            │    │  ATen Op                 │   │
│  │      │              │    │      │              │    │      │                   │   │
│  │      ▼              │    │      ▼              │    │      ▼                   │   │
│  │  CompiledSequenceV2 │    │  SequentialDispatch │    │  classification.py       │   │
│  │  + DtypeEngine      │    │  (dict-based)       │    │  (TRITON or METADATA)    │   │
│  │  (arena, AMP)       │    │                     │    │      │                   │   │
│  │      │              │    │      ▼              │    │      ▼                   │   │
│  │      ▼              │    │  torch.* native     │    │  Triton Kernels OR       │   │
│  │  Zero Python loop   │    │  calls              │    │  metadata_ops.py         │   │
│  └─────────────────────┘    └─────────────────────┘    └──────────────────────────┘   │
│                                                                                          │
└─────────────────────────────────────────────────────────────────────────────────────────┘
```

### Compiled Mode (DEFAULT)

**Status**: Production-ready (February 2026)

- **CompiledSequenceV2**: Pre-compiles all ops at load time, zero dict lookups
- **DtypeEngine**: PyTorch AMP autocast rules for numeric stability
- **TensorArena**: `__slots__`-based list for O(1) tensor access
- **Integer slot addressing**: All tensors by index, not string keys
- **Dead tensor analysis**: Frees memory during execution
- **Performance**: 95% GPU utilization, matches diffusers speed

**Key files**:
- `src/neurobrix/core/runtime/graph/compiled_sequence.py`
- `src/neurobrix/core/dtype/engine.py`

### Native Mode (`--seq_aten`)

**Status**: Debugging fallback

- **SequentialDispatcher**: Dict-based PyTorch ATen operations
- **Ad-hoc stability fixes**: Per-op workarounds for numeric issues
- **String-based dispatch**: Dict lookups for every op
- **Purpose**: Debugging when compiled mode has issues
- **Performance**: 80% GPU utilization (dict overhead)

**Key files**:
- `src/neurobrix/core/runtime/graph/sequential_dispatcher.py`

### Triton Mode (`--triton`)

**Status**: R&D / Experimental

- **Custom Triton kernels**: Hand-written optimizations
- **Classification system**: TRITON (36 ops) vs METADATA (128 ops)
- **Purpose**: Testing new kernel implementations
- **Stability**: Some fp16 precision drift issues
- **Performance**: Variable (depends on kernel quality)

**Key files**:
- `src/neurobrix/kernels/classification.py`
- `src/neurobrix/kernels/adapter.py`
- `src/neurobrix/kernels/metadata_ops.py`

### Files Involved

| File | Purpose |
|------|---------|
| `src/neurobrix/core/runtime/graph/compiled_sequence.py` | Compiled mode executor (DEFAULT) |
| `src/neurobrix/core/dtype/engine.py` | AMP dtype rules (compiled mode) |
| `src/neurobrix/core/runtime/graph/sequential_dispatcher.py` | Native ATen fallback (--seq_aten) |
| `src/neurobrix/kernels/classification.py` | Op classification (--triton mode) |
| `src/neurobrix/kernels/mapping.py` | ATen → kernel name mapping (--triton mode) |
| `src/neurobrix/kernels/metadata_ops.py` | Shape ops (--triton mode) |
| `src/neurobrix/core/runtime/graph_executor.py` | Mode routing logic |

### Usage

```bash
# Default: Compiled mode (production)
neurobrix run \
  --model PixArt-Sigma-XL-2-1024-MS \
  --hardware v100-32g \
  --prompt "A sunset"

# Dev installation (from repo root)
PYTHONPATH=src python -m neurobrix run \
  --model PixArt-Sigma-XL-2-1024-MS \
  --hardware v100-32g \
  --prompt "A sunset"

# Native mode (debugging fallback)
neurobrix run \
  --model PixArt-Sigma-XL-2-1024-MS \
  --hardware v100-32g \
  --prompt "A sunset" \
  --seq_aten

# Triton mode (R&D kernel testing)
neurobrix run \
  --model PixArt-Sigma-XL-2-1024-MS \
  --hardware v100-32g \
  --prompt "A sunset" \
  --triton
```

### Key Points

1. **Compiled is DEFAULT** - Production-ready, fastest execution
2. **Native is for debugging** - Use `--seq_aten` when troubleshooting
3. **Triton is experimental** - Use `--triton` for kernel R&D
4. **AMP in compiled mode** - Automatic numeric stability (no manual fixes)
5. **Zero Python overhead** - Compiled mode pre-computes all dispatch logic

---

## 9. CHANGELOG

| Version | Date | Changes |
|---------|------|---------|
| 6.0 | 2026-02-08 | Package restructure + compiled mode as default |
|  |  | - Moved kernels/ → src/neurobrix/kernels/ (pip installable) |
|  |  | - All imports: `from neurobrix.kernels.*` |
|  |  | - CompiledSequenceV2 + DtypeEngine is DEFAULT mode |
|  |  | - Native mode renamed: `--native` → `--seq_aten` |
|  |  | - Three modes: compiled (default), native (--seq_aten), triton (--triton) |
|  |  | - Updated paths: sequential_dispatcher.py (was: native_dispatcher.py) |
|  |  | - CLI: `neurobrix run` (was: python neurobrix.py run) |
| 5.0 | 2026-01-01 | Triton path stabilization |
|  |  | - Fixed inf/nan guard to skip METADATA ops (T5 GELU transient inf) |
|  |  | - Removed duplicate masked_fill entry |
|  |  | - Reorganized classification.py with clear section headers |
|  |  | - 36 TRITON ops, 128 METADATA ops |
|  |  | - Matrix/binary/unary ops routed to METADATA for fp16 precision |
| 4.0 | 2025-12-31 | BREAKING: Default changed to PyTorch, --triton opt-in |
|  |  | - Renamed `--native` → default behavior |
|  |  | - Added `--triton` flag for Triton kernels |
|  |  | - Moved kernels/aten/*.py → kernels/*.py |
| 3.0 | 2025-12-30 | Native mode documentation |
|  |  | - `--native` flag bypasses Triton kernels |
|  |  | - NativeATenDispatcher for PyTorch fallback |
| 2.3 | 2024-12-06 | CORRECTION: Interface torch.Tensor, définition Zero Compute PyTorch |
| 2.2 | 2024-12-06 | (Incorrect - remplacé) |
| 2.1 | 2024-12-06 | Renommage des tiers (`sm80+` -> `ampere`) pour clarté |
| 2.0 | 2024-12-06 | Introduction des Tiers de Compatibilité |
| 1.0 | 2024-12-06 | Initial specification |

---

## 9. REFERENCES

- **NeuroBrix Technical Specification** (`docs/TECHNICAL_SPEC.md`): Philosophy & Architecture
- **ZERO_FALLBACK_POLICY**: Politique stricte d'erreur
- **Vision 2026**: Performance & Multi-Hardware Strategy

---

*Kernels Operator System - NeuroBrix Infrastructure*
*"Zero Compute PyTorch, Full Triton Power"*