class MambaPool:
    def __init__(self, ...):
        # Conv State: 每层一个 tensor
        # qwen3-next是分别对qkv做卷积，这里把他们打包在conv_dim一起存储了，所以没有体现出来
        # 形状: [num_mamba_layers, size+1, conv_dim/tp, conv_kernel-1]
        conv_state = [
            torch.zeros(
                size=(num_mamba_layers, size + 1) + conv_shape,
                dtype=conv_dtype,  # 默认 bfloat16
                device=device,
            )
            for conv_shape in conv_state_shape  # 通常是 [(conv_dim/tp, conv_kernel-1)]
        ]

        # SSM State (Temporal State): 单个 tensor
        # 形状: [num_mamba_layers, size+1, num_heads/tp, head_dim, state_size]
        temporal_state = torch.zeros(
            size=(num_mamba_layers, size + 1) + temporal_state_shape,
            dtype=ssm_dtype,  # 默认 float32
            device=device,
        )

class TreeNode:
    key: RadixKey              # token 序列
    value: torch.Tensor        # KV cache indices (每个 token 对应的索引)
    mamba_value: torch.Tensor  # Mamba slot index (单个索引，指向 MambaPool)

def alloc(self, reqs: List["Req"]):
    for req in reqs:
        if req.mamba_pool_idx is not None:
            # 已有索引（来自 radix cache 前缀匹配）
            mid = req.mamba_pool_idx
        else:
            # 分配新 slot
            mid = self.mamba_pool.alloc(1)
            req.mamba_pool_idx = mid[0]

    # 建立映射
    self.req_index_to_mamba_index_mapping[select_index] = mamba_indices

def alloc(self, need_size: int):
    select_index = self.free_slots[:need_size]
    self.free_slots = self.free_slots[need_size:]

    # 分配时清零，避免脏数据
    for conv_tensor in self.mamba_cache.conv:
        conv_tensor[:, select_index] = 0
    self.mamba_cache.temporal[:, select_index] = 0

    return select_index

# 获取状态索引
state_indices_tensor = metadata.mamba_cache_indices  # [num_reqs]
conv_state = layer_cache.conv[0]
ssm_state = layer_cache.temporal

# Conv 计算，同时更新 conv_state
hidden_states_B_C = causal_conv1d_fn(
    x, conv_weights, bias,
    conv_states=conv_state,              # 被更新
    cache_indices=state_indices_tensor,  # 每个请求的 slot index
    query_start_loc=query_start_loc,
)

# SSM 计算，返回 final state
varlen_state = mamba_chunk_scan_combined(
    hidden_states, dt, A, B, C, ...
    initial_states=initial_states,  # 来自前缀匹配
    return_varlen_states=True,
)

# 更新 ssm_state
ssm_state[state_indices_tensor] = varlen_state

# Conv 增量更新（单 token）
hidden_states_B_C = causal_conv1d_update(
    hidden_states_B_C,
    conv_state,  # in-place 更新
    conv_weights, bias,
    conv_state_indices=state_indices_tensor,
)

# SSM 增量更新（单 token）
selective_state_update(
    ssm_state,  # in-place 更新
    hidden_states, dt, A, B, C, D,
    state_batch_indices=state_indices_tensor,
    out=output,
)

def _match_prefix_helper(self, key: RadixKey):
    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

        # 处理 split 情况（部分匹配）
        if prefix_len < len(child.key):
            new_node = self._split_node(child.key, child, prefix_len)
            if child.mamba_value is not None:
                best_value_len = len(value)
                best_last_node = child
            break

        value.append(child.value)
        node = child

if cow_mamba and last_node.mamba_value is not None:
    if req.mamba_pool_idx is None:
        dst_index = self.mamba_pool.alloc(1)
    else:
        dst_index = req.mamba_pool_idx.unsqueeze(0)

    src_index = last_node.mamba_value  # radix tree 中的 slot
    self.mamba_pool.copy_from(src_index, dst_index)  # 复制状态

