Repository navigation

feat: support kv cache reuse for MLA - #3571

Merged

kaiyux merged 15 commits into

NVIDIA:mainfrom

zhhuang-nv:user/zhhuang/mla-kv-cache-reuse

May 15, 2025

Merged

feat: support kv cache reuse for MLA#3571
kaiyux merged 15 commits into
NVIDIA:mainfrom
zhhuang-nv:user/zhhuang/mla-kv-cache-reuse

Commits on May 15, 2025
- support kv cache reuse for MLA
  zhhuang-nv
  committed
  cc0fffe
- add CI test
  zhhuang-nv
  committed
  576df53
- fix: set k_pe head_num to 1 for kernel 2 and kernel 2V2
  
  jmydurant
  authored and
  zhhuang-nv
  committed
  ee19565
- resolve comments
  zhhuang-nv
  committed
  779bd7c
- use GPTJ style RoPE for MLA
  zhhuang-nv
  committed
  3960d86
- fix rebase error and some docs
  zhhuang-nv
  committed
  2ca9012
- fix kv_lens
  zhhuang-nv
  committed
  dcc2196
- tiny fix
  zhhuang-nv
  committed
  6529621
- fix torch compile
  zhhuang-nv
  committed
  4e354c8
- fix: use normal device memory instead of pinned memory for unit test
  
  jmydurant
  authored and
  zhhuang-nv
  committed
  338fd78
- fix L0 tests
  zhhuang-nv
  committed
  743522d
- fix torch compile after rebase
  zhhuang-nv
  committed
  89bf494
- resolve comments
  zhhuang-nv
  committed
  229c218
- resolve comments again
  zhhuang-nv
  committed
  c918865
- Merge branch 'main' into user/zhhuang/mla-kv-cache-reuse
  zhhuang-nv
  authored
  399fa2f