Skip to content

Replace FlashAttention with xformers - #70

Merged
WoosukKwon merged 12 commits into
mainfrom
xformers
May 5, 2023
Merged

Replace FlashAttention with xformers#70
WoosukKwon merged 12 commits into
mainfrom
xformers

Conversation

@WoosukKwon

@WoosukKwon WoosukKwon commented May 4, 2023

Copy link
Copy Markdown
Collaborator

This PR replaces FlashAttention with xformers.

Pros:

  • Richer features & higher compatibility. xformers supports attention bias, FP32, head size 256, and old GPUs (such as V100) while FlashAttention does not.
  • xformers provides pre-compiled python wheels, while FlashAttention compiles the entire CUDA code during installation.
  • Future-proof, as the repository is maintained by many developers from Meta.

Cons:

  • xformers can be slower than FlashAttention for small inputs, because it incurs higher CPU overheads.
  • xformers internally creates a new tensor for the attention output. In our case, this leads to an extra copy overhead, because we concatenate the outputs of the two attention ops.

@WoosukKwon
WoosukKwon requested a review from zhuohan123 May 4, 2023 10:31
WoosukKwon
WoosukKwon commented May 4, 2023
View reviewed changes
Comment thread README.md
pip install sentencepiece # Required for LlamaTokenizer.
pip install ninja # To parallelize the compilation of flash-attn.
pip install flash-attn # This may take up to 10 mins.
pip install ninja psutil numpy sentencepiece ray torch transformers xformers

@WoosukKwon WoosukKwon May 4, 2023

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO (in the next PR): specify the exact dependencies in setup.py.

@zhisbug

zhisbug commented May 4, 2023

Copy link
Copy Markdown
Collaborator

is the memory footprint same with flashattention?

@zhisbug

zhisbug commented May 5, 2023

Copy link
Copy Markdown
Collaborator

I did a test myself and found the memory saving is almost the same.

@WoosukKwon

WoosukKwon commented May 5, 2023

Copy link
Copy Markdown
Collaborator Author

It seems the memory usage is comparable to FlashAttention's. @zhuohan123 Please review the PR.

zhuohan123
zhuohan123 approved these changes May 5, 2023
View reviewed changes

@zhuohan123 zhuohan123 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks!

Comment thread cacheflow/master/server.py
help='save a numpy copy of model weights for faster loading')
parser.add_argument('--use-dummy-weights', action='store_true', help='use dummy values for model weights')
# NOTE(woosuk): FlashAttention does not support float32.
# TODO(woosuk): Support FP32 for debugging.

@zhuohan123 zhuohan123 May 5, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does xformers support FP32?

@WoosukKwon WoosukKwon May 5, 2023

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it does. It is our attention kernel that does not support FP32. More precisely, our attention kernel currently does not support some block sizes when FP32 is used. I will fix this in the future.