def cache_finished_req(self, req: Req, is_insert: bool = True):
    # Fork: 分配新 slot 并复制状态
    mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone()
    mamba_value_forked = self.mamba_pool.fork_from(mamba_value)

    # 插入 radix tree
    result = self.insert(..., mamba_value=mamba_value_forked)

    # 释放请求的原始 slot
    self.req_to_token_pool.free_mamba_cache(req)

def cache_unfinished_req(self, req: Req, chunked=False):
    token_ids = req.fill_ids
    cache_len = len(token_ids)  # no_buffer 模式下缓存整个 chunk

    # 获取当前请求的 mamba slot index
    mamba_value = self.req_to_token_pool.get_mamba_indices(
        req.req_pool_idx
    ).unsqueeze(-1)

    # Fork 到新 slot（与 cache_finished_req 相同）
    mamba_value_forked = self.mamba_pool.fork_from(mamba_value)

    # 插入 radix tree
    result = self.insert(
        InsertParams(
            key=RadixKey(page_aligned_token_ids, req.extra_key),
            value=page_aligned_kv_indices,
            mamba_value=mamba_value_forked,
        )
    )

    # 关键区别：不释放原始 mamba slot！
    # 请求还在进行中，需要继续使用

    # 更新请求的 last_node 用于后续 chunk
    req.last_node = new_last_node

else:
    mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone()
    mamba_ping_pong_track_buffer_to_keep = None

result = self.insert(
    InsertParams(
        key=RadixKey(token_ids[:page_aligned_len], req.extra_key),
        value=page_aligned_kv_indices,
        mamba_value=mamba_value,
    )
)

else:
    mamba_value = self.req_to_token_pool.get_mamba_indices(
        req.req_pool_idx
    ).unsqueeze(-1)
# 关键：fork_from 会复制状态到新 slot
mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from(mamba_value)

result = self.insert(
    InsertParams(
        key=RadixKey(page_aligned_token_ids, req.extra_key),
        value=page_aligned_kv_indices,
        mamba_value=mamba_value_forked,  # 新 slot 的 index
    )
)

while len(key) > 0 and child_key in node.children.keys():
    child = node.children[child_key]
    # 更新 best_value_len
    if node.mamba_value is not None:
        best_value_len = len(value)
        best_last_node = node

    prefix_len = self.key_match_fn(child.key, key)
    if prefix_len < len(child.key):  # 1007 < 1008，触发 split
        new_node = self._split_node(child.key, child, prefix_len)
        value.append(new_node.value)
        node = new_node
        break  # 循环结束，但没有更新 best_value_len！

CHUNK_SIZE = 64

# cache_finished_req 第 508-512 行
cache_len = (
    req.mamba_last_track_seqlen  # ✅ extra_buffer: 对齐后的长度
    if self.enable_mamba_extra_buffer
    else len(token_ids)          # ❌ no_buffer: 完整长度，未对齐！
)

# schedule_batch.py - extra_buffer 模式的对齐计算
mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size  # 通常是 64

mamba_track_seqlen_aligned = (
    len(req.prefix_indices)
    + (req.extend_input_len // mamba_cache_chunk_size)
    * mamba_cache_chunk_size  # 向下对齐到 64 的倍数
)
req.mamba_last_track_seqlen = mamba_track_seqlen_aligned

# There are 3 cases for mamba_track_seqlen:
# 1) aligned with FLA_CHUNK_SIZE -> retrieve from last_recurrent_state
#    a) is the last position -> retrieve from last_recurrent_state
#    b) is NOT the last position -> retrieve from h
# 2) unaligned with FLA_CHUNK_SIZE -> retrieve from h

# cache_finished_req 修复
mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size

# no_buffer 模式也需要做 FLA 对齐
cache_len = (
    req.mamba_last_track_seqlen
    if self.enable_mamba_extra_buffer
    else (len(token_ids) // mamba_cache_chunk_size) * mamba_cache_chunk_size  # 向下对齐
)