[NPU] add NPU (Ascend) backend for chunk_gla #737
Conversation
noemotiovon
commented
Feb 5, 2026
noemotiovon
commented
- Auto-detect device and select NPU/GPU implementation
- Fix NaN from uninitialized memory in A/dA matrices
- Replace problematic triton kernel with PyTorch (NPU compiler limitation)
- Add exp overflow protection and varlen sequence support
- Auto-detect device and select NPU/GPU implementation - Fix NaN from uninitialized memory in A/dA matrices - Replace problematic triton kernel with PyTorch (NPU compiler limitation) - Add exp overflow protection and varlen sequence support
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the
✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @noemotiovon, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces NPU (Ascend) backend support for the chunk_gla operation. The changes include device-aware dynamic imports, NPU-specific Triton kernels, and several fixes for numerical stability and correctness on NPU hardware. Notably, a problematic Triton kernel in the backward pass is replaced with a pure PyTorch implementation to work around NPU compiler limitations, and explicit masking is added to prevent NaNs from uninitialized memory.
My review focuses on the new NPU implementation. I've suggested a performance improvement by vectorizing a loop that handles variable-length sequences. I've also pointed out a couple of type hint mismatches that should be corrected for code clarity and correctness. Overall, this is a well-structured and thoughtful implementation for adding new hardware support.
| output_final_state: bool, | ||
| cu_seqlens: torch.LongTensor | None = None, | ||
| chunk_size: int = 64, | ||
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
There was a problem hiding this comment.
The function chunk_gla_fwd_npu returns five tensors (g_cumsum, A, h, ht, o), but its return type hint indicates a tuple of three tensors. This should be corrected to match the actual return values for type consistency and clarity.
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| local_pos = torch.zeros(T_dim, dtype=torch.long, device=A.device) | ||
| for seq_idx in range(num_seqs): | ||
| seq_start = cu_seqlens_cpu[seq_idx].item() | ||
| seq_end = cu_seqlens_cpu[seq_idx + 1].item() | ||
| seq_len = seq_end - seq_start | ||
| if seq_len > 0: | ||
| # Local positions within this sequence, mod BT for chunk-local position | ||
| seq_positions = torch.arange(seq_len, device=A.device) % BT | ||
| local_pos[seq_start:seq_end] = seq_positions |
There was a problem hiding this comment.
The for-loop to compute local_pos for variable-length sequences can be a performance bottleneck when dealing with a large number of sequences. This can be vectorized using torch.repeat_interleave for better efficiency.
# Vectorized implementation for performance
lens = cu_seqlens_cpu[1:] - cu_seqlens_cpu[:-1]
# Create a tensor where each element is the start of its sequence
seq_starts = torch.repeat_interleave(cu_seqlens_cpu[:-1], lens).to(A.device)
# Create a tensor of global positions
global_pos = torch.arange(T_dim, device=A.device)
# Compute local positions
local_pos = (global_pos - seq_starts) % BT| local_pos = torch.zeros(T_dim, dtype=torch.long, device=dA.device) | ||
| for seq_idx in range(num_seqs): | ||
| seq_start = cu_seqlens_cpu[seq_idx].item() | ||
| seq_end = cu_seqlens_cpu[seq_idx + 1].item() | ||
| seq_len = seq_end - seq_start | ||
| if seq_len > 0: | ||
| seq_positions = torch.arange(seq_len, device=dA.device) % BT | ||
| local_pos[seq_start:seq_end] = seq_positions |
There was a problem hiding this comment.
The for-loop to compute local_pos for variable-length sequences can be a performance bottleneck when dealing with a large number of sequences. This can be vectorized using torch.repeat_interleave for better efficiency. This is similar to a suggestion in chunk_gla_fwd_npu.
# Vectorized implementation for performance
lens = cu_seqlens_cpu[1:] - cu_seqlens_cpu[:-1]
# Create a tensor where each element is the start of its sequence
seq_starts = torch.repeat_interleave(cu_seqlens_cpu[:-1], lens).to(dA.device)
# Create a tensor of global positions
global_pos = torch.arange(T_dim, device=dA.device)
# Compute local positions
local_pos = (global_pos - seq_starts) % BT| k: torch.Tensor, | ||
| v: torch.Tensor, | ||
| g: torch.Tensor, | ||
| scale: int | None = None, |
There was a problem hiding this comment.
The type hint for the scale parameter is int | None, but it is used as a float and the docstring also describes it as Optional[float]. The type hint should be changed to float | None to match its actual usage and documentation.
scale: float | None = None,
|
Thank you for your contribution. We'll be open-sourcing a different backend dispatch soon. At that time, it will be better adapted to different hardware |
That’s fantastic news! Since NPUs and GPUs differ significantly at the architectural level, Triton kernel implementations naturally diverge and really benefit from a proper dispatch mechanism. Could you share a bit more about the open-source plan on this front? Also, if there’s anything I can help with, I’d be very happy to contribute and get involved. |
We have introduced a new decorator that designates functions as dispatchable. Each backend targeted for dispatch includes a verifier to check whether the current input is compatible with that backend. This architecture allows substitution at arbitrary granularity (e.g., single operations or complete forward/backward passes) and enables mixing and matching different backends. Moving forward, we will also support registering multiple dispatch backends for the same function, enabling dispatch decisions based on varying performance characteristics and constraints. The core functionality is already implemented and will be open-sourced following code cleanup. Our primary concern is the lack of continuous integration for additional backends; currently, we are limited to a single H100. We would be extremely grateful if we could obtain on-demand NPU backends for CI gating purposes. |
|
That sounds really exciting! |
could we connect offline? |
2 participants
