nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

kimi delta attention (#17281)

Browse Source 
* kimi delta attention

* config

* test
This commit is contained in:
b1tg
2026-07-30 08:23:56 -07:00
committed by GitHub
parent b5a2a5666a
commit e25bf77ce9
3 changed files with 59 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_attention.py
+15 -1
View File
@@ -47,7 +47,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
def _make_config(self, **kwargs):
return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
"norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
"rope_dim":4, "v_head_dim":4, "max_context":4, "full_attention_interval":2,
"rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,),
"ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))
def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
@@ -176,6 +176,20 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
def test_kda_channel_decay(self):
config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
# f_b(f_a(x)) = [1, 2, 3, 4]
block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
block._init_state(x)
initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
block.recurrent_state.assign(initial_state).realize()
block.ssm_a = Tensor([[-1.], [-1.]])
block._attention(x, 0).realize()
alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)
class TestPairwiseTopk(unittest.TestCase):
def test_basic_topk(self):
x = Tensor([[[1.0, 3.0, 2.0, 5.0, 4.0]]])
tinygrad/llm/cli.py
+3 -2
View File
@@ -41,9 +41,10 @@ class SimpleTokenizer:
# https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
return SimpleTokenizer(dict(normal_tokens), dict(special_tokens), kv["tokenizer.ggml.pre"],
special_tokens_dict = dict(special_tokens)
return SimpleTokenizer(dict(normal_tokens), special_tokens_dict, kv["tokenizer.ggml.pre"],
bos_id=kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None,
eos_id=kv.get('tokenizer.ggml.eos_token_id', 0), eot_id=kv.get('tokenizer.ggml.eot_token_id'))
eos_id=kv.get('tokenizer.ggml.eos_token_id', 0), eot_id=kv.get('tokenizer.ggml.eot_token_id', special_tokens_dict.get('<|im_end|>')))
def _encode_word(self, word:bytes) -> list[int]:
if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
tinygrad/llm/model.py
+41 -15
View File
@@ -40,6 +40,7 @@ class SSMConfig:
group_count: int
time_step_rank: int
inner_size: int
kda: bool = False
@dataclass(frozen=True)
class TransformerConfig:
@@ -62,7 +63,7 @@ class TransformerConfig:
q_lora_rank: int = 0
kv_lora_rank: int = 0
shared_expert_dim: int = 0
full_attention_interval: int = 0
ssm_layers: tuple[bool, ...] = ()
attn_output_gate: bool = False
ssm: SSMConfig|None = None
shared_expert_gate: bool = True
@@ -211,13 +212,13 @@ class MLATransformerBlock(FFNBlock):
q_proj = self.attn_q_b(self.attn_q_a_norm(self.attn_q_a(x))) if self.config.q_lora_rank > 0 else self.attn_q(x)
q = q_proj.reshape(B, T, self.config.n_heads, self.config.head_dim).transpose(1, 2)
q_nope, q_rope = q[..., :q_nope_head_dim], q[..., q_nope_head_dim:]
q = (q_nope @ self.attn_k_b["weight"].transpose(-1, -2)).cat(apply_rope(q_rope, self.freqs_cis[start_pos:start_pos+T]), dim=-1)
if not self.config.ssm or not self.config.ssm.kda: q_rope = apply_rope(q_rope, self.freqs_cis[start_pos:start_pos+T])
q = (q_nope @ self.attn_k_b["weight"].transpose(-1, -2)).cat(q_rope, dim=-1)
kv_a = self.attn_kv_a_mqa(x)
c_kv = self.attn_kv_a_norm(kv_a[..., :self.config.kv_lora_rank])
k_rope = apply_rope(
kv_a[..., self.config.kv_lora_rank:].reshape(B, T, 1, self.config.rope_dim).transpose(1, 2),
self.freqs_cis[start_pos:start_pos+T])
k_rope = kv_a[..., self.config.kv_lora_rank:].reshape(B, T, 1, self.config.rope_dim).transpose(1, 2)
if not self.config.ssm or not self.config.ssm.kda: k_rope = apply_rope(k_rope, self.freqs_cis[start_pos:start_pos+T])
k_store = c_kv.reshape(B, 1, T, self.config.kv_lora_rank).cat(k_rope.reshape(B, 1, T, self.config.rope_dim), dim=-1)
k = Tensor(self.cache_k.uop.after(self.cache_k[:, :, start_pos:start_pos+T, :].uop.store(k_store.uop)))[:, :, 0:start_pos+T, :]
@@ -243,11 +244,17 @@ class GatedDeltaNetBlock(FFNBlock):
assert self.num_v_heads % self.num_k_heads == 0
self.head_v_dim, self.ssm_conv_kernel = ssm.inner_size // ssm.time_step_rank, ssm.conv_kernel
self.conv_channels, self.q_dim = ssm.inner_size + 2*ssm.group_count*ssm.state_size, ssm.state_size*ssm.group_count
self.attn_qkv, self.attn_gate = nn.Linear(config.dim, self.conv_channels, bias=False), nn.Linear(config.dim, ssm.inner_size, bias=False)
self.ssm_alpha, self.ssm_beta = nn.Linear(config.dim, self.num_v_heads, bias=False), nn.Linear(config.dim, self.num_v_heads, bias=False)
self.attn_qkv = nn.Linear(config.dim, self.conv_channels, bias=False)
if ssm.kda:
self.ssm_g_a, self.ssm_g_b = nn.Linear(config.dim, self.head_v_dim, bias=False), nn.Linear(self.head_v_dim, ssm.inner_size, bias=False)
self.ssm_f_a, self.ssm_f_b = nn.Linear(config.dim, self.head_k_dim, bias=False), nn.Linear(self.head_k_dim, ssm.inner_size, bias=False)
else:
self.attn_gate = nn.Linear(config.dim, ssm.inner_size, bias=False)
self.ssm_alpha = nn.Linear(config.dim, self.num_v_heads, bias=False)
self.ssm_beta = nn.Linear(config.dim, self.num_v_heads, bias=False)
self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)}
self.ssm_dt = {"bias": Tensor.zeros(self.num_v_heads)}
self.ssm_a = Tensor.zeros(self.num_v_heads)
self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)}
self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), nn.Linear(ssm.inner_size, config.dim, bias=False)
def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
@@ -256,9 +263,12 @@ class GatedDeltaNetBlock(FFNBlock):
# input processing
x = x.half()
out_gate = self.attn_gate(x).reshape(B, 1, self.num_v_heads, self.head_v_dim)
out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
alpha = ((self.ssm_alpha(x).float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, self.num_v_heads, 1, 1).exp()
alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
# qkv conv
conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
@@ -280,7 +290,8 @@ class GatedDeltaNetBlock(FFNBlock):
# output
core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
return self.ssm_out((core_attn_out * out_gate.silu()).reshape(B, 1, -1).cast(x.dtype))
out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
# recurrent state can't be partially reused after divergence, force a full rebuild
def _state_reset_ops(self):
@@ -298,7 +309,8 @@ class Transformer:
dense_config = replace(config, num_experts=0, num_experts_per_tok=0, shared_expert_dim=0, hidden_dim=config.dense_hidden_dim or config.hidden_dim)
if config.ssm: config = replace(config, qk_norm=config.head_dim)
block_cls = MLATransformerBlock if config.kv_lora_rank > 0 else TransformerBlock
self.blk:list[FFNBlock] = [GatedDeltaNetBlock(config, config.ssm) if config.ssm and (i+1) % config.full_attention_interval != 0 else
self.blk:list[FFNBlock] = [GatedDeltaNetBlock(dense_config if i < config.leading_dense_blocks else config, config.ssm)
if config.ssm and config.ssm_layers[i] else
block_cls(dense_config if i < config.leading_dense_blocks else config) for i in range(config.num_blocks)]
self.token_embd = nn.Embedding(config.vocab_size, config.dim)
self.output_norm = nn.RMSNorm(config.dim, config.norm_eps)
@@ -337,8 +349,21 @@ class Transformer:
n_heads, n_kv_heads = kv[f'{arch}.attention.head_count'], kv[f'{arch}.attention.head_count_kv']
ssm = None
ssm_layers: tuple[bool, ...] = ()
if arch in ('qwen35', 'qwen35moe'):
ssm = SSMConfig(**{k: kv[f'{arch}.ssm.{k}'] for k in ('conv_kernel','state_size','group_count','time_step_rank','inner_size')})
ssm_layers = tuple((i+1) % kv[f'{arch}.full_attention_interval'] != 0 for i in range(kv[f'{arch}.block_count']))
elif arch == 'kimi-linear':
ssm_layers = tuple(x == 0 for x in n_kv_heads)
n_kv_heads = max(n_kv_heads)
ssm = SSMConfig(kv[f'{arch}.ssm.conv_kernel'], kv[f'{arch}.kda.head_dim'], n_heads, n_heads, n_heads*kv[f'{arch}.kda.head_dim'], kda=True)
for i, is_ssm in enumerate(ssm_layers):
if not is_ssm: continue
state_dict[f"blk.{i}.attn_qkv.weight"] = state_dict.pop(f"blk.{i}.attn_q.weight").cat(
state_dict.pop(f"blk.{i}.attn_k.weight"), state_dict.pop(f"blk.{i}.attn_v.weight"), dim=0).contiguous()
state_dict[f"blk.{i}.ssm_conv1d.weight"] = state_dict.pop(f"blk.{i}.ssm_conv1d_q.weight").cat(
state_dict.pop(f"blk.{i}.ssm_conv1d_k.weight"), state_dict.pop(f"blk.{i}.ssm_conv1d_v.weight"), dim=0).squeeze(1).contiguous()
state_dict[f"blk.{i}.ssm_out.weight"] = state_dict.pop(f"blk.{i}.attn_output.weight")
if arch in ('qwen35', 'qwen35moe', 'glm4moe'):
state_dict = {k.replace('post_attention_norm', 'ffn_norm'):v for k,v in state_dict.items()}
@@ -348,6 +373,7 @@ class Transformer:
# Permute RoPE weights from interleaved to half-split layout.
for name in state_dict:
if arch == 'kimi-linear': continue
if ('attn_q.weight' in name or 'attn_q_b.weight' in name) and (arch == 'llama' or kv_lora_rank):
w = state_dict[name].reshape(n_heads, state_dict[name].shape[0]//n_heads, -1)
prefix = head_dim-rope_dim
@@ -369,7 +395,7 @@ class Transformer:
max_context=max_context,
qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0),
norm_topk_prob=kv.get(f'{arch}.expert_weights_norm', arch in ('qwen3moe', 'qwen35moe')),
norm_topk_prob=kv.get(f'{arch}.expert_weights_norm', arch in ('qwen3moe', 'qwen35moe', 'kimi-linear')),
kv_lora_rank=kv_lora_rank, q_lora_rank=kv.get(f'{arch}.attention.q_lora_rank', 0),
leading_dense_blocks=kv.get(f'{arch}.leading_dense_block_count', 0),
shared_expert_dim=kv.get(
@@ -378,7 +404,7 @@ class Transformer:
shared_expert_gate=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.ffn_gate_inp_shexp.weight" in state_dict,
dense_hidden_dim=kv.get(f'{arch}.feed_forward_length', 0) if kv.get(f'{arch}.leading_dense_block_count', 0) else 0,
routed_scaling_factor=kv.get(f'{arch}.expert_weights_scale', 1.0), attn_output_gate=arch in ('qwen35', 'qwen35moe'), ssm=ssm,
full_attention_interval=kv.get(f'{arch}.full_attention_interval', 0),
ssm_layers=ssm_layers,
qkv_bias='blk.0.attn_q.bias' in state_dict,
expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict)
model = Transformer(config)