# mlc-ai/tirx-kernels architecture

> Architecture diagram of the GitHub repository mlc-ai/tirx-kernels, 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/mlc-ai/tirx-kernels
- Repository: https://github.com/mlc-ai/tirx-kernels

## Overview

TIRx kernels is a GPU kernel library: callers select a workload implementation, which is compiled and run to produce a kernel result. Its distinct workload families cover dense and distributed GEMM, normalization, sparse and linear attention, and MoE; ports also preserve upstream kernel families. The README documents a unified runner for kernel testing and benchmarking. The graph groups implementations by workload and upstream, with the tirx-lite DSL as an internal authoring layer. Runner and benchmark-suite internals were not fully sampled, so uncertain wiring is omitted.

Main entry points: [`tirx_kernels/runner.py`](https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/runner.py) (Unified runner), [`tirx_kernels/tirx_lite`](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/tirx_lite) (tirx-lite DSL), [`tirx_kernels/gemm`](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/gemm) (GEMM kernels), [`tirx_kernels/norm`](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/norm) (Normalization)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Runtime"]
  node_runner["Unified runner<br/>[runner.py]"]
end

subgraph group_native["Native kernels"]
  node_gemm["GEMM kernels"]
  node_norm["Normalization"]
  node_kda["KDA kernels"]
  node_msa["MSA kernels"]
  node_sparse["Sparse attention"]
  node_vsa["VSA kernels"]
  node_moe["MoE kernels"]
end

subgraph group_ports["Kernel ports"]
  node_cudnn["cuDNN ports"]
  node_flashattention["FlashAttention ports"]
  node_flashinfer["FlashInfer ports"]
  node_flashmla["FlashMLA ports"]
  node_deepgemm["DeepGEMM ports"]
  node_msa_ports["MSA ports"]
  node_deepep["DeepEP ports"]
  node_fastcu["FastCU ports"]
end

subgraph group_authoring["Kernel authoring"]
  node_lite["tirx-lite DSL"]
end

subgraph group_benchmark["Benchmarking"]
  node_bench["Benchmark suite"]
end

node_caller(("Kernel caller"))
node_gpu(("CUDA GPU"))
node_tirx["TIRx compiler"]
node_result["Kernel output"]

node_caller -->|"invokes"| node_runner
node_runner -->|"dispatches"| node_gemm
node_runner -->|"dispatches"| node_norm
node_runner -->|"dispatches"| node_kda
node_runner -->|"dispatches"| node_msa
node_runner -->|"dispatches"| node_sparse
node_runner -->|"dispatches"| node_vsa
node_runner -->|"dispatches"| node_moe
node_runner -->|"dispatches"| node_cudnn
node_runner -->|"dispatches"| node_flashattention
node_runner -->|"dispatches"| node_flashinfer
node_runner -->|"dispatches"| node_flashmla
node_runner -->|"dispatches"| node_deepgemm
node_runner -->|"dispatches"| node_msa_ports
node_runner -->|"dispatches"| node_deepep
node_runner -->|"dispatches"| node_fastcu
node_gemm -->|"authored with"| node_lite
node_norm -->|"authored with"| node_lite
node_kda -->|"authored with"| node_lite
node_msa -->|"authored with"| node_lite
node_sparse -->|"authored with"| node_lite
node_vsa -->|"authored with"| node_lite
node_moe -->|"authored with"| node_lite
node_lite -->|"compiled through"| node_tirx
node_runner -->|"runs on"| node_gpu
node_runner -->|"returns"| node_result

click node_runner "https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/runner.py"
click node_gemm "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/gemm"
click node_norm "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/norm"
click node_kda "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/kda"
click node_msa "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/msa"
click node_sparse "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/mla"
click node_vsa "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/vsa"
click node_moe "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/moe"
click node_cudnn "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/cudnn"
click node_flashattention "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/flashattention"
click node_flashinfer "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/flashinfer"
click node_flashmla "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/flashmla"
click node_deepgemm "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/deepgemm"
click node_msa_ports "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/msa"
click node_deepep "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/deepep"
click node_fastcu "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/fastcu"
click node_lite "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/tirx_lite"
click node_bench "https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/bench_suite"

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_runner toneBlue
class node_gemm,node_norm,node_kda,node_msa,node_sparse,node_vsa,node_moe toneAmber
class node_cudnn,node_flashattention,node_flashinfer,node_flashmla,node_deepgemm,node_msa_ports,node_deepep,node_fastcu toneMint
class node_lite toneRose
class node_bench,node_caller,node_gpu,node_tirx,node_result toneIndigo
```

## Components

### Runtime

- [Unified runner](https://github.com/mlc-ai/tirx-kernels/blob/main/tirx_kernels/runner.py): `tirx_kernels/runner.py`

### Native kernels

- [GEMM kernels](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/gemm): `tirx_kernels/gemm`
- [Normalization](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/norm): `tirx_kernels/norm`
- [KDA kernels](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/kda): `tirx_kernels/kda`
- [MSA kernels](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/msa): `tirx_kernels/msa`
- [Sparse attention](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/mla): `tirx_kernels/mla`
- [VSA kernels](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/vsa): `tirx_kernels/vsa`
- [MoE kernels](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/moe): `tirx_kernels/moe`

### Kernel ports

- [cuDNN ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/cudnn): `tirx_kernels/ported/cudnn`
- [FlashAttention ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/flashattention): `tirx_kernels/ported/flashattention`
- [FlashInfer ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/flashinfer): `tirx_kernels/ported/flashinfer`
- [FlashMLA ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/flashmla): `tirx_kernels/ported/flashmla`
- [DeepGEMM ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/deepgemm): `tirx_kernels/ported/deepgemm`
- [MSA ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/msa): `tirx_kernels/ported/msa`
- [DeepEP ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/deepep): `tirx_kernels/ported/deepep`
- [FastCU ports](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/ported/fastcu): `tirx_kernels/ported/fastcu`

### Kernel authoring

- [tirx-lite DSL](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/tirx_lite): `tirx_kernels/tirx_lite`

### Benchmarking

- [Benchmark suite](https://github.com/mlc-ai/tirx-kernels/tree/main/tirx_kernels/bench_suite): `tirx_kernels/bench_suite`

### Other components

- **Kernel caller**
- **CUDA GPU**
- **TIRx compiler**
- **Kernel output**

## Connections

- Kernel caller → Unified runner: invokes
- Unified runner → GEMM kernels: dispatches
- Unified runner → Normalization: dispatches
- Unified runner → KDA kernels: dispatches
- Unified runner → MSA kernels: dispatches
- Unified runner → Sparse attention: dispatches
- Unified runner → VSA kernels: dispatches
- Unified runner → MoE kernels: dispatches
- Unified runner → cuDNN ports: dispatches
- Unified runner → FlashAttention ports: dispatches
- Unified runner → FlashInfer ports: dispatches
- Unified runner → FlashMLA ports: dispatches
- Unified runner → DeepGEMM ports: dispatches
- Unified runner → MSA ports: dispatches
- Unified runner → DeepEP ports: dispatches
- Unified runner → FastCU ports: dispatches
- GEMM kernels → tirx-lite DSL: authored with
- Normalization → tirx-lite DSL: authored with
- KDA kernels → tirx-lite DSL: authored with
- MSA kernels → tirx-lite DSL: authored with
- Sparse attention → tirx-lite DSL: authored with
- VSA kernels → tirx-lite DSL: authored with
- MoE kernels → tirx-lite DSL: authored with
- tirx-lite DSL → TIRx compiler: compiled through
- Unified runner → CUDA GPU: runs on
- Unified runner → Kernel output: returns