lightseekorg--tokenspeed/tokenspeed-kernel

TokenSpeed-kernel

TokenSpeed-kernel aims to provide a collection of the best portable and performant kernels for multi-silicon AI inference. It features:

• A clean layered API for maximal structured flexibility
• Kernel registration and selection logic to decouple complexity and increase reuse
• Plugin mechanism for multi-silicon extensibility
• A minimal list of curated dependencies for fast iteration

TokenSpeed-kernel is pip-installable on its own and can be directly used by others.

Design Goals

TokenSpeed-kernel is designed with the following functionality goals in mind:

• Support various kernels in AI models (attention, MoE, etc.)
• Support multiple silicon vendors and generations
• Marry default portability and performance solutions

In addition, to have a better devflow for fast iteration:

• Provide unified infra to verify and debug kernel numerics standalone
• Provide unified infra to run and benchmark kernels standalone
• Support tracing shapes and profiling workloads at runtime
• Stay forward-looking, with guardrails for agentic devflow

Overall Design

With the above goals in mind, we have made the following opinionated design choices (still evolving; subject to change):

Layered system

                       public API  (mha_prefill, mm, moe_fused, ...)
                                       │
                           ┌───────────┴───────────┐
                           │     select_kernel     │  (family, mode, format_signature, traits, ...)
                           └───────────┬───────────┘
                                       │ queries
                            ┌──────────┴──────────┐
                            │   KernelRegistry    │   ← @register_kernel(...) populates this
                            └──────────┬──────────┘
                                       │
       ┌──────────────┬────────────────┼────────────────┬───────────────┐
   attention         gemm             moe             norm     ...   (op family)
       │              │                │                │
  ┌────┼────┐    ┌────┼────┐      ┌────┼────┐      ┌────┼────┐
  triton         triton           triton           triton             ← in-tree portable JIT
  gluon          (...)            cute_dsl         (...)              ← in-tree perf JIT
  flash_mla      flashinfer       (...)            (...)              ← vendor library wrappers
                                ...
       │              │                │                │
       └──────────────┴────────────────┴────────────────┴── reference (PyTorch ground truth)

• Registration — backends register with @register_kernel(family, mode, ...), declaring supported format_signatures, arch capability requirements, non-format traits (head dim, GQA factor, ...), and a priority band.
• Auto-selection — select_kernel filters by capability and traits, ranks the survivors with an optional per-family SelectionOracle and priority, and returns a callable. Selection accepts an objective (latency, throughput, determinism, portability) and supports per-call override= plus config-file overrides for development.

Directory structure

tokenspeed_kernel/
  __init__.py            # Public API re-exports
  platform.py            # PlatformInfo, capability detection
  signature.py           # TensorFormat, ScaleFormat, FormatSignature
  registry.py            # KernelRegistry, register_kernel, Priority bands
  selection.py           # select_kernel, oracles, overrides
  profiling.py           # ShapeCapture, kernel_scope, Proton bootstrap
  _triton.py             # Single import point for the vendored Triton fork

  ops/
    attention/   { triton/, flash_attn/, ... }
    gemm/        { triton.py, trtllm.py, ... }
    moe/         { triton.py, deepep.py, triton_kernels.py, ... }
    ...

  numerics/              # Reference impls + tolerance + comparison + CLI
    reference/           # PyTorch ground-truth kernels
  benchmark/             # Unified runner, throughput model, report, CLI
  plugins/               # Out-of-tree backend discovery
  thirdparty/            # Vendored / wrapped third-party kernel sources

Each ops/<family>/ directory holds peer subdirectories — one per solution. A solution is either an in-tree JIT kernel (Triton/Gluon/CuteDSL), or a thin wrapper around an external library. All of them register through the same decorator and are scored by the same selection logic, so adding a backend is one new file in the right family folder.

Solution choices

• Triton — in-tree; default portable JIT path for various kernels
• Gluon / CuteDSL — in-tree; performant JIT path for key kernels
• Vendor libraries — wrapped (FlashAttention, TRT-LLM, etc.); no in-tree C++ build
• PyTorch reference — under numerics/reference/; never auto-selects over a real backend but always available as ground truth

Overall we carefully curate external dependencies and actively re-evaluate their inclusion, in order to maintain minimal dependencies and enable faster iteration.

Numerics, benchmarking, profiling

• python -m tokenspeed_kernel.numerics — dtype-aware tolerances, standard input generators, and a comparison/bisect flow that pits any registered kernel against the reference impl.
• python -m tokenspeed_kernel.benchmark — unified timing, throughput (FLOPs / bytes) per op family, tabular reports, and Proton integration.
• Runtime shape capture feeds replay and tuning workflows; kernel_scope scopes are visible in Proton/Chrome traces.

Plugins

python -m tokenspeed_kernel.plugins lists discovered out-of-tree backends. Plugins register via the same @register_kernel decorator from their own package, set their own priority, and participate in selection like in-tree backends. See tokenspeed_kernel/plugins/README.md.

Public API

from tokenspeed_kernel import (
    mha_prefill, mha_prefill_with_kvcache, mha_decode_with_kvcache,
    gdn_chunk_prefill,
    mm,
    moe_route, moe_dispatch, moe_experts, moe_combine, moe_fused,
    ...
)

Using the above platform and solution-agnostic public APIs can get the most value out of TokenSpeed-kernel; but one can also directly call into a specific solution under ops/<family>/, or manually select_kernel with targeted filters:

from tokenspeed_kernel.selection import select_kernel, kernel_override

For platform checks:

from tokenspeed_kernel.platform  import current_platform