`bridge.models.kimi.kimi_k3_ops`#

Numerical operators used by the Kimi K3 language backbone.

Module Contents#

Classes#

`SiTUAndMul`	Module wrapper used by MCore’s custom-activation path.
`KimiRMSNorm`	Kimi’s FP32-accumulating RMS normalization.

Functions#

`_patch_fla_kda_hopper_autotune`	Remove a KDA autotune choice that is invalid on Hopper.
`kda`	Run Kimi Delta Attention while preserving packed-sequence boundaries.
`situ_and_mul`	Apply Kimi K3’s SiTU gated activation.
`attn_res_aggregate`	Aggregate Kimi K3 AttnRes snapshots and the current prefix.

API#

bridge.models.kimi.kimi_k3_ops._patch_fla_kda_hopper_autotune() → None#: Remove a KDA autotune choice that is invalid on Hopper.

bridge.models.kimi.kimi_k3_ops.kda( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor, lower_bound: float, *, cu_seqlens: torch.Tensor | None = None, ) → torch.Tensor#: Run Kimi Delta Attention while preserving packed-sequence boundaries.

bridge.models.kimi.kimi_k3_ops.situ_and_mul( inputs: torch.Tensor, beta: float = 4.0, linear_beta: float = 25.0, ) → torch.Tensor#: Apply Kimi K3’s SiTU gated activation.

class bridge.models.kimi.kimi_k3_ops.SiTUAndMul(config)#

Bases: torch.nn.Module

Module wrapper used by MCore’s custom-activation path.

Initialization

forward(inputs: torch.Tensor) → torch.Tensor#: Apply SiTU to a fused gate/up projection.

class bridge.models.kimi.kimi_k3_ops.KimiRMSNorm( hidden_size: int, eps: float, device: torch.device | int | None = None, dtype: torch.dtype | None = None, )#