Skip to content

feat(linked): add FlashAttention FlashAttnVarlenFunc provider - #911

Merged
voltjia merged 1 commit into
masterfrom
feat/link-flash-attn-varlen
Aug 9, 2026
Merged

feat(linked): add FlashAttention FlashAttnVarlenFunc provider#911
voltjia merged 1 commit into
masterfrom
feat/link-flash-attn-varlen

Conversation

@voltjia

@voltjia voltjia commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add an NVIDIA linked provider for the existing FlashAttnVarlenFunc operator by resolving FlashAttention's exported flash::mha_varlen_fwd symbol.
  • Validate dense and paged KV layouts plus the existing optional block_table and alibi_slopes inputs.
  • Extend the existing operator test matrix with paged KV, ALiBi, auxiliary-output, and linked-provider coverage.

Motivation

FlashAttnVarlenFunc is the canonical replacement target for the deprecated PagedAttentionPrefillInfinilm interface, but its paged-KV path had no NVIDIA provider. The NVIDIA development image already ships FlashAttention and exports the required ABI, so the provider can be linked without adding another native kernel implementation or changing the public InfiniOps signature.

Type of Change

  • feat - new feature / new operator / new platform
  • fix - bug fix
  • perf - performance improvement (no behavioral change)
  • refactor - code restructuring without behavior change
  • test - adding or fixing tests only
  • docs - documentation only
  • build / ci - build system or CI configuration
  • chore - tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

Remote environment: ssh nvidia, image accelerator-dev/nvidia:latest, PyTorch 2.10.0a0+b4e4ee81d3.nv25.12, FlashAttention 2.7.4.post1+25.12.

$ CMAKE_BUILD_PARALLEL_LEVEL=8 python3 -m pip install . --no-build-isolation --no-deps \
    -C cmake.define.INFINI_RT_ROOT=/opt/infinirt \
    -C cmake.define.WITH_NVIDIA=ON \
    -C cmake.define.WITH_LINKED=ON \
    -C cmake.define.WITH_TORCH=OFF \
    -C cmake.define.INFINI_OPS_OPS=flash_attn_varlen_func \
    -C cmake.define.INFINI_OPS_TORCH_OPS=flash_attn_varlen_func
Successfully built infini-ops

$ python3 -m pytest tests/test_flash_attn_varlen_func.py -q --devices cuda
43 passed, 4 skipped in 2.78s

$ python3 -m pytest tests/test_resolve_linked_ops.py -q
24 passed in 1.14s

The four focused-test skips are the expected implementation-index 8 paged-KV/ALiBi cases; those arguments are supported by the new linked implementation-index 16 provider.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes focused linked build passed 43 passed, 4 skipped; resolver 24 passed
Iluvatar No N/A - not affected N/A - not affected
MetaX No N/A - not affected N/A - not affected
Cambricon No N/A - not affected N/A - not affected
Moore No N/A - not affected N/A - not affected
Ascend No N/A - not affected N/A - not affected

Formatting was checked with the CI-pinned tools: clang-format 21 and ruff 0.15.22 (ruff check and ruff format --check).

Benchmark / Performance Impact

N/A. This PR adds provider coverage and does not replace or tune a kernel.

Notes for Reviewers

API alignment

InfiniOps API Alignment target Evidence Mapping notes
FlashAttnVarlenFunc(q, k, v, cu_seqlens_q, cu_seqlens_k, alibi_slopes, block_table, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal, window_size, softcap, deterministic, return_attn_probs, out, softmax_lse, s_dmask) Dao-AILab FlashAttention flash_attn_varlen_func public Python interface and return contract, exported C++ ABI, paged-KV shape checks, C++ return tuple The operator name, argument set, optional inputs, and outputs match the public API. InfiniOps orders tensor inputs before attributes and output tensors last per CONTRIBUTING.md.

The linked symbol was verified directly in the image's FlashAttention DSO. A direct upstream paged-varlen smoke case also matched a materialized PyTorch reference (max_abs_error=0.000244140625, allclose=True).

The existing implementation-index 8 provider does not implement block_table or alibi_slopes; it now rejects those arguments explicitly instead of silently ignoring them. No public overload was added or removed.

@voltjia
voltjia requested a review from a team August 8, 2026 13:28
@voltjia voltjia changed the title feat(linked): add FlashAttention varlen provider feat(linked): add FlashAttention FlashAttnVarlenFunc provider Aug 9, 2026
@voltjia
voltjia merged commit 1c865ae into master Aug 9, 2026
20 checks passed
@voltjia
voltjia deleted the feat/link-flash-attn-varlen branch August 9, 2026 00:22
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