IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

kda

Graph-op bindings for the KDA (Kimi Delta Attention) kernels.

The kernel math and kda_chunk's host-side launch sequence live in //Kernels/lib/kda (kda.recurrent, kda.chunk_fwd, kda.chunk_launch); only the @extensibility.register wrappers live here, mirroring the msa.mojo binding for //Kernels/lib/msa. Registration has to be declared inside the built-in kernel library itself -- importing the struct from the standalone lib does not put the op in the graph compiler's registry, so a served graph resolving kda_decode / kda_chunk needs this declaration and not just the dependency.

Two ops, one per prefill regime:

  • kda_decode -- the token-sequential recurrence (kda.recurrent). Correct for any ragged batch (decode's 1-token sequences and a multi-token prefill alike), but O(T) sequential steps per sequence.
  • kda_chunk -- the chunk-parallel prefill pipeline (kda.chunk_fwd, driven by kda.chunk_launch): L1 kda_chunk_prepare_gpu (parallel per chunk) -> L2 kda_chunk_scan_gpu (sequential, but only O(T/CHUNK_SIZE) steps) -> L3 kda_chunk_output_gpu (parallel). This is what turns prefill from O(T) into O(T/CHUNK_SIZE) sequential depth; the Python KDA layer picks it for prefill and keeps kda_decode for single-token decode (see kda_attention.py).

Structsโ€‹

Was this page helpful?