Code

Add ragged paged attention#8659

Code

Open

vanbasten23 wants to merge 8 commits intomasterpytorch/xla:masterfrom xiowei/add_ragged_paged_attentionpytorch/xla:xiowei/add_ragged_paged_attentionCopy head branch name to clipboard

+1,359-1

Commits

Commits on Feb 3, 2025

• add the first version. All tests pass except for test_paged_attention_extreme_one_tokens_per_sequence.

  vanbasten23

  committed
  b129e9a

  Copy full SHA for b129e9a

• all tests passed.

  vanbasten23

  committed
  6c3bf73

  Copy full SHA for 6c3bf73

• all tests passed except for one test which oom'ed.

  vanbasten23

  committed
  dbc3fee

  Copy full SHA for dbc3fee

• Improved the tests and all tests passed except for the OOM one. Also added runtime check.

  vanbasten23

  committed
  e589c86

  Copy full SHA for e589c86

• clean up

  vanbasten23

  committed
  9b3cdab

  Copy full SHA for 9b3cdab

• linter

  vanbasten23

  committed
  ab69feb

  Copy full SHA for ab69feb

• address pr comments

  vanbasten23

  committed
  cd65cc4

  Copy full SHA for cd65cc4

• fix the rest of comments

  vanbasten23

  committed
  7fe5071

  Copy full SHA for 7fe5071