TIRx kernels

High-performance GPU kernels authored in tirx-lite and compiled through TIRx.

Kernels

Native kernels are grouped by workload under tirx_kernels/<workload>/. External ports live under tirx_kernels/ported/<upstream>/, with subpackages for upstream entry points. Unless annotated, a kernel runs on sm_100a, sm_103a and sm_107a, and its performance was measured and tuned on sm_100a; every listed architecture is gated by the registry correctness tests. Annotations mark the exceptions: ⟨sm_103a⟩ runs only on that architecture and was tuned there, and ⟨+sm_110a⟩ runs on that architecture in addition to the default three. Annotations follow KERNEL_META["runtime_cuda_archs"] and tests/lint/check_readme_kernels.py keeps them in sync.

Native TIRx

• GEMM: fp16_bf16_gemm ⟨+sm_110a⟩, nvfp4_gemm
• Normalization: rmsnorm
• Distributed: allgather_gemm ⟨sm_100a⟩, gemm_reduce_scatter ⟨sm_100a⟩
• KDA forward: curated_kda_forward_portfolio_multishape ⟨sm_100a⟩
• Qwen3-Next TP4 Alpha-MoE FP8: curated_alphamoe_fp8_blockscale_qwen3next ⟨sm_100a⟩
• KDA backward: curated_kda_backward_packed ⟨sm_100a⟩
• KDA decode: curated_kda_decode_multishape ⟨sm_100a⟩
• MSA prefill: curated_msa_prefill_multishape ⟨sm_100a⟩
• MSA decode: curated_msa_decode_multishape ⟨sm_100a⟩
• VSA forward: curated_vsa_multishape ⟨sm_100a⟩
• DeepSeek-V4 sparse MLA: curated_mla_dsv4_multishape ⟨sm_100a⟩

cuDNN Frontend ports

• Persistent GEMM: cudnn_sm100_dense_blockscaled_gemm_persistent_amax, cudnn_sm100_dense_gemm_persistent_swiglu, cudnn_sm100_dense_blockscaled_gemm_persistent_swiglu_interleaved_quant, cudnn_sm100_dense_blockscaled_gemm_persistent_srelu_quant, cudnn_sm100_dense_blockscaled_gemm_persistent_dsrelu_quant, cudnn_sm100_gemm_proj_rope_mxfp8_bf16in, cudnn_sm100_gemm_proj_rope_mxfp8_mxfp8in
• Grouped GEMM: cudnn_sm100_moe_blockscaled_grouped_gemm_dglu_dbias, cudnn_sm100_moe_grouped_gemm_dglu_dbias
• Linear attention: cudnn_sm100_kda_bprop_f16 ⟨+sm_110a⟩, cudnn_sm100_gdn_prefill_f16 ⟨+sm_110a⟩, cudnn_sm100_gdn_recompute_f16 ⟨+sm_110a⟩, cudnn_sm100_gdn2_prefill_f16 ⟨+sm_110a⟩, cudnn_sm100_gdn2_recompute_f16 ⟨+sm_110a⟩, cudnn_sm100_gdn2_bprop_f16, cudnn_sm100_gdn_bprop_f16
• CSA compression: cudnn_sm100_csa_compressor_fwd
• Flex attention: cudnn_sm100_flex_attention_forward_hd256 ⟨sm_100a sm_103a⟩, cudnn_sm103_flex_attention_forward ⟨sm_103a⟩, cudnn_sm100_flex_attention_backward ⟨sm_100a⟩
• Sparse attention: cudnn_sm100_dsa_sparse_attention_backward, cudnn_sm100_bsa_forward_blk128, cudnn_sm100_bsa_forward_blk64, cudnn_sm100_bsa_forward_combine_blk64, cudnn_sm100_bsa_backward_blk128, cudnn_sm100_bsa_backward_blk64

FlashAttention ports

• Forward: flash_attention4 ⟨+sm_110a⟩, flash_attention4_fp4 ⟨sm_103a⟩
• Backward: flash_attention_backward_sm100 ⟨+sm_110a⟩

FlashInfer ports

Grouped by the FlashInfer Python entry point each port backs.

