- in
- [B batch, 1 position, 7168 d] [B batch, 1 position, m stored sums, 7168 features]
- out
- [B batch, 1 position, 7168 d]
- parameters
- 30,228,720,256 · × 68 KDA layers = 2,055,552,977,408
- mode
- decode · one new token after P = 1,000 positions
- state
- kept between calls by a KDA layer: its memory, 96 matrices of 128 × 128, and the last 4 inputs, 12,288 wide, of three short convolutions · 1,720,320 values at any P
Parameters
30,228,720,256 parameters per KDA layer, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Attention input mix | 14,336 | |
| Attention RMSNorm | 7,168 | |
| Kimi Delta Attention (KDA) | 443,740,416 across its components | 443,740,416 |
| Feed-forward input mix | 14,336 | |
| MoE RMSNorm | 7,168 | |
| Latent mixture of experts | 29,784,936,832 across its components | 29,784,936,832 |
| Per-instance total | 30,228,720,256 | |
Across 68 independent KDA 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.
Codemodeling_kimi_linear.py 878–917, 973–1046 · configuration_kimi_k3.py 152–156
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.hidden_size = config.hidden_size
self.config = config
self.layer_idx = layer_idx
if config.is_kda_layer(layer_idx):
self.is_linear_attn = True
self.self_attn = KimiDeltaAttention(
config=config, layer_idx=layer_idx)
elif config.is_mla:
self.is_linear_attn = False
self.self_attn = KimiMLAAttention(
config=config, layer_idx=layer_idx)
else:
raise NotImplementedError
if (
config.num_experts is not None
and layer_idx >= config.first_k_dense_replace
and layer_idx % getattr(config, "moe_layer_freq", 1) == 0
):
self.block_sparse_moe = KimiSparseMoeBlock(config)
else:
self.mlp = KimiMLP(config)
self.input_layernorm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
# Attention residual
self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
if self.use_attn_residuals:
self.attn_res_block_size = config.attn_res_block_size
self.self_attention_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.mlp_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.self_attention_res_proj = nn.Linear(
config.hidden_size, 1, bias=False)
self.mlp_res_proj = nn.Linear(
config.hidden_size, 1, bias=False) def _forward_attn_residual(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: tuple[torch.Tensor] | None = None,
output_attentions: bool | None = False,
use_cache: bool | None = False,
block_residual: torch.Tensor | None = None,
**kwargs: Unpack[FlashAttentionKwargs],
):
batch_size, seq_len, hidden_size = hidden_states.shape
prefix_sum = hidden_states
if block_residual is not None and block_residual.shape[1] > 0:
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.self_attention_res_proj,
self.self_attention_res_norm,
).view(batch_size, seq_len, hidden_size)
if self.layer_idx % self.attn_res_block_size == 0:
block_residual = torch.cat(
[block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1)
prefix_sum = None
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
if self.is_linear_attn is False:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
else:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
cache_params=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
if prefix_sum is not None:
prefix_sum = prefix_sum + hidden_states
else:
prefix_sum = hidden_states
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.mlp_res_proj,
self.mlp_res_norm,
).view(batch_size, seq_len, hidden_size)
hidden_states = self.post_attention_layernorm(hidden_states)
if hasattr(self, "block_sparse_moe"):
hidden_states = self.block_sparse_moe(hidden_states)
else:
hidden_states = self.mlp(hidden_states)
if prefix_sum is None:
prefix_sum = hidden_states
else:
prefix_sum = prefix_sum + hidden_states
return prefix_sum, block_residual def is_kda_layer(self, layer_idx: int):
return (
self.linear_attn_config is not None
and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
)Sources3
Direct evidence
A layer has Kimi Delta Attention or latent attention by its index, a dense MLP in the first first_k_dense_replace layers and a MoE in the rest, two pre-norms, and for each of its two attention-residual mixes an RMSNorm and a one-row projection.
block-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 878–917
Read code excerpt
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.hidden_size = config.hidden_size
self.config = config
self.layer_idx = layer_idx
if config.is_kda_layer(layer_idx):
self.is_linear_attn = True
self.self_attn = KimiDeltaAttention(
config=config, layer_idx=layer_idx)
elif config.is_mla:
self.is_linear_attn = False
self.self_attn = KimiMLAAttention(
config=config, layer_idx=layer_idx)
else:
raise NotImplementedError
if (
config.num_experts is not None
and layer_idx >= config.first_k_dense_replace
and layer_idx % getattr(config, "moe_layer_freq", 1) == 0
):
self.block_sparse_moe = KimiSparseMoeBlock(config)
else:
self.mlp = KimiMLP(config)
self.input_layernorm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
# Attention residual
self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
if self.use_attn_residuals:
self.attn_res_block_size = config.attn_res_block_size
self.self_attention_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.mlp_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.self_attention_res_proj = nn.Linear(
config.hidden_size, 1, bias=False)
self.mlp_res_proj = nn.Linear(
config.hidden_size, 1, bias=False)Direct evidence
With attention residuals, a layer mixes its attention input from the stored sums and the running sum, stores the running sum and starts a new one when its index is divisible by the block size, adds the attention output to the running sum, mixes the feed-forward input the same way, and adds the feed-forward output.
block-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 973–1046
Read code excerpt
def _forward_attn_residual(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: tuple[torch.Tensor] | None = None,
output_attentions: bool | None = False,
use_cache: bool | None = False,
block_residual: torch.Tensor | None = None,
**kwargs: Unpack[FlashAttentionKwargs],
):
batch_size, seq_len, hidden_size = hidden_states.shape
prefix_sum = hidden_states
if block_residual is not None and block_residual.shape[1] > 0:
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.self_attention_res_proj,
self.self_attention_res_norm,
).view(batch_size, seq_len, hidden_size)
if self.layer_idx % self.attn_res_block_size == 0:
block_residual = torch.cat(
[block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1)
prefix_sum = None
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
if self.is_linear_attn is False:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
else:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
cache_params=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
if prefix_sum is not None:
prefix_sum = prefix_sum + hidden_states
else:
prefix_sum = hidden_states
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.mlp_res_proj,
self.mlp_res_norm,
).view(batch_size, seq_len, hidden_size)
hidden_states = self.post_attention_layernorm(hidden_states)
if hasattr(self, "block_sparse_moe"):
hidden_states = self.block_sparse_moe(hidden_states)
else:
hidden_states = self.mlp(hidden_states)
if prefix_sum is None:
prefix_sum = hidden_states
else:
prefix_sum = prefix_sum + hidden_states
return prefix_sum, block_residualDirect evidence
A layer uses Kimi Delta Attention when its 1-based index is in kda_layers; the lists are read 1-based.
layer-kinds · moonshotai/Kimi-K3 / configuration_kimi_k3.py
Pinned revision f831ab668142 ↗ · lines 152–156
Read code excerpt
def is_kda_layer(self, layer_idx: int):
return (
self.linear_attn_config is not None
and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
)