tokenspeed-mla

What it is and what it does

TokenSpeed-MLA provides hand-optimized CUDA kernels for Multi-head Latent Attention on NVIDIA Blackwell GPUs. It splits into two main paths: prefill (processing initial prompt tokens) and decode (generating one token at a time). The prefill path offers both a JIT-compiled CuTe DSL backend and an optional pre-compiled binary backend with NVIDIA-internal optimizations; it handles ragged variable-length sequences without padding and supports FP8 E4M3 quantization. The decode path uses a two-kernel split-KV strategy with runtime auto-sizing and compile caching, supporting FP16, BF16, and FP8 inputs while writing BF16 output for downstream stability.

The package targets latency-sensitive serving workloads—particularly coding agents and similar use cases with high request concurrency, short decode steps, and strict time-to-first-token requirements. It includes a fused Triton kernel for K/V packing and FP8 quantization that replaces separate cat and cast operations. Key optimizations include query-token folding for small head counts, reduced shared memory usage via 2CTA UTCMMA instructions, and multi-stage epilogue writes. The package is actively maintained and has no known security vulnerabilities.

Use it for:

• Accelerate token-by-token generation in LLM serving with strict latency budgets, especially for coding agents.
• Batch prefill processing of variable-length prompts without padding overhead using ragged attention.
• Reduce memory footprint and improve throughput by using FP8 quantization in both prefill and decode phases.
• Optimize KV cache utilization in paged attention scenarios with automatic workspace sizing and compile caching.
• Improve tile utilization in small-head decode scenarios by folding query tokens into the head dimension.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides optimized MLA (Multi-head Latent Attention) kernels for NVIDIA Blackwell GPUs, including prefill and decode operations with FP8 quantization and ragged sequence support.

Yes, if you are serving models on NVIDIA Blackwell GPUs and need low-latency token generation. The package is actively maintained, has no vulnerabilities, and offers measurable performance gains over TensorRT-LLM's native MLA implementation—particularly the AOT binary backend. Install friction is moderate due to compiled dependencies; ensure torch, apache-tvm-ffi, nvidia-cutlass-dsl, and tokenspeed-triton are available in your environment.

Install

tokenspeed-mla on PyPI

pip install tokenspeed-mla

uv add tokenspeed-mla

poetry add tokenspeed-mla

Installing tokenspeed-mla

Before you install

Medium install friction due to compiled wheel dependencies (manylinux_2_28 for aarch64 and x86_64). Requires torch, apache-tvm-ffi, nvidia-cutlass-dsl, and tokenspeed-triton. Package is actively maintained with recent releases and no known vulnerabilities.

License in practice

MIT License permits unrestricted use, modification, and distribution with minimal restrictions—suitable for both open-source and commercial projects.

Quickstart

pip install tokenspeed-mla

import torch
from tokenspeed_mla import tokenspeed_mla_decode

out = tokenspeed_mla_decode(
    query=query,
    kv_cache=kv_cache,
    workspace_buffer=workspace_buffer,
    kv_lora_rank=kv_lora_rank,
    qk_rope_head_dim=qk_rope_head_dim,
    block_tables=block_tables,
    seq_lens=seq_lens,
    max_seq_len=max_seq_len,
    softmax_scale=softmax_scale,
)

Requires NVIDIA Blackwell GPU (SM100/SM103); requires Python >=3.10 and torch, apache-tvm-ffi, nvidia-cutlass-dsl, tokenspeed-triton installed.

Verify before relying

• Whether the CuTe DSL JIT backend or AOT binary backend performs better for your specific workload and sequence lengths.
• Compatibility with non-Blackwell NVIDIA architectures or AMD GPUs.
• Memory overhead of workspace buffers and compile caching for different batch and sequence configurations.

Package facts

Evidence: tokenspeed_mla-0.2.5-py3-none-manylinux_2_28_aarch64.whl; tokenspeed_mla-0.2.5-py3-none-manylinux_2_28_x86_64.whl

Tags

More Artificial Intelligence packages

Further reading