# chunk_delta_h.py
CHUNK_SIZE = 64  # 固定chunk大小

# 序列被分成多个chunk:
# [0-63], [64-127], [128-191], ...

# Linear Attention的状态更新伪代码
def linear_attention(tokens):
    h = initial_state
    outputs = []

    for t in range(len(tokens)):
        # 每一步都依赖前一步的状态
        h = delta_rule(h, tokens[t])
        outputs.append(gate(h, tokens[t]))

    return outputs

# 如果要从位置 n 开始计算，必须有 h[n-1] 的状态

# chunk_delta_h.py
CHUNK_SIZE = 64  # FLA固定的chunk大小

# mamba_radix_cache.py
from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE

# mamba_radix_cache.py:955-963
# Calculate the branching point. It is defined as the last aligned position that
# does not have a mamba value.
if len(value) > best_value_len:
    fla_chunk_aligned_seqlen = (
        sum(len(v) for v in value) // FLA_CHUNK_SIZE
    ) * FLA_CHUNK_SIZE  # 向下对齐到64的倍数
    mamba_branching_seqlen = (
        fla_chunk_aligned_seqlen if fla_chunk_aligned_seqlen > 0 else None
    )

# chunk_gated_delta_rule_fwd_h 返回
h: [B, NT, H, K, V]  # NT = num_chunks
  # h[0]: chunk 0 后的状态
  # h[1]: chunk 1 后的状态
  # h[2]: chunk 2 后的状态
  # ...

last_recurrent_state: [N, H, K, V]  # N = batch_size
  # 整个序列的最终状态

# 缓存时只能保存:
# 1. h[i] - chunk边界的状态
# 2. last_recurrent_state - 序列末尾状态

class TreeNode:
    value: Optional[torch.Tensor]        # KV cache索引
    mamba_value: Optional[torch.Tensor]   # Mamba cache索引

    full_lock_ref: int     # KV cache锁计数
    mamba_lock_ref: int    # Mamba cache锁计数

# 示例树结构:
Root
├─ [0-95] (value=[...], mamba_value=[...])  # 有完整cache
├─ [96-150] (value=[...], mamba_value=None)  # Tombstone! 只有KV cache
└─ [96-191] (value=[...], mamba_value=[...]) # 另一条路径，有完整cache

# mamba_radix_cache.py

def evict_mamba(self, mamba_num: int):
    """只驱逐mamba cache，保留KV cache"""
    while need_evict:
        if len(x.children) > 0:
            # 内部节点: 只释放mamba，变成tombstone
            self.req_to_token_pool.mamba_pool.free(x.mamba_value)
            self._tombstone_internal_node(x)  # mamba_value = None
        else:
            # 叶子节点: 同时释放两者
            self._evict_leaf_node(x, True)

def evict(self, full_num_tokens: int):
    """驱逐KV cache (叶子节点)，同时释放mamba"""
    # 只能驱逐叶子节点
    # 同时释放 value 和 mamba_value

# mamba_radix_cache.py:900-967
def _match_prefix_helper(self, key: RadixKey) -> Tuple[List, TreeNode, Optional[int]]:
    """
    返回: (matched_value_list, last_node, mamba_branching_seqlen)

    核心逻辑:
    1. 遍历radix树，找到最长匹配路径
    2. 记录最后一个有 mamba_value 的节点
    3. 计算mamba_branching_seqlen
    """
    node = self.root_node
    value = []
    best_value_len = 0
    best_last_node = node

    # 遍历树
    while len(key) > 0 and child_key in node.children.keys():
        child = node.children[child_key]

        # 更新最佳节点: 当前节点必须有mamba_value
        if node.mamba_value is not None:
            best_value_len = len(value)
            best_last_node = node

        # ... 继续匹配

    # 计算branching point
    if len(value) > best_value_len:
        # 向下对齐到chunk边界
        fla_chunk_aligned_seqlen = (
            sum(len(v) for v in value) // FLA_CHUNK_SIZE
        ) * FLA_CHUNK_SIZE
        mamba_branching_seqlen = (
            fla_chunk_aligned_seqlen if fla_chunk_aligned_seqlen > 0 else None
        )
    else:
        mamba_branching_seqlen = None

    return value[:best_value_len], best_last_node, mamba_branching_seqlen

# ForwardBatch (schedule_batch.py)
mamba_track_indices: torch.Tensor    # [b] 追踪目标cache slot
mamba_track_mask: torch.Tensor       # [b] 是否需要追踪
mamba_track_seqlens: torch.Tensor    # [b] 追踪的序列长度

# Req (schedule_batch.py)
mamba_ping_pong_track_buffer: Tensor  # [2] ping-pong缓冲区
mamba_next_track_idx: int             # 下一个使用的buffer索引
mamba_last_track_seqlen: Optional[int] # 上次追踪的seq_len
mamba_branching_seqlen: Optional[int]  # 分支点seq_len

