Code

feat: support MLA decode#551

Code

Merged

yzh119 merged 17 commits intoflashinfer-ai:mainflashinfer-ai/flashinfer:mainfrom tsu-bin:mla_decode_devtsu-bin/flashinfer_dev:mla_decode_devCopy head branch name to clipboardNov 2, 2024

+1,887-102

Commits

Commits on Oct 30, 2024

• mv work_estimation_func outside from DecodePlan

  tsu-bin

  committed
  6edba14

  Copy full SHA for 6edba14

• mv work_estimation_func outside from DecodePlan, update AOT related code

  tsu-bin

  committed
  879a59d

  Copy full SHA for 879a59d

• add BatchDecodeWithPagedKVCacheKernelMLA BatchDecodeWithPagedKVCacheWorkEstimationDispatchedMLA

  tsu-bin

  committed
  386f162

  Copy full SHA for 386f162

• add BatchDecodeHandlerPlanMLA BatchDecodeWithPagedKVCacheWrapperMLA, now 512x128 is supported

  tsu-bin

  committed
  4084ae9

  Copy full SHA for 4084ae9

• jit support, add BatchDecodeMlaWithPagedKVCacheWrapper

  tsu-bin

  committed
  c64c2f3

  Copy full SHA for c64c2f3

• add paged_kv_mla_t and BatchDecodeParamsMLA

  tsu-bin

  committed
  4231bfc

  Copy full SHA for 4231bfc

• HEAD_DIM -> HEAD_DIM_CKV and HEAD_DIM_KPE

  tsu-bin

  committed
  04a3bb2

  Copy full SHA for 04a3bb2

• mla only use ROPE, and sm_scale should be set explicitly

  tsu-bin

  committed
  3fa9947

  Copy full SHA for 3fa9947

• add q_nope q_pe

  tsu-bin

  committed
  9ef5fd9

  Copy full SHA for 9ef5fd9

• implement MLA kernel

  tsu-bin

  committed
  f443bc1

  Copy full SHA for f443bc1

• some clean up

  tsu-bin

  committed
  3014b3f

  Copy full SHA for 3014b3f

• change some variable names of mla decode kernel

  tsu-bin

  committed
  cafa314

  Copy full SHA for cafa314

• add mla decode benchmark

  tsu-bin

  committed
  3cc1052

  Copy full SHA for 3cc1052

• add tile_size_qo_heads

  tsu-bin

  committed
  b14f7a1

  Copy full SHA for b14f7a1

• support num_qo_heads can be non-integral multiple of qo_heads_per_block to accommodate MiniCPM3-4B who has 40 num_qo_heads

  tsu-bin

  committed
  478c835

  Copy full SHA for 478c835

• update according to recent changes from upstream and some clean up

  tsu-bin

  committed
  d8f10ee

  Copy full SHA for d8f10ee

Commits on Nov 1, 2024

• update testcase, add WMAPE and MSE metrics

  tsu-bin

  committed
  89edcda

  Copy full SHA for 89edcda