Code

Add POD-Attention to FlashInfer#858

Code

Open

AKKamath wants to merge 21 commits intoflashinfer-ai:mainflashinfer-ai/flashinfer:mainfrom AKKamath:new_branchAKKamath/flashinfer:new_branchCopy head branch name to clipboard

+2,837-143

Commits

Commits on Feb 14, 2025

• Separate single prefill into device and kernel to prepare for POD

  

  Aditya K Kamath

  committed
  a2e8212

  Copy full SHA for a2e8212

• Add initial pod-attention code

  

  Aditya K Kamath

  committed
  e5fe841

  Copy full SHA for e5fe841

• Fix all compiler errors with newer version of repo

  

  Aditya K Kamath

  committed
  5bd1607

  Copy full SHA for 5bd1607

• Split BatchDecode into kernel and device function

  

  Aditya K Kamath

  committed
  99b2318

  Copy full SHA for 99b2318

• Rewrite pod params setup to explicit prefill params. TODO: Add decode params

  

  Aditya K Kamath

  committed
  658319c

  Copy full SHA for 658319c

• Add wrapper to POD for batched decode. TODO: Actually redirect wrapper to POD instead of batched decode.

  

  Aditya K Kamath

  committed
  3c2b817

  Copy full SHA for 3c2b817

• Add decode wrapper to pod test

  

  Aditya K Kamath

  committed
  08c323f

  Copy full SHA for 08c323f

• Add both prefill and decode params to POD-Attention methods. Currently decode kernel is manually called and needs to be merged into POD kernel

  

  Aditya K Kamath

  committed
  54daf81

  Copy full SHA for 54daf81

• Make threadId a function parameter

  

  Aditya K Kamath

  committed
  637a5b5

  Copy full SHA for 637a5b5

• Use more realistic parameters for pod test

  

  Aditya K Kamath

  committed
  e19d05c

  Copy full SHA for e19d05c

• Working merged kernel. TODO: SM-aware scheduling

  

  Aditya K Kamath

  committed
  55e5f71

  Copy full SHA for 55e5f71

• Add tensor-based decode version of POD. Currently decode output is mismatching!

  

  Aditya K Kamath

  committed
  3e10ac0

  Copy full SHA for 3e10ac0

• Add pod to compilation

  

  Aditya K Kamath

  committed
  a2602af

  Copy full SHA for a2602af

• Fix issues where POD's decode output didn't match.

  

  Aditya K Kamath

  committed
  558757c

  Copy full SHA for 558757c

• Add warning for POD's no_tensor version which performs worse

  

  Aditya K Kamath

  committed
  d382394

  Copy full SHA for d382394

Commits on Feb 15, 2025

• Fix compilation errors and cleanup POD

  

  Aditya K Kamath

  committed
  eb3dc4e

  Copy full SHA for eb3dc4e

Commits on Feb 17, 2025

• Working version of POD

  

  Aditya K Kamath

  committed
  e68c473

  Copy full SHA for e68c473

• Update setup

  

  Aditya K Kamath

  committed
  d6485d7

  Copy full SHA for d6485d7

• Speed up compilation for POD

  

  Aditya K Kamath

  committed
  bcb8ee6

  Copy full SHA for bcb8ee6

Commits on Feb 18, 2025

• Add JIT support for POD

  

  Aditya K Kamath

  committed
  547f14e

  Copy full SHA for 547f14e

Commits on Feb 19, 2025

• Fix POD performance bug (at the expense of larger compile time)

  

  Aditya K Kamath

  committed
  0c2cdd5

  Copy full SHA for 0c2cdd5