• flashinfer.activation: act_and_mul ⟨+sm_110a⟩, silu_and_mul_nvfp4_experts_quantize
• flashinfer.cascade: merge_state ⟨+sm_110a⟩
• flashinfer.quantization: nvfp4_quantize, nvfp4_quantize_per_token, mxfp4_quantize ⟨+sm_110a⟩, mxfp8_quantize
• flashinfer.norm: flashinfer_rmsnorm ⟨+sm_110a⟩, flashinfer_rmsnorm_quant, flashinfer_rmsnorm_fp4quant, flashinfer_add_rmsnorm_fp4quant, flashinfer_layernorm ⟨+sm_110a⟩, flashinfer_fused_add_rmsnorm ⟨sm_100a⟩, flashinfer_fused_add_rmsnorm_quant, flashinfer_fused_dit_layernorm, flashinfer_qk_rmsnorm ⟨+sm_110a⟩
• flashinfer.mamba: selective_state_update_stp_simple, selective_state_update_stp_vertical, selective_state_update_stp_horizontal, selective_state_update_mtp_simple ⟨+sm_110a⟩, selective_state_update_mtp_vertical ⟨+sm_110a⟩, selective_state_update_mtp_horizontal ⟨+sm_110a⟩
• flashinfer.kda: recurrent_kda_decode_one_warp ⟨+sm_110a⟩, recurrent_kda_decode_grouped
• flashinfer.gdn_decode: gdn_decode_bf16_ilp4, gdn_decode_bf16_wide_vec_t1, gdn_decode_bf16_wide_vec_mtp, gdn_decode_fp32_mtp_warp
• flashinfer.gdn_prefill: gdn_prefill_sm100, gdn_cp_prefill_sm100
• flashinfer.gemm: bmm_fp8_rubin ⟨sm_107a⟩, dense_blockscaled_gemm_sm107 ⟨sm_107a⟩, grouped_gemm_masked_rubin ⟨sm_107a⟩
• flashinfer.fused_moe: blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion_rubin ⟨sm_107a⟩
• flashinfer.topk: fast_topk_clusters ⟨+sm_110a⟩, filtered_topk ⟨+sm_110a⟩, radix_topk_multi_cta, radix_topk_single_cta, stable_sort_topk_by_value

FlashMLA ports

• Sparse prefill: sparse_flashmla_prefill_head64_phase1, sparse_flashmla_prefill_head128_phase1, sparse_flashmla_prefill_head128_small_topk_phase1
• Sparse decode: sparse_flashmla_decode_head64
• Sparse forward: flash_mla_sparse_fwd ⟨sm_100a⟩

DeepGEMM ports

• Dense and grouped GEMM: deepgemm_sm100_fp8_gemm_1d1d, deepgemm_sm100_m_grouped_fp8_gemm_contiguous, deepgemm_sm100_m_grouped_fp8_gemm_masked, deepgemm_sm100_k_grouped_fp8_gemm_contiguous, deepgemm_sm100_fp8_bmm, deepgemm_sm100_tf32_hc_prenorm_gemm
• MQA logits: deepgemm_sm100_fp4_mqa_logits, deepgemm_sm100_fp8_mqa_logits, deepgemm_sm100_fp4_paged_mqa_logits, deepgemm_sm100_fp8_paged_mqa_logits
• MoE: sm100_fp8_fp4_mega_moe

DeepEP ports

• Elastic communication: deepep_dispatch ⟨sm_100a⟩, deepep_combine ⟨sm_100a⟩

fast.cu ports

• NVFP4 GEMM: fastcu_nvfp4_gemm_gb300 ⟨sm_103a⟩

MSA ports

• Sparse-attention preparation: msa_sparse_prepare_flat_schedule_sm100 ⟨+sm_110a⟩, msa_sparse_prepare_fwd_split_atomic_sm100
• Sparse-attention forward: msa_sparse_atten_fwd_sm100, msa_sparse_atten_fwd_nvfp4_kv_sm100, msa_sparse_atten_fwd_combine_sm100

Performance

Per-workload numbers — our kernel time, every reference impl, and the ref/ours ratio (>1 means ours is faster) — are pinned in tirx_kernels/bench_suite/baseline.md, regenerated on every baseline promotion. See the bench-suite README for how the sweep runs and how to refresh the baseline.

Installation

pip install tirx-kernels==0.1.0rc0
# or, from a checkout:
pip install -e .

The kcoral client for remote benchmarking is installed automatically.

Pin 0.1.0rc0 explicitly: PyPI also contains older uploads numbered 0.1.1, which pip considers newer than this prerelease.

Packaging and releases

