tokenspeed-attention: add TokenSpeed MHA/MLA Triton kernels

tokenspeed-attention/CARD.md

@@ -0,0 +1,55 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Source code
+
+Source code of this kernel originally comes from {{ upstream }} and it was repurposed for compatibility with `kernels`.
+{% endif %}

tokenspeed-attention/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 LightSeek Foundation <contact@lightseek.org>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tokenspeed-attention/README.md

@@ -0,0 +1,75 @@
+---
+license: mit
+tags:
+- kernels
+- attention
+- triton
+- cuda
+- rocm
+---
+
+# TokenSpeed Attention
+
+Portable Triton attention kernels from
+[lightseekorg/tokenspeed](https://github.com/lightseekorg/tokenspeed), repackaged
+for the Hugging Face `kernels` library.
+
+This package exposes the Tokenspeed portable MHA and MLA attention paths without
+requiring the full Tokenspeed dispatcher:
+
+- `mha_prefill`
+- `mha_extend_with_kvcache`
+- `mha_decode_with_kvcache`
+- `mla_prefill`
+- `mla_decode_with_kvcache`
+- `attn_merge_state`
+
+The source was ported from `tokenspeed-kernel` commit
+`1492030a2a02d32bc7011645a74d2d691e99c2e6`.
+
+## Requirements
+
+- PyTorch with CUDA or ROCm support
+- Triton
+- NVIDIA Ampere or newer for the CUDA path, or AMD CDNA-class ROCm GPUs for the
+  ROCm path
+
+## Usage
+
+```python
+import torch
+from kernels import get_kernel
+
+attention = get_kernel("kernels-community/tokenspeed-attention")
+
+seqlens = [128, 256]
+cu_seqlens_cpu = [0]
+for seqlen in seqlens:
+    cu_seqlens_cpu.append(cu_seqlens_cpu[-1] + seqlen)
+
+cu_seqlens = torch.tensor(cu_seqlens_cpu, device="cuda", dtype=torch.int32)
+total = cu_seqlens_cpu[-1]
+
+q = torch.randn(total, 16, 128, device="cuda", dtype=torch.bfloat16)
+k = torch.randn(total, 4, 128, device="cuda", dtype=torch.bfloat16)
+v = torch.randn(total, 4, 128, device="cuda", dtype=torch.bfloat16)
+
+out = attention.mha_prefill(
+    q,
+    k,
+    v,
+    cu_seqlens,
+    cu_seqlens_cpu,
+    max(seqlens),
+)
+```
+
+## Notes
+
+This is a direct kernel-level port. It intentionally omits Tokenspeed's runtime
+registry, backend selection, profiling, and plugin system. Call the exported
+functions directly.
+
+## License
+
+MIT. See `LICENSE`.

tokenspeed-attention/build.toml

@@ -0,0 +1,12 @@
+[general]
+name = "tokenspeed-attention"
+version = 1
+license = "MIT"
+backends = ["cuda", "rocm"]
+
+[general.hub]
+repo-id = "kernels-community/tokenspeed-attention"
+
+[torch-noarch]
+
+[kernel]

tokenspeed-attention/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for tokenspeed-attention kernels";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}