Skip to content

[Refactor] NPU FLA: delegate to fla native operators, remove MindSpeed/twinkle-own kernels#251

Merged
kevssim merged 1 commit into
modelscope:mainfrom
ys2025-AI:main
Jul 22, 2026
Merged

[Refactor] NPU FLA: delegate to fla native operators, remove MindSpeed/twinkle-own kernels#251
kevssim merged 1 commit into
modelscope:mainfrom
ys2025-AI:main

Conversation

@ys2025-AI

@ys2025-AI ys2025-AI commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

PR type

  • Bug Fix
  • New Feature
  • Document Updates
  • More Models or Datasets Support

PR information

  • Remove twinkle-own Triton kernels (causal_conv1d.py, chunk_gated_delta_rule.py) that depended on MindSpeed; delegate to fla's native causal_conv1d and chunk_gated_delta_rule which ship their own Ascend-backend dispatch
  • Move npu_causal_conv1d_fn adapter into npu_impls/fla.py; drop redundant weight transpose (fla expects [D,W], not [W,D])
  • Simplify builtin.py FLA dependency check: only requires flash-linear-attention (no longer mindspeed>=0.12.1)
  • Update gdn_padding_free.py import path for npu_causal_conv1d_fn
  • Clean up unused Optional import and fix activation type hint
  • Update tests to stub fla package as unavailable

Experiment results

Test environment

Component Version
Hardware Ascend NPU A3
CANN 9.0.0
Python 3.11.15
torch 2.9.0+cpu
torch_npu 2.9.0.post2
triton 3.2.0
triton-ascend 3.2.1
transformers 5.9.0
flash-linear-attention 0.5.2

Unit tests

tests/kernel/ + tests/patch/: 66 passed, 5 skipped
test_qwen35_fla_bwd_precision.py (bf16 fla vs fp32 torch): 3 passed

fsdp2 e2e (Qwen3.5-4B, fsdp=8, dp=1, max_length=8192)

step FLA ON FLA OFF diff (ON−OFF)
1 1.5515 1.5520 −0.0005
2 1.3512 1.3517 −0.0005
3 1.6020 1.6018 +0.0002
4 1.6367 1.6358 +0.0009
5 1.6180 1.6192 −0.0012

max |diff| = 0.0012, mean |diff| = 0.0007

sp2 e2e (Qwen3.5-4B, fsdp=8, dp=1, ulysses=2, max_length=8192)

step FLA ON FLA OFF diff (ON−OFF)
0 1.5521 1.5538 −0.0017
1 1.3533 1.3536 −0.0003
2 1.6036 1.6035 +0.0001
3 1.6380 1.6368 +0.0012
4 1.6190 1.6195 −0.0005

max |diff| = 0.0017, mean |diff| = 0.0008

- Remove twinkle-own Triton kernels (causal_conv1d.py, chunk_gated_delta_rule.py)
  that depended on MindSpeed; delegate to fla's native causal_conv1d and
  chunk_gated_delta_rule which ship their own Ascend-backend dispatch
- Move npu_causal_conv1d_fn adapter into npu_impls/fla.py; drop redundant
  weight transpose (fla expects [D,W], not [W,D])
- Simplify builtin.py FLA dependency check: only requires flash-linear-attention
  (no longer mindspeed>=0.12.1)
- Update gdn_padding_free.py import path for npu_causal_conv1d_fn
- Clean up unused Optional import and fix activation type hint
- Update tests to stub fla package as unavailable
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@kevssim
kevssim merged commit 1a1d884 into modelscope:main Jul 22, 2026
2 of 3 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.

2 participants