TokenSpeed-MLA

Speed-of-light TokenSpeed MLA kernels for SM100/SM103/SM107 with:

• MLA prefill:
  • CuTe DSL JIT backend for ragged varlen FMHA (no padding)
  • BF16 output, optional LSE output, causal/non-causal modes, PDL support
• MLA decode:
  • CuTe DSL decode kernels for FP16/BF16/FP8 input paths
  • FP8 decode writes BF16 output for better downstream stability
  • Split-KV + workspace path with runtime auto-sizing and compile caching
• MLA K/V pack + FP8 quantize:
  • Fused Triton kernel replacing cat + cast + cast in chunked prefill
  • Supports strided views and optional pre-allocated output buffers

This package includes performance-oriented optimizations for latency-sensitive serving workloads, especially coding agent style use cases with high request concurrency, short decode steps, and strict time-to-first-token/next-token requirements. For MLA decode kernel, small q_len * num_heads configurations can fold a query-token group (fold_sq_factor) into heads for better tile utilization; remaining query groups are scheduled across the query sequence dimension.

SM107 support requires CuTe DSL 4.8.0 or newer and a compatible CUDA toolkit.

Performance Numbers

Prefill Performance

Where:

use case 1: batch_size = 1, seqlen_qo = 8 * 1024, seqlen_kv = 8 * 1024
use case 2: batch_size = 1, seqlen_qo = 8 * 1024, seqlen_kv = 32 * 1024
use case 3: batch_size = 1, seqlen_qo = 8 * 1024, seqlen_kv = 64 * 1024
use case 4: batch_size = 4, seqlen_qo = 512,      seqlen_kv = 80 * 1024
use case 5: batch_size = 4, seqlen_qo = 1024,     seqlen_kv = 80 * 1024

The prefill comparison above includes historical results from an AOT implementation. Current releases ship the public CuTe DSL JIT implementation; the historical AOT backend is not included in the package.

The performance numbers can be collected using the following command line:

python ./tokenspeed-mla/python/tokenspeed_mla/fmha.py \
  --is_causal \
  --bottom_right_align \
  --in_dtype Float8E4M3FN \
  --out_dtype Float8E4M3FN \
  --q_shape 1,8192,128,192 \
  --k_shape 1,8192,128,192 \
  --warmup_iterations 10 \
  --iterations 10 \
  --skip_ref_check

Decode Performance

In the above test cases, q_seqlen = 4 and kv_seqlen = 80K.

TensorRT-LLM uses a single kernel for MLA decode, which appears to adopt a swap-AB strategy in the tested cases. In contrast, TokenSpeed’s MLA decode kernel uses a two-kernel implementation: one kernel computes the MLA decode with split-KV, and a second kernel performs the reduction of the split-KV partial results.

Key Optimization of TokenSpeed MLA decode kernel: Group q_seqlen and num_heads into BMM1 M

In mla_decode.py, mla_decode_fp16.py, and mla_decode_fp8.py, decode uses fold_sq_factor to partially fold query tokens into the head axis when num_heads < 128. q_seqlen can be any positive length; the runtime chooses the largest factor F such that: q_seqlen % F == 0 and num_heads * F <= 128. If no factor greater than one divides q_seqlen, the kernel does not fold and schedules the full query sequence dimension directly.

The folded execution shape becomes:

• H_eff = num_heads * F
• q_seqlen_eff = q_seqlen / F

This improves BMM1 M-dimension utilization and reduces tile waste in small-head decode scenarios, especially token-by-token agent traffic. Example: num_heads=64, q_seqlen=4 chooses F=2, so two query tokens are folded into M (H_eff=128) and the remaining two query groups are scheduled on the scheduler second dimension (q_seqlen_eff=2).

The public tokenspeed_mla_decode also accepts enable_packed_q=True to opt into continuous query/head packing on the FP8 and FP16/BF16 M128 paths, adapted from FlashInfer PR #4178. The default is False, preserving the folded-query implementation. M64 and token-gapped Q/output views continue to use that implementation even when the option is enabled.

