Skip to content

refactor(ops): migrate paged attention prefill to canonical InfiniOps - #1482

Merged
voltjia merged 1 commit into
refactor/migrate-paged-attention-infiniopsfrom
refactor/migrate-paged-prefill-infiniops
Aug 10, 2026
Merged

refactor(ops): migrate paged attention prefill to canonical InfiniOps#1482
voltjia merged 1 commit into
refactor/migrate-paged-attention-infiniopsfrom
refactor/migrate-paged-prefill-infiniops

Conversation

@voltjia

@voltjia voltjia commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator

What

  • Route compatible NVIDIA mha_varlen paged-prefill calls through the canonical InfiniOps FlashAttnVarlenFunc API.
  • Preserve the existing ATen fallback for unsupported layouts, including dense MLA with different QK and V head dimensions.
  • Enable the linked varlen provider when either the legacy prefill operator or the canonical operator is selected, while keeping both wrappers available during migration.
  • Update the InfiniOps gitlink from 21b07ebcfdb0f993d2f3b672a4e38788e489fb80 to the merged [2025秋季][T1-1-49]zili2004  #911 revision 1c865aea58a6af8bbfdc67b76a4ed0ea8d1c167d.

Alignment

InfiniCore adapter Canonical InfiniOps call Open-source alignment
mha_varlen paged prefill FlashAttnVarlenFunc(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) InfiniOps API, linked provider #911, Dao-AILab public Python interface

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_varlen matrix already covers both paths.

This PR is stacked on #1481 and should be retargeted to main after #1481 merges. InfiniLM call-site changes are intentionally outside this PR.

Screenshots: N/A (backend adapter migration only).

Validation

Run on ssh nvidia in accelerator-dev/nvidia:latest:

  • clang-format 16.0.6: scripts/format.py --path src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc --check passed.
  • git diff --check passed.
  • Linked InfiniOps and full _infinicore builds passed with both FlashAttention linked providers enabled.
  • nm confirmed the built InfiniCore C++ API references the canonical FlashAttnVarlenFunc call.
  • 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_softmax because enabling its linked provider would require the image-missing vLLM _moe_C library. Python loading therefore used lazy symbol binding only for that unrelated wrapper; the canonical attention symbol was present and exercised.

@voltjia
voltjia merged commit bf24faa into refactor/migrate-paged-attention-infiniops Aug 10, 2026
9 of 10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant