[uc] add activation, norm, pos emb, quant, sdpa operators for uc backend. - #342
Open
shengw-bd wants to merge 5 commits into
Open
[uc] add activation, norm, pos emb, quant, sdpa operators for uc backend.#342shengw-bd wants to merge 5 commits into
shengw-bd wants to merge 5 commits into
Conversation
Claude Code ReviewVerdict: Request changes -- New UC backend has a few correctness/perf issues (sync calls on hot paths, broken empty-input scale shape, mis-scoped backend import) that should be fixed before merge. SummaryAdds a new "uc" (Unified Compiler) backend for Ascend NPU with operator implementations for activation, attention, gemm, normalization, position embedding, and quantization, plus README and registry updates. Also gates the ttx import in failures and adds an Must fix
SuggestionsSuggestions (6)
NitsNits (3)
Notes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a new Unified Compiler (UC) based backend for Ascend NPU kernels, adding support for various operators including activation, attention, GEMM, normalization, position embedding, and quantization, along with documentation and registry updates. The review feedback highlights several performance and robustness improvements, such as checking for uc_kernel availability at import time, eliminating redundant CPU-NPU synchronizations in quantization and position embedding operators, safer handling of DTensor operations in GEMM, and minor code cleanups like removing redundant code blocks and using idiomatic dictionary membership checks.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+1
to
+5
| from mojo_opset.utils.platform import get_impl_by_platform | ||
|
|
||
| _op_map = get_impl_by_platform() | ||
| globals().update(_op_map) | ||
| __all__ = list(_op_map.keys()) |
Contributor
There was a problem hiding this comment.
The uc backend will be imported and registered even if uc_kernel is not installed, because uc_kernel is only imported lazily at runtime. If the uc backend is selected or fell back to, it will crash at runtime with an ImportError. Checking for uc_kernel availability at import time allows the backend to be gracefully skipped and logged as a warning.
Suggested change
| from mojo_opset.utils.platform import get_impl_by_platform | |
| _op_map = get_impl_by_platform() | |
| globals().update(_op_map) | |
| __all__ = list(_op_map.keys()) | |
| import importlib.util | |
| if importlib.util.find_spec("uc_kernel") is None: | |
| raise ImportError("uc_kernel is not installed, which is required by the uc backend.") | |
| from mojo_opset.utils.platform import get_impl_by_platform | |
| _op_map = get_impl_by_platform() | |
| globals().update(_op_map) | |
| __all__ = list(_op_map.keys()) |
Comment on lines
+81
to
+97
| def _varlen_position_ids( | ||
| x: torch.Tensor, | ||
| cu_q_lens: torch.Tensor, | ||
| total_seq_lens: Optional[torch.Tensor], | ||
| ) -> torch.Tensor: | ||
| position_ids = torch.empty((x.shape[0],), device=x.device, dtype=torch.int32) | ||
| q_lens = cu_q_lens[1:] - cu_q_lens[:-1] | ||
| for i in range(q_lens.numel()): | ||
| q_len = q_lens[i].item() | ||
| context_len = 0 if total_seq_lens is None else total_seq_lens[i].item() - q_len | ||
| position_ids[cu_q_lens[i]:cu_q_lens[i + 1]] = torch.arange( | ||
| context_len, | ||
| context_len + q_len, | ||
| device=cu_q_lens.device, | ||
| dtype=torch.int32, | ||
| ) | ||
| return position_ids |
Contributor
There was a problem hiding this comment.
Calling .item() inside a loop over sequence lengths causes multiple CPU-GPU/NPU synchronizations, which is a major performance bottleneck in the hot path of LLM execution. Slicing with NPU tensors also adds overhead. Copying cu_q_lens and total_seq_lens to CPU once using .tolist() and performing the loop and slicing using CPU integers completely eliminates these synchronizations and slicing overhead.
@staticmethod
def _varlen_position_ids(
x: torch.Tensor,
cu_q_lens: torch.Tensor,
total_seq_lens: Optional[torch.Tensor],
) -> torch.Tensor:
position_ids = torch.empty((x.shape[0],), device=x.device, dtype=torch.int32)
cu_q_lens_cpu = cu_q_lens.cpu().tolist()
total_seq_lens_cpu = total_seq_lens.cpu().tolist() if total_seq_lens is not None else None
for i in range(len(cu_q_lens_cpu) - 1):
start, end = cu_q_lens_cpu[i], cu_q_lens_cpu[i + 1]
q_len = end - start
context_len = 0 if total_seq_lens_cpu is None else total_seq_lens_cpu[i] - q_len
position_ids[start:end] = torch.arange(
context_len,
context_len + q_len,
device=x.device,
dtype=torch.int32,
)
return position_ids
Comment on lines
+41
to
+42
| torch.npu.synchronize() | ||
| return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) |
Contributor
There was a problem hiding this comment.
Calling torch.npu.synchronize() forces CPU-NPU synchronization, which breaks the asynchronous execution pipeline of PyTorch and introduces massive CPU bubbles. Since PyTorch's stream semantics naturally guarantee that subsequent operations on the output tensors will wait for the kernel to finish, this synchronization is completely redundant and should be removed.
Suggested change
| torch.npu.synchronize() | |
| return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) | |
| return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) |
Comment on lines
+98
to
+99
| torch.npu.synchronize() | ||
| return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) |
Contributor
There was a problem hiding this comment.
Calling torch.npu.synchronize() forces CPU-NPU synchronization, which breaks the asynchronous execution pipeline of PyTorch and introduces massive CPU bubbles. Since PyTorch's stream semantics naturally guarantee that subsequent operations on the output tensors will wait for the kernel to finish, this synchronization is completely redundant and should be removed.
Suggested change
| torch.npu.synchronize() | |
| return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) | |
| return kernel_output.reshape(input.shape), kernel_scale.reshape(*input.shape[:-1], 1) |
Comment on lines
+44
to
+48
| if platform == "npu" and get_bool_env("MOJO_DETERMINISTIC", default=False): | ||
| import os | ||
|
|
||
| # special setting for npu deterministic matmul | ||
| os.environ["CLOSE_MATMUL_K_SHIFT"] = "1" |
Contributor
There was a problem hiding this comment.
This block of code is completely redundant. Setting CLOSE_MATMUL_K_SHIFT when MOJO_DETERMINISTIC is enabled is already handled at lines 11-17. Removing this redundant block keeps the code clean and maintainable.
Comment on lines
+35
to
+38
| if typed_api in kernels.keys(): | ||
| return typed_api | ||
| if dtype == torch.float16 and api in kernels.keys(): | ||
| return api |
Contributor
There was a problem hiding this comment.
In Python, checking membership directly in a dictionary (e.g., key in dict) is more idiomatic and significantly faster than calling .keys() (e.g., key in dict.keys()), which creates an unnecessary view object.
Suggested change
| if typed_api in kernels.keys(): | |
| return typed_api | |
| if dtype == torch.float16 and api in kernels.keys(): | |
| return api | |
| if typed_api in kernels: | |
| return typed_api | |
| if dtype == torch.float16 and api in kernels: | |
| return api |
Comment on lines
+43
to
+54
| if self.trans_weight: | ||
| weight = self.weight.t().contiguous() | ||
| else: | ||
| weight = self.weight | ||
| weight = _to_local_tensor(weight) | ||
| input = _to_local_tensor(input) | ||
| input_scale = _to_local_tensor(input_scale).flatten().float().contiguous() | ||
| weight_scale = _to_local_tensor(self.weight_scale).flatten().float().contiguous() | ||
| if not input.is_contiguous(): | ||
| input = input.contiguous() | ||
| if not weight.is_contiguous(): | ||
| weight = weight.contiguous() |
Contributor
There was a problem hiding this comment.
Calling .t().contiguous() on a DTensor before converting it to a local tensor can be extremely inefficient or fail depending on the PyTorch version. It is much safer and more efficient to convert the DTensor to a local tensor first using _to_local_tensor, and then perform transpose and contiguous operations locally. This also avoids redundant .contiguous() calls.
Suggested change
| if self.trans_weight: | |
| weight = self.weight.t().contiguous() | |
| else: | |
| weight = self.weight | |
| weight = _to_local_tensor(weight) | |
| input = _to_local_tensor(input) | |
| input_scale = _to_local_tensor(input_scale).flatten().float().contiguous() | |
| weight_scale = _to_local_tensor(self.weight_scale).flatten().float().contiguous() | |
| if not input.is_contiguous(): | |
| input = input.contiguous() | |
| if not weight.is_contiguous(): | |
| weight = weight.contiguous() | |
| input = _to_local_tensor(input) | |
| weight = _to_local_tensor(self.weight) | |
| if self.trans_weight: | |
| weight = weight.t() | |
| input_scale = _to_local_tensor(input_scale).flatten().float().contiguous() | |
| weight_scale = _to_local_tensor(self.weight_scale).flatten().float().contiguous() | |
| if not input.is_contiguous(): | |
| input = input.contiguous() | |
| if not weight.is_contiguous(): | |
| weight = weight.contiguous() |
Claude Code ReviewVerdict: Request changes -- New UC backend has several correctness/perf concerns that should be addressed before merging. SummaryThis PR introduces a new "uc" (Unified Compiler) backend for Ascend NPU, wiring it into the platform backend registry and providing initial UC implementations for activation, attention (SDPA), GEMM, normalization, RoPE, and dynamic quant ops. It also updates the README support matrix and gracefully handles import failures of optional backends. Must fix
SuggestionsSuggestions (6)
NitsNits (3)
Notes
|
Claude Code ReviewVerdict: Request changes -- New UC backend has several correctness concerns in MoE quant, ApplyRoPE shape handling, and registry priority placement. SummaryAdds a new "uc" (Unified Compiler) backend for Ascend NPU, registering it in the platform priority list and providing kernels for activation, attention (SDPA), gemm, normalization, position embedding, and quantization. Most ops route through a thin Must fix
SuggestionsSuggestions (6)
NitsNits (3)
Notes
|
Claude Code ReviewVerdict: Request changes -- New UC NPU backend has several correctness/UX issues that should be fixed before merge. SummaryAdds a new "uc" (Unified Compiler) backend for Ascend NPU, exposing kernels for activation, attention (SDPA), gemm, normalization, RoPE, and quantization, plus README/registry updates. Backend is registered with the lowest priority on npu and import failures are handled gracefully. Must fix
SuggestionsSuggestions (6)
NitsNits (4)
Notes
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.