nemo_automodel.components.models.kimi_linear.flops

Architecture FLOPs for Kimi Linear 48B A3B training.

Module Contents

Functions

Name	Description
`model_flops`	Estimate training work for independent, fully supervised sequences.

API

nemo_automodel.components.models.kimi_linear.flops.model_flops(
    config: nemo_automodel.components.models.kimi_linear.config.KimiLinear48BConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Estimate training work for independent, fully supervised sequences.

Counts linear projections, active SwiGLU experts, routers, the LM head, causal MLA attention (using the conventional S2 / 2 approximation), KDA convolutions and the chunkwise delta-rule kernel. KDA’s forward kernel count per head is 6Sd2 + 3SCd + SC**2, with C=64. Forward plus backward is estimated as three times forward work (six times MACs). Norms, activations, embedding lookup, optimizer work and activation recomputation are excluded. Packed document boundaries are unavailable through this interface: seq_len is treated as one document, so packed MLA work is overestimated. This is an architecture estimate, not executed FLOPs.

Parameters:

config

KimiLinear48BConfig

Complete model configuration before pipeline partitioning.

gbs

intDefaults to 1

Number of independent sequences in the global batch.

seq_len

int | NoneDefaults to None

Tokens per sequence; defaults to max_position_embeddings.

Returns: float

Estimated forward and backward FLOPs across the global batch.