# pearblossom-m/flash-swiglu-mlp architecture

> Architecture diagram of the GitHub repository pearblossom-m/flash-swiglu-mlp, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/pearblossom-m/flash-swiglu-mlp
- Repository: https://github.com/pearblossom-m/flash-swiglu-mlp

## Overview

Flash SwiGLU MLP is a GPU-optimized PyTorch/Triton library for the SwiGLU feed-forward block used in LLMs. Callers use either the functional API or a weight-owning module; both enter the same autograd operation. Forward fuses the gate and up projections with SiLU and gating, then uses PyTorch linear for the down projection. Training saves intermediates and computes gradients through fused backward kernels and matrix multiplies; inference skips saving gate/up activations. The README also documents Falcon H1 scaling and shape-bucketed autotuning. Some orchestration source was omitted, so the graph leaves uncertain internal dispatch details out.

Main entry points: [`src/flash_swiglu_mlp/ops.py`](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py) (Autograd function), [`src/flash_swiglu_mlp/module.py`](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/module.py) (FlashSwiGLUMLP), [`src/flash_swiglu_mlp/kernel.py`](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py) (Fused train forward)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_api["Public API"]
  node_module["FlashSwiGLUMLP<br/>[module.py]"]
  node_function["flash_swiglu<br/>[ops.py]"]
end

subgraph group_forward["Forward compute"]
  node_validate["Input preparation<br/>[ops.py]"]
  node_train_fwd["Fused train forward<br/>[kernel.py]"]
  node_infer_fwd["Fused inference forward<br/>[kernel.py]"]
  node_agu["AGU intermediates<br/>[ops.py]"]
  node_down["Down projection<br/>[ops.py]"]
end

subgraph group_backward["Autograd"]
  node_autograd["Autograd function<br/>[ops.py]"]
  node_bwd_dgu["Activation gradients<br/>[kernel.py]"]
  node_weight_grads["Weight gradients<br/>[ops.py]"]
  node_bwd_di["Input gradient<br/>[kernel.py]"]
end

node_caller(("Model / caller"))
node_torch["PyTorch / cuBLAS"]
node_result(("Output / gradients"))

node_caller -->|"calls"| node_module
node_caller -->|"calls"| node_function
node_module -->|"calls"| node_function
node_function -->|"applies"| node_autograd
node_autograd -->|"prepares tensors"| node_validate
node_autograd -->|"launches"| node_train_fwd
node_autograd -->|"launches"| node_infer_fwd
node_train_fwd -->|"stores activations"| node_agu
node_autograd -->|"projects activations"| node_down
node_down -->|"runs linear"| node_torch
node_torch -->|"returns output"| node_result
node_autograd -->|"launches backward"| node_bwd_dgu
node_autograd -->|"computes weight grads"| node_weight_grads
node_weight_grads -->|"uses matmul"| node_torch
node_autograd -->|"launches backward"| node_bwd_di
node_bwd_di -->|"returns input grad"| node_result
node_weight_grads -->|"returns weight grads"| node_result

click node_module "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/module.py"
click node_function "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py"
click node_autograd "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py"
click node_validate "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py"
click node_train_fwd "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py"
click node_infer_fwd "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py"
click node_agu "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py"
click node_down "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py"
click node_bwd_dgu "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py"
click node_weight_grads "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py"
click node_bwd_di "https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_module,node_function toneBlue
class node_validate,node_train_fwd,node_infer_fwd,node_agu,node_down toneAmber
class node_autograd,node_bwd_dgu,node_weight_grads,node_bwd_di toneMint
class node_caller,node_torch,node_result toneIndigo
```

## Components

### Public API

- [FlashSwiGLUMLP](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/module.py): `src/flash_swiglu_mlp/module.py`
- [flash_swiglu](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py): `src/flash_swiglu_mlp/ops.py`

### Forward compute

- [Input preparation](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py): `src/flash_swiglu_mlp/ops.py`
- [Fused train forward](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py): `src/flash_swiglu_mlp/kernel.py`
- [Fused inference forward](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py): `src/flash_swiglu_mlp/kernel.py`
- [AGU intermediates](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py): `src/flash_swiglu_mlp/ops.py`
- [Down projection](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py): `src/flash_swiglu_mlp/ops.py`

### Autograd

- [Autograd function](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py): `src/flash_swiglu_mlp/ops.py`
- [Activation gradients](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py): `src/flash_swiglu_mlp/kernel.py`
- [Weight gradients](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/ops.py): `src/flash_swiglu_mlp/ops.py`
- [Input gradient](https://github.com/pearblossom-m/flash-swiglu-mlp/blob/main/src/flash_swiglu_mlp/kernel.py): `src/flash_swiglu_mlp/kernel.py`

### Other components

- **Model / caller**
- **PyTorch / cuBLAS**
- **Output / gradients**

## Connections

- Model / caller → FlashSwiGLUMLP: calls
- Model / caller → flash_swiglu: calls
- FlashSwiGLUMLP → flash_swiglu: calls
- flash_swiglu → Autograd function: applies
- Autograd function → Input preparation: prepares tensors
- Autograd function → Fused train forward: launches
- Autograd function → Fused inference forward: launches
- Fused train forward → AGU intermediates: stores activations
- Autograd function → Down projection: projects activations
- Down projection → PyTorch / cuBLAS: runs linear
- PyTorch / cuBLAS → Output / gradients: returns output
- Autograd function → Activation gradients: launches backward
- Autograd function → Weight gradients: computes weight grads
- Weight gradients → PyTorch / cuBLAS: uses matmul
- Autograd function → Input gradient: launches backward
- Input gradient → Output / gradients: returns input grad
- Weight gradients → Output / gradients: returns weight grads