TokenSpeed Documentation

The SM100 FMHA (Fused Multi-Head Attention) with relative bias kernels are available in the tokenspeed_kernel package under the module tokenspeed_kernel.ops.attention.cute_dsl.rel_mha. This module provides device kernels, prepass helpers, a FA4 compatibility layer, and runtime-facing rel_* operator modules.

Runtime Components:

• flash_fwd_sm100_bias.py: Contains the device-kernel class, standalone callers/wrappers, and the direct runner.
• fmha_bias_helper.py: A compatibility layer over installed FA4 that provides APIs missing from standard FA4, such as the relative-bias block-range helper and MXFP8 block-scaled GEMM.

Requirements:

• tokenspeed-fa4 (specifically tokenspeed-fa4==4.0.0.post20260510 was verified)
• PyTorch with CUDA
• CUDA Python
• NVIDIA CuTe/CUTLASS DSL
• quack-kernels
• Hardware: Blackwell SM100/SM110 GPU

The tokenspeed-kernel repository contains several high-signal files for interacting with the sparse attention interface:

• interface.py: Public sparse attention interface.
• fp4_indexer_interface.py: Public FP4 indexer block-score interface.
• example.py: End-to-end CSR schedule + attention example with NVTX.
• sparse_index_utils.py: Public CSR build wrapper and reference helpers.
• src/sm100/prepare_k2q_csr.py: SM100 CUDA CSR builder dispatcher.
• src/sm100/fp4_indexer.py: SM100 FP4 indexer kernel classes.
• test_sparse_atten.py: Interface-level tests, benchmark CLI, and profile entrypoint.
• test_fp4_indexer.py: FP4 indexer correctness tests and benchmark CLI.
• Makefile: Setup, test, benchmark, and profiling shortcuts.
• src/sm100/fwd: Forward kernels (prefill).
• src/sm100/fwd_decode: Forward kernels (paged FP8 decode).

TokenSpeed-kernel uses a layered system to decouple high-level API calls from low-level kernel implementations:

1. Registration: Backends (like Triton, Gluon, or vendor libraries) register themselves using the @register_kernel(family, mode, ...) decorator. They declare their supported format_signatures, architecture requirements, non-format traits (e.g., head dim, GQA factor), and a priority band.
2. Registry: All registered kernels are stored in the KernelRegistry.
3. Auto-selection: When a public API is called, select_kernel queries the registry. It filters kernels by capability and traits, then ranks the remaining candidates using an optional SelectionOracle and priority levels. It returns a callable kernel that matches the requested objective (e.g., latency or throughput).

FP8 support in the MSA kernels is limited to identity-scale only.

Specifically, the implementation uses:

• BF16 Q (Query)
• FP8-E4M3 K/V (Key/Value) stages
• Conversion back to BF16 in-kernel.

There are currently no support for K/V descale parameters.

Paged mode is activated by specifying a page_size (8, 16, 32, or 64).

Storage Layouts:

• K and V: (physical_pages, page_size, kv_heads, head_dim).
• mPageTable: A flattened CUDA int32 logical-to-physical mapping.
• mSeqUsedK[b]: Logical KV length for batch b.
• Paged FP8 V: Uses a separate compact SFV tensor with ABI (physical_pages, ceil(page_size / 32), kv_heads, head_dim). Scale groups restart at every physical page.

Limitations:

• Paged MXFP8 Q/K and packed-Q with paged-KV are currently unsupported.
• Paged K remains BF16/FP16 even when V uses FP8 storage.

To improve BMM1 M-dimension utilization in small-head decode scenarios (where num_heads < 128), the kernel uses a fold_sq_factor to partially fold query tokens into the head axis.

The runtime selects the largest factor F such that:

1. q_seqlen % F == 0
2. num_heads * F <= 128

If no such factor F > 1 exists, the kernel schedules the full query sequence dimension directly.

Example: If num_heads=64 and q_seqlen=4, the runtime chooses F=2. This results in:

• H_eff = 128 (64 * 2)
• q_seqlen_eff = 2 (4 / 2) This allows two query tokens to be folded into the M dimension, improving tile utilization.

TokenSpeed allows stamping a weight_version into generation metadata via meta_info["weight_version"]. This is useful for RL trainers to identify which policy version produced a sample.

Updating Versions

• SGLang-compatible requests: update_weights_from_distributed, update_weights_from_tensor, and update_weights_from_disk accept an optional weight_version. The version only changes after a successful update.
• vLLM-compatible requests: finish_weight_update accepts an optional weight_version. The version is deferred until the update chunk is finished to ensure partially updated weights never advertise the new version.

API Endpoints

• GET /get_weight_version: Read the current version.
• POST /update_weight_version: Set the version directly using {"new_version": "..."}.
• GET /model_info: Read the model path and version together.

When deploying large Mixture-of-Experts (MoE) models for high performance, you must explicitly configure several key dimensions. These include:

• Model Identity: model path and revision.
• Memory & Context: context length and KV cache dtype.
• Scheduling: scheduler token and sequence budgets.
• Compute Backends: attention and MoE backends.
• Parallelism Strategies: tensor, data, and expert parallelism.
• Decoding & Parsing: reasoning, tool-call, and speculative decoding parsers.

The KernelRegistry selects kernels based on a priority system:

• Priority Range: An integer in [0, 20). Higher values win.
• Reference Implementation: Lives at priority 0.
• Built-in Optimized Kernels: Typically reside in the 10–18 range.
• Overriding Built-ins: To override a built-in kernel, a plugin must choose a priority strictly higher than the built-in it replaces (e.g., if FlashInfer decode is priority 18, use 19).
• Tie-breaking: If two registrations have the same priority for the same (family, mode), a warning is emitted and selection becomes dependent on load order. Always use distinct priorities to avoid this.
• Error Handling: If a plugin's register() function raises an exception, discovery continues; a UserWarning is emitted, but other plugins will still load.