# schedule_batch.py:1612-1664
def _mamba_radix_cache_v2_req_prepare_for_extend(self, req, ...):
    # 计算是否需要追踪
    mask = (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE > 0
    mamba_track_mask_cpu.append(mask)

    if mask:
        # 完整追踪长度
        mamba_track_seqlen = len(req.prefix_indices) + req.extend_input_len

        # 对齐后的追踪长度
        mamba_track_seqlen_aligned = (
            len(req.prefix_indices)
            + (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
        )

        # 处理分支点
        if req.mamba_branching_seqlen is not None:
            branching_seqlen_aligned_mask = (
                req.mamba_branching_seqlen - len(req.prefix_indices)
            ) % FLA_CHUNK_SIZE == 0

            if (req.mamba_branching_seqlen > len(req.prefix_indices)
                and req.mamba_branching_seqlen < mamba_track_seqlen
                and branching_seqlen_aligned_mask):
                # 需要追踪分支点
                mamba_track_seqlen = req.mamba_branching_seqlen + 1
                mamba_track_seqlen_aligned = req.mamba_branching_seqlen

        req.mamba_last_track_seqlen = mamba_track_seqlen_aligned

# hybrid_linear_attn_backend.py:259-298
def _init_track_conv_indices(self, query_start_loc, forward_batch):
    """
    计算需要追踪的卷积状态的位置

    Conv state存储最后 conv_kernel_size 个输入
    需要从输入序列中提取这些位置的数据
    """
    conv_state_len = self.conv_states_shape[-1]  # conv_kernel_size

    # 计算对齐长度
    lens_to_track = forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens
    aligned_len = (lens_to_track // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE

    # 起始位置 = 对齐长度 - conv_kernel_size
    start_indices = query_start_loc[:-1] + aligned_len - conv_state_len
    start_indices = start_indices[forward_batch.mamba_track_mask]

    # 创建索引: [num_tracked, conv_kernel_size]
    indices = start_indices.unsqueeze(-1) + torch.arange(conv_state_len, ...)

    return indices.clamp(0, query_start_loc[-1] - 1)

# hybrid_linear_attn_backend.py:300-378
def _init_track_ssm_indices(self, mamba_cache_indices, forward_batch):
    """
    计算SSM状态的追踪索引

    FLA kernel输出:
    - h: [NT, H, K, V] - 每个chunk边界的状态
    - last_recurrent_state: [N, H, K, V] - 序列末尾状态

    需要决定从h还是last_recurrent_state提取状态
    """
    # 计算chunk数量
    num_h_states = (extend_seq_lens - 1) // FLA_CHUNK_SIZE + 1

    # 计算偏移
    track_ssm_src_offset = torch.zeros_like(num_h_states)
    track_ssm_src_offset[1:] = torch.cumsum(num_h_states[:-1], dim=0)

    # 过滤需要追踪的请求
    lens_to_track = mamba_track_seqlens - prefix_lens
    lens_masked = lens_to_track[mamba_track_mask]
    offset_masked = track_ssm_src_offset[mamba_track_mask]
    dst_masked = mamba_track_indices[mamba_track_mask]

    # 判断是否对齐
    is_aligned = (lens_masked % FLA_CHUNK_SIZE) == 0

    # Case 1: 对齐 - 从 last_recurrent_state
    track_ssm_final_src = mamba_cache_indices[mamba_track_mask][is_aligned]
    track_ssm_final_dst = dst_masked[is_aligned]

    # Case 2: 不对齐 - 从 h (中间状态)
    not_aligned = ~is_aligned
    track_ssm_h_src = offset_masked[not_aligned] + (
        lens_masked[not_aligned] // FLA_CHUNK_SIZE
    )
    track_ssm_h_dst = dst_masked[not_aligned]

    return track_ssm_h_src, track_ssm_h_dst, track_ssm_final_src, track_ssm_final_dst

# mamba_radix_cache.py:440-459
def match_prefix(self, key, cow_mamba=False, req=None):
    value, last_node, mamba_branching_seqlen = self._match_prefix_helper(key)

    # Copy-on-Write for Mamba
    if cow_mamba and last_node.mamba_value is not None:
        if req.mamba_pool_idx is None:
            # 分配新的cache slot
            dst_index = self.req_to_token_pool.mamba_pool.alloc(1)
            if dst_index is None:
                # 空间不足，驱逐后重试
                self.inc_lock_ref(last_node)
                self.evict_mamba(1)
                dst_index = self.req_to_token_pool.mamba_pool.alloc(1)
                self.dec_lock_ref(last_node)

            # 复制状态
            src_index = last_node.mamba_value
            self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index)
            req.mamba_pool_idx = dst_index
        else:
            # 已有独立cache，更新
            src_index = last_node.mamba_value
            dst_index = req.mamba_pool_idx.unsqueeze(0)
            self.req_to_token_pool.mamba_pool.copy_from(src_index, dst_index)

    return MatchResult(device_indices=torch.cat(value), ...)

# memory_pool.py:304-309
def fork_from(self, src_index: torch.Tensor) -> Optional[torch.Tensor]:
    """Fork一个新的mamba cache副本"""
    dst_index = self.alloc(1)
    if dst_index == None:
        return None
    self.copy_from(src_index, dst_index)
    return dst_index

def copy_from(self, src_index, dst_index):
    """复制conv和temporal状态"""
    for i in range(len(self.mamba_cache.conv)):
        self.mamba_cache.conv[i][:, dst_index] = self.mamba_cache.conv[i][:, src_index]
    self.mamba_cache.temporal[:, dst_index] = self.mamba_cache.temporal[:, src_index]

# mamba_radix_cache.py:576-689
def cache_unfinished_req(self, req: Req):
    """
    缓存未完成的请求 (chunked prefill)
    """
    # Fork mamba cache
    mamba_value = self.req_to_token_pool.get_mamba_indices(req.req_pool_idx)
    mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from(mamba_value)

    # 如果fork失败，驱逐后重试
    if mamba_value_forked is None:
        self.evict_mamba(1)
        mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from(mamba_value)

    # 插入新的树节点
    new_prefix_len, mamba_exist = self.insert(
        RadixKey(page_aligned_token_ids, req.extra_key),
        page_aligned_kv_indices,
        mamba_value_forked,
    )

    # 如果已有mamba cache，释放fork的副本
    if mamba_exist:
        self.req_to_token_pool.mamba_pool.free(mamba_value_forked)

# mamba_radix_cache.py:480-572
def cache_finished_req(self, req: Req, is_insert: bool = True):
    # 获取已提交的KV长度
    kv_committed_len = req.pop_committed_kv_cache()

    token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
    kv_indices = self.req_to_token_pool.req_to_token[
        req.req_pool_idx, :kv_committed_len
    ]

    if is_insert:
        # 计算实际缓存长度
        cache_len = (
            req.mamba_last_track_seqlen
            if self.enable_mamba_extra_buffer
            else len(token_ids)
        )

        # 对齐到page_size
        if self.page_size != 1:
            page_aligned_len = len(kv_indices) // self.page_size * self.page_size
            page_aligned_kv_indices = kv_indices[:page_aligned_len]
            page_aligned_token_ids = token_ids[:page_aligned_len]

        # 从ping-pong buffer获取mamba cache
        if self.enable_mamba_extra_buffer:
            mamba_ping_pong_track_buffer_to_keep = (
                self.req_to_token_pool.get_mamba_ping_pong_other_idx(
                    req.mamba_ping_pong_track_buffer
                )
            )
            mamba_value = (
                req.mamba_ping_pong_track_buffer[
                    mamba_ping_pong_track_buffer_to_keep
                ].unsqueeze(-1).clone()
            )
        else:
            mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone()

        # 插入radix树
        new_prefix_len, mamba_exist = self.insert(
            RadixKey(page_aligned_token_ids, req.extra_key),
            page_aligned_kv_indices,
            mamba_value,
        )

        # 释放未使用的KV cache
        self.token_to_kv_pool_allocator.free(
            kv_indices[req.cache_protected_len : new_prefix_len]
        )

        # 释放mamba cache (如果已存在)
        if mamba_exist:
            mamba_ping_pong_track_buffer_to_keep = None

        # 释放请求slot
        self.req_to_token_pool.free(
            req.req_pool_idx,
            free_mamba_cache=free_mamba_cache,
            mamba_ping_pong_track_buffer_to_keep=mamba_ping_pong_track_buffer_to_keep,
        )

# hybrid_linear_attn_backend.py:585-617
def _track_mamba_state_extend(self, forward_batch, h, ssm_states, forward_metadata):
    """
    在extend阶段追踪SSM状态到缓存slot
    """
    if (forward_batch.mamba_track_mask is not None
        and forward_batch.mamba_track_mask.any()):

        h = h.squeeze(0)

        # Case 1: 从h获取 (不对齐)
        if forward_metadata.track_ssm_h_src.numel() > 0:
            ssm_states[forward_metadata.track_ssm_h_dst] = h[
                forward_metadata.track_ssm_h_src
            ].to(ssm_states.dtype, copy=False)

        # Case 2: 从last_recurrent_state获取 (对齐)
        if forward_metadata.track_ssm_final_src.numel() > 0:
            ssm_states[forward_metadata.track_ssm_final_dst] = ssm_states[
                forward_metadata.track_ssm_final_src
            ]

# hybrid_linear_attn_backend.py:997-1023
# 在forward_extend中

# 使用gather提取需要追踪的conv状态
if (forward_batch.mamba_track_mask is not None
    and forward_batch.mamba_track_mask.any()):

    conv_dst = forward_batch.mamba_track_indices

    # 从输入序列中提取conv window
    mixed_qkv_to_track = mixed_qkv[
        :, forward_metadata.track_conv_indices
    ].transpose(0, 1)

    # 写入缓存
    mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
    conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track

# hybrid_linear_attn_backend.py:555-583
def _track_mamba_state_decode(self, forward_batch, conv_states, ssm_states, cache_indices):
    """
    在decode阶段，将更新的状态复制到持久缓存slot
    """
    if forward_batch.mamba_track_mask is not None:
        track_mamba_states_if_needed(
            conv_states,      # 源: 工作slot
            ssm_states,       # 源: 工作slot
            cache_indices,    # 源索引
            forward_batch.mamba_track_mask,
            forward_batch.mamba_track_indices,  # 目标索引
            forward_batch.batch_size,
        )

# hybrid_linear_attn_backend.py:64-122
@triton.jit
def track_mamba_state_if_needed_kernel(
    conv_states_ptr,
    ssm_states_ptr,
    cache_indices_ptr,      # 源索引
    mamba_track_mask_ptr,   # 是否追踪
    mamba_track_indices_ptr,# 目标索引
    ...
):
    batch_idx = tl.program_id(0)

    # 检查mask
    track_mask = tl.load(mamba_track_mask_ptr + batch_idx)
    if not track_mask:
        return

    # 加载索引
    src_idx = tl.load(cache_indices_ptr + batch_idx)
    dst_idx = tl.load(mamba_track_indices_ptr + batch_idx)

    # 复制conv states
    for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE):
        src_ptr = conv_states_ptr + src_idx * conv_state_stride_0 + offset
        dst_ptr = conv_states_ptr + dst_idx * conv_state_stride_0 + offset
        data = tl.load(src_ptr, mask=mask, other=0.0)
        tl.store(dst_ptr, data, mask=mask)

    # 复制SSM states
    for offset in range(0, ssm_state_numel_per_row, BLOCK_SIZE):
        src_ptr = ssm_states_ptr + src_idx * ssm_state_stride_0 + offset
        dst_ptr = ssm_states_ptr + dst_idx * ssm_state_stride_0 + offset
        data = tl.load(src_ptr, mask=mask, other=0.0)
        tl.store(dst_ptr, data, mask=mask)

@dataclass
class ForwardBatch:
    # Mamba追踪
    mamba_track_indices: torch.Tensor      # [b] 目标cache slot
    mamba_track_mask: torch.Tensor         # [b] 是否需要追踪
    mamba_track_seqlens: torch.Tensor       # [b] 追踪到的seq_len

    # Extend相关
    extend_seq_lens: torch.Tensor          # [b] 每个请求的输入长度
    extend_prefix_lens: torch.Tensor        # [b] 每个请求的前缀长度
    extend_seq_lens_cpu: List[int]         # [b] CPU上的seq_len

@dataclass(kw_only=True)
class ForwardMetadata:
    query_start_loc: torch.Tensor         # 累积序列长度
    mamba_cache_indices: torch.Tensor      # Cache slot索引

    # 追踪相关
    track_conv_indices: Optional[torch.Tensor]  # Conv状态位置
    track_ssm_h_src: Optional[torch.Tensor]     # SSM h源索引
    track_ssm_h_dst: Optional[torch.Tensor]     # SSM h目标索引
    track_ssm_final_src: Optional[torch.Tensor]  # SSM final源索引
    track_ssm_final_dst: Optional[torch.Tensor]  # SSM final目标索引

class Req:
    # Mamba cache slot
    mamba_pool_idx: Optional[int]            # 当前使用的slot

    # Ping-pong buffer
    mamba_ping_pong_track_buffer: Tensor    # [2] 双缓冲
    mamba_next_track_idx: int                # 下一个buffer索引

    # 追踪状态
    mamba_last_track_seqlen: Optional[int]   # 上次追踪的seq_len
    mamba_branching_seqlen: Optional[int]    # 分支点seq_len

    # 前缀信息
    prefix_indices: torch.Tensor             # 匹配的token索引
    last_node: TreeNode                     # Radix树节点

class MambaPool:
    @dataclass(frozen=True, kw_only=True)
    class State:
        conv: List[torch.Tensor]             # [num_layers, pool_size+1, dim]
        temporal: torch.Tensor               # [num_layers, pool_size+1, H, K, V]

    @dataclass(frozen=True, kw_only=True)
    class SpeculativeState(State):
        intermediate_ssm: torch.Tensor        # [num_layers, pool_size+1, draft, H, K, V]
        intermediate_conv_window: List[torch.Tensor]  # [num_layers, pool_size+1, draft, dim, K-1]