Packed rows are ordered as query_token * num_heads + head. Each 2-CTA group owns 128 consecutive rows, including across query boundaries. Thus H96/Sq4 uses three query tiles instead of four, and H96/Sq8 uses six instead of eight. Only the final tile may contain padding. This is a tensor view transformation, with no additional packing kernel. The kernel uses per-row causal positions and predicates partial output/LSE rows.

For packed queries, auto split-KV uses ceil(H * q_len / 128) query tiles, and workspace needs B * 128 * ceil(H * q_len / 128) * split_kv * 513 * 4 bytes for D512/FP32 partials, or zero when split-KV is one. Callers enabling this option must provide sufficient workspace. Output shape, output dtype (BF16 for FP8; input dtype for FP16/BF16), and base-2 LSE semantics are unchanged. The FP16/BF16 implementation retains its existing split-KV heuristic, reducer capacity and PDL waits. The option is part of the compile cache key. Existing direct kernel callers retain the old layout; opting in through the public wrapper keeps tiling and workspace consistent.

Sliding-window and DCP masking remain supported; window boundaries use the packed row's original query-token position.

Regression coverage is in tests/test_mla_decode.py. It checks packed-query geometry, split-KV workspace sizing, FP8/FP16/BF16 outputs and LSE, reducer variants, CUDA-graph replay, sliding windows and DCP. From the repository root, select this checkout's sources explicitly:

PYTHONPATH=tokenspeed-mla/python python -m pytest -q tokenspeed-mla/tests/test_mla_decode.py

GPU cases require SM100, SM103 or SM107 and are skipped on other devices. Add -k 'not TestGPU and not TestCompile' for CPU checks, -k TestCompile for compilation checks across all three architectures, or -k TestGPU for GPU checks. Compilation checks require a CuTe DSL and CUDA toolchain that support the target architecture, even when no GPU is visible.

SM107 decode

FP8 and FP16/BF16 decode share the existing kernel classes, scheduler, split-KV reducer and masking logic across SM100, SM103 and SM107. The public wrapper selects the architecture from query.device and includes it in the compile cache key and compiler target. Direct kernel construction requires an explicit compute_capability=(10, 0), (10, 3) or (10, 7).

Other optimizations include:

• FP8 split-KV candidates are normalized to nonempty K partitions before workspace allocation and kernel launch. The reducer uses a 32/64-split capacity for the M128/M64 paths and selects 1/2/4 disjoint D512 output bands when the real output rows do not fill the GPU. These changes adapt the split-KV and reducer optimizations from FlashInfer PR #4178 to TokenSpeed's folded-query layout; both reducer settings are included in the compile cache.
• Using 2CTA UTCMMA instruction to reduce shared memory usage.
• Try to use as less mbarrier as possible.
• Split kv loading warp to get more latency hiding ability. After loading K, V is already in the L2 cache. Loading K of next tile will not have to wait for the completion of V loading.
• Using multiple stage (sub-tiling) for STG in epilogue.

The performance numbers can be collected using the following command line:

python ./tokenspeed-mla/python/tokenspeed_mla/mla_decode_fp8.py \
  --batch_size 4 \
  --softmax_scale 0.07216882 \
  --page_size 64 \
  --seq_len_k 81920 \
  --in_dtype Float8E4M3FN \
  --out_dtype Float8E4M3FN \
  --seq_len_q 4 \
  --warmup_iterations 1 \
  --iterations 10 \
  --num_heads 16 \
  --skip_ref_check

Kernel Capability Summary

MLA Prefill (tokenspeed_mla_prefill)

What it supports:

• Ragged varlen prefill without padding:
  • Q: [sum(q_lens), h_q, d_qk]
  • K: [sum(kv_lens), h_k, d_qk]
  • V: [sum(kv_lens), h_k, d_v]
• Different Q/KV sequence packs (cum_seq_lens_q and cum_seq_lens_kv can differ)
• Causal and non-causal execution
• Optional LSE return (return_lse=True)
• PDL enable/disable (enable_pdl)
• Kernel compile cache keyed by static config (dtype, d_qk, d_v, causal, LSE, PDL, etc.)
• Skip-correction is enabled in the wrapped FMHA path.
• ex2-emulation (disabled by default on B200, and not supported on B300)
• CuTe DSL JIT backend

Input/output dtype behavior:

• CuTe DSL backend accepts input dtypes supported :
  • torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2
  • MLA Prefill only support torch.float8_e4m3fn
