- in
- [B batch, T positions, 7168 d]
- out
- [B batch, 96 heads, T positions, 128 d_h]
- parameters
- 88,129,536
- mode
- prefill, nothing kept yet
- state
- written at the end of the call · fixed size: the last 4 inputs of this convolution, 49,152 values
Parameters
88,129,536 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Query projection | 88,080,384 | |
| Short convolution | 49,152 | |
| Per-instance total | 88,129,536 | |
Across 1 independent layers: parameters.
Derived from the pinned configuration and implementation, with the per-head log-scale at the checkpoint’s 128 elements where the code creates 96; equal to the checkpoint’s own tensors. The quantisation scales are not counted.
Explanation
The normalization to unit length runs inside the kernel, after the split into heads (modeling_kimi_linear.py 620).
Codemodeling_kimi_linear.py 478–649 · 2 more files
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.config = config
self.mode = "chunk"
self.hidden_size = config.hidden_size
self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
self.head_dim = config.linear_attn_config["head_dim"]
self.num_heads = config.linear_attn_config["num_heads"]
self.head_k_dim = self.head_dim
self.num_k_heads = self.num_heads
self.layer_idx = layer_idx
assert self.mode in [
'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`."
projection_k_size = self.head_k_dim * self.num_k_heads
projection_size = self.head_dim * self.num_heads
self.q_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.k_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
self.q_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.k_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.v_conv1d = ShortConvolution(
hidden_size=projection_size,
kernel_size=self.conv_size,
activation='silu',
)
self.A_log = torch.nn.Parameter(torch.log(torch.empty(
self.num_heads, dtype=torch.float32).uniform_(1, 16)))
self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.dt_bias = nn.Parameter(
torch.empty(projection_size, dtype=torch.float32))
self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False)
self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None)
if self.use_full_rank_gate:
self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
else:
self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.o_norm = FusedRMSNormGated(
self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) use_cache = cache_params is not None
batch_size, q_len, _ = hidden_states.shape
mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode conv_state_q, conv_state_k, conv_state_v = None, None, None
recurrent_state = None
if cache_params is not None:
if cache_params.conv_states[self.layer_idx] is not None:
conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
self.layer_idx]
recurrent_state = cache_params.recurrent_states[self.layer_idx]
q_proj_states = self.q_proj(hidden_states)
k_proj_states = self.k_proj(hidden_states)
v_proj_states = self.v_proj(hidden_states)
q, conv_state_q = self.q_conv1d(
x=q_proj_states,
cache=conv_state_q,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
k, conv_state_k = self.k_conv1d(
x=k_proj_states,
cache=conv_state_k,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
v, conv_state_v = self.v_conv1d(
x=v_proj_states,
cache=conv_state_v,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
g = self.f_b_proj(self.f_a_proj(hidden_states))
g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
beta = self.b_proj(hidden_states).float()
q, k = map(lambda x: rearrange(
x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
if mode == 'chunk':
o, recurrent_state = chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
safe_gate=self.gate_lower_bound is not None,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
else:
o, recurrent_state = fused_recurrent_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
if cache_params is not None:
cache_params.recurrent_states[self.layer_idx] = recurrent_state
cache_params.conv_states[self.layer_idx] = (
conv_state_q, conv_state_k, conv_state_v)class ShortConvolution(nn.Conv1d):
"""Short convolution layer for efficient causal convolution operations.
This class implements a depthwise 1D convolution with causal padding,
designed for efficient sequence processing. It supports multiple backends (Triton/CUDA)
and optional activation functions.
Args:
hidden_size (int): Number of input/output channels (must be equal for depthwise conv)
kernel_size (int): Size of the convolution kernel
bias (bool, optional): Whether to include learnable bias. Defaults to False.
activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'.
backend (Optional[str], optional): Backend implementation ('triton' or 'cuda'). Defaults to 'triton'.
device (Optional[torch.device], optional): Device to place the layer on. Defaults to None.
dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None.
**kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility)
Attributes:
hidden_size (int): Number of channels
activation (Optional[str]): Selected activation function
backend (str): Actual backend being used (may differ from input due to availability)
Note:
- Uses depthwise convolution (groups=hidden_size) for efficiency
- Applies causal padding (kernel_size-1) to ensure no future information leakage
- Falls back to Triton backend if CUDA backend is unavailable
"""
def __init__(
self,
hidden_size: int,
kernel_size: int,
bias: bool = False,
activation: str | None = 'silu',
backend: str | None = 'triton',
device: torch.device | None = None,
dtype: torch.dtype | None = None,
**kwargs,
):
super().__init__(
in_channels=hidden_size,
out_channels=hidden_size,
kernel_size=kernel_size,
groups=hidden_size,
bias=bias,
padding=kernel_size - 1,
device=device,
dtype=dtype,
)
self.hidden_size = hidden_size
self.activation = None
if activation is not None:
assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet."
self.activation = activation
def naive_recurrent_kda(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
):
r""" dtype = v.dtype
B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1]
G = HV // H
if scale is None:
scale = K ** -0.5
q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta])
q = q.repeat_interleave(G, dim=2) * scale # [B, T, HV, K]
k = k.repeat_interleave(G, dim=2) # [B, T, HV, K]
S = k.new_zeros(B, HV, K, V).to(q)
if initial_state is not None:
S += initial_state
o = torch.zeros_like(v)
for i in range(0, T):
q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
S = S * g_i[..., None].exp()
S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S)
if not output_final_state:
S = None
return o.to(dtype), SSources4
Direct evidence
Kimi Delta Attention owns query, key and value projections to 96 heads of 128 features, a causal short convolution on each, a per-head log-scale A_log, a rank-128 pair of projections and a bias for the decay logits, a per-head write-strength projection, a full-rank output gate projection, a gated RMSNorm over each head and an output projection.
kda-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 478–541
Read code excerpt
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.config = config
self.mode = "chunk"
self.hidden_size = config.hidden_size
self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
self.head_dim = config.linear_attn_config["head_dim"]
self.num_heads = config.linear_attn_config["num_heads"]
self.head_k_dim = self.head_dim
self.num_k_heads = self.num_heads
self.layer_idx = layer_idx
assert self.mode in [
'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`."
projection_k_size = self.head_k_dim * self.num_k_heads
projection_size = self.head_dim * self.num_heads
self.q_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.k_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
self.q_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.k_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.v_conv1d = ShortConvolution(
hidden_size=projection_size,
kernel_size=self.conv_size,
activation='silu',
)
self.A_log = torch.nn.Parameter(torch.log(torch.empty(
self.num_heads, dtype=torch.float32).uniform_(1, 16)))
self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.dt_bias = nn.Parameter(
torch.empty(projection_size, dtype=torch.float32))
self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False)
self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None)
if self.use_full_rank_gate:
self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
else:
self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.o_norm = FusedRMSNormGated(
self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)Direct evidence
Queries, keys and values are projected and convolved; the decay logits and the write strength are projected from the layer’s input; the kernel returns the output and the recurrent state, in the recurrent form for one new position with a cache and chunked otherwise; the state and the convolution states are kept per layer.
kda-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 559–649
Read code excerpt
use_cache = cache_params is not None
batch_size, q_len, _ = hidden_states.shape
mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode… lines 562–571 omitted …
conv_state_q, conv_state_k, conv_state_v = None, None, None
recurrent_state = None
if cache_params is not None:
if cache_params.conv_states[self.layer_idx] is not None:
conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
self.layer_idx]
recurrent_state = cache_params.recurrent_states[self.layer_idx]
q_proj_states = self.q_proj(hidden_states)
k_proj_states = self.k_proj(hidden_states)
v_proj_states = self.v_proj(hidden_states)
q, conv_state_q = self.q_conv1d(
x=q_proj_states,
cache=conv_state_q,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
k, conv_state_k = self.k_conv1d(
x=k_proj_states,
cache=conv_state_k,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
v, conv_state_v = self.v_conv1d(
x=v_proj_states,
cache=conv_state_v,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
g = self.f_b_proj(self.f_a_proj(hidden_states))
g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
beta = self.b_proj(hidden_states).float()
q, k = map(lambda x: rearrange(
x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
if mode == 'chunk':
o, recurrent_state = chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
safe_gate=self.gate_lower_bound is not None,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
else:
o, recurrent_state = fused_recurrent_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
if cache_params is not None:
cache_params.recurrent_states[self.layer_idx] = recurrent_state
cache_params.conv_states[self.layer_idx] = (
conv_state_q, conv_state_k, conv_state_v)Direct evidence
The short convolution is a depthwise, causal 1D convolution without bias, followed by SiLU.
short-conv · fla-org/flash-linear-attention / fla/modules/conv/short_conv.py
Pinned revision 07ca1e49ef6c ↗ · lines 24–80
Read code excerpt
class ShortConvolution(nn.Conv1d):
"""Short convolution layer for efficient causal convolution operations.
This class implements a depthwise 1D convolution with causal padding,
designed for efficient sequence processing. It supports multiple backends (Triton/CUDA)
and optional activation functions.
Args:
hidden_size (int): Number of input/output channels (must be equal for depthwise conv)
kernel_size (int): Size of the convolution kernel
bias (bool, optional): Whether to include learnable bias. Defaults to False.
activation (Optional[str], optional): Activation function ('silu' or 'swish'). Defaults to 'silu'.
backend (Optional[str], optional): Backend implementation ('triton' or 'cuda'). Defaults to 'triton'.
device (Optional[torch.device], optional): Device to place the layer on. Defaults to None.
dtype (Optional[torch.dtype], optional): Data type for layer parameters. Defaults to None.
**kwargs: Additional keyword arguments (deprecated 'use_fast_conv1d' supported for compatibility)
Attributes:
hidden_size (int): Number of channels
activation (Optional[str]): Selected activation function
backend (str): Actual backend being used (may differ from input due to availability)
Note:
- Uses depthwise convolution (groups=hidden_size) for efficiency
- Applies causal padding (kernel_size-1) to ensure no future information leakage
- Falls back to Triton backend if CUDA backend is unavailable
"""
def __init__(
self,
hidden_size: int,
kernel_size: int,
bias: bool = False,
activation: str | None = 'silu',
backend: str | None = 'triton',
device: torch.device | None = None,
dtype: torch.dtype | None = None,
**kwargs,
):
super().__init__(
in_channels=hidden_size,
out_channels=hidden_size,
kernel_size=kernel_size,
groups=hidden_size,
bias=bias,
padding=kernel_size - 1,
device=device,
dtype=dtype,
)
self.hidden_size = hidden_size
self.activation = None
if activation is not None:
assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet."
self.activation = activation
Direct evidence
The delta-rule recurrence per head: the state decays channel by channel, takes a write of the key times the difference between the value and what the state already returns for that key, scaled by the write strength, and the output reads the state with the scaled query.
kda-kernel · fla-org/flash-linear-attention / fla/ops/kda/naive.py
Pinned revision 07ca1e49ef6c ↗ · lines 12–66
Read code excerpt
def naive_recurrent_kda(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
):
r"""… lines 23–44 omitted …
dtype = v.dtype
B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1]
G = HV // H
if scale is None:
scale = K ** -0.5
q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta])
q = q.repeat_interleave(G, dim=2) * scale # [B, T, HV, K]
k = k.repeat_interleave(G, dim=2) # [B, T, HV, K]
S = k.new_zeros(B, HV, K, V).to(q)
if initial_state is not None:
S += initial_state
o = torch.zeros_like(v)
for i in range(0, T):
q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
S = S * g_i[..., None].exp()
S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S)
if not output_final_state:
S = None
return o.to(dtype), S