llmcompressor.modeling.kimi_k3.modeling_kimi_linear

Classes:

KimiDynamicCache –

Dynamic cache for Kimi model.
KimiLinearForCausalLM –
KimiLinearModel –
KimiMLAAttention –

Multi-Latent Attention adapted from deepseek-v3
KimiMoEGate –

MoEGate adapted from Deepseek-V3.
KimiSparseMoeBlock –

Adapted from Deepseek-V3's MOE implementation
SituAndMul –

SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up

Functions:

repeat_kv –

Expand the key/value heads from num_key_value_heads to num_attention_heads.

KimiDynamicCache

KimiDynamicCache(config: KimiLinearConfig)

Dynamic cache for Kimi model. Inspired by Qwen3-Next

Methods:

get_mask_sizes –

Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
get_seq_length –

Returns the sequence length of the cached states. A layer index can be optionally passed.
reorder_cache –

Reorders the cache for beam search, given the selected beam indices.

Attributes:

has_previous_state –

We have a previous state if the last linear (conv) layer was already updated.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, config: KimiLinearConfig):
    super().__init__()
    self.config = config

    if config.linear_attn_config is not None:
        self.layer_types = []
        for i in range(config.num_hidden_layers):
            if config.is_kda_layer(i):
                self.layer_types.append("linear_attention")
            else:
                self.layer_types.append("full_attention")
    else:
        self.layer_types = ["full_attention"] * config.num_hidden_layers

    self.transformer_layers = [
        i
        for i in range(config.num_hidden_layers)
        if self.layer_types[i] == "full_attention"
    ]

    linear_layers = [
        i
        for i in range(config.num_hidden_layers)
        if self.layer_types[i] == "linear_attention"
    ]
    self.last_linear_layer = linear_layers[-1] if linear_layers else -1

    self.conv_states = [None for _ in range(config.num_hidden_layers)]
    self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
    self.key_cache = [None for _ in range(config.num_hidden_layers)]
    self.value_cache = [None for _ in range(config.num_hidden_layers)]

has_previous_state `property`

has_previous_state

We have a previous state if the last linear (conv) layer was already updated.

get_mask_sizes

get_mask_sizes(
    cache_position, layer_idx: int
) -> tuple[int, int]

Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for the given layer at layer_idx. The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def get_mask_sizes(self, cache_position, layer_idx: int) -> tuple[int, int]:
    """
    Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for
    the given layer at `layer_idx`.
    The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer.
    """
    kv_offset = 0
    # cache_position may be an int (new API) or a 1-D tensor (old API)
    query_length = (
        cache_position
        if isinstance(cache_position, int)
        else cache_position.shape[0]
    )
    past_seen_tokens = self.get_seq_length(layer_idx)
    kv_length = query_length + past_seen_tokens
    return kv_length, kv_offset

get_seq_length

get_seq_length(layer_idx: int | None = 0) -> int

Returns the sequence length of the cached states. A layer index can be optionally passed.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def get_seq_length(self, layer_idx: int | None = 0) -> int:
    """Returns the sequence length of the cached states. A layer index can be optionally passed."""
    # take any layer that contains cache and not empty tensor
    layer_idx = (
        self.transformer_layers[0]
        if layer_idx not in self.transformer_layers
        else layer_idx
    )
    if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None:
        return 0
    return self.key_cache[layer_idx].shape[-2]

reorder_cache

reorder_cache(beam_idx: LongTensor)

Reorders the cache for beam search, given the selected beam indices.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def reorder_cache(self, beam_idx: torch.LongTensor):
    """Reorders the cache for beam search, given the selected beam indices."""
    for layer_idx in range(len(self.key_cache)):
        if self.key_cache[layer_idx] is not None:
            device = self.key_cache[layer_idx].device
            beam_idx = beam_idx.to(device)
            self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(
                0, beam_idx
            )
            self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(
                0, beam_idx
            )

        if self.conv_states[layer_idx] is not None:
            device = self.conv_states[layer_idx][0].device
            beam_idx = beam_idx.to(device)
            q_conv, k_conv, v_conv = self.conv_states[layer_idx]
            self.conv_states[layer_idx] = (
                q_conv.index_select(0, beam_idx),
                k_conv.index_select(0, beam_idx),
                v_conv.index_select(0, beam_idx),
            )
            self.recurrent_states[layer_idx] = self.recurrent_states[
                layer_idx
            ].index_select(0, beam_idx)

KimiLinearForCausalLM

KimiLinearForCausalLM(config)

Bases: KimiPreTrainedModel, GenerationMixin

Methods:

forward –

Args:

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, config):
    super().__init__(config)
    self.model = KimiLinearModel(config)
    self.vocab_size = config.vocab_size
    self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    # Initialize weights and apply final processing
    self.post_init()

forward

forward(
    input_ids: LongTensor = None,
    attention_mask: Tensor | None = None,
    position_ids: LongTensor | None = None,
    past_key_values: list[FloatTensor] | None = None,
    inputs_embeds: FloatTensor | None = None,
    labels: LongTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    generation_mode: bool | None = None,
    return_dict: bool | None = None,
    cache_position: LongTensor | None = None,
    **kwargs: Unpack[TransformersKwargs],
) -> tuple | CausalLMOutputWithPast

