> ## Documentation Index
> Fetch the complete documentation index at: https://snowflake-84d72a0d.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# Architecture

> The L1–L4 operator hierarchy, interface mirroring, and how to extend fastkernels.

fastkernels takes a **top-down, model-driven** approach: rather than collecting synthetic operators in isolation, benchmark tasks are dynamically instantiated from the configuration files of real production models. This ensures every operator is tested with realistic shapes, dtypes, and data flows. The layered hierarchy below makes it possible to benchmark a single kernel replacement while measuring its impact all the way up to full-model inference.

## The L1-L4 Hierarchy

All model operators live under [`tasks/baseline/`](https://github.com/sfc-gh-goliaro/fastkernels/tree/main/tasks/baseline) at four abstraction levels:

* **L1** — Single-kernel ops (e.g. `rms_norm`, `silu_and_mul`, `rotary_emb`, `linear`)
* **L2** — Multi-op blocks (e.g. `LlamaAttention`, `LlamaMLP`, `MixtralMoE`, `QKVParallelLinear`)
* **L3** — Decoder layers (e.g. `LlamaDecoderLayer`, `MixtralDecoderLayer`)
* **L4** — Full models (e.g. `LlamaForCausalLM`, `MixtralForCausalLM`)

Higher levels compose lower ones via standard Python imports:

```
L4/llama.py (LlamaForCausalLM)
  L3/llama_decoder.py (LlamaDecoderLayer)
    L2/attention.py (LlamaAttention)
      L1/store_kvcache.py, L1/flash_attn_*.py
      L2/parallel_linear.py (QKVParallelLinear, RowParallelLinear)
        L1/linear.py, L1/allreduce.py
    L2/llama_mlp.py (LlamaMLP)
      L1/silu_and_mul.py
      L2/parallel_linear.py
    L1/rms_norm.py
```

When you replace an L1 operator, the change propagates upward through every level that uses it. The bench suite traces these dependencies automatically via import graph analysis.

***

## Interface Mirroring

Every module's `__init__` and `forward` signatures are designed to mirror the corresponding vLLM module. This makes it easy to port optimizations between the two codebases and ensures candidate kernels can be validated against a well-known reference.

| fastkernels module           | vLLM equivalent              | `forward` signature                                        |
| ---------------------------- | ---------------------------- | ---------------------------------------------------------- |
| `LlamaAttention`             | `LlamaAttention`             | `(positions, hidden_states) -> Tensor`                     |
| `LlamaDecoderLayer`          | `LlamaDecoderLayer`          | `(positions, hidden_states, residual) -> (Tensor, Tensor)` |
| `QKVParallelLinear`          | `QKVParallelLinear`          | `(x) -> (output, bias)`                                    |
| `RowParallelLinear`          | `RowParallelLinear`          | `(x) -> (output, bias)`                                    |
| `MergedColumnParallelLinear` | `MergedColumnParallelLinear` | `(x) -> (output, bias)`                                    |
| `VocabParallelEmbedding`     | `VocabParallelEmbedding`     | `(input_ids) -> Tensor`                                    |
| `ParallelLMHead`             | `ParallelLMHead`             | `(hidden_states) -> Tensor`                                |
| `SiluAndMul`                 | `SiluAndMul`                 | `(x) -> Tensor`                                            |
| `RMSNorm`                    | `RMSNorm`                    | `(x, residual=None) -> Tensor or (Tensor, Tensor)`         |

All parallel linear layers return `(output, bias)` tuples to match vLLM. `LlamaAttention` stores `rotary_emb` as an `__init__` parameter (not a `forward` argument), matching vLLM's design where each `LlamaAttention` owns its rotary embedding.