Skip to content

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

Open
voltjia wants to merge 2 commits into
mainfrom
refactor/migrate-paged-attention-infiniops
Open

refactor(ops): migrate paged attention to canonical InfiniOps#1481
voltjia wants to merge 2 commits into
mainfrom
refactor/migrate-paged-attention-infiniops

Conversation

@voltjia

@voltjia voltjia commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator

What

  • Route compatible NVIDIA paged-attention decode calls through the canonical InfiniOps FlashAttnWithKvcache API.
  • Route compatible NVIDIA mha_varlen paged-prefill calls through the canonical InfiniOps FlashAttnVarlenFunc API.
  • Preserve the existing decode and ATen fallbacks for layouts that the linked FlashAttention providers cannot accept.
  • Enable the canonical linked providers together with the legacy wrappers 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
paged_attention decode FlashAttnWithKvcache(q, k_cache, v_cache, k=nullopt, v=nullopt, rotary_cos=nullopt, rotary_sin=nullopt, cache_seqlens, cache_batch_idx=nullopt, cache_leftpad=nullopt, block_table, alibi_slopes, softmax_scale, causal=true, window_size=(-1, -1), softcap=0, rotary_interleaved=true, num_splits=0, return_softmax_lse=false, out, softmax_lse=nullopt) Dao-AILab flash_attn_with_kvcache, InfiniOps API, linked provider #894
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) Dao-AILab flash_attn_varlen_func, InfiniOps API, linked provider #911

The adapters use metadata-only views to match the public FlashAttention layouts. 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. Unsupported decode layouts continue to use PagedAttentionInfinilm; unsupported mha_varlen layouts continue to use the existing ATen path. No native provider is added because both canonical implementations are loaded from the existing linked FlashAttention library.

This PR does not remove every InfiniLM-suffixed call from InfiniCore. The remaining call sites are tracked separately and include the paged-attention decode fallback plus the legacy KV caching, paged-attention prefill, random sampling, and rotary-embedding adapters.

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

Validation

Run on ssh nvidia in accelerator-dev/nvidia:latest on NVIDIA A100 GPUs:

  • CI-equivalent scripts/format.py --check --c clang-format-16 --py ruff passed with clang-format 16.0.6 and CI-pinned ruff 0.15.20 for the decode change; clang-format 16.0.6 also passed for the prefill source.
  • git diff --check passed for the combined branch.
  • Linked InfiniOps and full _infinicore builds passed with canonical and fallback wrappers enabled together.
  • Supported paged-attention decode matrix: 14/14 passed, covering 12 existing FP16/BF16 fallback cases and 2 canonical linked-provider cases.
  • python test/infinicore/ops/mha_varlen.py --nvidia --verbose: 22/22 passed, covering 20 FP16/BF16 paged cases and 2 dense MLA fallback cases.

@voltjia
voltjia requested a review from a team August 9, 2026 04:38
@voltjia voltjia changed the title refactor(ops): migrate paged attention decode to canonical InfiniOps refactor(ops): migrate paged attention to canonical InfiniOps Aug 10, 2026
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