Skip to content

[Fix] Cast Ascend kernel indices to int64 before stride multiply - #1216

Open
sunyi0505 wants to merge 1 commit into
fla-org:mainfrom
sunyi0505:fix_overflow
Open

[Fix] Cast Ascend kernel indices to int64 before stride multiply#1216
sunyi0505 wants to merge 1 commit into
fla-org:mainfrom
sunyi0505:fix_overflow

Conversation

@sunyi0505

Copy link
Copy Markdown
Contributor

Summary

Triton-Ascend kernels were overflowing int32 when building GM addresses on long sequences.
Program IDs and grid-derived indices (i_b, i_t, i_n, NT, bos) are int32. Products like i_b * T, i_tg * HV * K * V, or bos * HV wrap past 2³¹ before a trailing .to(tl.int64). That pattern also fails to compile when the index is specialized (constexpr has no .to()).
Typical wrap points:

  • state layout HV * K * V with K=V=128, HV=64, BT=64 → overflow at NT > 2048 (T > 131K)
  • packed offset * D with D=4096 → overflow at T > 524K
  • varlen cu_seqlens loaded as int32, then (bos * HV + i_h) * V (HV=32, V=4096 → safe bos ≈ 16K)
    This PR casts the index to int64 first, then multiplies:
  • tl.cast(i_b, tl.int64) * T instead of (i_b * T).to(tl.int64)
  • load cu_seqlens as tl.int64; keep T_cur = (eos - bos).to(tl.int32)
  • keep make_block_ptr offsets/block_shape as int32 (i_t * BT)
    Touched 20 triton_ascend files: modules (fused CE, layernorm, l2norm, rotary, fused_norm_gate) and ops (chunk_h / chunk_o / chunk_delta_h, GDN wy_fast/gate, GLA chunk, KDA, cumsum, solve_tril).

Test plan

hardware: Atlas 800T A3(X86)
FLA_NPU_XDIST=1 ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 pytest --exitfirst -n 8 --dist load tests/modules tests/ops/utils tests/ops/test_gdn_kernels.py tests/ops/test_gdn.py tests/ops/test_kda.py tests/ops/test_attnres.py tests/ops/test_solve_tril.py tests/ops/test_gla.py
image

Benchmark / NCU (kernel changes only)

Not involved

Breaking changes

None

Checklist

  • I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
  • I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
  • Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
  • Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
  • This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.

If you ticked the "minor" box above

Standalone minor PRs are normally not accepted (see No busywork PRs).
Justify here why yours is worth a maintainer's review time — minor PRs without a justification may be closed without review:

@zhiyuan1i zhiyuan1i added bug Something isn't working ascend-npu Ascend NPU (triton_ascend) related labels Sep 3, 2026
@sunyi0505
sunyi0505 marked this pull request as ready for review September 3, 2026 07:20
@sunyi0505

Copy link
Copy Markdown
Contributor Author

@zhiyuan1i This pr is ready for review. Could you please review this PR? Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ascend-npu Ascend NPU (triton_ascend) related bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants