Kimi-Linear W4A16 Megakernel

источник ↗ ·

Одноядерный CUDA-мегакернел (NBLK=188, NTHR=512) для декодирования Kimi-Linear W4A16: выполняет все 4 слоя (KDA×3 + MLA), int4-GEMV без материализации весов, gated-delta KDA, MLA с поглощением q_nope и двухпроходным вниманием, MoE top-8 + 9 экспертов, RMSNorm и residual за один запуск.

cuda-kernel int4-quantization moe mla kda single-launch w4a16 megakernel

privacy

Эта заметка участвует в semantic memory.

AI Summary

  1. Краткий обзор: один cooperative CUDA-мегакернел torch.load_inline выполняет полный per-token forward Kimi-Linear W4A16 (4 слоя: KDA×3 + MLA) за один запуск, включая fused int4 dequant-GEMV, short causal conv+silu, gated-delta state update, MLA absorbed-latent attention, MoE router+top-8+9 экспертов и все RMSNorm/residual.
  2. Ключевые тезисы: 14 grid-wide sense-reversing барьеров (~1 мкс каждый) + producer-consumer spin-счётчики; int4 dequant точен до bit-identical (LOP3 0x4300, HSUB2/HMUL2, single RN); Split-K с atomic done-счётчиком; MLA не материализует k/v, q_nope поглощается через kv_b k-half; Model владеет capacity-буфером для single-launch ingest старого кэша.
  3. Важные детали и нюансы: HID=2304, KDA_SCALE=0.08838834764831845, MLA_SCALE=0.07216878364870323, RSCALE=2.446, ROPE_THETA=10000, GROUP_SIZE=128, DSMEM=76800; точные размеры: NKD=4096, KVL=512, KVBN=8192, MI=1024, NQ=6144; 32 KDA-головы, 32 MLA-головы; 64 эксперта, top-8 + 1 shared; hardware: RTX_PRO_6000 (sm_120); cooperative launch требует ≥188 блоков на SMs.
  4. Практические выводы/следующие шаги: использовать mk_setup (48+15 тензоров) → mk_step с указателями состояния и кэша; state_dict-совместим с reference.py; при первом шаге ingest=1 копирует внешний кэш; последующие шаги append in-place; для профилирования доступны stamps[64].

Мои мысли

graph

Карта связей