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 the MLA decode kernel, small q_len * num_heads configurations can fold a query-token group (fold_sq_factor) into heads for better tile utilization; the scheduler places remaining query groups across the query sequence dimension.

SM107 support requires CuTe DSL 4.8.0 or newer and a compatible CUDA toolkit. Decode compilation and GPU regression tests also reject CuTe DSL warnings.

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 preceding prefill comparison 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.

Collect the performance numbers with this 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 these 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: with num_heads=64, q_seqlen=4, the runtime chooses F=2, so the kernel folds two query tokens into M (H_eff=128) and schedules the remaining two query groups 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.

The packing orders rows 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; pytest skips them 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).

On SM107, an out tensor with dtype torch.float8_e4m3fn selects mla_decode_fp8_sm107.py. The kernel packs query/head rows into M256 tiles and uses 4-CTA clusters to separate QK/softmax from PV, with a 2-CTA fallback. BF16 output keeps using the shared kernel.

The public wrapper uses FP32 for softmax exponentials and row sums. Packed FP16 softmax remains an explicit option for direct kernel construction.

FP8 output supports up to four M256 tiles per request. By default, H must divide 256 and q_len must be a multiple of 256 / H; enable_packed_q=True accepts other shapes such as H96 and partial tiles. The kernel supports causal/non-causal attention, split-KV, LSE, PDL and CUDA graphs; it does not support sliding windows, DCP, or local_visible_lens. Rows with no visible keys produce zero output and -inf LSE.

SM107 FP8-output compilation and GPU regression tests reject CuTe DSL warnings. The mixed-cluster kernel issues CLC queries directly and owns its unswizzled grid and response decoding, avoiding the deprecated example scheduler classes.

H96/Q4 fills 75% of M256 tiles; H96/Q8 fills them completely. M128 packing fills three and six tiles respectively, so relative performance depends on tile utilization as well as batch size and split-KV scheduling.

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 few mbarriers as possible.
• Split the KV loading warp to get more latency-hiding ability. After loading K, V is already in the L2 cache. Loading K of the next tile does not have to wait for the completion of V loading.
• Using multiple stages (sub-tiling) for STG in the epilogue.

Collect the performance numbers with this 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 (such as dtype, d_qk, d_v, causal, LSE, and PDL)
• The wrapped FMHA path enables skip-correction.
• ex2-emulation (disabled by default on B200, and not supported on B300)
• CuTe DSL JIT backend

Input/output dtype behavior:

• The CuTe DSL backend accepts these input dtypes:
  • torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2
  • MLA Prefill supports only 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, or FP8 output on SM107 when out is FP8 (see SM107 decode).
• Supports H <= 128 and 1 <= q_len <= 4; for example, H=64, q_len=4 is supported.
• The wrapper computes and caches split_kv and workspace_size 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

FP8 split-KV partial storage

FP8 decode keeps split-KV partial outputs in FP32 by default. Set TOKENSPEED_MLA_FP16_PARTIALS=1 before importing the package to opt into FP16 partial storage. Accumulation and log-sum-exp remain FP32; the reducer applies output_scale after combining normalized partials. Direct FP8 kernel callers must explicitly specify partial_fp16 and a reducer_max_splits capacity that covers every split in their launches.