> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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

* **[`nemo_automodel.components.models.kimi_k3.kda_fused.chunk_kda_bwd_triton`](/nemo-automodel/nemo_automodel/components/models/kimi_k3/kda_fused/chunk_kda_bwd_triton)**
* **[`nemo_automodel.components.models.kimi_k3.kda_fused.chunk_kda_fwd_cuda`](/nemo-automodel/nemo_automodel/components/models/kimi_k3/kda_fused/chunk_kda_fwd_cuda)**

## Package Contents

### Classes

| Name                                                                                                   | Description |
| ------------------------------------------------------------------------------------------------------ | ----------- |
| [`_FusedChunkKDAFunction`](#nemo_automodel-components-models-kimi_k3-kda_fused-_FusedChunkKDAFunction) | -           |

### Functions

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

### Data

[`_EXT_NAME`](#nemo_automodel-components-models-kimi_k3-kda_fused-_EXT_NAME)

[`_HEAD_DIM`](#nemo_automodel-components-models-kimi_k3-kda_fused-_HEAD_DIM)

[`_STALE_LOCK_SECONDS`](#nemo_automodel-components-models-kimi_k3-kda_fused-_STALE_LOCK_SECONDS)

[`__all__`](#nemo_automodel-components-models-kimi_k3-kda_fused-__all__)

[`_ext_lock`](#nemo_automodel-components-models-kimi_k3-kda_fused-_ext_lock)

[`_fwd_ext`](#nemo_automodel-components-models-kimi_k3-kda_fused-_fwd_ext)

### API

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

**Bases:** `Function`

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

staticmethod

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

staticmethod

```python
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).

```python
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.

```python
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).

```python
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.

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

```python
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.

```python
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.

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

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

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

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

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

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