Args: labels (torch.LongTensor of shape (batch_size, sequence_length), optional): Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

@can_return_tuple
# @auto_docstring
def forward(
    self,
    input_ids: torch.LongTensor = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: list[torch.FloatTensor] | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    generation_mode: bool | None = None,
    return_dict: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    **kwargs: Unpack[TransformersKwargs],
) -> tuple | CausalLMOutputWithPast:
    r"""
    Args:
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
    """

    output_attentions = (
        output_attentions
        if output_attentions is not None
        else self.config.output_attentions
    )
    output_hidden_states = (
        output_hidden_states
        if output_hidden_states is not None
        else self.config.output_hidden_states
    )
    return_dict = (
        return_dict if return_dict is not None else self.config.use_return_dict
    )

    outputs = self.model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_values=past_key_values,
        inputs_embeds=inputs_embeds,
        use_cache=use_cache,
        output_attentions=output_attentions,
        output_hidden_states=output_hidden_states,
        return_dict=return_dict,
        cache_position=cache_position,
    )

    logits = outputs[0]
    if generation_mode:
        logits = logits[:, -1:]
    logits = self.lm_head(logits)

    loss = None
    if labels is not None:
        loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)

    return CausalLMOutputWithPast(
        loss=loss,
        logits=logits,
        past_key_values=outputs.past_key_values,
        hidden_states=outputs.hidden_states,
        attentions=outputs.attentions,
    )

KimiLinearModel

KimiLinearModel(config: KimiLinearConfig)

Bases: KimiPreTrainedModel

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, config: KimiLinearConfig):
    super().__init__(config)
    self.padding_idx = config.pad_token_id
    self.vocab_size = config.vocab_size

    self.embed_tokens = nn.Embedding(
        config.vocab_size, config.hidden_size, self.padding_idx
    )
    self.layers = nn.ModuleList(
        [
            KimiDecoderLayer(config, layer_idx)
            for layer_idx in range(config.num_hidden_layers)
        ]
    )
    self.norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    self.use_attn_residuals = (
        getattr(config, "attn_res_block_size", None) is not None
    )
    if self.use_attn_residuals:
        self.output_attn_res_norm = KimiRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.output_attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False)

    from transformers.utils import is_flash_attn_2_available as _fa2_avail

    _requested = getattr(config, "_attn_implementation", None)
    if _requested not in (None, "flash_attention_2") or not _fa2_avail():
        # Fall back gracefully when flash-attn2 is unavailable or a different impl is requested
        if _requested == "flash_attention_2" and not _fa2_avail():
            logger.warning_once(
                "flash_attention_2 requested but not available; falling back to sdpa."
            )
        config._attn_implementation = (
            _requested if _requested not in (None, "flash_attention_2") else "eager"
        )
    else:
        config._attn_implementation = "flash_attention_2"

    self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
    self.gradient_checkpointing = False
    # Initialize weights and apply final processing
    self.post_init()

KimiMLAAttention

KimiMLAAttention(config: KimiLinearConfig, layer_idx: int)

Bases: Module

Multi-Latent Attention adapted from deepseek-v3

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, config: KimiLinearConfig, layer_idx: int):
    nn.Module.__init__(self)
    self.config = config
    self.layer_idx = layer_idx
    self.hidden_size = config.hidden_size
    self.num_heads = config.num_attention_heads
    self.num_key_value_heads = config.num_key_value_heads
    self.num_key_value_groups = self.num_heads // self.num_key_value_heads

    self.attention_dropout = getattr(config, "attention_dropout", 0.0)

    try:
        self.q_lora_rank = config.q_lora_rank
        self.qk_rope_head_dim = config.qk_rope_head_dim
        self.kv_lora_rank = config.kv_lora_rank
        self.v_head_dim = config.v_head_dim
        self.qk_nope_head_dim = config.qk_nope_head_dim
        self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
        self.use_nope = config.mla_use_nope
        self.scaling = self.q_head_dim ** (-0.5)
    except Exception as e:
        raise ValueError(
            f"Kimi MLA config is not found or not properly formatted: {e}"
        )

    if self.q_lora_rank is not None:
        self.q_a_proj = nn.Linear(
            self.hidden_size,
            self.q_lora_rank,
            bias=False,
        )
        self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank)
        self.q_b_proj = nn.Linear(
            self.q_lora_rank,
            self.num_heads * self.q_head_dim,
            bias=False,
        )
    else:
        self.q_proj = nn.Linear(
            self.hidden_size,
            self.num_heads * self.q_head_dim,
            bias=False,
        )
    self.kv_a_proj_with_mqa = nn.Linear(
        self.hidden_size,
        self.kv_lora_rank + self.qk_rope_head_dim,
        bias=False,
    )
    self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank)
    self.kv_b_proj = nn.Linear(
        self.kv_lora_rank,
        self.num_heads
        * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
        bias=False,
    )
    self.o_proj = nn.Linear(
        self.num_heads * self.v_head_dim,
        self.hidden_size,
        bias=False,
    )
    self.is_causal = True
    assert self.use_nope

    self.use_output_gate = getattr(config, "mla_use_output_gate", False)
    if self.use_output_gate:
        projection_size = self.num_heads * self.v_head_dim
        self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)

    self.rotary_emb = None

