# Kernel Inventory - Complete Audit

## Date: 2026-02-08 (Last Updated)

---

## Execution Modes

NeuroBrix supports two execution modes:

### Default Mode (Compiled)
- Uses `CompiledSequenceV2` with `DtypeEngine`
- All ATen ops executed via PyTorch native implementations
- DtypeEngine applies AMP (Automatic Mixed Precision) autocast rules
- Zero Python overhead, arena-based memory management
- **This is the default and recommended mode**

### Triton Mode (Experimental)
- Enabled with `--triton` flag
- Uses custom Triton kernels from `src/neurobrix/kernels/ops/`
- Op classification in `kernels/classification.py` determines execution:
  - `METADATA`: PyTorch native (cuBLAS, cuDNN)
  - `TRITON`: Custom Triton kernel
- Tier cascade: hopper → ampere → volta → common
- **Currently experimental, use for testing only**

---

## Current Kernel Structure

**Location:** All kernel paths below are relative to `src/neurobrix/kernels/ops/`

```
src/neurobrix/kernels/ops/
├── image/                  # Image/Diffusion models
│   ├── nvidia/
│   │   ├── common/         # Compatible Volta, Ampere, Hopper
│   │   ├── volta/
│   │   ├── ampere/
│   │   └── hopper/
│   └── amd/
│       └── common/
├── llm/                    # Language models
│   ├── nvidia/
│   │   ├── common/
│   │   ├── volta/
│   │   ├── ampere/         # flash_attention.py, paged_attention.py
│   │   └── hopper/
│   └── amd/
│       └── common/
├── video/                  # Video generation models
│   ├── nvidia/...
│   └── amd/...
└── audio/                  # Audio models
    ├── nvidia/...
    └── amd/...
```

**Total kernel files:** 314

---

## Implemented Kernels by Family

### All Families (image, llm, video, audio) - Common Tier

| Kernel | File | Status |
|--------|------|--------|
| add | add.py | ✅ Done |
| sub | sub.py | ✅ Done |
| mul | mul.py | ✅ Done |
| div | div.py | ✅ Done |
| softmax | softmax.py | ✅ Done |
| layernorm | layernorm.py | ✅ Done |
| gelu | gelu.py | ✅ Done |
| matmul | matmul.py | ✅ Done |
| gemm | gemm.py | ✅ Done |
| abs | abs.py | ✅ Done |
| neg | neg.py | ✅ Done |
| exp | exp.py | ✅ Done |
| log | log.py | ✅ Done |
| sqrt | sqrt.py | ✅ Done |
| rsqrt | rsqrt.py | ✅ Done |
| sin | sin.py | ✅ Done |
| cos | cos.py | ✅ Done |
| sigmoid | sigmoid.py | ✅ Done |
| silu | silu.py | ✅ Done |
| relu | relu.py | ✅ Done |
| tanh | tanh.py | ✅ Done |
| transpose | transpose.py | ✅ Done |
| reshape | reshape.py | ✅ Done |
| squeeze | squeeze.py | ✅ Done |
| unsqueeze | unsqueeze.py | ✅ Done |
| cast | cast.py | ✅ Done |
| constant | constant.py | ✅ Done |
| expand | expand.py | ✅ Done |
| flatten | flatten.py | ✅ Done |
| split | split.py | ✅ Done |
| greater | greater.py | ✅ Done |
| less | less.py | ✅ Done |
| equal | equal.py | ✅ Done |
| maximum | maximum.py | ✅ Done |
| minimum | minimum.py | ✅ Done |

### LLM-Specific Kernels

| Kernel | Path | Tier | Status |
|--------|------|------|--------|
| **rmsnorm** | `llm/nvidia/common/rmsnorm.py` | common | ✅ Done |
| **rope** | `llm/nvidia/common/rope.py` | common | ✅ Done |
| **swiglu** | `llm/nvidia/common/swiglu.py` | common | ✅ Done |
| **cross_entropy** | `llm/nvidia/common/cross_entropy.py` | common | ✅ Done |
| **flash_attention** | `llm/nvidia/ampere/flash_attention.py` | ampere | ✅ Done |
| **paged_attention** | `llm/nvidia/ampere/paged_attention.py` | ampere | ✅ Done |

### Image-Specific Kernels

| Kernel | Path | Tier | Status |
|--------|------|------|--------|
| groupnorm | `image/nvidia/common/groupnorm.py` | common | ✅ Done |
| instancenorm | `image/nvidia/common/instancenormalization.py` | common | ✅ Done |
| resize | `image/nvidia/common/resize.py` | common | ✅ Done |
| einsum | `image/nvidia/common/einsum.py` | common | ✅ Done |
| fused_ops | `image/nvidia/common/fused_ops.py` | common | ✅ Done |

---

## Additional Kernels (Common across families)

| Kernel | Tier | Description |
|--------|------|-------------|
| argmax | common | Argmax reduction |
| argmin | common | Argmin reduction |
| batchmatmul | common | Batched matmul |
| bmm | common | Batch matrix multiply |
| ceil | common | Ceiling function |
| clip | common | Clamp/clip values |
| concat | common | Tensor concatenation |
| constant_of_shape | common | Create constant tensor |
| conv2d | common | 2D convolution |
| celu | common | CELU activation |
| elu | common | ELU activation |
| embedding | common | Embedding lookup |
| eq | common | Equal comparison |
| erf | common | Error function |
| floor | common | Floor function |
| gather | common | Gather along axis |
| ge | common | Greater or equal |
| gt | common | Greater than |
| hardsigmoid | common | Hard sigmoid |
| hardswish | common | Hard swish |
| hardtanh | common | Hard tanh |
| index_select | common | Index selection |
| instance_norm | common | Instance normalization |
| le | common | Less or equal |
| leaky_relu | common | Leaky ReLU |
| lt | common | Less than |
| max | common | Max reduction |
| min | common | Min reduction |
| mish | common | Mish activation |
| mm | common | Matrix multiply |
| ne | common | Not equal |
| pow | common | Power function |
| prod | common | Product reduction |
| range_op | common | Range tensor |
| reduce_max | common | Max reduction |
| reduce_mean | common | Mean reduction |
| reduce_min | common | Min reduction |
| reduce_prod | common | Product reduction |
| reduce_sum | common | Sum reduction |
| scaled_dot_product_attention | common | SDPA |
| scatter | common | Scatter operation |
| selu | common | SELU activation |
| shape_op | common | Shape operation |
| slice_op | common | Slice operation |
| softplus | common | Softplus |
| softsign | common | Softsign |
| tile | common | Tile tensor |
| where | common | Where (conditional) |

---

## Notes

1. **Family Organization:** Kernels are organized by model family (image, llm, video, audio) and vendor (nvidia, amd).
2. **Triton Mode Only:** Kernels are ONLY used when running with `--triton` flag. Default execution uses compiled mode (CompiledSequenceV2 + DtypeEngine).
3. **Tier Cascade (Triton Mode):** When `--triton` is enabled, runtime resolves kernels from most specific tier (hopper) to most general (common).
4. **LLM-Specific:** Some kernels (rmsnorm, rope, swiglu, cross_entropy) exist only in LLM family.
5. **Ampere Features:** flash_attention and paged_attention require Ampere or newer GPUs.
6. **Default Execution:** Without `--triton`, all ops are executed via CompiledSequenceV2 with DtypeEngine (AMP autocast rules, no custom kernels).

---

*Kernel Inventory - NeuroBrix Infrastructure*
