perf(npu): vectorize ring attention LSE extractio

[Ruihan11]

PR type

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

PR information

Summary

Vectorize second-half LSE extraction in the Ascend NPU ring-attention backward path.

Motivation

_get_second_half_lse previously iterated over every packed sequence and called .item() for each sequence boundary. When cu_seqlens is an NPU tensor, these per-sequence scalar reads introduce repeated NPU-to-CPU synchronizations.

The overhead grows with the number of packed sequences and becomes significant for padding-free batches containing many short sequences.

Changes

• Replace per-sequence slicing and copying with device-side index construction.
• Use repeat_interleave to build source and destination offsets.
• Use a single index_select to gather the second-half LSE values.
• Preserve the existing function signature, output layout, dtype, device, and autograd behavior.
• Retain a bounded-loop fast path for up to four packed sequences to avoid regressions for single long sequences.
• Keep the CUDA implementation, ring schedule, HCCL communication, and native NPU attention APIs unchanged.

Tests

Added:

• CPU reference-based correctness tests.
• Ascend NPU correctness tests for:
  • float32
  • bfloat16
  • variable-length packed sequences
  • empty segments
  • int32 and int64 cu_seqlens
  • supported LSE layouts
  • backward gradient propagation
• Opt-in NPU microbenchmark.
• Opt-in 2-rank HCCL end-to-end ring-attention test.

Validation:

22 passed, 2 skipped
python -m compileall -q swift: passed
flake8: passed
isort: passed
yapf: passed
git diff --check: passed
2-rank HCCL E2E: passed

Experiment results

Environment:

Ascend 910B3
torch==2.10.0+cpu
torch_npu==2.10.0
bfloat16
8 attention heads
total length: 4096 tokens
warmup: 10
measurement: 30

Helper-level benchmark:

Packed sequences	Tokens/sequence	Baseline	Optimized	Speedup
1	4096	229.3 us	229.6 us	1.00x
8	512	1072.8 us	853.4 us	1.26x
32	128	3930.1 us	864.4 us	4.55x
128	32	15265.7 us	874.0 us	17.47x
512	8	60203.3 us	971.1 us	62.00x
1024	4	119459.4 us	1099.4 us	108.66x

Synthetic 2-rank HCCL ring-attention forward and backward benchmark:

Packed sequences	Baseline	Optimized	Speedup
1	6.771 ms	6.531 ms	1.04x
32	30.466 ms	27.596 ms	1.10x
128	99.209 ms	79.222 ms	1.25x

The helper-level improvement is larger than the end-to-end improvement because the complete step also includes native attention kernels, HCCL communication, and other backward operations.

Reproduction

Run the NPU helper benchmark:

SWIFT_RUN_NPU_BENCHMARK=1 \
python -m pytest -q -s \
    tests/sequence_parallel/test_zigzag_ring_attn_npu.py \
    -k benchmark

Run the distributed NPU E2E test:

SWIFT_RUN_NPU_E2E=1 \
ASCEND_RT_VISIBLE_DEVICES=0,1 \
python -m torch.distributed.run \
    --standalone \
    --nproc_per_node=2 \
    tests/sequence_parallel/test_zigzag_ring_attn_npu_e2e.py

Limitations

• This change targets the Ascend NPU ring-attention backward path only.
• Existing sequence-length divisibility requirements are unchanged.
• For a small number of packed sequences, the bounded-loop path is retained because index construction can cost more than the original loop.
• The CUDA implementation is intentionally unchanged.