@WoosukKwon WoosukKwon mentioned this pull request May 5, 2023
Closed
@WoosukKwon
WoosukKwon merged commit c9d5b6d into main May 5, 2023
@WoosukKwon
WoosukKwon deleted the xformers branch May 5, 2023 09:01
@tmm1 tmm1 mentioned this pull request Aug 3, 2023
Closed
hongxiayang pushed a commit to hongxiayang/vllm that referenced this pull request Feb 13, 2024
@WoosukKwon
Replace FlashAttention with xformers (vllm-project#70)
b43acbb
yukavio pushed a commit to yukavio/vllm that referenced this pull request Jul 3, 2024
@tlrmchlsmth
Mark files that have been modified by Neural Magic (vllm-project#70)
13787e6 
SUMMARY:
for Apache 4(b) -- "You must cause any modified files to carry prominent
notices stating that You changed the files"
https://www.apache.org/licenses/LICENSE-2.0 

TEST PLAN:
GHA
dllehr-amd pushed a commit to dllehr-amd/vllm that referenced this pull request Jul 22, 2024
@mawong-amd @Alexei-V-Ivanov-AMD
[2/2] Using xfail instead of skip for ROCm 6.2 tests (vllm-project#70)
cce6281 
* Enabling some basic tests for ROCm 6.2

Use strict xfail for ROCm 6.2 test repairs

* Use lenient xfail instead

---------

Co-authored-by: Alexei V. Ivanov <alexei.ivanov@amd.com>
@yazdanbakhsh yazdanbakhsh mentioned this pull request Sep 13, 2024
Closed
wuhuikx pushed a commit to wuhuikx/vllm that referenced this pull request Mar 27, 2025
@MengqingCao
[doc] fix feature support (vllm-project#70)
c935b70 
Check and update the feature support table.

- both multi-step and speculative decoding require adaptation of corresponding workers
- prompt adapter (finetune method) require adaption in worker.py and model_runner.py

Signed-off-by: MengqingCao <cmq0113@163.com>
robertgshaw2-redhat added a commit to robertgshaw2-redhat/vllm that referenced this pull request May 6, 2025
@tlrmchlsmth @ApostaC
@robertgshaw2-redhat Robert Shaw
[P/D Disagg] NIXL MLA (vllm-project#70)
2c50275 
* [Update] LMcache connector v1 implementation

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* [Add] examples for disaggregated prefill

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* [add] extra information about evns

Signed-off-by: ApostaC <yihua98@uchicago.edu>

* Initial stubs for P/D scheduling changes

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Updates

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Rs branch (#3)

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Rs branch (#5)

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Remove Unneeded Arguments (#7)

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* stash

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* cleanup

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

---------

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Improve disagg-example.sh (#8)

- fix spelling
- CUDA_VISIBLE_DEVICES should be set externally

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* added connector

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* update

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* remove

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* seems to load properly

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Revert "updated"

This reverts commit 97316d9.

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* stash

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* added

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* diffs for local dev on macos

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* update

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updaed

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* Checkpoint.

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* Cleanup

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* WIP

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated on scheduler side

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* Hacking away

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* cleanup

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* ensure request removed from running list

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* Runs E2E. Garbage output. Crashes on 2nd request

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* update

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* rename files

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* updated

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* update

Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>

* Second request no longer crashes

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Remove gpu_model_runner hacks

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Clean up Justfile

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* [Bugfix] Stale finished requests in EMPTY_MODEL_RUNNER_OUTPUT

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* update

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* justfile edits

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Update

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Fixes - lm_eval gsm8k has correctness

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* "just delete the assert"

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* fixup precommit issues

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Fixes

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* updated (#12)

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Add Accuracy Test (#13)

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

---------

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Preemption Bugfixes (#15)

* stash fixed double free issue

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* fixed issue

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatrd

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatrd

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatrd

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatrd

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatrd

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatrd

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

---------

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated (#16)

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Fix Bad Merge | Fix Memory Leak in Upstream (#18)

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* fix merge

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

---------

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* cleanup code

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* cleanup code

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* stash

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updatted

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* revert

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* more spurious changes

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* updated

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>

* Support MLA in NIXL connector

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* WIP adding tests

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* wip

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

* Fixes

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

---------

Signed-off-by: ApostaC <yihua98@uchicago.edu>
Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Signed-off-by: Robert Shaw <rshaw@neuralmagic.com>
Co-authored-by: ApostaC <yihua98@uchicago.edu>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Co-authored-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Co-authored-by: Robert Shaw <rshaw@neuralmagic.com>
GirasoleY added a commit to GirasoleY/vllm that referenced this pull request Jul 28, 2026
@GirasoleY @claude girasoley
[Attention] Direct symmetric-memory DCP A2A for MLA
e20238b 
Port of upstream vllm-project#48897 (branch
proto/upstream-direct-dcp-mla, commits 6569754, 1691322, c15e0f1,
983f562). The CUDA kernel is byte-identical to the upstream head; the
Python workspace differs only in hoisting function-local imports to
module scope (repo style).

- csrc/libtorch_stable/attention/dcp_direct_a2a.cu: single-kernel
  symmetric-memory output/LSE exchange + LSE-weighted combine with
  epoch-based signaling; persistent workspace, CUDA-graph safe, one
  staging slot per ubatch (min 1 -- num_ubatches is 0 without DBO).
- vllm/v1/attention/ops/dcp_direct_a2a.py: DirectDCPA2AWorkspace and a
  cached per-(group, geometry) accessor.
- VLLM_USE_DIRECT_DCP_A2A tri-state: unset = auto on single-node CUDA
  bf16/fp16; 1 = force (multi-node NVLink domains); 0 = off.
- Generic MLAAttention: init-time dcp_combine dispatch (direct A2A when
  available, else packed dcp_a2a_lse_reduce for dcp_comm_backend="a2a",
  else cp_lse_ag_out_rs); dcp_world_size derived statically from
  parallel_config.

Test: tests/distributed/test_dcp_direct_a2a.py -- fp16/bf16, with and
without CUDA graphs, including Kimi-K3 decode geometry (6 heads/rank,
D=512, world size 4 matching gathered H=24, non-contiguous 128-head
backing storage). 8 passed on 2 GPUs; 12 passed on a 4-GPU GB200 node.
No regression in the existing packed-A2A suite (1 pre-existing ordering
failure reproduces on the unmodified base).

docs/design/attention_backends.md is regenerated by the pre-commit hook;
its one-line delta (FLASHINFER_MLA Non-Causal) is pre-existing drift
from vllm-project#70, not a change of this commit.

Amended: the combine kernel now masks empty-KV-shard rows itself --
the dispatch kernel takes the decode seq_lens (optional Tensor in the op
schema), publishes LSE = -inf and skips the payload copy for empty
shards, and the combine skips zero-weight sources; lse_reduce grows a
seq_lens kwarg and direct_dcp_a2a_active() is exposed; the generic
MLAAttention passes decode seq_lens to the direct combine. Folded from
dcp-all-bucket 91fc3dba8 (empty-shard combine masking).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: girasoley <girasoley@inferact.ai>
GirasoleY added a commit to GirasoleY/vllm that referenced this pull request Jul 30, 2026
@GirasoleY @claude
[Attention] Direct symmetric-memory DCP A2A for MLA
b1913c0 
Port of upstream vllm-project#48897 (branch
proto/upstream-direct-dcp-mla, commits 6569754, 1691322, c15e0f1,
983f562). The CUDA kernel is byte-identical to the upstream head; the
Python workspace differs only in hoisting function-local imports to
module scope (repo style).

- csrc/libtorch_stable/attention/dcp_direct_a2a.cu: single-kernel
  symmetric-memory output/LSE exchange + LSE-weighted combine with
  epoch-based signaling; persistent workspace, CUDA-graph safe, one
  staging slot per ubatch (min 1 -- num_ubatches is 0 without DBO).
- vllm/v1/attention/ops/dcp_direct_a2a.py: DirectDCPA2AWorkspace and a
  cached per-(group, geometry) accessor.
- VLLM_USE_DIRECT_DCP_A2A tri-state: unset = auto on single-node CUDA
  bf16/fp16; 1 = force (multi-node NVLink domains); 0 = off.
- Generic MLAAttention: init-time dcp_combine dispatch (direct A2A when
  available, else packed dcp_a2a_lse_reduce for dcp_comm_backend="a2a",
  else cp_lse_ag_out_rs); dcp_world_size derived statically from
  parallel_config.

Test: tests/distributed/test_dcp_direct_a2a.py -- fp16/bf16, with and
without CUDA graphs, including Kimi-K3 decode geometry (6 heads/rank,
D=512, world size 4 matching gathered H=24, non-contiguous 128-head
backing storage). 8 passed on 2 GPUs; 12 passed on a 4-GPU GB200 node.
No regression in the existing packed-A2A suite (1 pre-existing ordering
failure reproduces on the unmodified base).

docs/design/attention_backends.md is regenerated by the pre-commit hook;
its one-line delta (FLASHINFER_MLA Non-Causal) is pre-existing drift
from vllm-project#70, not a change of this commit.

Amended: the combine kernel now masks empty-KV-shard rows itself --
the dispatch kernel takes the decode seq_lens (optional Tensor in the op
schema), publishes LSE = -inf and skips the payload copy for empty
shards, and the combine skips zero-weight sources; lse_reduce grows a
seq_lens kwarg and direct_dcp_a2a_active() is exposed; the generic
MLAAttention passes decode seq_lens to the direct combine. Folded from
dcp-all-bucket 91fc3dba8 (empty-shard combine masking).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: girasoley <girasoley@inferact.ai>
GirasoleY added a commit to GirasoleY/vllm that referenced this pull request Aug 6, 2026
@GirasoleY @claude
[Attention] Direct symmetric-memory DCP A2A for MLA
b91e26c 
Port of upstream vllm-project#48897 (branch
proto/upstream-direct-dcp-mla, commits 6569754, 1691322, c15e0f1,
983f562). The CUDA kernel is byte-identical to the upstream head; the
Python workspace differs only in hoisting function-local imports to
module scope (repo style).

- csrc/libtorch_stable/attention/dcp_direct_a2a.cu: single-kernel
  symmetric-memory output/LSE exchange + LSE-weighted combine with
  epoch-based signaling; persistent workspace, CUDA-graph safe, one
  staging slot per ubatch (min 1 -- num_ubatches is 0 without DBO).
- vllm/v1/attention/ops/dcp_direct_a2a.py: DirectDCPA2AWorkspace and a
  cached per-(group, geometry) accessor.
- VLLM_USE_DIRECT_DCP_A2A tri-state: unset = auto on single-node CUDA
  bf16/fp16; 1 = force (multi-node NVLink domains); 0 = off.
- Generic MLAAttention: init-time dcp_combine dispatch (direct A2A when
  available, else packed dcp_a2a_lse_reduce for dcp_comm_backend="a2a",
  else cp_lse_ag_out_rs); dcp_world_size derived statically from
  parallel_config.

Test: tests/distributed/test_dcp_direct_a2a.py -- fp16/bf16, with and
without CUDA graphs, including Kimi-K3 decode geometry (6 heads/rank,
D=512, world size 4 matching gathered H=24, non-contiguous 128-head
backing storage). 8 passed on 2 GPUs; 12 passed on a 4-GPU GB200 node.
No regression in the existing packed-A2A suite (1 pre-existing ordering
failure reproduces on the unmodified base).

docs/design/attention_backends.md is regenerated by the pre-commit hook;
its one-line delta (FLASHINFER_MLA Non-Causal) is pre-existing drift
from vllm-project#70, not a change of this commit.

Amended: the combine kernel now masks empty-KV-shard rows itself --
the dispatch kernel takes the decode seq_lens (optional Tensor in the op
schema), publishes LSE = -inf and skips the payload copy for empty
shards, and the combine skips zero-weight sources; lse_reduce grows a
seq_lens kwarg and direct_dcp_a2a_active() is exposed; the generic
MLAAttention passes decode seq_lens to the direct combine. Folded from
dcp-all-bucket 91fc3dba8 (empty-shard combine masking).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: girasoley <girasoley@inferact.ai>
Fangzhou-Ai pushed a commit to Fangzhou-Ai/vllm that referenced this pull request Sep 9, 2026
@xudonlyu
Merge pull request vllm-project#70 from AMD-DCGPU-AICE-CN/flydsl_gemm…
96a689b 
…_mxfp8

support blockscale 128 by flydsl mxfp8 gemm
Isotr0py pushed a commit that referenced this pull request Sep 10, 2026
@WoosukKwon @Isotr0py
[Perf] Reorder only local SP tokens after Engram all-gather (#70)
2bc62cc 
* Reorder only local SP tokens after Engram all-gather

Select each rank token slice in the gathered rank-major view before materializing checkpoint head order. Add TP/SP sharding microbenchmarks.

Co-authored-by: OpenAI Codex <noreply@openai.com>

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

* Fuse Engram SP selection and head reorder into one copy

Keep the microbenchmark as a local experiment artifact rather than in the PR. Copy only valid local tokens and heads, zeroing padded tokens directly in the output.

Co-authored-by: OpenAI Codex <noreply@openai.com>

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

---------

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
triangleXIV pushed a commit to tokmo-infra/vllm that referenced this pull request Nov 9, 2026
@wzhao18 @claude @ZJY0516
[Kimi K3.1] FlashInfer prims_ts sparse MLA for DSA (fp8 + nvfp4 KV) (v…
ad32438 
…llm-project#70)

* [Kimi K3.1] Gather nvfp4_ds_mla prefill context under DCP (vllm-project#54)

With --kv-cache-dtype nvfp4_ds_mla and DCP > 1, the DCP chunked-context
gather fell through to gather_and_maybe_dequant_cache, which rejects
nvfp4_ds_mla. Kimi-K3.1's sparse layers take this path for any prefill
with context whose sequence fits the top-k (dense MHA), which the
startup warmup's mixed spec / non-spec decode step already produces, so
the server failed to start.

Gather with the NVFP4 upconvert kernel, as the non-DCP path does, over
the rank's padded local rows. That kernel takes no seq_starts, so a
chunk that resumes a request mid-context moves the request's workspace
start back by its local start instead; only a chunk's first request can
resume, so the workspace starts stay sorted.

Functional dependency: vllm-project#54 as of 6c1a209569 (Python only, so it runs on an
existing extension build). vllm-project#54 has since been revised to gather from
seq_starts in the CUDA kernel (17fc65982d), which needs an extension rebuild;
that revision replaces this commit when rebasing.

Signed-off-by: Wei Zhao <weizha@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Kimi K3.1] NIXL: keep manager blocks as the transfer unit of packed caches (from vllm-project#65)

Functional dependency, taken unchanged from the NIXL hunk of vllm-project#65
(yifan/k31-packed-dsa-followups). On mum-tracking a 1P1D NIXL deployment of
Kimi K3.1 cannot start: get_current_attn_backends_and_specs asserts on the
packed DSA caches' kernel-paged specs. In a packed (block-outermost) layout
the kernel blocks are views inside the transferred manager blocks, so no
(backend, spec) pairs are returned and the manager block stays the
transfer unit.

Signed-off-by: Wei Zhao <weizha@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Kimi K3.1] Keep per-tensor fp8 KV for FLASHINFER_MLA_SPARSE (from vllm-project#65)

Functional dependency, taken unchanged from the config hunk of vllm-project#65
(yifan/k31-packed-dsa-followups). On mum-tracking, KimiK3 config
verification rewrites every "fp8*" kv-cache dtype of a DSA model to
fp8_ds_mla, so `--kv-cache-dtype fp8` with FLASHINFER_MLA_SPARSE (which
cannot read fp8_ds_mla) is impossible: the per-tensor fp8 K3.1 path cannot
be selected at all. With FLASHINFER_MLA_SPARSE the explicit dtype is kept and
"auto" becomes per-tensor fp8_e4m3. (The vllm-project#52 description announced this
default, but the code did not land with it.)

Signed-off-by: Wei Zhao <weizha@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Kimi K3.1] FLASHINFER_MLA_SPARSE on FlashInfer prims_ts sparse MLA (fp8 + nvfp4_ds_mla)

The SM100 FLASHINFER_MLA_SPARSE backend can now run FlashInfer's
task-scheduled sparse MLA decode (flashinfer.attention.prims_ts.
BatchSparseMLADecodePagedTSWrapper, experimental public API) instead of the
trtllm-gen kernel, for D576 per-tensor fp8 caches and, new for this backend,
nvfp4_ds_mla caches (byte-identical to prims_ts "nvfp4.fp8rope").

- Selection: VLLM_FLASHINFER_SPARSE_MLA_KERNEL=auto|trtllm|prims_ts. auto
  picks prims_ts for (kv_lora_rank, rope) = (512, 64) fp8 / nvfp4_ds_mla on
  SM100/103 without HiSparse when FlashInfer provides it, else trtllm-gen.
  nvfp4_ds_mla requires prims_ts; an explicit prims_ts request that cannot
  run raises. The FlashInfer check probes the public workspace-size query
  for the format, so a FlashInfer without flashinfer-ai/flashinfer#7410
  fails with a clear message and moving to a nightly with it is an install
  change only.
- Every row (decode, MTP verify, sparse prefill) attends its own storage-row
  indices over the flat cache rows (packed block strides are whole rows);
  lengths are the converter's valid counts (compacted under DCP), -1 holes
  masked. Natural-log LSE; rows without entries give (0, -inf), the DCP merge
  identity, so no fix-up pass.
- 65..127 query heads (96 = DCP8 gathered heads of K3.1) are zero-padded to
  128: only H128 selects the 2-CTA schedule, 2-3x faster than the two M64
  tiles of H96.
- Q is the layer's static per-tensor E4M3 query (q_scale); the descales fold
  into the host-scalar softmax/output scales, which CUDA graphs capture.
- Plans are cached per (rows, heads, format, top-k width) and shared by all
  layers and one workspace; capture-time plan misses raise. Long prefill
  steps run in 4096-row chunks; reserve_workspace claims their peak. The
  trtllm-gen DCP softmax slab (GBs at 16K tokens) is not allocated when
  prims_ts runs.

The sparse MLA mock layer of the backend tests now quantizes the query for
nvfp4_ds_mla, as MLAAttention.forward_impl does, and the backend correctness
test checks an E4M3-query nvfp4_ds_mla backend at the fp8 tolerance.

Signed-off-by: Wei Zhao <weizha@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Kimi K3.1] prims_ts: plan every row capacity at startup; fp8 opt-in

Review follow-ups for the FlashInfer prims_ts sparse MLA path:

- Plans were created per exact row count, and FlashInfer compiles one
  kernel per dispatch schedule on first use (7-11 schedules over 1..4096
  rows, several seconds each). Eager row counts outside the CUDA graph
  sizes (small prefills, chunk tails) therefore compiled mid-serving.
  Now each geometry plans a fixed set of row capacities at once, from
  the profile run's reserve_workspace: every CUDA graph size (graph
  replays keep the schedule of their exact row count) plus the powers of
  two up to the largest call. A call runs on the smallest capacity that
  holds its rows as a packed-query plan of one-token requests
  (qo_indptr = [0, 1, ..., rows, ..., rows]), which keeps FlashInfer's
  decode (Q = 1) schedule selection. No planning or JIT while serving,
  and the plan count is bounded.
- VLLM_FLASHINFER_SPARSE_MLA_KERNEL=auto no longer moves per-tensor fp8
  models (e.g. DeepSeek V3.2) from trtllm-gen to prims_ts when FlashInfer
  is upgraded; fp8 needs =prims_ts. nvfp4_ds_mla always uses prims_ts.
- Refuse prims_ts with micro-batching (DBO / ubatching): its plans share
  one workspace, which FlashInfer allows for one stream only.

Signed-off-by: Wei Zhao <weizha@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* [Kimi K3.1] Select the FlashInfer sparse MLA kernel with kernel_config

Replace VLLM_FLASHINFER_SPARSE_MLA_KERNEL with
KernelConfig.flashinfer_sparse_mla_kernel ("auto" | "trtllm" | "prims_ts"),
e.g. --kernel-config '{"flashinfer_sparse_mla_kernel": "prims_ts"}'.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

* [Kimi K3.1] prims_ts: fix pre-commit (mypy, torch.accelerator)

- Assert the head-padded prims_ts LSE chunk before slicing it (mypy).
- Use torch.accelerator for the current device and synchronize/empty_cache
  in the backend, the FlashInfer probe, the test and the benchmark.
- Unpack the test reference before _check (mypy).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

* [Kimi K3.1] prims_ts: default use_prims_ts on the impl class

Impls built without __init__ (tests/v1/attention/test_sparse_mla_backends.py
on mum-tracking) reach the trtllm-gen decode, which now reads use_prims_ts.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

---------

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zhuohan123 zhuohan123 zhuohan123 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@WoosukKwon @zhisbug @zhuohan123