lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/attention/configs/mha.py

# 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 dataclasses import dataclass

import torch

from tokenspeed.runtime.configs.model_config import ModelConfig
from tokenspeed.runtime.configs.paged_cache_spec import (
    STATE_LAYER_TYPES,
    scheduler_ext_flat_kvcache,
)
from tokenspeed.runtime.layers.attention.configs.base import (
    BaseAttnConfig,
    resolve_dtype,
)
from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool
from tokenspeed.runtime.utils.server_args import ServerArgs


@dataclass
class MHAConfig(BaseAttnConfig):
    # Per-layer attention-type labels + window, forwarded to the KV pool for
    # paged_cache_group_specs publication (empty -> single full-history group).
    layer_types: tuple[str, ...] = ()
    sliding_window_tokens: int | tuple[int | None, ...] | None = None
    max_scheduled_tokens: int = 0
    # True iff server_args.disaggregation_mode != "null"; the pool's slab
    # guards consume it.
    pd_disaggregation_enabled: bool = False
    # Mamba2/GDN per-state-layer shapes and dtypes (the configs'
    # mamba2_cache_params), forwarded to the pool's state slabs. Populated
    # only on a flat-built scheduler ext — the radix path keeps its
    # SimpleMambaPool state ownership byte-identical (None here means the
    # pool neither allocates state slabs nor runs the page-geometry check).
    conv_state_shape: tuple[int, ...] | None = None
    temporal_state_shape: tuple[int, ...] | None = None
    conv_dtype: torch.dtype | None = None
    ssm_dtype: torch.dtype | None = None

    @classmethod
    def generate(
        cls, server_args: ServerArgs, model_config: ModelConfig, is_draft: bool = False
    ):
        kwargs = {}
        if server_args.speculative_algorithm is not None:
            kwargs.update(
                speculative_num_steps=server_args.speculative_num_steps,
                speculative_num_draft_tokens=server_args.speculative_num_draft_tokens,
            )
        kv_cache_dtype = server_args.kv_cache_dtype
        draft_block_decode = bool(
            is_draft and server_args.speculative_algorithm == "DFLASH"
        )
        if draft_block_decode:
            kv_cache_dtype = "bfloat16"

        hf_config = getattr(model_config, "hf_config", None)
        layer_types = tuple(getattr(hf_config, "layer_types", None) or ())
        sliding_window_tokens = getattr(hf_config, "sliding_window", None)
        conv_state_shape = temporal_state_shape = None
        conv_dtype = ssm_dtype = None
        if (
            any(label in STATE_LAYER_TYPES for label in layer_types)
            and scheduler_ext_flat_kvcache()
        ):
            # GDN hybrid on the flat ext: the KV pool owns the recurrent
            # state (state slabs), so it needs the mamba2 shapes/dtypes.
            # Radix branch untouched: SimpleMambaPool owns the state there.
            text_config = getattr(hf_config, "text_config", hf_config)
            (
                conv_state_shape,
                temporal_state_shape,
                conv_dtype,
                ssm_dtype,
                _,
            ) = text_config.mamba2_cache_params
        return cls(
            device=server_args.device,
            context_len=model_config.context_len,
            backend_name=(
                server_args.attention_backend
                if not is_draft
                else server_args.drafter_attention_backend
            ),
            num_attention_heads=model_config.num_attention_heads,
            num_kv_heads=model_config.num_key_value_heads,
            head_dim=model_config.head_dim,
            attn_tp_size=server_args.attn_tp_size or server_args.mapping.attn.tp_size,
            dtype=model_config.dtype,
            kv_cache_dtype=resolve_dtype(kv_cache_dtype),
            page_size=server_args.block_size,
            max_bs=server_args.max_num_seqs
            // (server_args.data_parallel_size or server_args.mapping.attn.dp_size),
            max_graph_bs=server_args.max_cudagraph_capture_size,
            kv_cache_quant_method=server_args.kv_cache_quant_method,
            is_draft=is_draft,
            draft_block_decode=draft_block_decode,
            layer_types=layer_types,
            sliding_window_tokens=sliding_window_tokens,
            max_scheduled_tokens=getattr(server_args, "chunked_prefill_size", 8192),
            pd_disaggregation_enabled=getattr(
                server_args, "disaggregation_mode", "null"
            )
            != "null",
            conv_state_shape=conv_state_shape,
            temporal_state_shape=temporal_state_shape,
            conv_dtype=conv_dtype,
            ssm_dtype=ssm_dtype,
            **kwargs,
        )

    def cache_cell_size(self) -> int:
        return (
            max(self.num_kv_heads // self.attn_tp_size, 1)
            * self.head_dim
            * 2
            * torch._utils._element_size(self.kv_cache_dtype)
        )

    def create_pool(
        self,
        num_layers: int,
        max_total_num_tokens: int,
        rank: int,
        enable_memory_saver: bool,
    ) -> BaseTokenToKVPool:
        from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool

        return MHATokenToKVPool(
            size=max_total_num_tokens,
            dtype=self.kv_cache_dtype,
            head_num=max(self.num_kv_heads // self.attn_tp_size, 1),
            head_dim=self.head_dim,
            layer_num=num_layers,
            device=self.device,
            enable_memory_saver=enable_memory_saver,
            max_batch_size=self.max_bs,
            max_context_len=self.context_len,
            page_size=self.page_size,
            rank=rank,
            layer_types=self.layer_types,
            sliding_window_tokens=self.sliding_window_tokens,
            max_scheduled_tokens=self.max_scheduled_tokens,
            pd_disaggregation_enabled=self.pd_disaggregation_enabled,
            conv_state_shape=self.conv_state_shape,
            temporal_state_shape=self.temporal_state_shape,
            conv_dtype=self.conv_dtype,
            ssm_dtype=self.ssm_dtype,
        )