tokenspeed-mla
Speed-of-light TokenSpeed MLA kernels for Blackwell SM100 and SM103.
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires NVIDIA Blackwell GPU (SM100/SM103); requires Python >=3.10 and torch, apache-tvm-ffi, nvidia-cutlass-dsl, tokenspeed-triton installed.
- 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.
License · maintenance · safety
permissive license (permissive) — MIT License permits unrestricted use, modification, and distribution with minimal restrictions—suitable for both open-source and commercial projects.
last release 2026-08-09 (5 days) · last repo commit 2026-08-14 · 1,897 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 1,972,060 downloads/mo, #3,395 on PyPI
Alternatives
Verify before relying
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,
)- 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.
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 on it.
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
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.
Requires NVIDIA Blackwell GPU (SM100/SM103); requires Python >=3.10 and torch, apache-tvm-ffi, nvidia-cutlass-dsl, tokenspeed-triton installed.
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,
)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
| License | permissive license permissive |
| Python support | Supports the current Python release >=3.10 |
| Install friction | Medium. Platform-specific wheel |
| Runtime dependencies | 4 packagesapache-tvm-ffinvidia-cutlass-dsltokenspeed-tritontorch |
| Maintenance | Actively maintained 5 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 1,972,060 / month, #3,395 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
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
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “MLA attention kernels blackwell”
- tokenspeed-mlaProvides optimized MLA (Multi-head Latent Attention) kernels for…
- flashinfer-pythonFlashInfer provides optimized GPU kernels for LLM inference…
- nvidia-cudnn-frontendProvides Python and C++ APIs to NVIDIA's cuDNN library, exposing…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also humming-kernels · nvidia-cudnn-frontend · flash-attn · flashinfer-python · flash-attn-4 · CoLT5-attention · transformer-engine-cu12 · sglang-kernel · flashinfer-cubin · genai-perf