# 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.ops.attention.flash_attn import flash_attn_varlen_func from tokenspeed_kernel.ops.attention.flash_mla import ( flash_mla_with_kvcache, get_mla_metadata, ) from tokenspeed_kernel.ops.attention.flashinfer import ( BatchMLAPagedAttentionWrapper, BatchPrefillWithRaggedKVCacheWrapper, ) 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 ( create_flashinfer_kv_indices_triton, ) from tokenspeed.runtime.spec_decode.eagle import ( EagleDraftInput, generate_attn_arg_prefill, ) from tokenspeed.runtime.utils.env import global_server_args_dict from tokenspeed.runtime.utils.flashinfer_config import get_flashinfer_workspace_size PAGE_SIZE = 64 if TYPE_CHECKING: from tokenspeed.runtime.layers.paged_attention import PagedAttention @dataclass class FlashMLADecodeMetadata: num_extends: int = 0 flashmla_metadata: tuple | None = None num_splits: torch.Tensor | None = None block_table: torch.Tensor | None = None @dataclass class _PrefillMetadata: prefill_wrapper: BatchMLAPagedAttentionWrapper use_ragged: bool @dataclass class _ChunkedPrefillMetadata: extend_prefix_lens: torch.Tensor extend_prefix_lens_cpu: torch.Tensor extend_seq_lens: torch.Tensor extend_seq_lens_cpu: torch.Tensor req_pool_indices: torch.Tensor cum_extend_seq_lens: torch.Tensor max_extend_seq_len: int chunked_loop_num: int chunk_kv_indices_list: list chunked_seq_len: torch.Tensor cu_chunked_seq_len: torch.Tensor max_chunk_len_per_loop: list # Shared across all flashinfer prefill wrappers used by FlashMLABackend. _global_workspace_buffer = None class FlashMLABackend(AttentionBackend): """FlashMLA attention backend for TokenSpeed scheduling. Uses the FlashMLA kernel for decode (any q_len); uses FlashInfer's MLA prefill wrappers for the EXTEND path. """ def __init__(self, config: MLAConfig): super().__init__(config) # Parse constants self.max_context_len = config.context_len self.kv_cache_quant_method = config.kv_cache_quant_method self.cache_dtype = config.kv_cache_dtype # MLA-specific dimensions 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_lora_rank + config.qk_rope_head_dim self.scaling = config.scaling self.softmax_scale = 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.num_q_heads = config.num_attention_heads // config.attn_tp_size # FlashMLA-specific self.draft_token_num = 0 if self.kv_cache_quant_method == "per_token_head": raise NotImplementedError( "FlashMLABackend no longer supports " "kv_cache_quant_method='per_token_head'." ) if self.cache_dtype == torch.float8_e4m3fn: raise NotImplementedError( "FlashMLABackend no longer supports dense FP8 KV cache. " "Use a non-FP8 KV cache." ) # Workspace buffer + flashinfer prefill wrappers (EXTEND path only). global _global_workspace_buffer if _global_workspace_buffer is None: _global_workspace_buffer = torch.empty( get_flashinfer_workspace_size(), dtype=torch.uint8, device=config.device, ) self.workspace_buffer = _global_workspace_buffer max_bs = config.max_bs self.kv_indptr = torch.zeros( (max_bs + 1,), dtype=torch.int32, device=config.device ) self.qo_indptr = torch.zeros( (max_bs + 1,), dtype=torch.int32, device=config.device ) self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper( self.workspace_buffer, "NHD" ) self.prefill_wrapper_paged = BatchMLAPagedAttentionWrapper( self.workspace_buffer, backend="auto", ) self.indices_updater_prefill = _PrefillIndicesUpdater(config, self) # Metadata state. Decode and prefill metadata are split so MIXED batches # can carry both simultaneously (decode-half + prefill-half sub-contexts # dispatch to their respective metadata). self.forward_decode_metadata: FlashMLADecodeMetadata | None = None self.forward_prefill_metadata: _PrefillMetadata | None = None self.chunked_prefill_metadata: _ChunkedPrefillMetadata | None = None self.last_seq_lens_sum: int | None = None # ------------------------------------------------------------------ # Metadata init # ------------------------------------------------------------------ def init_forward_metadata( self, bs: int, num_extends: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, forward_mode: ForwardMode, req_to_page: torch.Tensor = None, extend_with_prefix: bool = False, extend_prefix_lens: torch.Tensor | None = None, spec_info=None, **kwargs, ): if forward_mode.is_extend_or_mixed(): self._init_prefill_metadata( req_pool_indices=req_pool_indices[:num_extends], seq_lens=seq_lens[:num_extends], req_to_page=req_to_page, extend_with_prefix=extend_with_prefix, extend_prefix_lens=extend_prefix_lens, extend_prefix_lens_cpu=kwargs.pop("extend_prefix_lens_cpu"), extend_seq_lens=kwargs.pop("extend_seq_lens"), extend_seq_lens_cpu=kwargs.pop("extend_seq_lens_cpu"), ) # Under is_draft, also fill decode_metadata under any forward_mode so # the drafter's multi-step loop has metadata. Wrapper pre-writes # draft_seq_lens before calling here, so `seq_lens` aliases the # drafter's live buffer for step-1+ advances. if ( forward_mode.is_decode_or_idle() or forward_mode.is_mixed() or (forward_mode.is_extend() and self.is_draft) ): self._init_decode_metadata( bs, num_extends, req_pool_indices, seq_lens, 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_decode_metadata( self, bs: int, num_extends: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, req_to_page: torch.Tensor, ): if req_to_page is not None: block_table = req_to_page[req_pool_indices] else: block_table = None # When spec-dec is active (self.spec_num_tokens > 1), advance per-row # seq_lens by the worst-case verify width so the tile planner covers # the longest path. if self.spec_num_tokens > 1: plan_seq_lens = seq_lens + self.draft_token_num num_heads_plan = self.draft_token_num * self.num_q_heads else: plan_seq_lens = seq_lens num_heads_plan = self.num_q_heads mla_metadata, num_splits = get_mla_metadata( plan_seq_lens.to(torch.int32), num_heads_plan, 1, ) self.forward_decode_metadata = FlashMLADecodeMetadata( num_extends=num_extends, flashmla_metadata=mla_metadata, num_splits=num_splits, block_table=block_table, ) def _init_prefill_metadata( self, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, req_to_page: torch.Tensor, extend_with_prefix: bool, extend_prefix_lens: torch.Tensor | None, extend_prefix_lens_cpu: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_cpu: torch.Tensor, ): # EXTEND path — flashinfer ragged/paged prefill. if extend_prefix_lens is None: raise RuntimeError( "FlashMLABackend.init_forward_metadata requires " "extend_prefix_lens in extend mode." ) seq_lens_cpu = seq_lens.cpu() seq_lens_sum = seq_lens_cpu.sum().item() self.last_seq_lens_sum = seq_lens_sum extend_no_prefix = not extend_with_prefix use_ragged = ( not global_server_args_dict["mla_disable_ragged"] and extend_no_prefix ) self.indices_updater_prefill.update( req_pool_indices, seq_lens, seq_lens_sum, extend_prefix_lens, req_to_page=req_to_page, prefill_wrapper_paged=self.prefill_wrapper_paged, use_ragged=use_ragged, ) self.forward_prefill_metadata = _PrefillMetadata( self.prefill_wrapper_paged, use_ragged ) num_extends = extend_seq_lens.shape[0] cum_extend_seq_lens = torch.zeros( num_extends + 1, device=self.device, dtype=torch.int32 ) torch.cumsum(extend_seq_lens, dim=0, out=cum_extend_seq_lens[1:]) max_extend_seq_len = extend_seq_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, PAGE_SIZE, ) self.chunked_prefill_metadata = _ChunkedPrefillMetadata( extend_prefix_lens=extend_prefix_lens, extend_prefix_lens_cpu=extend_prefix_lens_cpu, extend_seq_lens=extend_seq_lens, extend_seq_lens_cpu=extend_seq_lens_cpu, req_pool_indices=req_pool_indices, cum_extend_seq_lens=cum_extend_seq_lens, max_extend_seq_len=max_extend_seq_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, ) # ------------------------------------------------------------------ # CUDA graph (decode only, any q_len) # ------------------------------------------------------------------ def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): del seq_lens_buf # flashmla allocates its own buffers. max_context_len = self.max_context_len + PAGE_SIZE - 1 # 4 PAGES are reserved for speculation cuda_graph_kv_indices = torch.full( (max_bs, (max_context_len + 4 * PAGE_SIZE) // PAGE_SIZE), 1, dtype=torch.int32, device="cuda", ) if self.draft_token_num: ( self.cuda_graph_mla_metadata, self.cuda_graph_num_splits, ) = get_mla_metadata( torch.ones( max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device ), self.draft_token_num * self.num_q_heads, 1, ) else: ( self.cuda_graph_mla_metadata, self.cuda_graph_num_splits, ) = get_mla_metadata( torch.ones( max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device ), self.num_q_heads, 1, ) self.cuda_graph_kv_indices = cuda_graph_kv_indices def init_forward_metadata_capture_cuda_graph( self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, forward_mode: ForwardMode, ): block_table = self.cuda_graph_kv_indices[:bs] is_target_verify = ( forward_mode.is_decode_or_idle() and not self.is_draft and self.spec_num_tokens > 1 ) is_draft_extend = ( forward_mode.is_decode_or_idle() and self.is_draft and self.spec_num_tokens > 1 ) if forward_mode.is_decode_or_idle() and self.spec_num_tokens == 1: mla_metadata, num_splits = get_mla_metadata( seq_lens.to(torch.int32), self.num_q_heads, 1, ) self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) self.cuda_graph_kv_indices[:bs].copy_(block_table) self.forward_decode_metadata = FlashMLADecodeMetadata( num_extends=0, flashmla_metadata=self.cuda_graph_mla_metadata, num_splits=self.cuda_graph_num_splits[: bs + 1], block_table=self.cuda_graph_kv_indices[:bs, :], ) elif is_target_verify or is_draft_extend: seq_lens = seq_lens + self.draft_token_num mla_metadata, num_splits = get_mla_metadata( seq_lens.to(torch.int32), self.draft_token_num * self.num_q_heads, 1, ) self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) self.cuda_graph_kv_indices[:bs].copy_(block_table) self.forward_decode_metadata = FlashMLADecodeMetadata( num_extends=0, flashmla_metadata=self.cuda_graph_mla_metadata, num_splits=self.cuda_graph_num_splits[: bs + 1], block_table=self.cuda_graph_kv_indices[:bs], ) else: raise RuntimeError(f"Not supported forward mode: {forward_mode}") 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 None or not forward_mode.is_decode_or_idle(): raise RuntimeError(f"Not supported forward mode: {forward_mode}") req_pool_indices = req_pool_indices[:bs] if req_to_page is not None: block_table = req_to_page[req_pool_indices] else: block_table = self.cuda_graph_kv_indices[:bs] seq_lens = seq_lens[:bs] is_target_verify = not self.is_draft and self.spec_num_tokens > 1 is_draft_extend = self.is_draft and self.spec_num_tokens > 1 if self.spec_num_tokens == 1: mla_metadata, num_splits = get_mla_metadata( seq_lens.to(torch.int32), self.num_q_heads, 1, ) elif is_target_verify or is_draft_extend: seq_lens = seq_lens + self.draft_token_num mla_metadata, num_splits = get_mla_metadata( seq_lens.to(torch.int32), self.draft_token_num * self.num_q_heads, 1, ) else: raise RuntimeError(f"Not supported forward mode: {forward_mode}") self.cuda_graph_mla_metadata.copy_(mla_metadata) self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) self.cuda_graph_kv_indices[:bs].copy_(block_table) self.forward_decode_metadata.num_extends = 0 self.forward_decode_metadata.flashmla_metadata = self.cuda_graph_mla_metadata self.forward_decode_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1] self.forward_decode_metadata.block_table = self.cuda_graph_kv_indices[:bs] def get_cuda_graph_seq_len_fill_value(self): return 1 # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ 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, seq_lens: torch.Tensor | None = None, forward_mode: ForwardMode | None = None, **kwargs, ): q_len_per_req = q.shape[0] // bs if bs > 0 else 1 is_target_verify = ( forward_mode is not None and forward_mode.is_decode_or_idle() and not self.is_draft and q_len_per_req > 1 ) is_draft_extend = ( forward_mode is not None and forward_mode.is_decode_or_idle() and self.is_draft and q_len_per_req > 1 ) if forward_mode is None or forward_mode.is_extend(): # Prefill: dispatch to ragged (MHA-style) or absorbed (MQA) path. if self.forward_prefill_metadata.use_ragged: return self._forward_normal_extend(q, k, v, layer, save_kv_cache) else: return self._forward_absorbed_extend( q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache, ) assert is_target_verify or is_draft_extend if k is not None: assert v is not None if save_kv_cache: token_to_kv_pool.set_kv_buffer(layer, out_cache_loc, k, v) metadata = self.forward_decode_metadata num_extends = metadata.num_extends bs = ( q.shape[0] if is_draft_extend else metadata.block_table.shape[0] - num_extends ) k_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) assert ( layer.tp_q_head_num == self.num_q_heads ), f"{layer.tp_q_head_num=} != {self.num_q_heads=}" reshape_q = q.view(bs, -1, self.num_q_heads, layer.head_dim) o, _ = flash_mla_with_kvcache( q=reshape_q, k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), block_table=metadata.block_table[num_extends : num_extends + bs], cache_seqlens=seq_lens.to(torch.int32) + self.draft_token_num, head_dim_v=self.kv_lora_rank, tile_scheduler_metadata=metadata.flashmla_metadata, num_splits=metadata.num_splits, softmax_scale=layer.scaling, causal=True, ) return o.view(-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 # flash_attn_varlen_func has no `out=` parameter; copy into the # caller-provided buffer at the end when requested. output, lse, *_ = flash_attn_varlen_func( q=q.view(-1, self.num_local_heads, head_dim), k=k.view(-1, self.num_local_heads, head_dim).to(q.dtype), v=v.view(-1, self.num_local_heads, self.v_head_dim).to(q.dtype), cu_seqlens_q=cum_seq_lens_q, cu_seqlens_k=cum_seq_lens_kv, max_seqlen_q=max_q_len, max_seqlen_k=max_kv_len, softmax_scale=scaling, causal=causal, return_attn_probs=True, ) if out is not None: out.copy_(output.view(out.shape)) output = out # lse must be transposed when using fa3. return output, lse.T.contiguous() 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, seq_lens: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: # Multi-token decode (target verify or drafter compound) reuses # the multi-token kernel path in forward_extend. q_len_per_req = q.shape[0] // bs if bs > 0 else 1 if q_len_per_req > 1: return self.forward_extend( q, k, v, layer, out_cache_loc, token_to_kv_pool, bs, save_kv_cache=save_kv_cache, seq_lens=seq_lens, forward_mode=ForwardMode.DECODE, **kwargs, ) if k is not None: assert v is not None if save_kv_cache: token_to_kv_pool.set_kv_buffer( layer, out_cache_loc, k, v, ) bs = q.shape[0] metadata = self.forward_decode_metadata num_extends = metadata.num_extends k_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) assert ( layer.tp_q_head_num == self.num_q_heads ), f"{layer.tp_q_head_num=} != {self.num_q_heads=}" reshape_q = q.view(bs, -1, self.num_q_heads, layer.head_dim) cache_lens = seq_lens o, _ = flash_mla_with_kvcache( q=reshape_q, k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), block_table=metadata.block_table[num_extends : num_extends + bs], cache_seqlens=cache_lens.to(torch.int32), head_dim_v=self.kv_lora_rank, tile_scheduler_metadata=metadata.flashmla_metadata, num_splits=metadata.num_splits, softmax_scale=layer.scaling, causal=True, ) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) # ------------------------------------------------------------------ # EXTEND prefill helpers # ------------------------------------------------------------------ def _forward_normal_extend( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: PagedAttention, save_kv_cache: bool = True, ): assert not save_kv_cache o = self.prefill_wrapper_ragged.forward( q, k.view(-1, layer.tp_k_head_num, layer.head_dim), v.view(-1, layer.tp_k_head_num, layer.v_head_dim), causal=True, sm_scale=layer.scaling, logits_soft_cap=layer.logit_cap, ) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) def _forward_absorbed_extend( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: PagedAttention, out_cache_loc: torch.Tensor, token_to_kv_pool, save_kv_cache: bool = True, ): # q is whole Q [T, H, head_dim]; k is whole latent [T, 1, head_dim]. # flashinfer prefill_wrapper.run() requires q_nope / q_pe split, so # slice views here (free) before handing off to the kernel. assert k is not None if save_kv_cache: token_to_kv_pool.set_mla_kv_buffer( layer, out_cache_loc, k[..., : layer.v_head_dim], k[..., layer.v_head_dim :], ) q = q.view(-1, layer.tp_q_head_num, layer.head_dim) q_nope = q[..., : layer.v_head_dim] q_pe = q[..., layer.v_head_dim :] o = q_nope.new_empty(q_nope.shape) k_buf = token_to_kv_pool.get_key_buffer(layer.layer_id).to(q_nope.dtype) o = self.forward_prefill_metadata.prefill_wrapper.run( q_nope, q_pe, k_buf[:, :, : layer.v_head_dim], k_buf[:, :, layer.v_head_dim :], out=o, ) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) class _PrefillIndicesUpdater: """Plans FlashInfer MLA prefill wrappers for the EXTEND path.""" def __init__(self, config: MLAConfig, attn_backend: FlashMLABackend): self.num_local_heads = config.num_attention_heads // config.attn_tp_size self.kv_cache_quant_method = config.kv_cache_quant_method 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.scaling = config.scaling self.data_type = config.kv_cache_dtype self.q_data_type = config.dtype self.attn_backend = attn_backend self.kv_indptr = attn_backend.kv_indptr self.qo_indptr = attn_backend.qo_indptr self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged def update( self, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, seq_lens_sum: int, prefix_lens: torch.Tensor, req_to_page: torch.Tensor = None, prefill_wrapper_paged: BatchMLAPagedAttentionWrapper = None, use_ragged: bool = False, spec_info: EagleDraftInput | None = None, ): if use_ragged: paged_kernel_lens = prefix_lens paged_kernel_lens_sum = 0 else: paged_kernel_lens = seq_lens paged_kernel_lens_sum = seq_lens_sum self._call_begin_forward( self.prefill_wrapper_ragged, prefill_wrapper_paged, req_pool_indices, paged_kernel_lens, paged_kernel_lens_sum, seq_lens, prefix_lens, self.kv_indptr, self.qo_indptr, use_ragged, req_to_page=req_to_page, spec_info=spec_info, ) def _call_begin_forward( self, wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper, wrapper_paged: BatchMLAPagedAttentionWrapper, req_pool_indices: torch.Tensor, paged_kernel_lens: torch.Tensor, paged_kernel_lens_sum: int, seq_lens: torch.Tensor, prefix_lens: torch.Tensor, kv_indptr: torch.Tensor, qo_indptr: torch.Tensor, use_ragged: bool, req_to_page: torch.Tensor = None, spec_info: EagleDraftInput | None = None, ): bs = len(seq_lens) sm_scale = self.scaling if spec_info is None: assert len(seq_lens) == len(req_pool_indices) torch.cumsum(paged_kernel_lens, dim=0, out=kv_indptr[1 : bs + 1]) kv_indptr = kv_indptr[: bs + 1] if wrapper_paged._use_cuda_graph: kv_indices = wrapper_paged._kv_indices_buf else: kv_indices = torch.empty( paged_kernel_lens_sum, dtype=torch.int32, device=req_pool_indices.device, ) if req_to_page is not None: create_flashinfer_kv_indices_triton[(bs,)]( req_to_page, req_pool_indices, paged_kernel_lens, kv_indptr, None, kv_indices, req_to_page.shape[1], ) torch.cumsum(seq_lens - prefix_lens, dim=0, out=qo_indptr[1 : bs + 1]) qo_indptr = qo_indptr[: bs + 1] else: kv_indices, kv_indptr, qo_indptr, _ = generate_attn_arg_prefill( spec_info.draft_token_num, req_pool_indices, paged_kernel_lens, req_to_page, ) if use_ragged: wrapper_ragged.begin_forward( qo_indptr=qo_indptr, kv_indptr=qo_indptr, num_qo_heads=self.num_local_heads, num_kv_heads=self.num_local_heads, head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, head_dim_vo=self.v_head_dim, q_data_type=self.q_data_type, ) else: kv_len_arr = kv_indptr[1:] - kv_indptr[:-1] wrapper_paged.plan( qo_indptr, kv_indptr, kv_indices, kv_len_arr, self.num_local_heads, self.kv_lora_rank, self.qk_rope_head_dim, 1, True, sm_scale, self.q_data_type, self.data_type, ) register_backend("flashmla", {AttentionArch.MLA}, FlashMLABackend)