TIRx kernels

Curated high-performance TIRx GPU kernels, with support for tirx-harness.

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: kda_forward_portfolio_multishape ⟨sm_100a⟩
• Qwen3-Next TP4 Alpha-MoE FP8: alphamoe_fp8_blockscale_qwen3next ⟨sm_100a⟩
• KDA backward: kda_backward_packed ⟨sm_100a⟩
• KDA decode: kda_decode_multishape ⟨sm_100a⟩
• MSA prefill: msa_prefill_multishape ⟨sm_100a⟩
• MSA decode: msa_decode_multishape ⟨sm_100a⟩
• VSA forward: vsa_multishape ⟨sm_100a⟩
• DeepSeek-V4 sparse MLA: 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
# or, from a checkout:
pip install -e .

The kcoral client for remote benchmarking is installed automatically.

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.