Module mla_cuda

Module mla_cuda 

Source
Expand description

2026-09-25: Host launch of the Kimi K3 MLA decode kernels (PTX module MODULE) for one token.

Two launches per token: ROPE_ENTRY rotates the rope slice of q and k (it returns without rotating when mla_use_nope is set), then SDPA_ENTRY attends over the cached KV and applies the sigmoid output gate when mla_use_output_gate is set. The CPU reference is mla_decode_token. Decode uses launch_k3_mla_decode_token_on_device unless K3_CUDA_MLA=0.

Owner: model-arch, Kimi K3. Invariants:

  • Every temporary device buffer a launch allocates gets a free call before it returns, on success and on error; a failed free is ignored.
  • MlaDeviceKv::alloc and MlaDeviceKv::alloc_and_upload, when they fail, call free on what they allocated, ignoring a failed free.

Structs§

K3MlaDecodeKernels
MlaDeviceKv
2026-09-25: One sequence’s MLA K and V on the device, allocated for cap tokens. A decode token appends one row (K device to device after rope, V from the host) and attention reads the buffers in place; download copies the history to a host MlaKv.

Constants§

MODULE
2026-09-25: Built from kernels/gb10/kimi-k3/bf16/mla_decode.cu.
ROPE_ENTRY
SDPA_ENTRY

Functions§

launch_k3_mla_decode_token
2026-09-25: One decode token against a host MlaKv: writes the roped q/k back, appends to kv, and uploads the whole KV history.
launch_k3_mla_decode_token_on_device
2026-09-25: One decode token against an MlaDeviceKv: rope stays on the device, the append is one row, and only the output is copied back.