Skip to content

[New features] Add HySparse MQA/MLA sparse attention integration #1453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ForFishes wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

[New features] Add HySparse MQA/MLA sparse attention integration #1453

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b9f3145
HySparse MQA sparse attention (FA4 block scoring + SWA MQA + block-sp…
ForFishes Jul 15, 2026
da832e0
Add HySparse MQA sparse-attention test coverage
ForFishes Jul 15, 2026
2b6e928
HySparse: thread origin_input_ids through HySparseTransformerLayer; r…
ForFishes Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/paddlefleet/cudnn_ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@
"""cuDNN frontend ops bridged into PaddleFleet."""

from .attn.csa_sparse_attn_bwd_cudnn import csa_sparse_attn_bwd_cudnn
from .block_sparse_mqa_dsa import (
block_sparse_mqa_attention_dsa,
is_dsa_available,
)
from .indexer.csa_indexer_bwd_cudnn import csa_indexer_bwd
from .indexer.csa_indexer_fwd_cudnn import (
cudnn_indexer_forward,
Expand All @@ -23,9 +27,11 @@
)

__all__ = [
"block_sparse_mqa_attention_dsa",
"csa_indexer_bwd",
"csa_sparse_attn_bwd_cudnn",
"cudnn_indexer_forward",
"cudnn_indexer_topk",
"cudnn_indexer_topk_fwd",
"is_dsa_available",
]
14 changes: 11 additions & 3 deletions src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,14 +47,18 @@ def _get_topk_alignment() -> int:


def flash_mla_sparse_attn(
q, kv, attn_sink, topk_idxs, sm_scale=None, indexer_topk: int = 0
q, kv, attn_sink, topk_idxs, sm_scale=None, indexer_topk: int = 0, d_v=None
):
if _flash_mla_sparse_fwd is None:
raise RuntimeError("flash_mla is not available")

b, sq, h, d = q.shape
_, skv, _ = kv.shape
topk = topk_idxs.shape[-1]
# Value dim may be smaller than the query/key dim (absorbed MLA MQA uses
# d_qk=576 / d_v=512). Default to a symmetric d_v=d.
if d_v is None:
d_v = d

q_flat = q.reshape([b * sq, h, d])
kv_flat = kv.reshape([b * skv, d])
Expand All @@ -70,7 +74,7 @@ def flash_mla_sparse_attn(
kv_flat.unsqueeze(1),
global_idxs.unsqueeze(1),
sm_scale,
d_v=d,
d_v=d_v,
attn_sink=attn_sink,
topk_length=None,
indexer_topk=indexer_topk,
Expand All @@ -81,4 +85,8 @@ def flash_mla_sparse_attn(
else:
out_flat, _max_logits, lse = res
lse_indexer = None
return out_flat.reshape([b, sq, h, d]), lse.reshape([b, sq, h]), lse_indexer
return (
out_flat.reshape([b, sq, h, d_v]),
lse.reshape([b, sq, h]),
lse_indexer,
)
Loading
Loading