• Prefill output tensor is BF16 (torch.bfloat16)
• Optional LSE output is FP32

MLA Decode (tokenspeed_mla_decode)

What it supports:

• Query shape: [B, q_len, H, kv_lora_rank + qk_rope_head_dim]
• KV cache shape:
  • 3D: [num_pages, page_size, D_total]
  • 4D accepted and normalized internally
• Auto split_kv + workspace sizing and caching
• Supports FP16/BF16/FP8; FP8 path writes BF16 output.
• Supports H <= 128 and 1 <= q_len <= 4; for example, H=64, q_len=4 is supported.
• split_kv and workspace_size are computed and cached from runtime shape/device info.
• is_var_seq, is_persistent, and enable_pdl affect scheduling/compile variants.
• causal_mask supports causal and non-causal execution on FP16/BF16/FP8 paths.
• window_left bounds each block row's history: row i sees keys [max(0, K - q_len - window_left + i), k_bound), the whole q_len block plus window_left tokens of context. The kernel starts its KV walk at the window rather than at key 0, so cost tracks the window, not the cache. -1 (the default) is full history and compiles the same kernel it always did.
• Optional out tensor reuse
• is_var_seq and enable_pdl controls

Minimal Usage

1) Decode

import torch
from tokenspeed_mla import tokenspeed_mla_decode

# query: [B, q_len, H, D_qk]
# kv_cache: [num_pages, page_size, D_total]
out = tokenspeed_mla_decode(
    query=query,
    kv_cache=kv_cache,
    workspace_buffer=workspace_buffer,  # torch.int8, 1D
    kv_lora_rank=kv_lora_rank,
    qk_rope_head_dim=qk_rope_head_dim,
    block_tables=block_tables,          # [B, max_pages]
    seq_lens=seq_lens,                  # [B]
    max_seq_len=max_seq_len,
    softmax_scale=softmax_scale,
    enable_pdl=False,
)

2) Prefill

import torch
from tokenspeed_mla import tokenspeed_mla_prefill

# query: [sum(q_lens), h_q, d_qk]
# key:   [sum(kv_lens), h_k, d_qk]
# value: [sum(kv_lens), h_k, d_v]
out, lse = tokenspeed_mla_prefill(
    query=query,
    key=key,
    value=value,
    seq_lens=seq_lens,
    cum_seq_lens=cum_seq_lens_kv,
    max_seq_len=max_kv_len,
    batch_size=batch_size,
    softmax_scale=softmax_scale,
    is_causal=True,
    return_lse=True,
    cum_seq_lens_q=cum_seq_lens_q,  # optional, when Q/KV lengths differ
    max_seq_len_q=max_q_len,        # optional
    enable_pdl=False,
)

Releases

The release-tokenspeed-mla workflow builds a source-only py3-none-any wheel from this repository and publishes it to PyPI. Packaging does not require a GPU or CUDA compiler; the kernels compile with CuTe DSL and Triton at runtime.

Before the first release through this workflow, configure a PyPI Trusted Publisher for the existing tokenspeed-mla project:

• Owner: lightseekorg
• Repository: tokenspeed
• Workflow filename: release-tokenspeed-mla.yml
• Environment: pypi

Release steps:

1. Merge kernel changes, then update [project].version in tokenspeed-mla/pyproject.toml when a release is needed. Prefer a separate version-bump PR; multiple code changes can share one release.
2. After merging the version bump, dispatch from main: gh workflow run release-tokenspeed-mla.yml -R lightseekorg/tokenspeed --ref main. The workflow refuses versions already present on PyPI and checks the wheel's metadata, JIT sources, and license notices before publishing.
3. Wait for PyPI publication, then use update-tokenspeed-kernel-mla.yml with mla_version=<version> to open the dependency-update PR.

For local packaging checks, use Python 3.12 and install build, packaging, pytest, and twine in a virtual environment. From the repository root:

(cd tokenspeed-mla && python -m pytest tests/test_release.py -q)
python -m build tokenspeed-mla --wheel --outdir dist
python -m twine check --strict dist/*
python tokenspeed-mla/scripts/check_release.py --package-dir tokenspeed-mla --dist-dir dist