add split_qk_rmsnorm fusion op for forward_mla_prepare_npu#570
add split_qk_rmsnorm fusion op for forward_mla_prepare_npu#570Estrella-xx wants to merge 1 commit into
Conversation
![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 a new Triton-based split QK RMSNorm kernel and its wrapper function. The review feedback highlights several improvement opportunities: removing the unused and duplicated get_device_properties helper along with its unused imports, adding other=0.0 and casting to tl.float32 when loading q_weight and k_weight to prevent potential NaN propagation, and updating the wrapper function to robustly handle both 2D and 3D input tensors by flattening and reshaping them appropriately.
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.
| from functools import cache | ||
| from typing import Any, Dict, Tuple | ||
|
|
||
|
|
||
| @cache | ||
| def get_device_properties() -> Tuple[int, int]: | ||
| device = torch.npu.current_device() | ||
| device_properties: Dict[str, Any] = ( | ||
| triton.runtime.driver.active.utils.get_device_properties(device) | ||
| ) | ||
|
|
||
| num_aicore = device_properties.get("num_aicore", -1) | ||
| num_vectorcore = device_properties.get("num_vectorcore", -1) | ||
|
|
||
| assert num_aicore > 0 and num_vectorcore > 0, "Failed to detect device properties." | ||
| return num_aicore, num_vectorcore | ||
|
|
||
|
|
||
| # before |
There was a problem hiding this comment.
The get_device_properties function is duplicated from sgl_kernel_npu.utils.triton_utils and is completely unused in this file. We should remove it along with the unused imports (cache, Any, Dict, Tuple) and the leftover # before comment to keep the codebase clean.
|
|
||
| q_norm = q * q_rstd | ||
|
|
||
| q_weight = tl.load(q_weight_ptr + offsets, mask=offsets < q_lora_rank) |
There was a problem hiding this comment.
Specify other=0.0 and cast to tl.float32 when loading q_weight. This ensures consistent precision during multiplication and avoids potential NaN propagation from undefined register values at masked-out offsets.
| q_weight = tl.load(q_weight_ptr + offsets, mask=offsets < q_lora_rank) | |
| q_weight = tl.load(q_weight_ptr + offsets, mask=offsets < q_lora_rank, other=0.0).to(tl.float32) |
|
|
||
| k_norm = k * k_rstd | ||
|
|
||
| k_weight = tl.load(k_weight_ptr + offsets, mask=offsets < kv_lora_rank) |
There was a problem hiding this comment.
Specify other=0.0 and cast to tl.float32 when loading k_weight. This ensures consistent precision during multiplication and avoids potential NaN propagation from undefined register values at masked-out offsets.
| k_weight = tl.load(k_weight_ptr + offsets, mask=offsets < kv_lora_rank) | |
| k_weight = tl.load(k_weight_ptr + offsets, mask=offsets < kv_lora_rank, other=0.0).to(tl.float32) |
| batch_size = input.shape[0] | ||
| # total_hidden_size = q_lora_rank + kv_lora_rank | ||
| total_hidden_size = input.shape[-1] | ||
| input = input.contiguous() | ||
|
|
||
| q_output = torch.empty( | ||
| batch_size, q_lora_rank, device=input.device, dtype=input.dtype | ||
| ) | ||
| k_nope_output = torch.empty( | ||
| batch_size, kv_lora_rank, device=input.device, dtype=input.dtype | ||
| ) | ||
| k_pe_output = torch.empty( | ||
| batch_size, qk_rope_head_dim, device=input.device, dtype=input.dtype | ||
| ) | ||
|
|
||
| BLOCK_SIZE = triton.next_power_of_2(max(q_lora_rank, kv_lora_rank, qk_rope_head_dim)) | ||
| split_qk_rmsnorm_kernel[(batch_size,)]( | ||
| input, | ||
| q_output, | ||
| k_nope_output, | ||
| k_pe_output, | ||
| q_weight, | ||
| k_weight, | ||
| q_lora_rank, | ||
| kv_lora_rank, | ||
| qk_rope_head_dim, | ||
| total_hidden_size, | ||
| eps, | ||
| BLOCK_SIZE, | ||
| ) | ||
|
|
||
| return q_output, k_nope_output, k_pe_output |
There was a problem hiding this comment.
The current implementation assumes a 2D input tensor where input.shape[0] represents the batch size (or total number of tokens). If a 3D tensor of shape (batch, seq_len, hidden) is passed, input.shape[0] will only capture the batch dimension, leading to incorrect grid launching and output shapes.
We should flatten the input to 2D before running the kernel, and then reshape the outputs back to 3D if the input was 3D. This makes the function robust to both 2D and 3D inputs.
input_shape = input.shape
if input.ndim == 3:
input = input.view(-1, input_shape[-1])
num_rows = input.shape[0]
total_hidden_size = input.shape[-1]
input = input.contiguous()
q_output = torch.empty(
num_rows, q_lora_rank, device=input.device, dtype=input.dtype
)
k_nope_output = torch.empty(
num_rows, kv_lora_rank, device=input.device, dtype=input.dtype
)
k_pe_output = torch.empty(
num_rows, qk_rope_head_dim, device=input.device, dtype=input.dtype
)
BLOCK_SIZE = triton.next_power_of_2(max(q_lora_rank, kv_lora_rank, qk_rope_head_dim))
split_qk_rmsnorm_kernel[(num_rows,)](
input,
q_output,
k_nope_output,
k_pe_output,
q_weight,
k_weight,
q_lora_rank,
kv_lora_rank,
qk_rope_head_dim,
total_hidden_size,
eps,
BLOCK_SIZE,
)
if len(input_shape) == 3:
q_output = q_output.view(input_shape[0], input_shape[1], q_lora_rank)
k_nope_output = k_nope_output.view(input_shape[0], input_shape[1], kv_lora_rank)
k_pe_output = k_pe_output.view(input_shape[0], input_shape[1], qk_rope_head_dim)
return q_output, k_nope_output, k_pe_output1 participant
No description provided.