Skip to content

bugfix: disable multibuffer for fused recurrent state writes. - #4

Open
jarvis666666 wants to merge 1 commit into
xLLM-AI:mainfrom
jarvis666666:main
Open

bugfix: disable multibuffer for fused recurrent state writes.#4
jarvis666666 wants to merge 1 commit into
xLLM-AI:mainfrom
jarvis666666:main

Conversation

@jarvis666666

@jarvis666666 jarvis666666 commented Jul 29, 2026

Copy link
Copy Markdown

1.PR 总结

本PR是对 fused_recurrent_gated_delta_rule_spec_fwd_kernel 这一个 launch显式设置 multibuffer=False,修复新版 Triton Ascend auto multi-buffer 下,带accepted-token offset 的循环内动态 state slot 写回错误。修改不改变数学公式、shape、精度阈值或 CANN API,已在 CANN 9.0 和 9.1 环境完成测试;

2.问题与根因

CANN 9.0 升级到 CANN 9.1 的配套软件过程中,Triton Ascend 升级改变了 auto multi-buffer 的 lowering/调度行为,暴露了 fused recurrent spec kernel 对动态逐 token 状态写回的兼容性问题。当前修复属于 torch-npu-ops 对 Triton Ascend 编译器优化行为的约束,不是 CANN 9.1 API 变更适配。

CANN 9.1 配套环境构建 torch-npu-ops Triton npubin 时,主输出精度断言通过,但是per-token state 断言失败:

AssertionError: Spec per-token state mismatch with accepted offsets
1 failed, 209 passed, 2 skipped

3. 代码变更

文件:
triton_npu/triton_src/test_fused_recurrent_gated_delta_rule.py

fused_recurrent_gated_delta_rule_spec_fwd_kernel[grid](
    # existing arguments
    IS_KDA=False,
    multibuffer=False,
)

虽然路径名以 test_ 开头,该文件同时承载 kernel 定义、launch wrapper、精度测试和npubin 生成入口,因此这不是“只让测试跳过失败”的修改;launch option 会改变生成的 fused_recurrent_gated_delta_rule_spec_fwd_kernel.npubin

4. 测试验证结果

已完成的测试:

  • git diff d324c3c^ d324c3c --check:通过;
  • CANN 9.0 验证fused recurrent Spec Decode 单算子精度:通过;
  • CANN 9.0 torch-npu-ops Triton UT:通过 210 passed, 2skipped;
  • CANN 9.1 验证fused recurrent Spec Decode 单算子精度:通过。

  Disable multibuffer scheduling for the fused recurrent kernel because
  newer Triton Ascend versions produce incorrect per-token state updates
  when accepted-token offsets are used.

  The conservative scheduling mode is compatible with both CANN 9.0 and
  CANN 9.1.
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