refactor(ops): migrate paged attention prefill to canonical InfiniOps - #1482
Merged
voltjia merged 1 commit intoAug 10, 2026
Conversation
voltjia
merged commit Aug 10, 2026
bf24faa
into
refactor/migrate-paged-attention-infiniops
9 of 10 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
mha_varlenpaged-prefill calls through the canonical InfiniOpsFlashAttnVarlenFuncAPI.21b07ebcfdb0f993d2f3b672a4e38788e489fb80to the merged [2025秋季][T1-1-49]zili2004 #911 revision1c865aea58a6af8bbfdc67b76a4ed0ea8d1c167d.Alignment
mha_varlenpaged prefillFlashAttnVarlenFunc(q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table, max_seqlen_q, max_seqlen_k, dropout_p=0, softmax_scale=scale, causal=true, window_size=(-1, -1), softcap=0, deterministic=false, return_attn_probs=false, out, softmax_lse=nullopt, s_dmask=nullopt)The canonical API matches Dao-AILab FlashAttention's operator name, argument set, optional inputs, and return contract. InfiniOps orders tensor inputs before attributes and outputs last according to its contribution rules.
Scope
No public InfiniCore Python or C++ API changes are introduced. Compatible NVIDIA FP16/BF16 paged layouts use linked implementation index 16; incompatible layouts keep the existing fallback. No native provider or new test function is added because the existing
mha_varlenmatrix already covers both paths.This PR is stacked on #1481 and should be retargeted to
mainafter #1481 merges. InfiniLM call-site changes are intentionally outside this PR.Screenshots: N/A (backend adapter migration only).
Validation
Run on
ssh nvidiainaccelerator-dev/nvidia:latest:scripts/format.py --path src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc --checkpassed.git diff --checkpassed._infinicorebuilds passed with both FlashAttention linked providers enabled.nmconfirmed the built InfiniCore C++ API references the canonicalFlashAttnVarlenFunccall.python test/infinicore/ops/mha_varlen.py --nvidia --verbose: 22/22 passed, covering 20 FP16/BF16 paged cases and 2 dense MLA fallback cases.The focused operator allowlist omits unrelated
topk_softmaxbecause enabling its linked provider would require the image-missing vLLM_moe_Clibrary. Python loading therefore used lazy symbol binding only for that unrelated wrapper; the canonical attention symbol was present and exercised.