nemo_automodel.components.models.kimi_k3.kda_fused

Opt-in fused kernels for the Kimi-K3 KDA chunked delta rule (KimiK3TextConfig.kda_chunk_impl = "fused").

Forward: a fused CUDA kernel (source text in chunk_kda_fwd_cuda.py, JIT-built once per build directory through torch.utils.cpp_extension.load_inline and imported straight from the built .so afterwards; build dir $NEMO_KDA_FUSED_BUILD_DIR or ~/.cache/nemo_automodel/kda_fused, keyed by sources, torch version and device arch). Backward: a Triton kernel (chunk_kda_bwd_triton.py) that recomputes every intermediate from the raw forward inputs, so the autograd function saves only q, k, v, g, beta and cu_seqlens — no chunk states or WY factors.

Both kernels are specialised to the K3 production call: head dims K = V = 128, packed layout (batch 1 with int32 cu_seqlens; a dense batch of equal-length sequences is viewed as B packed documents), q/k l2-normalised inside the kernel, a bounded log gate (safe_gate), no initial or final state, no context parallelism, kda_mode = "chunk". fused_kda_unsupported_reason names the first violated condition; the model raises instead of silently falling back, so a benchmark that asks for the fused kernels either runs them or stops. FLA’s chunk_kda remains the default (kda_chunk_impl = "fla"); kda_transpose_state_layout and kda_disable_recompute are FLA-only knobs and have no effect under "fused".

Submodules

Package Contents

Classes

Name	Description
`_FusedChunkKDAFunction`	-

Functions

Name	Description
`_build_dir`	One build directory per (sources, torch, device arch); shared file systems are fine (torch takes a file lock).
`_clear_stale_lock`	Drop a build lock nobody can be holding (older than the longest build by far), with a warning.
`_forward_ext`	Return the CUDA forward extension: prebuilt .so if present, else JIT-build it once (file-locked by torch).
`_import_prebuilt`	Import the extension straight from its built `.so` when it exists and no build is in progress.
`_lock_age`	-
`fused_chunk_kda`	Drop-in for FLA’s `chunk_kda(..., output_final_state=False)` return convention: `(o, None)`.
`fused_kda_unsupported_reason`	Return why this call cannot run on the fused kernels, or None when it can.

Data

_EXT_NAME

_HEAD_DIM

_STALE_LOCK_SECONDS

__all__

_ext_lock

_fwd_ext

API

class nemo_automodel.components.models.kimi_k3.kda_fused._FusedChunkKDAFunction()

Bases: Function

nemo_automodel.components.models.kimi_k3.kda_fused._FusedChunkKDAFunction.backward(
    ctx,
    do
)

staticmethod

nemo_automodel.components.models.kimi_k3.kda_fused._FusedChunkKDAFunction.forward(
    ctx,
    q,
    k,
    v,
    g,
    beta,
    cu_seqlens,
    cu_seqlens_cpu = None
)

staticmethod

nemo_automodel.components.models.kimi_k3.kda_fused._build_dir() -> str

One build directory per (sources, torch, device arch); shared file systems are fine (torch takes a file lock).

nemo_automodel.components.models.kimi_k3.kda_fused._clear_stale_lock(
    build_dir: str
) -> None

Drop a build lock nobody can be holding (older than the longest build by far), with a warning.

nemo_automodel.components.models.kimi_k3.kda_fused._forward_ext() -> typing.Any

Return the CUDA forward extension: prebuilt .so if present, else JIT-build it once (file-locked by torch).

nemo_automodel.components.models.kimi_k3.kda_fused._import_prebuilt(
    build_dir: str
) -> typing.Any | None

Import the extension straight from its built .so when it exists and no build is in progress.

torch.utils.cpp_extension.load_inline takes a file lock in the build directory on EVERY call, even when nothing needs rebuilding, and a process killed while holding it (a cancelled job) leaves a lock that hangs every later caller for good. Importing the finished artifact directly needs no lock; only a genuine build goes through load_inline. A fresh lock means another process may still be linking the .so: return None and let the caller wait in load_inline instead of importing a half-written file.

nemo_automodel.components.models.kimi_k3.kda_fused._lock_age(
    build_dir: str
) -> float | None

nemo_automodel.components.models.kimi_k3.kda_fused.fused_chunk_kda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    cu_seqlens: torch.Tensor | None
) -> tuple[torch.Tensor, None]

Drop-in for FLA’s chunk_kda(..., output_final_state=False) return convention: (o, None).

Packed input (cu_seqlens given, batch 1) goes straight to the kernels. Dense input (cu_seqlens None, batch B of equal-length sequences — the fixed-length benchmark data) is viewed as one packed batch of B documents with cu_seqlens = [0, T, 2T, ..., BT]: the recurrent state resets at every document boundary, so this is the same computation FLA performs on the batched layout, the views cost no copies, and the offsets are handed to the backward as host ints (no device->host sync, matching FLA’s sync-free dense path). Packed input syncs once in the backward, as FLA’s varlen path does.

nemo_automodel.components.models.kimi_k3.kda_fused.fused_kda_unsupported_reason(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    cu_seqlens: torch.Tensor | None,
    mode: str = 'chunk',
    cp_context: typing.Any = None,
    use_qk_l2norm_in_kernel: bool = True,
    safe_gate: bool = True
) -> str | None

Return why this call cannot run on the fused kernels, or None when it can.

nemo_automodel.components.models.kimi_k3.kda_fused._EXT_NAME = 'nemo_fused_chunk_kda_fwd'

nemo_automodel.components.models.kimi_k3.kda_fused._HEAD_DIM = 128

nemo_automodel.components.models.kimi_k3.kda_fused._STALE_LOCK_SECONDS = 600.0

nemo_automodel.components.models.kimi_k3.kda_fused.__all__ = ['fused_chunk_kda', 'fused_kda_unsupported_reason']

nemo_automodel.components.models.kimi_k3.kda_fused._ext_lock = threading.Lock()

nemo_automodel.components.models.kimi_k3.kda_fused._fwd_ext: Any = None