Skip to content

feat(iluvatar): add native canonical attention providers - #975

Open
gongchensu wants to merge 2 commits into
InfiniTensor:masterfrom
gongchensu:feat/iluvatar-canonical-attention
Open

feat(iluvatar): add native canonical attention providers#975
gongchensu wants to merge 2 commits into
InfiniTensor:masterfrom
gongchensu:feat/iluvatar-canonical-attention

Conversation

@gongchensu

@gongchensu gongchensu commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add native Iluvatar implementation-index 0 providers for the canonical FlashAttnVarlenFunc and FlashAttnWithKvcache APIs, reusing the existing native paged prefill and decode kernels.
  • Adapt canonical paged-cache layouts and cumulative KV sequence offsets without allocating temporary device tensors.
  • Keep existing CUDA-compatible providers on their original per-sequence-length and launch configurations while applying an Iluvatar-only 128-thread launch cap to ReshapeAndCacheFlash.
  • Fix the standalone Iluvatar CMake link configuration by removing the unavailable CUDA::cuda_driver imported target; the backend only uses the CUDA runtime and cuBLAS APIs.
  • Add Iluvatar canonical-attention smoke selection and an independent PyTorch paged-decode reference.

Motivation

Modern InfiniLM dispatches paged prefill and decode through the canonical
InfiniOps FlashAttnVarlenFunc and FlashAttnWithKvcache APIs. Before this
change, Iluvatar only had the legacy _infinilm compatibility providers, so
the canonical calls had no active Iluvatar implementation.

The canonical varlen API supplies cumulative KV offsets such as
[0, 130, 430], while the reused native paged-prefill kernel expects one
length per sequence, such as [130, 300]. This change lets the kernel derive
adjacent differences directly, avoiding an additional device allocation and
conversion kernel during InfiniRT graph execution.

The Iluvatar CMake configuration also referenced CUDA::cuda_driver, which is
not provided by the CoreX FindCUDAToolkit package. InfiniOps does not use the
CUDA Driver API on this backend, so only CUDA::cudart and CUDA::cublas are
required.

Validated with static, paged-attention plus InfiniRT graph, and explicit
FlashAttention plus InfiniRT graph inference using the 9G-8B model on
Iluvatar BI-V150.

N/A - no linked issue.

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
  • Hygon (WITH_HYGON) — shared ReshapeAndCacheFlash template
  • Mars (WITH_MARS) — shared paged-attention and cache templates
  • T-Head (WITH_THEAD) — shared ReshapeAndCacheFlash template

Smoke Test Result

paste smoke test output here

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA
Iluvatar
MetaX
Cambricon
Moore
Ascend
Full `pytest` output (optional)
paste here

Benchmark / Performance Impact

Notes for Reviewers

@gongchensu gongchensu self-assigned this Sep 3, 2026
@gongchensu
gongchensu requested a review from a team September 3, 2026 07:56
@gongchensu
gongchensu force-pushed the feat/iluvatar-canonical-attention branch from 6f93119 to 8f0770d Compare September 4, 2026 01:43
voltjia
voltjia previously approved these changes Sep 4, 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.

2 participants