diff --git a/tokenspeed-attention/CARD.md b/tokenspeed-attention/CARD.md new file mode 100644 index 00000000..3b7169c5 --- /dev/null +++ b/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 %} diff --git a/tokenspeed-attention/LICENSE b/tokenspeed-attention/LICENSE new file mode 100644 index 00000000..e95f89c0 --- /dev/null +++ b/tokenspeed-attention/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 LightSeek Foundation + +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. diff --git a/tokenspeed-attention/README.md b/tokenspeed-attention/README.md new file mode 100644 index 00000000..760ba1a8 --- /dev/null +++ b/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`. diff --git a/tokenspeed-attention/build.toml b/tokenspeed-attention/build.toml new file mode 100644 index 00000000..532ecab5 --- /dev/null +++ b/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] diff --git a/tokenspeed-attention/flake.lock b/tokenspeed-attention/flake.lock new file mode 100644 index 00000000..e81f2c91 --- /dev/null +++ b/tokenspeed-attention/flake.lock @@ -0,0 +1,117 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1767039857, + "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs", + "rust-overlay": "rust-overlay" + }, + "locked": { + "lastModified": 1781689680, + "narHash": "sha256-h1SSwX195HdbFHvvPl1nBVu5794YYoyXiYXJOCWGVsA=", + "owner": "huggingface", + "repo": "kernels", + "rev": "19a568f64aa291880427d4ab56ee48bfeb26234b", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernels", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1776927958, + "narHash": "sha256-XOzEtft7E0P6TgQViLUOQeGHlEYiQ0+FY24BPEksj6s=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "fec2c46cca5bf9767486a290abae51200b656d69", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable-small", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "rust-overlay": { + "inputs": { + "nixpkgs": [ + "kernel-builder", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1776914043, + "narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=", + "owner": "oxalica", + "repo": "rust-overlay", + "rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3", + "type": "github" + }, + "original": { + "owner": "oxalica", + "repo": "rust-overlay", + "type": "github" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/tokenspeed-attention/flake.nix b/tokenspeed-attention/flake.nix new file mode 100644 index 00000000..d0c46058 --- /dev/null +++ b/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 = ./.; + }; +} diff --git a/tokenspeed-attention/tests/test_attention.py b/tokenspeed-attention/tests/test_attention.py new file mode 100644 index 00000000..3b8398ac --- /dev/null +++ b/tokenspeed-attention/tests/test_attention.py @@ -0,0 +1,165 @@ +import math + +import pytest +import torch + + +def _make_cu_seqlens(seqlens, device): + cu_cpu = [0] + for seqlen in seqlens: + cu_cpu.append(cu_cpu[-1] + seqlen) + return cu_cpu, torch.tensor(cu_cpu, device=device, dtype=torch.int32) + + +def _varlen_attention_reference( + q, + k, + v, + cu_q, + cu_kv, + *, + causal, + scale, +): + out = torch.empty((q.shape[0], q.shape[1], v.shape[-1]), device=q.device) + kv_group = q.shape[1] // k.shape[1] + for batch in range(len(cu_q) - 1): + q_start, q_end = cu_q[batch], cu_q[batch + 1] + kv_start, kv_end = cu_kv[batch], cu_kv[batch + 1] + sq = q_end - q_start + sk = kv_end - kv_start + for head in range(q.shape[1]): + kv_head = head // kv_group + scores = ( + q[q_start:q_end, head].float() + @ k[kv_start:kv_end, kv_head].float().T + ) * scale + if causal: + q_idx = torch.arange(sq, device=q.device).view(-1, 1) + k_idx = torch.arange(sk, device=q.device).view(1, -1) + scores = scores.masked_fill(k_idx > q_idx + (sk - sq), float("-inf")) + probs = torch.softmax(scores, dim=-1) + out[q_start:q_end, head] = probs @ v[kv_start:kv_end, kv_head].float() + return out.to(q.dtype) + + +def _paged_mla_reference( + q, + kv_cache, + page_table, + cache_seqlens, + *, + kv_lora_rank, + qk_rope_head_dim, + softmax_scale, +): + out = torch.empty(q.shape[:-1] + (kv_lora_rank,), device=q.device) + lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32) + page_size = kv_cache.shape[1] + for batch in range(q.shape[0]): + tokens = [] + for pos in range(int(cache_seqlens[batch].item())): + page = int(page_table[batch, pos // page_size].item()) + tokens.append(kv_cache[page, pos % page_size, 0]) + kv = torch.stack(tokens).float() + k_latent = kv[:, :kv_lora_rank] + k_rope = kv[:, kv_lora_rank : kv_lora_rank + qk_rope_head_dim] + for head in range(q.shape[2]): + q_head = q[batch, 0, head].float() + scores = ( + k_latent @ q_head[:kv_lora_rank] + + k_rope @ q_head[kv_lora_rank : kv_lora_rank + qk_rope_head_dim] + ) * softmax_scale + probs = torch.softmax(scores, dim=0) + out[batch, 0, head] = probs @ k_latent + lse[batch, 0, head] = torch.logsumexp(scores, dim=0) + return out.to(q.dtype), lse + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/ROCm device") +def test_mha_prefill_matches_reference(): + from tokenspeed_attention import mha_prefill + + torch.manual_seed(0) + device = "cuda" + dtype = torch.bfloat16 + seqlens = [8, 5] + cu_cpu, cu = _make_cu_seqlens(seqlens, device) + total = cu_cpu[-1] + q = torch.randn(total, 4, 32, device=device, dtype=dtype) + k = torch.randn(total, 2, 32, device=device, dtype=dtype) + v = torch.randn(total, 2, 32, device=device, dtype=dtype) + + out = mha_prefill(q, k, v, cu, cu_cpu, max(seqlens)) + ref = _varlen_attention_reference( + q, + k, + v, + cu_cpu, + cu_cpu, + causal=True, + scale=1.0 / math.sqrt(q.shape[-1]), + ) + + torch.testing.assert_close(out, ref, atol=8e-2, rtol=8e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/ROCm device") +def test_mla_decode_with_kvcache_matches_reference(): + from tokenspeed_attention import mla_decode_with_kvcache + + torch.manual_seed(1) + device = "cuda" + dtype = torch.bfloat16 + batch, q_len, heads = 2, 1, 4 + rank, rope_dim, page_size = 16, 8, 8 + cache_lens = torch.tensor([7, 11], device=device, dtype=torch.int32) + pages_per_batch = (cache_lens + page_size - 1) // page_size + total_pages = int(pages_per_batch.sum().item()) + max_pages = int(pages_per_batch.max().item()) + + q = torch.randn(batch, q_len, heads, rank + rope_dim, device=device, dtype=dtype) + kv_cache = torch.randn( + total_pages, + page_size, + 1, + rank + rope_dim, + device=device, + dtype=dtype, + ) + page_table = torch.zeros(batch, max_pages, device=device, dtype=torch.int32) + next_page = 0 + for batch_idx, num_pages in enumerate(pages_per_batch.tolist()): + page_table[batch_idx, :num_pages] = torch.arange( + next_page, + next_page + num_pages, + device=device, + dtype=torch.int32, + ) + next_page += num_pages + + scale = 1.0 / math.sqrt(rank + rope_dim) + out, lse = mla_decode_with_kvcache( + q, + kv_cache, + page_table, + cache_lens, + int(cache_lens.max().item()), + rank, + rank, + rope_dim, + scale, + return_lse=True, + ) + ref, ref_lse = _paged_mla_reference( + q, + kv_cache, + page_table, + cache_lens, + kv_lora_rank=rank, + qk_rope_head_dim=rope_dim, + softmax_scale=scale, + ) + + torch.testing.assert_close(out, ref, atol=8e-2, rtol=8e-2) + torch.testing.assert_close(lse, ref_lse, atol=8e-2, rtol=8e-2) diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/__init__.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/__init__.py new file mode 100644 index 00000000..3bf4abdb --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/__init__.py @@ -0,0 +1,285 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import math + +import torch +from ._triton import tl, triton +from .mha_decode import decode_attention_fwd +from .mha_prefill import prefill_attention_fwd +from .mla_decode import ( # noqa: F401 + triton_mla_decode_with_kvcache, +) +from .mla_prefill import ( # noqa: F401 + triton_mla_prefill, +) + +__version__ = "0.1.0" + + +@triton.jit +def attn_merge_state_kernel( + OutA, + LseA, + OutB, + LseB, + Out, + Lse, + head_dim: tl.constexpr, + lse_scale_log2: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row = tl.program_id(0) + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < head_dim + value_offsets = row * head_dim + offs_d + + lse_a = tl.load(LseA + row).to(tl.float32) + lse_b = tl.load(LseB + row).to(tl.float32) + lse_a_log2 = lse_a * lse_scale_log2 + lse_b_log2 = lse_b * lse_scale_log2 + lse_max_log2 = tl.maximum(lse_a_log2, lse_b_log2) + + weight_a = tl.exp2(lse_a_log2 - lse_max_log2) + weight_b = tl.exp2(lse_b_log2 - lse_max_log2) + denom = weight_a + weight_b + + out_a = tl.load(OutA + value_offsets, mask=mask_d, other=0.0).to(tl.float32) + out_b = tl.load(OutB + value_offsets, mask=mask_d, other=0.0).to(tl.float32) + out = (out_a * weight_a + out_b * weight_b) / denom + merged_lse = (lse_max_log2 + tl.log2(denom)) / lse_scale_log2 + + tl.store(Out + value_offsets, out, mask=mask_d) + tl.store(Lse + row, merged_lse) + + +def triton_mha_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_cpu: list[int], + max_seqlen: int, + window_left: int = -1, + logit_cap: float = 0.0, + sinks: torch.Tensor | None = None, + return_lse: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + out = torch.empty_like(q) + lse = ( + torch.empty((q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device) + if return_lse + else None + ) + cache_seqlens = torch.empty((0,), dtype=torch.int32, device=q.device) + empty_k = torch.empty((0, k.shape[1], k.shape[2]), dtype=k.dtype, device=k.device) + empty_v = torch.empty((0, v.shape[1], v.shape[2]), dtype=v.dtype, device=v.device) + prefill_attention_fwd( + q, + k, + v, + out, + empty_k, + empty_v, + cu_seqlens, + cache_seqlens, + None, + True, + max_seqlen, + sm_scale=1.0 / math.sqrt(q.shape[-1]), + logit_cap=logit_cap, + sliding_window_size=window_left, + sinks=sinks, + has_kv_cache=False, + lse_extend=lse, + ) + if return_lse: + return out, lse + return out + + +def triton_mha_extend_with_kvcache( + q: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + is_causal: bool = False, + window_left: int = -1, + logit_cap: float = 0.0, + sinks: torch.Tensor | None = None, + return_lse: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + k = torch.empty( + (0, k_cache.shape[2], k_cache.shape[3]), + dtype=k_cache.dtype, + device=k_cache.device, + ) + v = torch.empty( + (0, v_cache.shape[2], v_cache.shape[3]), + dtype=v_cache.dtype, + device=v_cache.device, + ) + + out = torch.empty_like(q) + lse = ( + torch.empty((q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device) + if return_lse + else None + ) + prefill_attention_fwd( + q, + k, + v, + out, + k_cache.view(-1, k_cache.shape[2], k_cache.shape[3]), + v_cache.view(-1, v_cache.shape[2], v_cache.shape[3]), + cu_seqlens_q, + cache_seqlens, + None, + is_causal, + max_seqlen_q, + sm_scale=1.0 / math.sqrt(q.shape[-1]), + logit_cap=logit_cap, + sliding_window_size=window_left, + sinks=sinks, + page_table=page_table, + page_table_stride_b=page_table.stride(0), + page_size=k_cache.shape[1], + has_kv_cache=True, + lse_extend=lse, + ) + if return_lse: + return out, lse + return out + + +def triton_mha_decode_with_kvcache( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + max_seqlen_k: int, + max_seqlen_q: int = 1, + window_left: int = -1, + logit_cap: float = 0.0, + sinks: torch.Tensor | None = None, + return_lse: bool = False, +) -> torch.Tensor: + if return_lse: + raise NotImplementedError("Triton MHA decode does not return LSE") + out = torch.empty_like(q) + max_kv_splits = 4 + attn_logits = torch.empty( + q.shape[0], + q.shape[1], + max_kv_splits, + q.shape[2], + dtype=torch.float32, + device=q.device, + ) + attn_lse = torch.empty( + q.shape[0], + q.shape[1], + max_kv_splits, + dtype=torch.float32, + device=q.device, + ) + num_kv_splits = torch.ones( + (cache_seqlens.shape[0],), dtype=torch.int32, device=q.device + ) + decode_attention_fwd( + q, + k_cache.view(-1, k_cache.shape[2], k_cache.shape[3]), + v_cache.view(-1, v_cache.shape[2], v_cache.shape[3]), + out, + page_table, + cache_seqlens, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + page_table.stride(0), + k_cache.shape[1], + window_left, + sm_scale=1.0 / math.sqrt(q.shape[-1]), + logit_cap=logit_cap, + sinks=sinks, + ) + return out + + +def triton_attn_merge_state( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, + lse_scale_log2: float, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.empty_like(out_a) + lse = torch.empty_like(lse_a) + total_rows = out_a.shape[0] * out_a.shape[1] + head_dim = out_a.shape[2] + block_d = triton.next_power_of_2(head_dim) + attn_merge_state_kernel[(total_rows,)]( + out_a, + lse_a, + out_b, + lse_b, + out, + lse, + head_dim, + float(lse_scale_log2), + BLOCK_D=block_d, + ) + return out, lse + + +mha_prefill = triton_mha_prefill +mha_extend_with_kvcache = triton_mha_extend_with_kvcache +mha_decode_with_kvcache = triton_mha_decode_with_kvcache +mla_prefill = triton_mla_prefill +mla_decode_with_kvcache = triton_mla_decode_with_kvcache +attn_merge_state = triton_attn_merge_state + + +__all__ = [ + "__version__", + "mha_prefill", + "mha_extend_with_kvcache", + "mha_decode_with_kvcache", + "mla_prefill", + "mla_decode_with_kvcache", + "attn_merge_state", + "triton_mha_prefill", + "triton_mha_extend_with_kvcache", + "triton_mha_decode_with_kvcache", + "triton_mla_prefill", + "triton_mla_decode_with_kvcache", + "triton_attn_merge_state", +] diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/_platform.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/_platform.py new file mode 100644 index 00000000..c279d9b9 --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/_platform.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache + + +@dataclass(frozen=True) +class ArchVersion: + major: int + minor: int + + def __ge__(self, other: "ArchVersion") -> bool: + return (self.major, self.minor) >= (other.major, other.minor) + + def __gt__(self, other: "ArchVersion") -> bool: + return (self.major, self.minor) > (other.major, other.minor) + + def __le__(self, other: "ArchVersion") -> bool: + return (self.major, self.minor) <= (other.major, other.minor) + + def __lt__(self, other: "ArchVersion") -> bool: + return (self.major, self.minor) < (other.major, other.minor) + + +@dataclass(frozen=True) +class PlatformInfo: + vendor: str + arch_version: ArchVersion + + @property + def is_nvidia(self) -> bool: + return self.vendor == "nvidia" + + @property + def is_amd(self) -> bool: + return self.vendor == "amd" + + @property + def is_ampere_plus(self) -> bool: + return self.is_nvidia and self.arch_version >= ArchVersion(8, 0) + + @property + def is_hopper_plus(self) -> bool: + return self.is_nvidia and self.arch_version >= ArchVersion(9, 0) + + +@lru_cache(maxsize=1) +def current_platform() -> PlatformInfo: + import torch + + if not torch.cuda.is_available(): + raise RuntimeError( + "tokenspeed_attention requires an NVIDIA CUDA or AMD ROCm GPU" + ) + + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + if getattr(torch.version, "hip", None): + arch = getattr(props, "gcnArchName", "").split(":")[0] + arch_map = { + "gfx942": ArchVersion(9, 4), + "gfx950": ArchVersion(9, 5), + } + return PlatformInfo("amd", arch_map.get(arch, ArchVersion(9, 0))) + return PlatformInfo("nvidia", ArchVersion(props.major, props.minor)) diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/_triton.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/_triton.py new file mode 100644 index 00000000..35f85feb --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/_triton.py @@ -0,0 +1,16 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import triton +import triton.language as tl + +__all__ = ["tl", "triton"] diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/mha_decode.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/mha_decode.py new file mode 100644 index 00000000..e8dc55f3 --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/mha_decode.py @@ -0,0 +1,822 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from ._triton import tl, triton +from ._platform import current_platform + +_MIN_BLOCK_KV = 32 + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _fwd_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + page_table, + cache_seqlens, + Att_Out, + Att_Lse, + num_kv_splits, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + page_table_stride_b: tl.constexpr, + PAGE_SIZE: tl.constexpr, + MAX_SEQLEN_Q: tl.constexpr, + WINDOW_LEFT: tl.constexpr, + kv_group_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_q = tl.program_id(0) + cur_batch = cur_q // MAX_SEQLEN_Q + q_pos = cur_q - cur_batch * MAX_SEQLEN_Q + cur_head = tl.program_id(1) + split_kv_id = tl.program_id(2) + + cur_kv_head = cur_head // kv_group_num + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + + cache_len = tl.load(cache_seqlens + cur_batch) + cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos) + cache_len = tl.maximum(cache_len, 0) + cur_batch_seq_len = ( + tl.minimum(cache_len, WINDOW_LEFT) if WINDOW_LEFT >= 0 else cache_len + ) + kv_start_offset = cache_len - cur_batch_seq_len + kv_splits = tl.load(num_kv_splits + cur_batch) + + off_q = cur_q * stride_qbs + cur_head * stride_qh + offs_d + + kv_len_per_split = ( + tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + ) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + e_max = -float("inf") + e_sum = 0.0 + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + q = tl.load(Q + off_q, mask=mask_d, other=0.0) + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + token_indices = kv_start_offset + offs_n + page_indices = token_indices // PAGE_SIZE + page_offsets = token_indices - page_indices * PAGE_SIZE + physical_pages = tl.load( + page_table + cur_batch * page_table_stride_b + page_indices, + mask=offs_n < split_kv_end, + other=0, + ) + kv_loc = physical_pages * PAGE_SIZE + page_offsets + offs_buf_k = ( + kv_loc[:, None] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[None, :] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]), + other=0.0, + ) + qk = tl.sum(q[None, :] * k, 1) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(offs_n < split_kv_end, qk, float("-inf")) + + offs_buf_v = ( + kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 0), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max) + acc *= re_scale + acc += tl.sum(p[:, None] * v, 0) + + e_sum = e_sum * re_scale + tl.sum(p, 0) + e_max = n_e_max + + offs_mid_o = ( + cur_q * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_dv + ) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum, + mask=(mask_dv), + ) + + offs_mid_o_1 = ( + cur_q * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + ) // Lv + + tl.store( + Att_Lse + offs_mid_o_1, + e_max + tl.log(e_sum), + ) + + +def _decode_att_m_fwd( + q, + k_buffer, + v_buffer, + att_out, + att_lse, + page_table, + cache_seqlens, + num_kv_splits, + max_kv_splits, + page_table_stride_b, + page_size, + max_seqlen_q, + window_left, + sm_scale, + logit_cap, +): + platform = current_platform() + + BLOCK = 64 + # MI3xx uses a smaller block size to stay within the SGPR limit. + if platform.is_amd: + BLOCK = 8 + MAX_KV_SPLITS = max_kv_splits + Lk = k_buffer.shape[-1] + Lv = v_buffer.shape[-1] + + total_q, head_num = q.shape[0], q.shape[1] + + grid = (total_q, head_num, MAX_KV_SPLITS) + kv_group_num = q.shape[1] // k_buffer.shape[1] + + if kv_group_num == 1: + num_warps = 4 + else: + num_warps = 2 + if platform.is_amd: + num_warps = 1 + + BLOCK_DMODEL = triton.next_power_of_2(Lk) + BLOCK_DV = triton.next_power_of_2(Lv) + + _fwd_kernel_stage1[grid]( + q, + k_buffer, + v_buffer, + sm_scale, + page_table, + cache_seqlens, + att_out, + att_lse, + num_kv_splits, + q.stride(0), + q.stride(1), + k_buffer.stride(0), + k_buffer.stride(1), + v_buffer.stride(0), + v_buffer.stride(1), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + page_table_stride_b, + page_size, + MAX_SEQLEN_Q=max_seqlen_q, + WINDOW_LEFT=window_left, + kv_group_num=kv_group_num, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DV=BLOCK_DV, + BLOCK_N=BLOCK, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + logit_cap=logit_cap, + num_warps=num_warps, + num_stages=2, + Lk=Lk, + Lv=Lv, + ) + + +@triton.jit +def _fwd_grouped_kernel_stage1( + Q, + K_Buffer, + V_Buffer, + sm_scale, + page_table, + cache_seqlens, + Att_Out, + Att_Lse, + num_kv_splits, + stride_qbs, + stride_qh, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + page_table_stride_b: tl.constexpr, + PAGE_SIZE: tl.constexpr, + MAX_SEQLEN_Q: tl.constexpr, + WINDOW_LEFT: tl.constexpr, + kv_group_num: tl.constexpr, + q_head_num: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + logit_cap: tl.constexpr, + Lk: tl.constexpr, + Lv: tl.constexpr, +): + cur_q = tl.program_id(0) + cur_batch = cur_q // MAX_SEQLEN_Q + q_pos = cur_q - cur_batch * MAX_SEQLEN_Q + cur_head_id = tl.program_id(1) + cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H) + split_kv_id = tl.program_id(2) + + if BLOCK_H < kv_group_num: + VALID_BLOCK_H: tl.constexpr = BLOCK_H + else: + VALID_BLOCK_H: tl.constexpr = kv_group_num + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H + mask_h = mask_h & (cur_head < q_head_num) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lk + mask_dv = offs_dv < Lv + + cache_len = tl.load(cache_seqlens + cur_batch) + cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos) + cache_len = tl.maximum(cache_len, 0) + cur_batch_seq_len = ( + tl.minimum(cache_len, WINDOW_LEFT) if WINDOW_LEFT >= 0 else cache_len + ) + kv_start_offset = cache_len - cur_batch_seq_len + kv_splits = tl.load(num_kv_splits + cur_batch) + + offs_q = cur_q * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :] + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + mask_dpe = offs_dpe < Lk + off_qpe = cur_q * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :] + + kv_len_per_split = ( + tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + ) + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) + + if split_kv_end > split_kv_start: + q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0) + if BLOCK_DPE > 0: + qpe = tl.load( + Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0 + ) + for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + token_indices = kv_start_offset + offs_n + page_indices = token_indices // PAGE_SIZE + page_offsets = token_indices - page_indices * PAGE_SIZE + physical_pages = tl.load( + page_table + cur_batch * page_table_stride_b + page_indices, + mask=offs_n < split_kv_end, + other=0, + ) + kv_loc = physical_pages * PAGE_SIZE + page_offsets + offs_buf_k = ( + kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[:, None] + ) + k = tl.load( + K_Buffer + offs_buf_k, + mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]), + other=0.0, + ) + qk = tl.dot(q, k.to(q.dtype)) + if BLOCK_DPE > 0: + offs_buf_kpe = ( + kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Buffer + offs_buf_kpe, + mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]), + other=0.0, + ) + qk += tl.dot(qpe, kpe.to(qpe.dtype)) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where( + mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf") + ) + + offs_buf_v = ( + kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + + n_e_max = tl.maximum(tl.max(qk, 1), e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + acc *= re_scale[:, None] + acc += tl.dot(p.to(v.dtype), v) + + e_sum = e_sum * re_scale + tl.sum(p, 1) + e_max = n_e_max + + offs_mid_o = ( + cur_q * stride_mid_ob + + cur_head[:, None] * stride_mid_oh + + split_kv_id * stride_mid_os + + offs_dv[None, :] + ) + + tl.store( + Att_Out + offs_mid_o, + acc / e_sum[:, None], + mask=(mask_h[:, None]) & (mask_dv[None, :]), + ) + + offs_mid_o_1 = ( + cur_q * stride_mid_ob + + cur_head * stride_mid_oh + + split_kv_id * stride_mid_os + ) // Lv + + tl.store( + Att_Lse + offs_mid_o_1, + e_max + tl.log(e_sum), + mask=mask_h, + ) + + +def _decode_grouped_att_m_fwd( + q, + k_buffer, + v_buffer, + att_out, + att_lse, + page_table, + cache_seqlens, + num_kv_splits, + max_kv_splits, + page_table_stride_b, + page_size, + max_seqlen_q, + window_left, + sm_scale, + logit_cap, +): + platform = current_platform() + + BLOCK = 32 + Lk = k_buffer.shape[-1] + Lv = v_buffer.shape[-1] + + # MI3xx uses a smaller block size for large heads to stay within shmem limits. + if platform.is_amd and Lk >= 576: + BLOCK = 16 + + if Lk == 576: + BLOCK_DMODEL = 512 + BLOCK_DPE = 64 + elif Lk == 288: + BLOCK_DMODEL = 256 + BLOCK_DPE = 32 + else: + BLOCK_DMODEL = triton.next_power_of_2(Lk) + BLOCK_DPE = 0 + BLOCK_DV = triton.next_power_of_2(Lv) + + total_q, head_num = q.shape[0], q.shape[1] + kv_group_num = q.shape[1] // k_buffer.shape[1] + + BLOCK_H = 16 + MAX_KV_SPLITS = max_kv_splits + grid = ( + total_q, + triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), + MAX_KV_SPLITS, + ) + + extra_kargs = {} + num_stages = 2 + if platform.is_amd: + extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16} + num_stages = 1 + + _fwd_grouped_kernel_stage1[grid]( + q, + k_buffer, + v_buffer, + sm_scale, + page_table, + cache_seqlens, + att_out, + att_lse, + num_kv_splits, + q.stride(0), + q.stride(1), + k_buffer.stride(0), + k_buffer.stride(1), + v_buffer.stride(0), + v_buffer.stride(1), + att_out.stride(0), + att_out.stride(1), + att_out.stride(2), + page_table_stride_b, + page_size, + MAX_SEQLEN_Q=max_seqlen_q, + WINDOW_LEFT=window_left, + kv_group_num=kv_group_num, + q_head_num=head_num, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DPE=BLOCK_DPE, + BLOCK_DV=BLOCK_DV, + BLOCK_N=BLOCK, + BLOCK_H=BLOCK_H, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + logit_cap=logit_cap, + num_warps=4, + num_stages=num_stages, + Lk=Lk, + Lv=Lv, + **extra_kargs, + ) + + +@triton.jit +def _fwd_kernel_stage2( + Mid_O, + Mid_O_1, + O, + cache_seqlens, + num_kv_splits, + sink_ptr, + stride_mid_ob, + stride_mid_oh, + stride_mid_os, + stride_obs, + stride_oh, + MAX_SEQLEN_Q: tl.constexpr, + MAX_KV_SPLITS: tl.constexpr, + MIN_BLOCK_KV: tl.constexpr, + BLOCK_DV: tl.constexpr, + Lv: tl.constexpr, + WINDOW_LEFT: tl.constexpr, + HAS_SINK: tl.constexpr, +): + cur_q = tl.program_id(0) + cur_batch = cur_q // MAX_SEQLEN_Q + q_pos = cur_q - cur_batch * MAX_SEQLEN_Q + cur_head = tl.program_id(1) + + cache_len = tl.load(cache_seqlens + cur_batch) + cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos) + cache_len = tl.maximum(cache_len, 0) + cur_batch_seq_len = ( + tl.minimum(cache_len, WINDOW_LEFT) if WINDOW_LEFT >= 0 else cache_len + ) + kv_splits = tl.load(num_kv_splits + cur_batch) + + offs_d = tl.arange(0, BLOCK_DV) + mask_d = offs_d < Lv + + e_sum = 0.0 + e_max = -float("inf") + acc = tl.zeros([BLOCK_DV], dtype=tl.float32) + + offs_v = cur_q * stride_mid_ob + cur_head * stride_mid_oh + offs_d + offs_logic = (cur_q * stride_mid_ob + cur_head * stride_mid_oh) // Lv + kv_len_per_split = ( + tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV + ) + + for split_kv_id in range(0, MAX_KV_SPLITS): + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + + if split_kv_end > split_kv_start: + tv = tl.load( + Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0 + ) + tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv) + n_e_max = tl.maximum(tlogic, e_max) + + old_scale = tl.exp(e_max - n_e_max) + acc *= old_scale + exp_logic = tl.exp(tlogic - n_e_max) + acc += exp_logic * tv + + e_sum = e_sum * old_scale + exp_logic + e_max = n_e_max + + if HAS_SINK: + cur_sink = tl.load(sink_ptr + cur_head) + e_sum += tl.exp(cur_sink - e_max) + + tl.store( + O + cur_q * stride_obs + cur_head * stride_oh + offs_d, + acc / e_sum, + mask=mask_d, + ) + + +def _decode_softmax_reducev_fwd( + logits, + lse, + q, + o, + v_buffer, + cache_seqlens, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + window_left, + sinks=None, +): + platform = current_platform() + total_q, head_num = q.shape[0], q.shape[1] + Lv = v_buffer.shape[-1] + BLOCK_DV = triton.next_power_of_2(Lv) + + MAX_KV_SPLITS = max_kv_splits + HAS_SINK = sinks is not None + + extra_kargs = {} + if platform.is_amd: + extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16} + + grid = (total_q, head_num) + _fwd_kernel_stage2[grid]( + logits, + lse, + o, + cache_seqlens, + num_kv_splits, + sinks, + logits.stride(0), + logits.stride(1), + logits.stride(2), + o.stride(0), + o.stride(1), + MAX_SEQLEN_Q=max_seqlen_q, + MAX_KV_SPLITS=MAX_KV_SPLITS, + MIN_BLOCK_KV=_MIN_BLOCK_KV, + BLOCK_DV=BLOCK_DV, + Lv=Lv, + WINDOW_LEFT=window_left, + HAS_SINK=HAS_SINK, + num_warps=4, + num_stages=2, + **extra_kargs, + ) + + +def decode_attention_fwd_normal( + q, + k_buffer, + v_buffer, + o, + page_table, + cache_seqlens, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + page_table_stride_b, + page_size, + window_left, + sm_scale, + logit_cap=0.0, + sinks=None, +): + _decode_att_m_fwd( + q, + k_buffer, + v_buffer, + attn_logits, + attn_lse, + page_table, + cache_seqlens, + num_kv_splits, + max_kv_splits, + page_table_stride_b, + page_size, + max_seqlen_q, + window_left, + sm_scale, + logit_cap, + ) + _decode_softmax_reducev_fwd( + attn_logits, + attn_lse, + q, + o, + v_buffer, + cache_seqlens, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + window_left, + sinks, + ) + + +def decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + page_table, + cache_seqlens, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + page_table_stride_b, + page_size, + window_left, + sm_scale, + logit_cap=0.0, + sinks=None, +): + _decode_grouped_att_m_fwd( + q, + k_buffer, + v_buffer, + attn_logits, + attn_lse, + page_table, + cache_seqlens, + num_kv_splits, + max_kv_splits, + page_table_stride_b, + page_size, + max_seqlen_q, + window_left, + sm_scale, + logit_cap, + ) + _decode_softmax_reducev_fwd( + attn_logits, + attn_lse, + q, + o, + v_buffer, + cache_seqlens, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + window_left, + sinks, + ) + + +def decode_attention_fwd( + q, + k_buffer, + v_buffer, + o, + page_table, + cache_seqlens, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + page_table_stride_b, + page_size, + window_left, + sm_scale, + logit_cap=0.0, + sinks=None, +): + assert max_kv_splits == attn_logits.shape[2] + assert q.shape[0] == cache_seqlens.shape[0] * max_seqlen_q + assert q.shape[0] <= attn_logits.shape[0] + + kv_group_num = q.shape[1] // v_buffer.shape[1] + + if kv_group_num == 1: + # MHA + decode_attention_fwd_normal( + q, + k_buffer, + v_buffer, + o, + page_table, + cache_seqlens, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + page_table_stride_b, + page_size, + window_left, + sm_scale, + logit_cap=logit_cap, + sinks=sinks, + ) + else: + # GQA/MQA/MLA + decode_attention_fwd_grouped( + q, + k_buffer, + v_buffer, + o, + page_table, + cache_seqlens, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + max_seqlen_q, + page_table_stride_b, + page_size, + window_left, + sm_scale, + logit_cap=logit_cap, + sinks=sinks, + ) diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/mha_prefill.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/mha_prefill.py new file mode 100644 index 00000000..739e9b18 --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/mha_prefill.py @@ -0,0 +1,448 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from ._triton import tl, triton +from ._platform import current_platform + + +@triton.jit +def tanh(x): + # Tanh is just a scaled sigmoid + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _fwd_kernel( + Q_Extend, + K_Extend, + V_Extend, + O_Extend, + LSE_Extend, + K_Buffer, + V_Buffer, + cu_seqlens_q, + cache_seqlens, + mask_ptr, + sink_ptr, + page_table, + sm_scale, + kv_group_num, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + stride_lse_bs, + stride_lse_h, + stride_buf_kbs, + stride_buf_kh, + stride_buf_vbs, + stride_buf_vh, + page_table_stride_b: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SLIDING_WINDOW_SIZE: tl.constexpr, + logit_cap: tl.constexpr, + Lq: tl.constexpr, + Lv: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + USE_CUSTOM_MASK: tl.constexpr, + IS_CAUSAL: tl.constexpr, + SKIP_PREFIX_CUSTOM_MASK: tl.constexpr, + HAS_KV_CACHE: tl.constexpr, + STORE_TRANSPOSE: tl.constexpr, + HAS_SINK: tl.constexpr, + HAS_LSE: tl.constexpr, +): + cur_seq = tl.program_id(0) + cur_head = tl.program_id(1) + cur_block_m = tl.program_id(2) + cur_kv_head = cur_head // kv_group_num + + cur_seq_extend_start_idx = tl.load(cu_seqlens_q + cur_seq) + cur_seq_len_extend = tl.load(cu_seqlens_q + cur_seq + 1) - cur_seq_extend_start_idx + if HAS_KV_CACHE: + cur_seq_len = tl.load(cache_seqlens + cur_seq) + else: + cur_seq_len = cur_seq_len_extend + cur_q_start = tl.maximum(cur_seq_len - cur_seq_len_extend, 0) + + if USE_CUSTOM_MASK: + cur_seq_mask_start_idx = tl.load(cu_seqlens_q + cur_seq) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + offs_m = tl.arange(0, BLOCK_M) + mask_m = (cur_block_m * BLOCK_M + offs_m) < cur_seq_len_extend + + mask_d = offs_d < Lq + mask_dv = offs_dv < Lv + + offs_q = ( + (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] + ) + q = tl.load( + Q_Extend + offs_q, mask=(mask_m[:, None]) & (mask_d[None, :]), other=0.0 + ) + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + offs_qpe = ( + (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_qbs + + cur_head * stride_qh + + offs_dpe[None, :] + ) + qpe = tl.load(Q_Extend + offs_qpe, mask=mask_m[:, None], other=0.0) + + # Compute attention over the full visible KV range. For causal cached + # prefill, query tokens are treated as the suffix of the KV sequence. + offs_n = tl.arange(0, BLOCK_N) + + acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32) + deno = tl.zeros([BLOCK_M], dtype=tl.float32) + e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + + for start_n in range(0, cur_seq_len, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + mask_n = (start_n + offs_n) < cur_seq_len + + final_mask = mask_m[:, None] & mask_n[None, :] + if USE_CUSTOM_MASK: + custom_mask = tl.load( + mask_ptr + + cur_seq_mask_start_idx + + (cur_block_m * BLOCK_M + offs_m[:, None]) * cur_seq_len + + start_n + + offs_n[None, :], + mask=(mask_m[:, None] & mask_n[None, :]), + other=0, + ) + final_mask &= custom_mask + elif IS_CAUSAL: + query_positions = cur_q_start + cur_block_m * BLOCK_M + offs_m[:, None] + key_positions = start_n + offs_n[None, :] + final_mask &= query_positions >= key_positions + if SLIDING_WINDOW_SIZE > 0: + # Add mask where q_id <= kv_id + sliding_window_size + query_positions = cur_q_start + cur_block_m * BLOCK_M + offs_m[:, None] + key_positions = start_n + offs_n[None, :] + window_mask = query_positions <= key_positions + SLIDING_WINDOW_SIZE + final_mask &= window_mask + + SKIP_TILE = False + if USE_CUSTOM_MASK or SLIDING_WINDOW_SIZE > 0: + SKIP_TILE = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 + + if not SKIP_TILE: + if HAS_KV_CACHE: + cache_token_indices = start_n + offs_n + page_indices = cache_token_indices // PAGE_SIZE + page_offsets = cache_token_indices - page_indices * PAGE_SIZE + physical_pages = tl.load( + page_table + cur_seq * page_table_stride_b + page_indices, + mask=mask_n, + other=0, + ) + offs_kv_loc = physical_pages * PAGE_SIZE + page_offsets + offs_k = ( + offs_kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_d[:, None] + ) + k = tl.load( + K_Buffer + offs_k, + mask=(mask_n[None, :]) & (mask_d[:, None]), + other=0.0, + ) + else: + offs_k = ( + (cur_seq_extend_start_idx + start_n + offs_n[None, :]) * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] + ) + k = tl.load( + K_Extend + offs_k, + mask=(mask_n[None, :]) & (mask_d[:, None]), + other=0.0, + ) + + qk = tl.dot(q.to(k.dtype), k) + if BLOCK_DPE > 0: + if HAS_KV_CACHE: + offs_kpe = ( + offs_kv_loc[None, :] * stride_buf_kbs + + cur_kv_head * stride_buf_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Buffer + offs_kpe, + mask=mask_n[None, :], + other=0.0, + ) + else: + offs_kpe = ( + (cur_seq_extend_start_idx + start_n + offs_n[None, :]) + * stride_kbs + + cur_kv_head * stride_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K_Extend + offs_kpe, + mask=mask_n[None, :], + other=0.0, + ) + qk += tl.dot(qpe.to(kpe.dtype), kpe) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(final_mask, qk, float("-inf")) + + row_max = tl.max(qk, 1) + row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max) + n_e_max = tl.maximum(row_max_fixed, e_max) + + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + deno = deno * re_scale + tl.sum(p, 1) + + if HAS_KV_CACHE: + offs_v = ( + offs_kv_loc[:, None] * stride_buf_vbs + + cur_kv_head * stride_buf_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Buffer + offs_v, + mask=mask_n[:, None] & mask_dv[None, :], + other=0.0, + ) + else: + offs_v = ( + (cur_seq_extend_start_idx + start_n + offs_n[:, None]) * stride_vbs + + cur_kv_head * stride_vh + + offs_dv[None, :] + ) + v = tl.load( + V_Extend + offs_v, + mask=mask_n[:, None] & mask_dv[None, :], + other=0.0, + ) + p = p.to(v.dtype) + acc = acc * re_scale[:, None] + tl.dot(p, v) + + e_max = n_e_max + + if HAS_SINK: + cur_sink = tl.load(sink_ptr + cur_head) + deno += tl.exp(cur_sink - e_max) + + safe_deno = tl.where(deno > 0.0, deno, 1.0) + + offs_o = ( + (cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m[:, None]) + * stride_obs + + cur_head * stride_oh + + offs_dv[None, :] + ) + if STORE_TRANSPOSE: + tl.store( + O_Extend + offs_o.T, + (acc / safe_deno[:, None]).T, + mask=(mask_m[:, None] & mask_dv[None, :]).T, + ) + else: + tl.store( + O_Extend + offs_o, + acc / safe_deno[:, None], + mask=mask_m[:, None] & mask_dv[None, :], + ) + + if HAS_LSE: + offs_lse = ( + cur_seq_extend_start_idx + cur_block_m * BLOCK_M + offs_m + ) * stride_lse_bs + cur_head * stride_lse_h + lse = tl.where(deno > 0.0, tl.log(deno) + e_max, float("-inf")) + tl.store(LSE_Extend + offs_lse, lse, mask=mask_m) + + +def prefill_attention_fwd( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + cu_seqlens_q, + cache_seqlens, + custom_mask, + is_causal, + max_len_extend, + sm_scale=None, + logit_cap=0.0, + skip_prefix_custom_mask=True, + sliding_window_size=-1, + sinks=None, + page_table=None, + page_table_stride_b=0, + page_size=1, + has_kv_cache=False, + lse_extend=None, +): + """ + q_extend, k_extend, v_extend, o_extend: contiguous tensors + + k_buffer, v_buffer: (prefix + extend) tensors in mem_manager + """ + platform = current_platform() + + Lq, Lk, Lv = ( + q_extend.shape[-1], + k_extend.shape[-1], + v_extend.shape[-1], + ) + + if Lq == 576: + BLOCK_DMODEL = 512 + BLOCK_DPE = 64 + elif Lq == 288: + BLOCK_DMODEL = 256 + BLOCK_DPE = 32 + elif Lq == 192: + BLOCK_DMODEL = 128 + BLOCK_DPE = 64 + else: + BLOCK_DMODEL = triton.next_power_of_2(Lq) + BLOCK_DPE = 0 + BLOCK_DV = triton.next_power_of_2(Lv) + + if platform.is_amd: + BLOCK_M, BLOCK_N = (64, 64) + num_warps = 4 + + else: + if platform.is_hopper_plus: + if Lq <= 256: + BLOCK_M, BLOCK_N = (128, 64) + else: + BLOCK_M, BLOCK_N = (32, 64) + elif platform.is_ampere_plus: + # sm86/sm89 has a much smaller shared memory size (100K) than sm80 (160K) + if platform.arch_version.minor == 9 or platform.arch_version.minor == 6: + if Lq <= 128: + BLOCK_M, BLOCK_N = (64, 128) + elif Lq <= 256: + BLOCK_M, BLOCK_N = (64, 64) + else: + BLOCK_M, BLOCK_N = (32, 32) + else: + if Lq <= 128: + BLOCK_M, BLOCK_N = (128, 128) + elif Lq <= 256: + BLOCK_M, BLOCK_N = (64, 64) + else: + BLOCK_M, BLOCK_N = (32, 64) + else: + BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32) + + num_warps = 4 if Lk <= 64 else 8 + + sm_scale = sm_scale or 1.0 / (Lq**0.5) + batch_size, head_num = cu_seqlens_q.shape[0] - 1, q_extend.shape[1] + kv_group_num = q_extend.shape[1] // k_extend.shape[1] + + USE_CUSTOM_MASK = custom_mask is not None + # Skip custom mask for prefix part + SKIP_PREFIX_CUSTOM_MASK = skip_prefix_custom_mask + + HAS_SINK = sinks is not None + HAS_LSE = lse_extend is not None + lse_arg = lse_extend if lse_extend is not None else o_extend + page_table_arg = page_table if page_table is not None else cache_seqlens + + grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M)) + num_stages = 1 + + extra_kargs = {} + if platform.is_amd: + extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16} + + _fwd_kernel[grid]( + q_extend, + k_extend, + v_extend, + o_extend, + lse_arg, + k_buffer, + v_buffer, + cu_seqlens_q, + cache_seqlens, + custom_mask, + sinks, + page_table_arg, + sm_scale, + kv_group_num, + q_extend.stride(0), + q_extend.stride(1), + k_extend.stride(0), + k_extend.stride(1), + v_extend.stride(0), + v_extend.stride(1), + o_extend.stride(0), + o_extend.stride(1), + lse_arg.stride(0), + lse_arg.stride(1), + k_buffer.stride(0), + k_buffer.stride(1), + v_buffer.stride(0), + v_buffer.stride(1), + page_table_stride_b, + page_size, + SLIDING_WINDOW_SIZE=sliding_window_size, + logit_cap=logit_cap, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_DPE=BLOCK_DPE, + BLOCK_DV=BLOCK_DV, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + Lq=Lq, + Lv=Lv, + USE_CUSTOM_MASK=USE_CUSTOM_MASK, + IS_CAUSAL=is_causal, + SKIP_PREFIX_CUSTOM_MASK=SKIP_PREFIX_CUSTOM_MASK, + HAS_KV_CACHE=has_kv_cache, + HAS_SINK=HAS_SINK, + HAS_LSE=HAS_LSE, + STORE_TRANSPOSE=platform.is_amd, + num_warps=num_warps, + num_stages=num_stages, + **extra_kargs, + ) diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/mla_decode.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/mla_decode.py new file mode 100644 index 00000000..867e14c5 --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/mla_decode.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from ._triton import tl, triton + +_FP8_DTYPES = frozenset( + getattr(torch, name) + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz") + if hasattr(torch, name) +) + + +@triton.jit +def tanh(x): + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _mla_decode_kernel( + Q, + KV_Cache, + O, + LSE, + page_table, + cache_seqlens, + sm_scale, + stride_qb, + stride_qq, + stride_qh, + stride_kv_page, + stride_kv_token, + stride_kv_head, + stride_ob, + stride_oq, + stride_oh, + stride_lse_b, + stride_lse_q, + stride_lse_h, + page_table_stride_b: tl.constexpr, + PAGE_SIZE: tl.constexpr, + MAX_SEQLEN_K: tl.constexpr, + logit_cap: tl.constexpr, + KV_LORA_RANK: tl.constexpr, + QK_ROPE_HEAD_DIM: tl.constexpr, + BLOCK_R: tl.constexpr, + BLOCK_ROPE: tl.constexpr, + BLOCK_N: tl.constexpr, + HAS_LSE: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_q = tl.program_id(1) + cur_head = tl.program_id(2) + + offs_r = tl.arange(0, BLOCK_R) + offs_rope = tl.arange(0, BLOCK_ROPE) + mask_r = offs_r < KV_LORA_RANK + mask_rope = offs_rope < QK_ROPE_HEAD_DIM + + q_base = cur_batch * stride_qb + cur_q * stride_qq + cur_head * stride_qh + q_latent = tl.load(Q + q_base + offs_r, mask=mask_r, other=0.0) + q_rope = tl.load( + Q + q_base + KV_LORA_RANK + offs_rope, + mask=mask_rope, + other=0.0, + ) + + cache_len = tl.load(cache_seqlens + cur_batch) + offs_n = tl.arange(0, BLOCK_N) + acc = tl.zeros([BLOCK_R], dtype=tl.float32) + e_sum = 0.0 + e_max = -float("inf") + + for start_n in tl.range(0, cache_len, BLOCK_N, num_stages=2): + start_n = tl.multiple_of(start_n, BLOCK_N) + token_offsets = start_n + offs_n + mask_n = token_offsets < cache_len + page_indices = token_offsets // PAGE_SIZE + page_offsets = token_offsets - page_indices * PAGE_SIZE + physical_pages = tl.load( + page_table + cur_batch * page_table_stride_b + page_indices, + mask=mask_n, + other=0, + ) + cache_base = ( + physical_pages * stride_kv_page + + page_offsets * stride_kv_token + + 0 * stride_kv_head + ) + + k_latent = tl.load( + KV_Cache + cache_base[:, None] + offs_r[None, :], + mask=mask_n[:, None] & mask_r[None, :], + other=0.0, + ) + k_rope = tl.load( + KV_Cache + cache_base[:, None] + KV_LORA_RANK + offs_rope[None, :], + mask=mask_n[:, None] & mask_rope[None, :], + other=0.0, + ) + + qk = tl.sum(k_latent.to(tl.float32) * q_latent[None, :].to(tl.float32), axis=1) + qk += tl.sum(k_rope.to(tl.float32) * q_rope[None, :].to(tl.float32), axis=1) + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(mask_n, qk, float("-inf")) + block_max = tl.max(qk, axis=0) + block_max_fixed = tl.where(block_max == float("-inf"), -1e20, block_max) + n_e_max = tl.maximum(block_max_fixed, e_max) + old_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max) + acc = acc * old_scale + tl.sum(p[:, None] * k_latent.to(tl.float32), axis=0) + e_sum = e_sum * old_scale + tl.sum(p, axis=0) + e_max = n_e_max + + safe_sum = tl.where(e_sum > 0.0, e_sum, 1.0) + out_base = cur_batch * stride_ob + cur_q * stride_oq + cur_head * stride_oh + tl.store(O + out_base + offs_r, acc / safe_sum, mask=mask_r) + + if HAS_LSE: + lse = tl.where(e_sum > 0.0, tl.log(e_sum) + e_max, float("-inf")) + tl.store( + LSE + + cur_batch * stride_lse_b + + cur_q * stride_lse_q + + cur_head * stride_lse_h, + lse, + ) + + +def _normalize_kv_cache(kv_cache: torch.Tensor) -> torch.Tensor: + if kv_cache.dim() == 3: + return kv_cache.unsqueeze(2) + if kv_cache.dim() != 4: + raise ValueError(f"kv_cache must be 3D or 4D, got {kv_cache.dim()}D") + if kv_cache.shape[2] == 1: + return kv_cache + if kv_cache.shape[1] == 1: + return kv_cache[:, 0].unsqueeze(2) + raise ValueError(f"unsupported kv_cache shape {tuple(kv_cache.shape)}") + + +def mla_decode_fwd( + q: torch.Tensor, + kv_cache: torch.Tensor, + out: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + max_seqlen_k: int, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + softmax_scale: float, + *, + logit_cap: float = 0.0, + lse: torch.Tensor | None = None, +) -> None: + if q.dim() != 4: + raise ValueError( + f"q must have shape [B, q_len, H, R + rope], got {tuple(q.shape)}" + ) + if q.shape[1] != 1: + raise NotImplementedError("Triton MLA decode currently supports q_len == 1") + if q.shape[-1] != kv_lora_rank + qk_rope_head_dim: + raise ValueError( + f"q head dim must be {kv_lora_rank + qk_rope_head_dim}, " + f"got {q.shape[-1]}" + ) + if out.shape != q.shape[:-1] + (kv_lora_rank,): + raise ValueError( + f"out shape must be {q.shape[:-1] + (kv_lora_rank,)}, " + f"got {tuple(out.shape)}" + ) + if q.stride(-1) != 1 or out.stride(-1) != 1: + raise ValueError("q and out must have contiguous last dimension") + if lse is not None and lse.shape != q.shape[:-1]: + raise ValueError(f"lse shape must be {q.shape[:-1]}, got {tuple(lse.shape)}") + + kv_cache = _normalize_kv_cache(kv_cache) + if kv_cache.shape[2] != 1: + raise ValueError(f"MLA kv_cache must have one KV head, got {kv_cache.shape[2]}") + if kv_cache.shape[-1] != kv_lora_rank + qk_rope_head_dim: + raise ValueError( + f"kv_cache head dim must be {kv_lora_rank + qk_rope_head_dim}, " + f"got {kv_cache.shape[-1]}" + ) + if kv_cache.stride(-1) != 1: + raise ValueError("kv_cache must have contiguous last dimension") + + block_n = 16 + block_r = triton.next_power_of_2(kv_lora_rank) + block_rope = triton.next_power_of_2(qk_rope_head_dim) + grid = (q.shape[0], q.shape[1], q.shape[2]) + lse_arg = lse if lse is not None else out + + _mla_decode_kernel[grid]( + q, + kv_cache, + out, + lse_arg, + page_table, + cache_seqlens, + softmax_scale, + q.stride(0), + q.stride(1), + q.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + lse_arg.stride(0), + lse_arg.stride(1), + lse_arg.stride(2), + page_table.stride(0), + kv_cache.shape[1], + max_seqlen_k, + logit_cap=logit_cap, + KV_LORA_RANK=kv_lora_rank, + QK_ROPE_HEAD_DIM=qk_rope_head_dim, + BLOCK_R=block_r, + BLOCK_ROPE=block_rope, + BLOCK_N=block_n, + HAS_LSE=lse is not None, + num_warps=8, + num_stages=2, + ) + + +def triton_mla_decode_with_kvcache( + q: torch.Tensor, + kv_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + max_seqlen_k: int, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + softmax_scale: float, + *, + logit_cap: float = 0.0, + return_lse: bool = False, + out: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if out is None: + out_dtype = torch.bfloat16 if q.dtype in _FP8_DTYPES else q.dtype + out = torch.empty( + q.shape[:-1] + (kv_lora_rank,), dtype=out_dtype, device=q.device + ) + + lse = ( + torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device) + if return_lse + else None + ) + mla_decode_fwd( + q, + kv_cache, + out, + page_table, + cache_seqlens, + max_seqlen_k, + qk_nope_head_dim, + kv_lora_rank, + qk_rope_head_dim, + softmax_scale, + logit_cap=logit_cap, + lse=lse, + ) + if return_lse: + return out, lse + return out diff --git a/tokenspeed-attention/torch-ext/tokenspeed_attention/mla_prefill.py b/tokenspeed-attention/torch-ext/tokenspeed_attention/mla_prefill.py new file mode 100644 index 00000000..f8cb7568 --- /dev/null +++ b/tokenspeed-attention/torch-ext/tokenspeed_attention/mla_prefill.py @@ -0,0 +1,324 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from ._triton import tl, triton + +_FP8_DTYPES = frozenset( + getattr(torch, name) + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz") + if hasattr(torch, name) +) + + +@triton.jit +def tanh(x): + return 2 * tl.sigmoid(2 * x) - 1 + + +@triton.jit +def _mla_prefill_kernel( + Q, + K, + V, + O, + LSE, + cu_seqlens_q, + cu_seqlens_kv, + sm_scale, + kv_group_num, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + stride_lse_bs, + stride_lse_h, + logit_cap: tl.constexpr, + Lq: tl.constexpr, + Lv: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_DPE: tl.constexpr, + BLOCK_DV: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, + HAS_LSE: tl.constexpr, +): + cur_seq = tl.program_id(0) + cur_head = tl.program_id(1) + cur_block_m = tl.program_id(2) + cur_kv_head = cur_head // kv_group_num + + q_start = tl.load(cu_seqlens_q + cur_seq) + q_len = tl.load(cu_seqlens_q + cur_seq + 1) - q_start + kv_start = tl.load(cu_seqlens_kv + cur_seq) + kv_len = tl.load(cu_seqlens_kv + cur_seq + 1) - kv_start + q_causal_start = tl.maximum(kv_len - q_len, 0) + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_dv = tl.arange(0, BLOCK_DV) + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + q_offsets_m = cur_block_m * BLOCK_M + offs_m + mask_m = q_offsets_m < q_len + mask_d = offs_d < Lq + mask_dv = offs_dv < Lv + + offs_q = ( + (q_start + q_offsets_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_d[None, :] + ) + q = tl.load(Q + offs_q, mask=mask_m[:, None] & mask_d[None, :], other=0.0) + + if BLOCK_DPE > 0: + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) + mask_dpe = offs_dpe < Lq + offs_qpe = ( + (q_start + q_offsets_m[:, None]) * stride_qbs + + cur_head * stride_qh + + offs_dpe[None, :] + ) + qpe = tl.load(Q + offs_qpe, mask=mask_m[:, None] & mask_dpe[None, :], other=0.0) + + acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32) + deno = tl.zeros([BLOCK_M], dtype=tl.float32) + e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + + for start_n in range(0, kv_len, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + kv_offsets_n = start_n + offs_n + mask_n = kv_offsets_n < kv_len + final_mask = mask_m[:, None] & mask_n[None, :] + + if IS_CAUSAL: + query_positions = q_causal_start + q_offsets_m[:, None] + key_positions = kv_offsets_n[None, :] + final_mask &= query_positions >= key_positions + + offs_k = ( + (kv_start + kv_offsets_n[None, :]) * stride_kbs + + cur_kv_head * stride_kh + + offs_d[:, None] + ) + k = tl.load(K + offs_k, mask=mask_n[None, :] & mask_d[:, None], other=0.0) + + qk = tl.dot(q.to(k.dtype), k) + + if BLOCK_DPE > 0: + offs_kpe = ( + (kv_start + kv_offsets_n[None, :]) * stride_kbs + + cur_kv_head * stride_kh + + offs_dpe[:, None] + ) + kpe = tl.load( + K + offs_kpe, mask=mask_n[None, :] & mask_dpe[:, None], other=0.0 + ) + qk += tl.dot(qpe.to(kpe.dtype), kpe) + + qk *= sm_scale + + if logit_cap > 0: + qk = logit_cap * tanh(qk / logit_cap) + + qk = tl.where(final_mask, qk, float("-inf")) + + row_max = tl.max(qk, 1) + row_max_fixed = tl.where(row_max == float("-inf"), -1e20, row_max) + n_e_max = tl.maximum(row_max_fixed, e_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + deno = deno * re_scale + tl.sum(p, 1) + + offs_v = ( + (kv_start + kv_offsets_n[:, None]) * stride_vbs + + cur_kv_head * stride_vh + + offs_dv[None, :] + ) + v = tl.load(V + offs_v, mask=mask_n[:, None] & mask_dv[None, :], other=0.0) + p = p.to(v.dtype) + acc = acc * re_scale[:, None] + tl.dot(p, v) + e_max = n_e_max + + safe_deno = tl.where(deno > 0.0, deno, 1.0) + offs_o = ( + (q_start + q_offsets_m[:, None]) * stride_obs + + cur_head * stride_oh + + offs_dv[None, :] + ) + tl.store( + O + offs_o, + acc / safe_deno[:, None], + mask=mask_m[:, None] & mask_dv[None, :], + ) + + if HAS_LSE: + offs_lse = (q_start + q_offsets_m) * stride_lse_bs + cur_head * stride_lse_h + lse = tl.where(deno > 0.0, tl.log(deno) + e_max, float("-inf")) + tl.store(LSE + offs_lse, lse, mask=mask_m) + + +def mla_prefill_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + max_seqlen_q: int, + max_seqlen_kv: int, + softmax_scale: float, + *, + is_causal: bool, + logit_cap: float = 0.0, + lse: torch.Tensor | None = None, +) -> None: + if q.shape[-1] != k.shape[-1]: + raise ValueError( + f"q/k head dims must match, got {q.shape[-1]} and {k.shape[-1]}" + ) + if q.shape[1] % k.shape[1] != 0: + raise ValueError( + "num_q_heads must be divisible by num_kv_heads, " + f"got {q.shape[1]} and {k.shape[1]}" + ) + if out.shape != (q.shape[0], q.shape[1], v.shape[-1]): + raise ValueError( + f"out shape must be {(q.shape[0], q.shape[1], v.shape[-1])}, " + f"got {tuple(out.shape)}" + ) + for name, tensor in (("q", q), ("k", k), ("v", v), ("out", out)): + if tensor.stride(-1) != 1: + raise ValueError(f"{name} must have contiguous last dimension") + if lse is not None and lse.shape != (q.shape[0], q.shape[1]): + raise ValueError( + f"lse shape must be {(q.shape[0], q.shape[1])}, got {tuple(lse.shape)}" + ) + + q_head_dim = q.shape[-1] + v_head_dim = v.shape[-1] + + if q_head_dim == 576: + block_dmodel = 512 + block_dpe = 64 + elif q_head_dim == 288: + block_dmodel = 256 + block_dpe = 32 + elif q_head_dim == 192: + block_dmodel = 128 + block_dpe = 64 + else: + block_dmodel = triton.next_power_of_2(q_head_dim) + block_dpe = 0 + block_dv = triton.next_power_of_2(v_head_dim) + block_m, block_n = (64, 64) + num_warps = 4 + + lse_arg = lse if lse is not None else out + grid = (cu_seqlens_q.shape[0] - 1, q.shape[1], triton.cdiv(max_seqlen_q, block_m)) + + _mla_prefill_kernel[grid]( + q, + k, + v, + out, + lse_arg, + cu_seqlens_q, + cu_seqlens_kv, + softmax_scale, + q.shape[1] // k.shape[1], + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + out.stride(0), + out.stride(1), + lse_arg.stride(0), + lse_arg.stride(1), + logit_cap=logit_cap, + Lq=q_head_dim, + Lv=v_head_dim, + BLOCK_DMODEL=block_dmodel, + BLOCK_DPE=block_dpe, + BLOCK_DV=block_dv, + BLOCK_M=block_m, + BLOCK_N=block_n, + IS_CAUSAL=is_causal, + HAS_LSE=lse is not None, + num_warps=num_warps, + num_stages=1, + ) + + +def triton_mla_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + max_seqlen_q: int, + max_seqlen_kv: int, + softmax_scale: float, + *, + is_causal: bool = True, + logit_cap: float = 0.0, + return_lse: bool = False, + out: torch.Tensor | None = None, + seq_lens_kv: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if out is None: + out_dtype = torch.bfloat16 if q.dtype in _FP8_DTYPES else q.dtype + out = torch.empty( + (q.shape[0], q.shape[1], v.shape[-1]), + dtype=out_dtype, + device=q.device, + ) + + lse = ( + torch.empty((q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device) + if return_lse + else None + ) + mla_prefill_fwd( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + softmax_scale, + is_causal=is_causal, + logit_cap=logit_cap, + lse=lse, + ) + if return_lse: + return out, lse + return out