Versions come from Git tags through setuptools-scm, as in kcoral. A clean checkout of v0.1.0rc0 builds version 0.1.0rc0; commits after a tag get a development version automatically. Build both the source distribution and a pure-Python wheel with:

python -m pip install build twine
python -m build
python -m twine check --strict dist/*

The build workflow checks the installed wheel on Python 3.10 and 3.12, including kernel discovery and packaged benchmark configurations, baselines, and CUDA reference sources. The wheel does not compile GPU code during installation.

To publish, tag the release commit (for example, v0.1.0rc0) and publish a GitHub Release for that tag, marking release candidates as prereleases. .github/workflows/publish_pypi.yml builds and validates that exact tag, then uploads the wheel and source distribution using PyPI Trusted Publishing. The PyPI publisher for tirx-kernels must use owner mlc-ai, repository TIRx-kernels, workflow publish_pypi.yml, and environment pypi. Configure it in the PyPI publishing settings. Alternatively, as in kcoral, set a PYPI_API_TOKEN Actions secret on the repository or its pypi environment; when present, the workflow uses that token and disables OIDC attestations. The token must have permission to publish tirx-kernels.

The workflow can also be run manually with an existing release tag. For example, after configuring publishing credentials:

gh workflow run publish_pypi.yml --repo mlc-ai/TIRx-kernels --ref main -f ref=v0.1.0rc0

External dependencies

Correctness uses the original upstream implementations. Install the exact, mutually compatible revisions from the repository lock:

python scripts/install_reference_dependencies.py

reference-dependencies.json is the single source of truth for reference revisions and the shared CUTLASS DSL version. The same command installs the pinned pytest/xdist runner. torch and tvm.tirx remain externally managed runtime/compiler dependencies.

Correctness tests import and run these upstream implementations. The bench suite does not launch or time benchmark reference implementations by default (kernel data-preparation helpers may still import their upstream package). Pass --with-references to enable reference launches; a missing enabled reference fails its workload. See tirx_kernels/bench_suite/README.md for the prerequisites and workarounds.

Usage

Command line

# List discovered kernels (with their config labels)
python -m tirx_kernels.registry --format json

# Run correctness tests (optionally filter by kernel / config label)
pytest -n 16 tests/test_correctness.py

# Benchmark
python -m tirx_kernels.bench --kernel nvfp4_gemm
python -m tirx_kernels.bench --kernel nvfp4_gemm --with-references

# Pre-commit regression benchmark sweep on the kcoral benchmark server
# (see tirx_kernels/bench_suite/README.md)
python -m tirx_kernels.bench_suite --server http://127.0.0.1:8901

Programmatic API

Every kernel module exposes a small, uniform interface (see tirx_kernels/_protocol.py):

from tirx_kernels.registry import discover_kernels

kernels = discover_kernels()          # {name: module}
mod = kernels["fp16_bf16_gemm"]

mod.run_test(M=1024, N=1024, K=1024)  # compile + run + correctness check
mod.run_bench(M=1024, N=1024, K=1024) # profile (needs a GPU)

func = mod.get_kernel(M=1024, N=1024, K=1024)  # the TIRx PrimFunc

A native module uses its workload as KERNEL_META["category"] (e.g. kda); a port uses ported.<upstream> (e.g. ported.flashinfer). Registry category filters use those same names. CLI kernel names are independent of module paths.

Each module also provides KERNEL_META: its name, category, exact runtime_cuda_archs, and optional correctness-only reference_requirements. The registry and test harness reject unsupported architectures before compile, and skip correctness before GPU work when a declared reference package, version, or Git source identity is unavailable. CONFIGS contains the test parameter sweep.

License

Except where otherwise noted, this project is licensed under the Apache License 2.0; see LICENSE. Required Apache attribution notices are collected in NOTICE.

Every Python source file carries SPDX tags. Kernel ports derived from third-party projects (cuDNN Frontend, DeepGEMM, DeepEP, fast.cu, FlashMLA, flash-attention, flash-attention-fp4, FlashInfer, MSA) additionally cite the upstream project and the exact commit ported, retain the upstream copyright notice, and declare the combined terms — for example Apache-2.0 AND MIT. Where an upstream license requires its conditions text to travel with the source, that text is kept in the file verbatim. The third-party section at the end of LICENSE lists which components fall under which license, and licenses/ holds the corresponding license texts.