Skip to content

feat: add LightningAttentionInfinilm for MiniMax linear attention - #994

Closed
y258dd wants to merge 3 commits into
InfiniTensor:masterfrom
y258dd:archive/lightning-attention-infinilm
Closed

y258dd wants to merge 3 commits into
InfiniTensor:masterfrom
y258dd:archive/lightning-attention-infinilm

Conversation

@y258dd

@y258dd y258dd commented Sep 19, 2026

Copy link
Copy Markdown

Summary

  • Add LightningAttentionInfinilm, an InfiniLM-specific indexed-pool lightning attention operator.

  • Add the shared operator interface and metadata in src/base/lightning_attention_infinilm.h.

  • Add a CPU reference implementation in src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h.

  • Add CUDA device kernel, launcher, and NVIDIA vendor binding under src/native/cuda/ and src/native/cuda/nvidia/.

  • Add an Ascend ACLNN implementation in src/native/ascend/ops/lightning_attention_infinilm/kernel.h.

  • Add pytest coverage in tests/test_lightning_attention_infinilm.py.

Motivation

MiniMax-style linear attention uses a recurrent lightning-attention state with a per-head decay and an indexed state pool:

ratio = exp(-slope)
state = ratio * state + outer(k_t, v_t)
out_t = q_t @ state

PyTorch has no directly corresponding operator. The closest public reference is Flash-Linear-Attention's fused_recurrent_lightning_attn, but this implementation also has the InfiniLM-specific indexed state-pool contract.

The operator is needed by the MiniMax model support work in InfiniLM.
Closes #

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

Affected-platform builds succeeded.

Operator-specific smoke commands and results:

# CPU
python3 -m pytest tests/test_lightning_attention_infinilm.py --devices cpu -v
24 passed

# NVIDIA RTX 5090, CUDA 12.8
pytest tests/test_lightning_attention_infinilm.py -v
48 passed  # CPU + CUDA

# Ascend 910B1, CANN 9.0.0, torch_npu 2.10.0
python3 -m pytest tests/test_lightning_attention_infinilm.py --devices ascend -v
24 passed```

## Test Results on Supported Platforms

<!--
Per `CONTRIBUTING.md` §Pull Requests, build and run the smoke test set on every
affected platform. Use `smoke passed`, `full passed`, or `N/A - not affected` in
the result columns. Run the full suite for high-risk changes, release prep,
maintainer spot checks, or changes affecting shared build, dispatch, wrapper
generation, or cross-platform behavior.

If an affected platform was not tested, state the reason and tag a reviewer or
owner with access. Reviewers may request full-suite or full-platform validation
when the risk profile justifies it.
-->

| Platform   | Affected | Build / Smoke Result | Full Result / Notes |
| ---------- | :------: | ------------------- | ------------------- |
| CPU        |   Yes    | build passed        | operator tests passed: 24 |
| NVIDIA     |   Yes    | build passed        | operator tests passed: 48 (24 CPU + 24 CUDA) |
| Iluvatar   |   No     | N/A                 | N/A |
| MetaX      |   No     | N/A                 | N/A |
| Cambricon  |   No     | N/A                 | N/A |
| Moore      |   No     | N/A                 | N/A |
| Ascend     |   Yes    | build passed        | operator tests passed: 24 on 910B1 / CANN 9.0.0 |
<details>
<summary>Full `pytest` output (optional)</summary>

```text
CPU:    24 passed
NVIDIA: 48 passed
Ascend: 24 passed```

</details>

## Benchmark / Performance Impact

N/A. This PR is a correctness-first implementation. The Ascend path composes ACLNN operators and does not claim fused-kernel performance; an AscendC fused kernel is a possible follow-up.
<!--
Required for `perf` PRs; optional otherwise. Describe the benchmark harness,
shapes, dtypes, hardware, and include baseline vs. new numbers. If the PR is
not performance-sensitive, write "N/A".
-->

## Notes for Reviewers
- `initial_state` is intentionally passed as a mutable `Tensor` because the destination row is updated.
- The row selected by `initial_state_indices` is never modified; the source row is copied into a workspace state first and the destination row is written only after all tokens have been processed.
- Requests are independent. A destination row must not also be the source row of another request in the same call.
- The current pytest matrix covers `int32` state indices. The operator interface implements `int64` as well, but that path is not currently parameterized in the test matrix.
- Ascend uses `aclnnMul`, `aclnnMatmul`, and `aclnnAdd` with the existing `WorkspacePool`. FP32 matmul uses `cube_math_type=0` to avoid FP32 downcast on 910B.
- The CUDA implementation assumes `head_dim` fits within one block.
<!--
Anything reviewers should focus on: subtle invariants, known trade-offs,
follow-up work intentionally left out of scope, etc.
-->

…ttention

Recurrent lightning attention with an indexed state pool: S <- exp(-slope) * S + k^T v ; out = q @ S. Includes the operator interface, a CPU reference implementation, an NVIDIA CUDA kernel, and pytest coverage for float32/float16/bfloat16 on both backends (validated on sm_120).
Compose aclnn Mul, Matmul and Add over a per-request workspace row. The source state row is staged before the recurrence and written back only after all tokens, preserving the read-only source-row contract for MiniMax lightning attention on Ascend.
Use KEEP_DTYPE for float32 aclnnMatmul so the 910B output projection does not silently downcast FP32 and exceed the operator tolerance.
@y258dd
y258dd requested a review from a team September 19, 2026 05:32
@y258dd

y258dd commented Sep 19, 2026

Copy link
Copy Markdown
Author

Closing this PR because the source branch does not follow the required
feat/... branch naming convention in CONTRIBUTING.md.

Reopening from feat/lightning-attention-infinilm.

@y258dd y258dd closed this Sep 19, 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