KimiMoEGate

KimiMoEGate(config: KimiLinearConfig)

Bases: Module

MoEGate adapted from Deepseek-V3. Parameter correspondences: num_experts -> n_routed_experts num_experts_per_token -> num_experts_per_tok num_expert_group -> n_group moe_router_activation_func -> scoring_func

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, config: KimiLinearConfig):
    super().__init__()
    self.config = config
    self.top_k = config.num_experts_per_token
    self.num_experts = config.num_experts
    self.routed_scaling_factor = config.routed_scaling_factor
    self.moe_router_activation_func = config.moe_router_activation_func
    self.num_expert_group = getattr(config, "num_expert_group", 1)
    self.topk_group = getattr(config, "topk_group", 1)

    # topk selection algorithm
    self.moe_renormalize = config.moe_renormalize
    self.gating_dim = config.hidden_size
    self.weight = nn.Parameter(
        torch.empty((self.num_experts, self.gating_dim)),
    )

    self.e_score_correction_bias = nn.Parameter(
        torch.empty(self.num_experts),
    )
    self.reset_parameters()

KimiSparseMoeBlock

KimiSparseMoeBlock(config: KimiLinearConfig)

Bases: Module

Adapted from Deepseek-V3's MOE implementation The namings are consistent with Kimi's version.

Methods:

moe_train –

Calibration-compatible MoE dispatch with gradient flow.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, config: KimiLinearConfig):
    super().__init__()
    self.config = config
    self.hidden_dim = config.hidden_size
    self.num_experts = config.num_experts
    self.top_k = config.num_experts_per_token
    self.moe_renormalize = config.moe_renormalize

    self.use_latent_moe = (
        getattr(config, "routed_expert_hidden_size", None) is not None
    )
    self.moe_hidden_size = (
        config.routed_expert_hidden_size
        if self.use_latent_moe
        else config.hidden_size
    )
    self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False)

    self.ep_size = 1
    self.experts_per_rank = config.num_experts
    self.ep_rank = 0
    self.experts = nn.ModuleList(
        [
            KimiBlockSparseMLP(
                config,
                hidden_size=self.moe_hidden_size,
                intermediate_size=config.moe_intermediate_size,
            )
            for _ in range(config.num_experts)
        ],
    )
    self.gate = KimiMoEGate(config)
    if config.num_shared_experts is not None:
        intermediate_size = config.moe_intermediate_size * config.num_shared_experts
        self.shared_experts = KimiMLP(
            config=config,
            intermediate_size=intermediate_size,
        )

    if self.use_latent_moe:
        self.routed_expert_down_proj = nn.Linear(
            config.hidden_size,
            self.moe_hidden_size,
            bias=False,
        )
        self.routed_expert_up_proj = nn.Linear(
            self.moe_hidden_size,
            config.hidden_size,
            bias=False,
        )
        if self.latent_moe_use_norm:
            self.routed_expert_norm = KimiRMSNorm(
                self.moe_hidden_size,
                eps=config.rms_norm_eps,
            )

moe_train

moe_train(x, topk_ids, topk_weight)

Calibration-compatible MoE dispatch with gradient flow.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def moe_train(self, x, topk_ids, topk_weight):
    """Calibration-compatible MoE dispatch with gradient flow."""
    y = torch.zeros_like(x)

    with torch.no_grad():
        expert_mask = F.one_hot(topk_ids, self.num_experts).permute(2, 1, 0)

    for expert_idx, expert in enumerate(self.experts):
        top_k_pos, token_indices = torch.where(expert_mask[expert_idx])

        expert_out = expert(x)[token_indices]

        expert_weights = topk_weight[token_indices, top_k_pos, None]
        y.index_add_(0, token_indices, (expert_out * expert_weights).to(y.dtype))

    return y

SituAndMul

SituAndMul(
    beta: float = 1.0, linear_beta: float | None = None
)

Bases: Module

SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta).

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def __init__(self, beta: float = 1.0, linear_beta: float | None = None):
    super().__init__()
    self.beta = beta
    self.linear_beta = linear_beta

repeat_kv

repeat_kv(
    hidden_states: Tensor, n_rep: int
) -> torch.Tensor

Expand the key/value heads from num_key_value_heads to num_attention_heads.

Source code in src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Expand the key/value heads from `num_key_value_heads` to `num_attention_heads`."""
    if n_rep == 1:
        return hidden_states
    return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep)

llmcompressor.modeling.kimi_k3.modeling_kimi_linear

KimiDynamicCache

has_previous_state property

get_mask_sizes

get_seq_length

reorder_cache

KimiLinearForCausalLM

forward

KimiLinearModel

KimiMLAAttention

KimiMoEGate

KimiSparseMoeBlock

moe_train

SituAndMul

repeat_kv

has_previous_state `property`