# 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 from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING import torch from tokenspeed_kernel import mla_decode_with_kvcache, mla_prefill from tokenspeed.runtime.configs.model_config import AttentionArch from tokenspeed.runtime.execution.forward_batch_info import ForwardMode from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend from tokenspeed.runtime.layers.attention.chunk import ( build_chunked_prefill_metadata_arrays, ) from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig from tokenspeed.runtime.layers.attention.registry import register_backend from tokenspeed.runtime.layers.attention.utils import build_page_table from tokenspeed.runtime.utils.common import ceil_div if TYPE_CHECKING: from tokenspeed.runtime.layers.paged_attention import PagedAttention @dataclass(kw_only=True) class MLAPrefillMetadata: # Device-side metadata for explicit Q/K/V MLA prefill and prefix replay. seq_lens: torch.Tensor req_pool_indices: torch.Tensor extend_prefix_lens: torch.Tensor extend_seq_lens: torch.Tensor cum_extend_seq_lens: torch.Tensor # Host-side metadata. extend_seq_lens_cpu: list[int] max_extend_seq_len: int max_extend_prefix_len: int # Per-prefix-chunk arrays consumed by DeepSeek's chunked prefix replay. chunked_loop_num: int chunk_kv_indices_list: list[torch.Tensor] chunked_seq_len: torch.Tensor cu_chunked_seq_len: torch.Tensor max_chunk_len_per_loop: list[int] @dataclass(kw_only=True) class MLADecodeMetadata: # num_extends lets mixed batches slice decode requests after extend requests. num_extends: int page_table: torch.Tensor seq_lens: torch.Tensor @property def block_kv_indices(self) -> torch.Tensor: return self.page_table @property def seq_lens_k(self) -> torch.Tensor: return self.seq_lens class MLAAttnBackend(AttentionBackend): """Unified MLA backend routed through tokenspeed_kernel MLA APIs.""" def __init__(self, config: MLAConfig): super().__init__(config) self.max_context_len = config.context_len self.page_size = config.page_size self.max_num_pages = ceil_div(self.max_context_len, self.page_size) self.kv_lora_rank = config.kv_lora_rank self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_rope_head_dim = config.qk_rope_head_dim self.v_head_dim = config.v_head_dim self.kv_cache_dim = config.kv_cache_dim self.scaling = config.scaling self.data_type = config.kv_cache_dtype self.q_data_type = config.dtype self.num_local_heads = config.num_attention_heads // config.attn_tp_size self.kernel_solution = None self.forward_decode_metadata: MLADecodeMetadata | None = None self.forward_prefill_metadata: MLAPrefillMetadata | None = None self.chunked_prefill_metadata: MLAPrefillMetadata | None = None self.decode_cuda_graph_metadata: dict[int, MLADecodeMetadata] = {} self.cuda_graph_page_table: torch.Tensor | None = None self.cuda_graph_seq_lens: torch.Tensor | None = None def init_forward_metadata( self, bs: int, num_extends: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, req_to_page: torch.Tensor, forward_mode: ForwardMode, extend_seq_lens: torch.Tensor | None = None, extend_seq_lens_cpu: torch.Tensor | None = None, extend_prefix_lens: torch.Tensor | None = None, extend_prefix_lens_cpu: torch.Tensor | None = None, **kwargs, ): if forward_mode.is_extend_or_mixed(): self._init_prefill_metadata( seq_lens=seq_lens[:num_extends], req_pool_indices=req_pool_indices[:num_extends], req_to_page=req_to_page, extend_prefix_lens=extend_prefix_lens[:num_extends], extend_prefix_lens_cpu=extend_prefix_lens_cpu[:num_extends], extend_seq_lens=extend_seq_lens[:num_extends], extend_seq_lens_cpu=extend_seq_lens_cpu[:num_extends], ) if ( forward_mode.is_decode() or forward_mode.is_mixed() or (forward_mode.is_extend() and self.is_draft) ): self._init_decode_metadata( bs=bs, num_extends=num_extends, req_pool_indices=req_pool_indices, seq_lens=seq_lens, req_to_page=req_to_page, ) @contextmanager def override_num_extends(self, num_extends: int): assert self.forward_decode_metadata is not None prev = self.forward_decode_metadata.num_extends self.forward_decode_metadata.num_extends = num_extends try: yield finally: self.forward_decode_metadata.num_extends = prev def _init_prefill_metadata( self, seq_lens: torch.Tensor, req_pool_indices: torch.Tensor, req_to_page: torch.Tensor, extend_prefix_lens: torch.Tensor, extend_prefix_lens_cpu: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_cpu: torch.Tensor, ): extend_seq_lens_cpu_list = [int(x) for x in extend_seq_lens_cpu.tolist()] cum_extend_seq_lens = torch.zeros( extend_seq_lens.shape[0] + 1, device=self.device, dtype=torch.int32, ) torch.cumsum(extend_seq_lens, dim=0, out=cum_extend_seq_lens[1:]) max_extend_seq_len = max(extend_seq_lens_cpu_list, default=0) max_extend_prefix_len = int(extend_prefix_lens_cpu.max().item()) ( chunked_loop_num, chunk_kv_indices_list, chunked_seq_len, cu_chunked_seq_len, max_chunk_len_per_loop, ) = build_chunked_prefill_metadata_arrays( extend_prefix_lens, extend_prefix_lens_cpu, req_to_page, req_pool_indices, self.page_size, ) metadata = MLAPrefillMetadata( seq_lens=seq_lens, req_pool_indices=req_pool_indices, extend_prefix_lens=extend_prefix_lens, extend_seq_lens=extend_seq_lens, cum_extend_seq_lens=cum_extend_seq_lens, extend_seq_lens_cpu=extend_seq_lens_cpu_list, max_extend_seq_len=max_extend_seq_len, max_extend_prefix_len=max_extend_prefix_len, chunked_loop_num=chunked_loop_num, chunk_kv_indices_list=chunk_kv_indices_list, chunked_seq_len=chunked_seq_len, cu_chunked_seq_len=cu_chunked_seq_len, max_chunk_len_per_loop=max_chunk_len_per_loop, ) self.forward_prefill_metadata = metadata self.chunked_prefill_metadata = metadata def _init_decode_metadata( self, bs: int, num_extends: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, req_to_page: torch.Tensor, ): page_table = build_page_table( req_pool_indices[:bs], req_to_page, self.page_size, self.max_context_len, ) self.forward_decode_metadata = MLADecodeMetadata( num_extends=num_extends, page_table=page_table, seq_lens=seq_lens[:bs], ) def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): assert ( seq_lens_buf.dtype == torch.int32 and seq_lens_buf.dim() == 1 and seq_lens_buf.shape[0] >= max_bs ), ( f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, " f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}" ) self.cuda_graph_page_table = torch.zeros( (max_bs, self.max_num_pages), dtype=torch.int32, device=self.device ) self.cuda_graph_seq_lens = seq_lens_buf self.decode_cuda_graph_metadata = {} def init_forward_metadata_capture_cuda_graph( self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, forward_mode: ForwardMode, ): if forward_mode.is_extend_or_mixed(): raise NotImplementedError( f"mla CUDA graph capture not supported for {forward_mode}" ) metadata = MLADecodeMetadata( num_extends=0, page_table=self.cuda_graph_page_table[:bs, :], seq_lens=self.cuda_graph_seq_lens[:bs], ) self.decode_cuda_graph_metadata[bs] = metadata self.forward_decode_metadata = metadata def init_forward_metadata_replay_cuda_graph( self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, forward_mode: ForwardMode = None, req_to_page: torch.Tensor = None, **kwargs, ): if forward_mode is not None and forward_mode.is_extend_or_mixed(): raise NotImplementedError( f"mla CUDA graph replay not supported for {forward_mode}" ) self.cuda_graph_page_table[:bs, : self.max_num_pages].copy_( req_to_page[req_pool_indices[:bs], : self.max_num_pages] ) self.forward_decode_metadata = self.decode_cuda_graph_metadata[bs] def get_cuda_graph_seq_len_fill_value(self): return 1 def forward_decode( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: PagedAttention, out_cache_loc: torch.Tensor, token_to_kv_pool, bs: int, save_kv_cache: bool = True, **kwargs, ) -> torch.Tensor: # q is absorbed MLA query [T, H, R + D_rope]; k is compressed KV # [T, 1, R + D_rope]. DeepSeek normally writes cache before this call. if save_kv_cache: assert k is not None token_to_kv_pool.set_mla_kv_buffer( layer, out_cache_loc, k[..., : self.kv_lora_rank], k[..., self.kv_lora_rank :], ) metadata = self.forward_decode_metadata assert metadata is not None num_extends = metadata.num_extends q_len_per_req = q.shape[0] // bs if bs > 0 else 1 if q_len_per_req > 1: query = q.view(-1, layer.tp_q_head_num, layer.head_dim).unsqueeze(1) page_table = metadata.page_table[num_extends:].repeat_interleave( q_len_per_req, dim=0 ) cache_seqlens = metadata.seq_lens[num_extends:].repeat_interleave( q_len_per_req ) # Draft catch-up starts from the current draft KV length; target # verify starts from the final target KV length and backs up. offset_start = 0 if self.is_draft else 1 - q_len_per_req offsets = torch.arange( offset_start, offset_start + q_len_per_req, device=cache_seqlens.device, dtype=cache_seqlens.dtype, ).repeat(bs) cache_seqlens = cache_seqlens + offsets max_seqlen_k = self.max_context_len else: query = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) page_table = metadata.page_table[num_extends:] cache_seqlens = metadata.seq_lens[num_extends:] max_seqlen_k = self.max_context_len softmax_scale = layer.scaling if self.data_type == torch.float8_e4m3fn: query = query.to(self.data_type) k_scale = ( layer.k_scale_float if getattr(layer, "k_scale_float", None) is not None else 1.0 ) softmax_scale = k_scale * softmax_scale kv_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) if self.data_type != kv_cache.dtype: kv_cache = kv_cache.to(self.data_type) kv_cache = kv_cache.view(-1, self.page_size, 1, self.kv_cache_dim) result = mla_decode_with_kvcache( q=query, kv_cache=kv_cache, page_table=page_table, cache_seqlens=cache_seqlens, max_seqlen_k=max_seqlen_k, qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, softmax_scale=softmax_scale, logit_cap=layer.logit_cap, solution=self.kernel_solution, ) output = self._unwrap_output(result) return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) def forward_extend( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: PagedAttention, out_cache_loc: torch.Tensor, token_to_kv_pool, bs: int, save_kv_cache: bool = True, **kwargs, ) -> torch.Tensor: if save_kv_cache: raise NotImplementedError( "MLA forward_extend cannot derive compressed cache rows from " "materialized K/V; DeepSeek writes MLA cache in the model path" ) metadata = self.forward_prefill_metadata assert metadata is not None if metadata.max_extend_prefix_len > 0: raise NotImplementedError( "MLA prefix-cache extend is handled by DeepSeek's chunked " "prefix replay path via forward_extend_chunked" ) q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) k = k.view(-1, layer.tp_k_head_num, layer.qk_head_dim) v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) result = mla_prefill( q=q, k=k, v=v, cu_seqlens_q=metadata.cum_extend_seq_lens, cu_seqlens_kv=metadata.cum_extend_seq_lens, max_seqlen_q=metadata.max_extend_seq_len, max_seqlen_kv=metadata.max_extend_seq_len, softmax_scale=layer.scaling, seq_lens_kv=metadata.extend_seq_lens, is_causal=True, logit_cap=layer.logit_cap, solution=self.kernel_solution, ) output = self._unwrap_output(result) return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) def forward_extend_chunked( self, q, k, v, scaling, logits_soft_cap=None, *, cum_seq_lens_q, cum_seq_lens_kv, max_q_len, max_kv_len, seq_lens, batch_size, causal, out: torch.Tensor | None = None, ): if causal: step_counter = getattr(self, "step_counter", None) if step_counter is not None: step_counter.record_cache() head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim q = q.reshape(-1, self.num_local_heads, head_dim) k = k.reshape(-1, self.num_local_heads, head_dim) v = v.reshape(-1, self.num_local_heads, self.v_head_dim) if q.dtype == torch.float8_e4m3fn: k = k.to(torch.float8_e4m3fn) v = v.to(torch.float8_e4m3fn) result = mla_prefill( q=q, k=k, v=v, cu_seqlens_q=cum_seq_lens_q, cu_seqlens_kv=cum_seq_lens_kv, max_seqlen_q=max_q_len, max_seqlen_kv=max_kv_len, softmax_scale=scaling, seq_lens_kv=seq_lens, is_causal=causal, logit_cap=logits_soft_cap or 0.0, return_lse=True, out=out, solution=self.kernel_solution, ) if isinstance(result, tuple): return result[0], result[1] return result, None def _unwrap_output(self, result): if isinstance(result, tuple): return result[0] return result register_backend("mla", {AttentionArch.MLA}, MLAAttnBackend)