refactor: 分页 KV cache 替换固定 slot，删除 PrefixCache 及相关死代码
- 用 PagedCache + CacheView 替换固定 slot 式 KV cache，attention 层只通过 page_table 间接索引 - 删除 PrefixCache（radix tree）及 scheduler 中所有 prefix cache 命中/插入/释放逻辑 - 删除无用函数：pin、version、free_count、_mark_seq_mask 及 seq_mask 分配 - 修复 write 在多页 prefill 时 offset 为负导致 chunk 计算错误 - _make_page_table_tensor 改用 list 拼接一次 tensor，去掉逐元素赋值 - 清理 model 接口参数：kv_cache, slot_indices → paged_cache（CacheView） - 精简 docstring 为单行，删除冗余 section 注释和旧代码 - 修复 test_scheduler_concurrency.py 缺少 import pytest
This commit is contained in:
+23
-72
@@ -5,17 +5,11 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.inference.cache import CacheView
|
||||
|
||||
|
||||
def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
|
||||
"""
|
||||
Repeat k times along the dimension for attention heads.
|
||||
Args:
|
||||
x (Tensor): The input tensor.
|
||||
n_rep (int): The number of repetitions.
|
||||
Returns:
|
||||
Tensor: The repeated tensor.
|
||||
"""
|
||||
|
||||
"""Repeat KV heads n_rep times for GQA."""
|
||||
bs, slen, n_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
@@ -32,49 +26,25 @@ def get_rotary_emb(
|
||||
base: float = 10000,
|
||||
device: Optional[torch.device] = None,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Get the rotary embedding for the given dimension and maximum length.
|
||||
Args:
|
||||
dim (int): The dimension of the input.
|
||||
max_len (int): The maximum length of the input.
|
||||
base (float, optional): The base for the frequency. Defaults to 10000.
|
||||
device (optional): The device to create tensors on. Defaults to None.
|
||||
Returns:
|
||||
Tensor: The rotary embedding tensor.
|
||||
"""
|
||||
|
||||
"""Precompute cos/sin for RoPE."""
|
||||
theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
|
||||
t = torch.arange(0, max_len, dtype=torch.float64, device=device)
|
||||
freqs = torch.outer(t, theta)
|
||||
|
||||
return torch.cos(freqs).float(), torch.sin(freqs).float()
|
||||
|
||||
|
||||
def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tensor:
|
||||
"""
|
||||
Apply rotary embedding to the input tensor using cos/sin form.
|
||||
Args:
|
||||
x (Tensor): The input tensor (shape [..., seq_len, dim]).
|
||||
rotary_emb (Tuple[Tensor, Tensor]): The rotary embedding (shape [seq_len, dim//2]).
|
||||
Returns:
|
||||
Tensor: The output tensor (rotated, same shape as input).
|
||||
"""
|
||||
|
||||
"""Apply rotary embedding via cos/sin (shape-preserving)."""
|
||||
dtype = x.dtype
|
||||
cos, sin = rotary_emb
|
||||
|
||||
cos = cos.unsqueeze(0).unsqueeze(2) # [1, seq_len, 1, dim//2]
|
||||
sin = sin.unsqueeze(0).unsqueeze(2) # [1, seq_len, 1, dim//2]
|
||||
|
||||
x_real = x[..., 0::2] # [batch, seq_len, dim//2]
|
||||
x_imag = x[..., 1::2] # [batch, seq_len, dim//2]
|
||||
|
||||
cos = cos.unsqueeze(0).unsqueeze(2)
|
||||
sin = sin.unsqueeze(0).unsqueeze(2)
|
||||
x_real = x[..., 0::2]
|
||||
x_imag = x[..., 1::2]
|
||||
x_real_rot = x_real * cos - x_imag * sin
|
||||
x_imag_rot = x_real * sin + x_imag * cos
|
||||
|
||||
x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1) # [batch, seq_len, dim//2, 2]
|
||||
x_out = x_out.view(*x_out.shape[:-2], -1) # [batch, seq_len, dim]
|
||||
|
||||
x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)
|
||||
x_out = x_out.view(*x_out.shape[:-2], -1)
|
||||
return x_out.to(dtype)
|
||||
|
||||
|
||||
@@ -95,13 +65,10 @@ class RotaryEmbedding(nn.Module):
|
||||
|
||||
def forward(self, x: Tensor, start_pos: int = 0) -> Tuple[Tensor, Tensor]:
|
||||
seq_len = x.size(1)
|
||||
|
||||
if self.max_len_cached < seq_len + start_pos:
|
||||
self._set_rotary_buffer(self.max_len_cached * 2, x.device)
|
||||
|
||||
cos = self.cos_cached[start_pos : start_pos + seq_len]
|
||||
sin = self.sin_cached[start_pos : start_pos + seq_len]
|
||||
|
||||
return (cos, sin)
|
||||
|
||||
|
||||
@@ -185,14 +152,13 @@ class GQA(nn.Module):
|
||||
x: Tensor,
|
||||
rotary_emb: Tuple[Tensor, Tensor],
|
||||
mask: Tensor = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
paged_cache: Optional[CacheView] = None,
|
||||
start_pos: int = 0,
|
||||
slot_indices: Optional[Tensor] = None,
|
||||
) -> Tensor:
|
||||
bsz, seq_len, _ = x.size()
|
||||
is_causal = mask is None
|
||||
|
||||
# x(bsz, seq_len, n_heads * head_dim) -> (bsz, seq_len, n_heads, head_dim)
|
||||
# (bsz, seq_len, dim) -> (bsz, seq_len, n_heads, head_dim)
|
||||
q = self._split_heads(self.q_proj(x), self.n_heads)
|
||||
k = self._split_heads(self.k_proj(x), self.n_kv_heads)
|
||||
v = self._split_heads(self.v_proj(x), self.n_kv_heads)
|
||||
@@ -201,18 +167,14 @@ class GQA(nn.Module):
|
||||
if self.use_qk_norm:
|
||||
q, k = self.q_norm(q), self.k_norm(k)
|
||||
|
||||
if kv_cache is not None:
|
||||
k_cache, v_cache = kv_cache
|
||||
k_cache[slot_indices, start_pos : start_pos + seq_len, self.layer_id] = k
|
||||
v_cache[slot_indices, start_pos : start_pos + seq_len, self.layer_id] = v
|
||||
k = k_cache[slot_indices, : start_pos + seq_len, self.layer_id]
|
||||
v = v_cache[slot_indices, : start_pos + seq_len, self.layer_id]
|
||||
if paged_cache is not None:
|
||||
paged_cache.write(self.layer_id, start_pos, k, v)
|
||||
k, v = paged_cache.gather(self.layer_id)
|
||||
|
||||
k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
|
||||
|
||||
# (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
|
||||
q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
|
||||
# (bsz, n_heads, seq_len, head_dim) - > (bsz, seq_len, n_heads*head_dim)
|
||||
sdqa_out = (
|
||||
F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
|
||||
.permute(0, 2, 1, 3)
|
||||
@@ -224,7 +186,6 @@ class GQA(nn.Module):
|
||||
sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
|
||||
|
||||
out = self.o_proj(sdqa_out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
@@ -257,7 +218,7 @@ class MLA(nn.Module):
|
||||
self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
|
||||
self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
|
||||
|
||||
# KV (k_nope, k_rope, v)
|
||||
# fused KV: (k_nope, k_rope, v)
|
||||
self.kv_b_proj = Linear(
|
||||
kv_lora_rank,
|
||||
n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
|
||||
@@ -273,9 +234,8 @@ class MLA(nn.Module):
|
||||
x: Tensor,
|
||||
rotary_emb: Tuple[Tensor, Tensor],
|
||||
mask: Tensor = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
paged_cache: Optional[CacheView] = None,
|
||||
start_pos: int = 0,
|
||||
slot_indices: Optional[Tensor] = None,
|
||||
) -> Tensor:
|
||||
bsz, seq_len, _ = x.size()
|
||||
is_causal = mask is None
|
||||
@@ -303,12 +263,9 @@ class MLA(nn.Module):
|
||||
q = torch.cat([q_nope, q_rope], dim=-1)
|
||||
k = torch.cat([k_nope, k_rope], dim=-1)
|
||||
|
||||
if kv_cache is not None:
|
||||
k_cache, v_cache = kv_cache
|
||||
k_cache[slot_indices, start_pos : start_pos + seq_len, self.layer_id] = k
|
||||
v_cache[slot_indices, start_pos : start_pos + seq_len, self.layer_id] = v
|
||||
k = k_cache[slot_indices, : start_pos + seq_len, self.layer_id]
|
||||
v = v_cache[slot_indices, : start_pos + seq_len, self.layer_id]
|
||||
if paged_cache is not None:
|
||||
paged_cache.write(self.layer_id, start_pos, k, v)
|
||||
k, v = paged_cache.gather(self.layer_id)
|
||||
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
@@ -321,7 +278,6 @@ class MLA(nn.Module):
|
||||
attn_out = attn_out * F.sigmoid(self.gate(x))
|
||||
|
||||
out = self.o_proj(attn_out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
@@ -356,24 +312,19 @@ class DecoderBlock(nn.Module):
|
||||
x: Tensor,
|
||||
rotary_emb: Tuple[Tensor, Tensor],
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
paged_cache: Optional[CacheView] = None,
|
||||
start_pos: int = 0,
|
||||
slot_indices: Optional[Tensor] = None,
|
||||
) -> Tensor:
|
||||
# attention
|
||||
attn_output = self.attention(
|
||||
self.input_norm(x),
|
||||
rotary_emb,
|
||||
attention_mask,
|
||||
kv_cache,
|
||||
paged_cache,
|
||||
start_pos,
|
||||
slot_indices,
|
||||
)
|
||||
x = attn_output + x
|
||||
|
||||
# feed forward
|
||||
x = self.mlp(self.post_attention_norm(x)) + x
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
from typing import Any, Mapping, Optional, Tuple
|
||||
from typing import Any, Mapping, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.config.model_config import ModelConfig
|
||||
from astrai.inference.cache import CacheView
|
||||
from astrai.model.automodel import AutoModel
|
||||
from astrai.model.module import (
|
||||
DecoderBlock,
|
||||
@@ -21,39 +22,25 @@ def process_attention_mask(
|
||||
start_pos: int = 0,
|
||||
is_causal: bool = False,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Create attention mask for GQA
|
||||
Args:
|
||||
seq_mask (Tensor): A tensor indicating whether each position is valid or not.
|
||||
input_tensor (Tensor): The input tensor.
|
||||
start_pos (int): The starting position of the sequence.
|
||||
is_causal (bool): Whether the attention is causal or not.
|
||||
Returns:
|
||||
Tensor: The attention mask tensor.
|
||||
"""
|
||||
"""Build 4D attention mask from 2D seq_mask, with optional causal masking."""
|
||||
device = input_tensor.device
|
||||
dtype = input_tensor.dtype
|
||||
seq_len = input_tensor.size(1)
|
||||
|
||||
if seq_mask is None:
|
||||
if start_pos != 0:
|
||||
# for single prompt chat
|
||||
seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
|
||||
else:
|
||||
return None
|
||||
|
||||
if seq_mask.dim() > 2:
|
||||
# shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)
|
||||
# if ndim > 2, it's 4D tensor
|
||||
return seq_mask
|
||||
|
||||
batch_size = seq_mask.size(0)
|
||||
seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
|
||||
# (bsz, start_pos + seq_len)
|
||||
expanded_mask = seq_mask.unsqueeze(1).expand(
|
||||
batch_size, seq_len, start_pos + seq_len
|
||||
)
|
||||
# (bsz, seq_len, start_pos + seq_len)
|
||||
|
||||
if is_causal:
|
||||
expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
|
||||
@@ -62,16 +49,13 @@ def process_attention_mask(
|
||||
attention_mask = attention_mask.masked_fill_(
|
||||
~expanded_mask, -torch.finfo(dtype).max / 2
|
||||
).unsqueeze(1)
|
||||
# (bsz, 1, seq_len, seq_len + start_pos)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
@AutoModel.register("transformer")
|
||||
class Transformer(AutoModel):
|
||||
"""
|
||||
Transformer language model.
|
||||
"""
|
||||
"""Transformer language model with paged KV cache."""
|
||||
|
||||
def __init__(self, config: ModelConfig):
|
||||
super().__init__(config)
|
||||
@@ -114,18 +98,15 @@ class Transformer(AutoModel):
|
||||
lm_head_key = "lm_head.weight"
|
||||
embed_key = "embed_tokens.weight"
|
||||
|
||||
# Make a copy to avoid modifying the original state_dict
|
||||
state_dict = dict(state_dict)
|
||||
|
||||
if self.config.tie_weight:
|
||||
# same tensor
|
||||
# same tensor for embed and lm_head
|
||||
if embed_key in state_dict:
|
||||
state_dict[lm_head_key] = state_dict[embed_key]
|
||||
else:
|
||||
# If lm_head.weight exists in checkpoint, use it directly
|
||||
# If not, copy from embed_tokens.weight
|
||||
if lm_head_key not in state_dict and embed_key in state_dict:
|
||||
# use clone to avoid sharing the same tensor
|
||||
# clone to avoid sharing gradients
|
||||
state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
|
||||
|
||||
return super().load_state_dict(state_dict, strict, assign)
|
||||
@@ -146,9 +127,8 @@ class Transformer(AutoModel):
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
input_mask: Optional[Tensor] = None,
|
||||
persistent_key_values: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
paged_cache: Optional[CacheView] = None,
|
||||
start_pos: int = 0,
|
||||
slot_indices: Optional[Tensor] = None,
|
||||
) -> Tensor:
|
||||
assert input_ids.ndim == 2
|
||||
|
||||
@@ -157,13 +137,8 @@ class Transformer(AutoModel):
|
||||
|
||||
attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)
|
||||
|
||||
if slot_indices is None:
|
||||
slot_indices = slice(input_ids.size(0))
|
||||
|
||||
for layer in self.layers:
|
||||
x = layer(
|
||||
x, rotary_emb, attn_mask, persistent_key_values, start_pos, slot_indices
|
||||
)
|
||||
x = layer(x, rotary_emb, attn_mask, paged_cache, start_pos)
|
||||
|
||||
hidden_states = self.norm(x)
|
||||
logits = self.lm_head(hidden_states)
|
||||
|
||||
Reference in New Issue
Block a user