How mlc-ai/tirx-kernels works
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: runner.py Unified runner · tirx_lite tirx-lite DSL · gemm GEMM kernels · norm Normalization
Components 22
Runtime
- Unified runner
tirx_kernels/runner.py
Native kernels
- GEMM kernels
tirx_kernels/gemm - Normalization
tirx_kernels/norm - KDA kernels
tirx_kernels/kda - MSA kernels
tirx_kernels/msa - Sparse attention
tirx_kernels/mla - VSA kernels
tirx_kernels/vsa - MoE kernels
tirx_kernels/moe
Kernel ports
- cuDNN ports
tirx_kernels/ported/cudnn - FlashAttention ports
tirx_kernels/ported/flashattention - FlashInfer ports
tirx_kernels/ported/flashinfer - FlashMLA ports
tirx_kernels/ported/flashmla - DeepGEMM ports
tirx_kernels/ported/deepgemm - MSA ports
tirx_kernels/ported/msa - DeepEP ports
tirx_kernels/ported/deepep - FastCU ports
tirx_kernels/ported/fastcu
Kernel authoring
- tirx-lite DSL
tirx_kernels/tirx_lite
Benchmarking
- Benchmark suite
tirx_kernels/bench_suite
Other components
- Kernel caller
- CUDA GPU
- TIRx compiler
- Kernel output
Connections 26
- 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