Star137
242
代码
代码
Issues72
Pull Requests182
项目讨论
Wiki
分析
Star137
242
已合并
[feature] triton算子for qwen3.5 part5 #460
stanzzzzz创建于 3月2日
[feature] triton算子for qwen3.5 part5 #460
已合并
stanzzzzz创建于 3月2日
从stanzzzzz/MindIE-LLM_newtriton_p5合入到Ascend/MindIE-LLMdev
共 3 个文件变更+760-0
Amindie_llm/runtime/ops/triton/layernorm_gated.py+185-0
@@ -0,0 +1,185 @@
1+# SPDX-FileCopyrightText: Copyright contributors to the sgl-project
2+# SPDX-License-Identifier: Apache-2.0
3+# Part of this file implemented based on sgl-project.
4+#
5+# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
6+# MindIE is licensed under Mulan PSL v2.
7+# You can use this software according to the terms and conditions of the Mulan PSL v2.
8+# You may obtain a copy of Mulan PSL v2 at:
9+# http://license.coscl.org.cn/MulanPSL2
10+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
11+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
12+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
13+# See the Mulan PSL v2 for more details.
14+ 
15+import torch
16+import triton
17+import triton.language as tl
18+ 
19+ 
20+@triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None})
21+@triton.heuristics({"HAS_Z": lambda args: args["Z"] is not None})
22+@triton.jit
23+def _layer_norm_fwd_1pass_kernel_npu(
24+ X, # pointer to the input
25+ Y, # pointer to the output
26+ W, # pointer to the weights
27+ B, # pointer to the biases
28+ Z, # pointer to the other branch
29+ Mean, # pointer to the mean
30+ Rstd, # pointer to the 1/std
31+ stride_x_row, # how much to increase the pointer when moving by 1 row
32+ stride_y_row,
33+ stride_z_row,
34+ M, # number of rows in X
35+ N, # number of columns in X
36+ eps, # epsilon to avoid division by zero
37+ BLOCK_M: tl.constexpr,
38+ BLOCK_N: tl.constexpr,
39+ HAS_BIAS: tl.constexpr,
40+ HAS_Z: tl.constexpr,
41+ NORM_BEFORE_GATE: tl.constexpr,
42+ IS_RMS_NORM: tl.constexpr,
43+):
44+ # Map the program id to the row of X and Y it should compute.
45+ pid_m = tl.program_id(0)
46+ group = tl.program_id(1)
47+ if not IS_RMS_NORM:
48+ Mean += group * M
49+ Rstd += group * M
50+ W += group * N
51+ if HAS_BIAS:
52+ B += group * N
53+ 
54+ # Compute row indices for this program
55+ rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
56+ cols = tl.arange(0, BLOCK_N)
57+ 
58+ # Mask for valid rows and cols
59+ row_mask = rows < M
60+ col_mask = cols < N
61+ 
62+ # Load weight once (broadcasted over rows)
63+ w = tl.load(W + cols, mask=col_mask).to(tl.float32)
64+ if HAS_BIAS:
65+ b = tl.load(B + cols, mask=col_mask).to(tl.float32)
66+ 
67+ # Load X: shape [BLOCK_M, BLOCK_N]
68+ x_ptrs = X + rows[:, None] * stride_x_row + cols[None, :] + group * N
69+ x = tl.load(x_ptrs, mask=row_mask[:, None] & col_mask[None, :]).to(tl.float32)
70+ 
71+ # Load Z if needed
72+ if HAS_Z:
73+ z_ptrs = Z + rows[:, None] * stride_z_row + cols[None, :] + group * N
74+ z = tl.load(z_ptrs, mask=row_mask[:, None] & col_mask[None, :]).to(tl.float32)
75+ if not NORM_BEFORE_GATE:
76+ x *= z * tl.sigmoid(z)
77+ 
78+ # Compute statistics per row
79+ if not IS_RMS_NORM:
80+ mean = tl.sum(x, axis=1) / N # [BLOCK_M]
81+ xbar = tl.where(col_mask[None, :], x - mean[:, None], 0.0)
82+ var = tl.sum(xbar * xbar, axis=1) / N
83+ tl.store(Mean + rows, mean, mask=row_mask)
84+ else:
85+ xbar = tl.where(col_mask[None, :], x, 0.0)
86+ var = tl.sum(xbar * xbar, axis=1) / N
87+ 
88+ rstd = 1.0 / tl.sqrt(var + eps) # [BLOCK_M]
89+ tl.store(Rstd + rows, rstd, mask=row_mask)
90+ 
91+ # Normalize
92+ if not IS_RMS_NORM:
93+ x_hat = (x - mean[:, None]) * rstd[:, None]
94+ else:
95+ x_hat = x * rstd[:, None]
96+ 
97+ y = x_hat * w[None, :]
98+ if HAS_BIAS:
99+ y += b[None, :]
100+ 
101+ # Post-gate
102+ if HAS_Z and NORM_BEFORE_GATE:
103+ y *= z * tl.sigmoid(z)
104+ 
105+ # Store output
106+ y_ptrs = Y + rows[:, None] * stride_y_row + cols[None, :] + group * N
107+ tl.store(y_ptrs, y, mask=row_mask[:, None] & col_mask[None, :])
108+ 
109+ 
110+def layer_norm_fwd_npu(
111+ x,
112+ weight,
113+ bias,
114+ eps,
115+ z=None,
116+ out=None,
117+ group_size=None,
118+ norm_before_gate=True,
119+ is_rms_norm=False,
120+):
121+ M, N = x.shape
122+ if group_size is None:
123+ group_size = N
124+ if N % group_size != 0:
125+ raise ValueError(f"N ({N}) must be divisible by group_size ({group_size})")
126+ ngroups = N // group_size
127+ 
128+ if x.stride(-1) != 1:
129+ raise ValueError("x must be contiguous in the last dimension (stride(-1) == 1)")
130+ if z is not None:
131+ if z.stride(-1) != 1:
132+ raise ValueError("z must be contiguous in the last dimension (stride(-1) == 1)")
133+ if z.shape != (M, N):
134+ raise ValueError(f"z.shape must be (M, N) = ({M}, {N}), got {z.shape}")
135+ if weight.shape != (N,):
136+ raise ValueError(f"weight.shape must be (N,) = ({N},), got {weight.shape}")
137+ if weight.stride(-1) != 1:
138+ raise ValueError("weight must be contiguous in the last dimension (stride(-1) == 1)")
139+ if bias is not None:
140+ if bias.stride(-1) != 1:
141+ raise ValueError("bias must be contiguous in the last dimension (stride(-1) == 1)")
142+ if bias.shape != (N,):
143+ raise ValueError(f"bias.shape must be (N,) = ({N},), got {bias.shape}")
144+ # allocate output
145+ if out is not None:
146+ if out.shape != x.shape:
147+ raise ValueError(f"out.shape must match x.shape {x.shape}, got {out.shape}")
148+ else:
149+ out = torch.empty_like(x)
150+ if out.stride(-1) != 1:
151+ raise ValueError("out must be contiguous in the last dimension (stride(-1) == 1)")
152+ mean = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None
153+ rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)
154+ 
155+ MAX_FUSED_SIZE = 65536 // x.element_size()
156+ BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size))
157+ if group_size > BLOCK_N:
158+ raise RuntimeError("Feature dim too large.")
159+ 
160+ # Choose BLOCK_M: e.g., 16, 32, 64 — depends on NPU vector core capacity
161+ BLOCK_M = 64 # Tune this based on your NPU's register/shared memory
162+ 
163+ # Now grid is (num blocks over M, num groups)
164+ grid = (triton.cdiv(M, BLOCK_M), ngroups)
165+ _layer_norm_fwd_1pass_kernel_npu[grid](
166+ x,
167+ out,
168+ weight,
169+ bias,
170+ z,
171+ mean,
172+ rstd,
173+ x.stride(0),
174+ out.stride(0),
175+ z.stride(0) if z is not None else 0,
176+ M,
177+ group_size,
178+ eps,
179+ BLOCK_M=BLOCK_M,
180+ BLOCK_N=BLOCK_N,
181+ NORM_BEFORE_GATE=norm_before_gate,
182+ IS_RMS_NORM=is_rms_norm,
183+ # Remove multibuffer if not needed
184+ )
185+ return out, mean, rstd
Amindie_llm/runtime/ops/triton/rope.py+212-0
@@ -0,0 +1,212 @@
1+# SPDX-FileCopyrightText: Copyright contributors to the vllm-project
2+# SPDX-License-Identifier: Apache-2.0
3+# Part of this file implemented based on vllm-project.
4+#
5+# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved.
6+# MindIE is licensed under Mulan PSL v2.
7+# You can use this software according to the terms and conditions of the Mulan PSL v2.
8+# You may obtain a copy of Mulan PSL v2 at:
9+# http://license.coscl.org.cn/MulanPSL2
10+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
11+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
12+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
13+# See the Mulan PSL v2 for more details.
14+ 
15+import torch
16+import triton
17+import triton.language as tl
18+ 
19+ 
20+@triton.jit
21+def _triton_rope(
22+ q_ptr,
23+ q_row_stride,
24+ k_ptr,
25+ k_row_stride,
26+ cos,
27+ cos_row_stride,
28+ sin,
29+ sin_row_stride,
30+ num_tokens,
31+ n_qh,
32+ n_kh,
33+ hd,
34+ rope_dim,
35+ pad_rope_dim: tl.constexpr,
36+ BLOCK_HEADS: tl.constexpr, # 新增：每次处理的 Head 数量
37+ IS_NEOX_STYLE: tl.constexpr,
38+):
39+ pid = tl.program_id(0).to(tl.int64)
40+ row_block_size = tl.num_programs(0)
41+ 
42+ half_rope_dim = rope_dim // 2
43+ rope_offs = tl.arange(0, pad_rope_dim // 2)
44+ rope_mask = rope_offs < half_rope_dim
45+ 
46+ for row_idx in tl.range(pid, num_tokens, row_block_size):
47+ q_start_ptr = q_ptr + row_idx * q_row_stride
48+ k_start_ptr = k_ptr + row_idx * k_row_stride
49+ cos_start_ptr = cos + row_idx * cos_row_stride
50+ sin_start_ptr = sin + row_idx * sin_row_stride
51+ cos_row = tl.load(cos_start_ptr + rope_offs, mask=rope_mask, other=0).to(tl.float32)
52+ sin_row = tl.load(sin_start_ptr + rope_offs, mask=rope_mask, other=0).to(tl.float32)
53+ for h_start in range(0, n_qh, BLOCK_HEADS):
54+ 
55+ h_offs = h_start + tl.arange(0, BLOCK_HEADS)
56+ h_mask = h_offs < n_qh
57+ head_base_offsets = h_offs[:, None] * hd
58+ if IS_NEOX_STYLE:
59+ 
60+ offs_1 = head_base_offsets + rope_offs[None, :]
61+ offs_2 = offs_1 + half_rope_dim
62+ else:
63+ 
64+ offs_1 = head_base_offsets + (2 * rope_offs[None, :])
65+ offs_2 = offs_1 + 1
66+ 
67+ 
68+ mask = h_mask[:, None] & rope_mask[None, :]
69+ 
70+ q1 = tl.load(q_start_ptr + offs_1, mask=mask, other=0).to(tl.float32)
71+ q2 = tl.load(q_start_ptr + offs_2, mask=mask, other=0).to(tl.float32)
72+ 
73+ out_q1 = q1 * cos_row - q2 * sin_row
74+ out_q2 = q2 * cos_row + q1 * sin_row
75+ 
76+ tl.store(q_start_ptr + offs_1, out_q1, mask=mask)
77+ tl.store(q_start_ptr + offs_2, out_q2, mask=mask)
78+ 
79+ for h_start in range(0, n_kh, BLOCK_HEADS):
80+ h_offs = h_start + tl.arange(0, BLOCK_HEADS)
81+ h_mask = h_offs < n_kh
82+ 
83+ head_base_offsets = h_offs[:, None] * hd
84+ 
85+ if IS_NEOX_STYLE:
86+ offs_1 = head_base_offsets + rope_offs[None, :]
87+ offs_2 = offs_1 + half_rope_dim
88+ else:
89+ offs_1 = head_base_offsets + (2 * rope_offs[None, :])
90+ offs_2 = offs_1 + 1
91+ 
92+ mask = h_mask[:, None] & rope_mask[None, :]
93+ 
94+ k1 = tl.load(k_start_ptr + offs_1, mask=mask, other=0).to(tl.float32)
95+ k2 = tl.load(k_start_ptr + offs_2, mask=mask, other=0).to(tl.float32)
96+ 
97+ out_k1 = k1 * cos_row - k2 * sin_row
98+ out_k2 = k2 * cos_row + k1 * sin_row
99+ 
100+ tl.store(k_start_ptr + offs_1, out_k1, mask=mask)
101+ tl.store(k_start_ptr + offs_2, out_k2, mask=mask)
102+ 
103+ 
104+def rope_forward_triton(
105+ q: torch.Tensor,
106+ k: torch.Tensor,
107+ cos: torch.Tensor,
108+ sin: torch.Tensor,
109+ rope_dim: int = -1,
110+ is_neox_style: bool = True,
111+) -> tuple[torch.Tensor, torch.Tensor]:
112+ if not q.is_contiguous():
113+ q = q.contiguous()
114+ if not k.is_contiguous():
115+ k = k.contiguous()
116+ 
117+ num_tokens, n_q_head, head_dim = q.shape
118+ n_kv_head = k.shape[1]
119+
120+ # Flatten cos/sin to [num_tokens, dim]
121+ cos = cos.view(num_tokens, -1)
122+ sin = sin.view(num_tokens, -1)
123+
124+ if rope_dim == -1:
125+ rope_dim = cos.shape[-1] * 2
126+ 
127+ pad_rope_dim = triton.next_power_of_2(rope_dim)
128+
129+ BLOCK_HEADS = 16
130+
131+ grid = (num_tokens,)
132+ 
133+ _triton_rope[grid](
134+ q,
135+ q.stride(0),
136+ k,
137+ k.stride(0),
138+ cos,
139+ cos.stride(0),
140+ sin,
141+ sin.stride(0),
142+ num_tokens,
143+ n_q_head,
144+ n_kv_head,
145+ head_dim,
146+ rope_dim,
147+ pad_rope_dim=pad_rope_dim,
148+ BLOCK_HEADS=BLOCK_HEADS,
149+ IS_NEOX_STYLE=is_neox_style,
150+ num_stages=1,
151+ )
152+ return q, k
153+ 
154+ 
155+@triton.jit
156+def triton_rope_kernel_in_place(
157+ x_ptr,
158+ sin_ptr,
159+ cos_ptr,
160+ x_stride,
161+ cos_stride,
162+ hidden_size: tl.constexpr,
163+ rope_dim: tl.constexpr,
164+ head_num: tl.constexpr,
165+):
166+ cur_b = tl.program_id(0)
167+ dim_start = hidden_size - rope_dim
168+ # load x
169+ offset_x = cur_b * x_stride + dim_start + tl.arange(0, rope_dim)
170+ x = tl.load(x_ptr + offset_x).to(tl.float32)
171+ # load sin cos
172+ offset_sin_cos = cur_b // head_num * cos_stride + tl.arange(0, rope_dim)
173+ sin = tl.load(sin_ptr + offset_sin_cos).to(tl.float32)
174+ cos = tl.load(cos_ptr + offset_sin_cos).to(tl.float32)
175+ 
176+ even = tl.extract_slice(x, [0], [rope_dim // 2], [2])
177+ odd = tl.extract_slice(x, [1], [rope_dim // 2], [2])
178+ odd = -odd
179+ 
180+ x_rotate = tl.zeros([rope_dim], dtype=tl.float32)
181+ x_rotate = tl.insert_slice(x_rotate, odd, [0], [rope_dim // 2], [2])
182+ x_rotate = tl.insert_slice(x_rotate, even, [1], [rope_dim // 2], [2])
183+ 
184+ out = x * cos + x_rotate * sin
185+ tl.store(x_ptr + offset_x, out.to(tl.bfloat16))
186+ 
187+ 
188+def triton_apply_rope_partial_in_place(x, sin, cos):
189+ rope_dim = sin.shape[-1]
190+ org_shape = x.shape
191+ if x.dim() == 2:
192+ bsz, hidden_size = x.shape
193+ head_num = 1
194+ elif x.dim() == 3:
195+ bsz, head_num, hidden_size = x.shape
196+ x = x.view(-1, hidden_size)
197+ else:
198+ raise NotImplementedError(f"x_shape={x.shape} not supported")
199+ cores = bsz * head_num
200+ if cores >= 65535:
201+ raise ValueError(f"cores ({cores}) must be less than 65535 (uint16 limit for triton grid)")
202+ triton_rope_kernel_in_place[(cores,)](
203+ x,
204+ sin,
205+ cos,
206+ x.stride(0),
207+ sin.stride(0),
208+ hidden_size,
209+ rope_dim,
210+ head_num,
211+ )
212+ return x.view(org_shape)
Amindie_llm/runtime/ops/triton/split_qkv_rmsnorm_mrope.py+363-0
@@ -0,0 +1,363 @@
1+#
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
3+#
4+# Licensed under the Apache License, Version 2.0 (the "License");
5+# you may not use this file except in compliance with the License.
6+# You may obtain a copy of the License at
7+#
8+# http://www.apache.org/licenses/LICENSE-2.0
9+#
10+# Unless required by applicable law or agreed to in writing, software
11+# distributed under the License is distributed on an "AS IS" BASIS,
12+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+# See the License for the specific language governing permissions and
14+# limitations under the License.
15+#
16+ 
17+ 
18+import torch
19+import triton # type: ignore
20+import triton.language as tl # type: ignore
21+ 
22+from .triton_utils import get_vectorcore_num
23+ 
24+ 
25+@triton.jit(
26+ do_not_specialize=["num_tokens", "front_core_num", "num_tokens_each_front_core", "num_tokens_each_tail_core"]
27+)
28+def split_qkv_rmsnorm_mrope_kernel(
29+ in_qkv_ptr: torch.Tensor,
30+ q_weight_ptr: torch.Tensor,
31+ q_bias_ptr: torch.Tensor,
32+ k_weight_ptr: torch.Tensor,
33+ k_bias_ptr: torch.Tensor,
34+ cos_sin_ptr: torch.Tensor,
35+ out_q_ptr: torch.Tensor,
36+ out_k_ptr: torch.Tensor,
37+ out_v_ptr: torch.Tensor,
38+ num_tokens,
39+ front_core_num,
40+ num_tokens_each_front_core,
41+ num_tokens_each_tail_core,
42+ num_q_heads: tl.constexpr,
43+ num_kv_heads: tl.constexpr,
44+ head_size: tl.constexpr,
45+ q_size: tl.constexpr,
46+ kv_size: tl.constexpr,
47+ eps: tl.constexpr,
48+ mrope_section_t: tl.constexpr,
49+ mrope_section_h: tl.constexpr,
50+ mrope_section_w: tl.constexpr,
51+ half_head_size: tl.constexpr,
52+ has_bias: tl.constexpr,
53+ is_interleaved: tl.constexpr,
54+ rope_dim: tl.constexpr,
55+ half_rope_dim: tl.constexpr,
56+ IS_PARTIAL_ROPE: tl.constexpr,
57+):
58+ block_idx = tl.program_id(0)
59+ 
60+ loop_num = num_tokens_each_front_core
61+ if block_idx >= front_core_num:
62+ loop_num = num_tokens_each_tail_core
63+ 
64+ block_offset = num_tokens_each_front_core * block_idx
65+ if block_idx >= front_core_num:
66+ block_offset = (
67+ num_tokens_each_front_core * front_core_num + (block_idx - front_core_num) * num_tokens_each_tail_core
68+ )
69+ 
70+ q_rmsnorm_weight = tl.load(q_weight_ptr + tl.arange(0, head_size))
71+ k_rmsnorm_weight = tl.load(k_weight_ptr + tl.arange(0, head_size))
72+ 
73+ if has_bias:
74+ q_bias = tl.load(q_bias_ptr + tl.arange(0, head_size))
75+ k_bias = tl.load(k_bias_ptr + tl.arange(0, head_size))
76+ 
77+ for index in range(loop_num):
78+ ## load ##
79+ # q
80+ in_q_offset = in_qkv_ptr + (block_offset + index) * (q_size + 2 * kv_size)
81+ in_q_tensor = tl.load(in_q_offset + tl.arange(0, q_size)).to(tl.float32).reshape(num_q_heads, head_size)
82+ 
83+ # k
84+ in_k_offset = in_qkv_ptr + (block_offset + index) * (q_size + 2 * kv_size) + q_size
85+ in_k_tensor = tl.load(in_k_offset + tl.arange(0, kv_size)).to(tl.float32).reshape(num_kv_heads, head_size)
86+ # v
87+ in_v_offset = in_qkv_ptr + (block_offset + index) * (q_size + 2 * kv_size) + q_size + kv_size
88+ in_v_tensor = tl.load(in_v_offset + tl.arange(0, kv_size))
89+ 
90+ # cos, sin
91+ cos_offsets = tl.arange(0, half_rope_dim)
92+ if is_interleaved:
93+ h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h)
94+ w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w)
95+ t_mask = ~(h_mask | w_mask)
96+ else:
97+ t_mask = cos_offsets < mrope_section_t
98+ h_mask = (mrope_section_t - 1 < cos_offsets) & (cos_offsets < mrope_section_t + mrope_section_h)
99+ w_mask = (mrope_section_t + mrope_section_h - 1 < cos_offsets) & (
100+ cos_offsets < mrope_section_t + mrope_section_h + mrope_section_w
101+ )
102+ 
103+ t_cos_offset = cos_sin_ptr + (block_offset + index) * rope_dim
104+ h_cos_offset = t_cos_offset + num_tokens * rope_dim
105+ w_cos_offset = h_cos_offset + num_tokens * rope_dim
106+ 
107+ t_sin_offset = cos_sin_ptr + (block_offset + index) * rope_dim + half_rope_dim
108+ h_sin_offset = t_sin_offset + num_tokens * rope_dim
109+ w_sin_offset = h_sin_offset + num_tokens * rope_dim
110+ 
111+ t_cos_tensor = tl.load(t_cos_offset + cos_offsets, mask=t_mask, other=0)
112+ h_cos_tensor = tl.load(h_cos_offset + cos_offsets, mask=h_mask, other=0)
113+ w_cos_tensor = tl.load(w_cos_offset + cos_offsets, mask=w_mask, other=0)
114+ t_sin_tensor = tl.load(t_sin_offset + cos_offsets, mask=t_mask, other=0)
115+ h_sin_tensor = tl.load(h_sin_offset + cos_offsets, mask=h_mask, other=0)
116+ w_sin_tensor = tl.load(w_sin_offset + cos_offsets, mask=w_mask, other=0)
117+ 
118+ cos_tensor = (t_cos_tensor + h_cos_tensor + w_cos_tensor).to(tl.float32).reshape(1, half_rope_dim)
119+ cos_tensor = tl.broadcast_to(cos_tensor, (2, half_rope_dim)).reshape(1, rope_dim)
120+ 
121+ sin_tensor = (t_sin_tensor + h_sin_tensor + w_sin_tensor).to(tl.float32).reshape(1, half_rope_dim)
122+ sin_tensor = tl.broadcast_to(sin_tensor, (2, half_rope_dim)).reshape(1, rope_dim)
123+ 
124+ ## compute ##
125+ # q-rmsnorm
126+ squares = in_q_tensor * in_q_tensor
127+ variances = tl.sum(squares, axis=1) / head_size
128+ reciprocal_std = (1 / tl.sqrt(variances + eps)).reshape(num_q_heads, 1)
129+ q_normalized = in_q_tensor * reciprocal_std
130+ q_normalized = q_normalized * q_rmsnorm_weight
131+ if has_bias:
132+ q_normalized = q_normalized + q_bias
133+ 
134+ # k-rmsnorm
135+ squares = in_k_tensor * in_k_tensor
136+ variances = tl.sum(squares, axis=1) / head_size
137+ reciprocal_std = (1 / tl.sqrt(variances + eps)).reshape(num_kv_heads, 1)
138+ k_normalized = in_k_tensor * reciprocal_std
139+ k_normalized = k_normalized * k_rmsnorm_weight
140+ if has_bias:
141+ k_normalized = k_normalized + k_bias
142+ 
143+ # q-mrope
144+ x1 = tl.extract_slice(
145+ q_normalized,
146+ offsets=(0, 0),
147+ sizes=(num_q_heads, half_rope_dim),
148+ strides=(1, 1),
149+ )
150+ x2 = tl.extract_slice(
151+ q_normalized,
152+ offsets=(0, half_rope_dim),
153+ sizes=(num_q_heads, half_rope_dim),
154+ strides=(1, 1),
155+ )
156+ cat_x = tl.zeros((num_q_heads, rope_dim), dtype=tl.float32)
157+ cat_x = tl.insert_slice(
158+ cat_x,
159+ -x2,
160+ offsets=(0, 0),
161+ sizes=(num_q_heads, half_rope_dim),
162+ strides=(1, 1),
163+ )
164+ cat_x = tl.insert_slice(
165+ cat_x,
166+ x1,
167+ offsets=(0, half_rope_dim),
168+ sizes=(num_q_heads, half_rope_dim),
169+ strides=(1, 1),
170+ )
171+ if IS_PARTIAL_ROPE:
172+ orig_qk = tl.extract_slice(q_normalized, offsets=(0, 0),
173+ sizes=(num_q_heads, rope_dim),
174+ strides=(1, 1))
175+ else:
176+ orig_qk = q_normalized
177+ roped_q = cat_x * sin_tensor + orig_qk * cos_tensor
178+ 
179+ # k-mrope
180+ y1 = tl.extract_slice(
181+ k_normalized,
182+ offsets=(0, 0),
183+ sizes=(num_kv_heads, half_rope_dim),
184+ strides=(1, 1),
185+ )
186+ y2 = tl.extract_slice(
187+ k_normalized,
188+ offsets=(0, half_rope_dim),
189+ sizes=(num_kv_heads, half_rope_dim),
190+ strides=(1, 1),
191+ )
192+ cat_y = tl.zeros((num_kv_heads, rope_dim), dtype=tl.float32)
193+ cat_y = tl.insert_slice(
194+ cat_y,
195+ -y2,
196+ offsets=(0, 0),
197+ sizes=(num_kv_heads, half_rope_dim),
198+ strides=(1, 1),
199+ )
200+ cat_y = tl.insert_slice(
201+ cat_y,
202+ y1,
203+ offsets=(0, half_rope_dim),
204+ sizes=(num_kv_heads, half_rope_dim),
205+ strides=(1, 1),
206+ )
207+ if IS_PARTIAL_ROPE:
208+ orig_qk = tl.extract_slice(k_normalized, offsets=(0, 0), sizes=(num_kv_heads, rope_dim),
209+ strides=(1, 1))
210+ else:
211+ orig_qk = k_normalized
212+ roped_k = cat_y * sin_tensor + orig_qk * cos_tensor
213+ if IS_PARTIAL_ROPE:
214+ q_normalized = tl.insert_slice(q_normalized, roped_q,
215+ offsets=(0, 0), sizes=(num_q_heads, rope_dim),
216+ strides=(1, 1)).to(tl.bfloat16)
217+ k_normalized = tl.insert_slice(k_normalized, roped_k,
218+ offsets=(0, 0), sizes=(num_kv_heads, rope_dim),
219+ strides=(1, 1)).to(tl.bfloat16)
220+ else:
221+ q_normalized = roped_q.to(tl.bfloat16)
222+ k_normalized = roped_k.to(tl.bfloat16)
223+ 
224+ ## store ##
225+ # out_q
226+ out_q_offset = out_q_ptr + (block_offset + index) * q_size
227+ out_q_indices = tl.arange(0, q_size)
228+ tl.store(out_q_offset + out_q_indices, q_normalized.reshape(q_size))
229+ 
230+ # out_k
231+ out_k_offset = out_k_ptr + (block_offset + index) * kv_size
232+ out_k_indices = tl.arange(0, kv_size)
233+ tl.store(out_k_offset + out_k_indices, k_normalized.reshape(kv_size))
234+ 
235+ # out_v
236+ out_v_offset = out_v_ptr + (block_offset + index) * kv_size
237+ tl.store(out_v_offset + tl.arange(0, kv_size), in_v_tensor)
238+ 
239+ 
240+def triton_split_qkv_rmsnorm_mrope(
241+ qkv: torch.Tensor,
242+ q_weight: torch.Tensor,
243+ k_weight: torch.Tensor,
244+ cos_sin: torch.Tensor,
245+ num_q_heads: int,
246+ num_kv_heads: int,
247+ head_size: int,
248+ eps: float,
249+ mrope_section: list[int],
250+ is_interleaved: bool,
251+ rope_dim: int | None = None,
252+ q_bias: torch.Tensor | None = None,
253+ k_bias: torch.Tensor | None = None,
254+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
255+ core_num = get_vectorcore_num()
256+ 
257+ q_size = num_q_heads * head_size
258+ kv_size = num_kv_heads * head_size
259+ num_tokens = qkv.shape[0]
260+ 
261+ if rope_dim is None:
262+ rope_dim = head_size
263+ IS_PARTIAL_ROPE = rope_dim != head_size
264+ 
265+ front_core_num = core_num
266+ if num_tokens % core_num != 0:
267+ front_core_num = num_tokens % core_num
268+ 
269+ num_tokens_each_front_core = (num_tokens + core_num - 1) // core_num
270+ 
271+ tail_core_num = 0
272+ if num_tokens > core_num:
273+ tail_core_num = core_num - front_core_num
274+ 
275+ num_tokens_each_tail_core = num_tokens // core_num
276+ 
277+ q_output = torch.empty(num_tokens, q_size, device=qkv.device, dtype=qkv.dtype)
278+ k_output = torch.empty(num_tokens, kv_size, device=qkv.device, dtype=qkv.dtype)
279+ v_output = torch.empty(num_tokens, kv_size, device=qkv.device, dtype=qkv.dtype)
280+ 
281+ total_core = front_core_num + tail_core_num
282+ block_dim = core_num
283+ if total_core < core_num:
284+ block_dim = total_core
285+ 
286+ half_head_size = head_size // 2
287+ 
288+ has_bias = q_bias is not None
289+ 
290+ split_qkv_rmsnorm_mrope_kernel[(block_dim,)](
291+ qkv,
292+ q_weight,
293+ q_bias,
294+ k_weight,
295+ k_bias,
296+ cos_sin,
297+ q_output,
298+ k_output,
299+ v_output,
300+ num_tokens,
301+ front_core_num,
302+ num_tokens_each_front_core,
303+ num_tokens_each_tail_core,
304+ num_q_heads,
305+ num_kv_heads,
306+ head_size,
307+ q_size,
308+ kv_size,
309+ eps,
310+ mrope_section[0],
311+ mrope_section[1],
312+ mrope_section[2],
313+ half_head_size,
314+ has_bias,
315+ is_interleaved,
316+ rope_dim,
317+ rope_dim // 2,
318+ IS_PARTIAL_ROPE
319+ )
320+ 
321+ return q_output, k_output, v_output
322+ 
323+ 
324+def triton_split_qkv_rmsnorm_mrope_fake(
325+ qkv: torch.Tensor,
326+ q_weight: torch.Tensor,
327+ k_weight: torch.Tensor,
328+ cos_sin: torch.Tensor,
329+ num_q_heads: int,
330+ num_kv_heads: int,
331+ head_size: int,
332+ eps: float,
333+ mrope_section: list[int],
334+ q_bias: torch.Tensor | None = None,
335+ k_bias: torch.Tensor | None = None,
336+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
337+ num_tokens = qkv.shape[0]
338+ q_size = num_q_heads * head_size
339+ kv_size = num_kv_heads * head_size
340+ 
341+ q_output = torch.empty(
342+ num_tokens,
343+ q_size,
344+ device=qkv.device,
345+ dtype=qkv.dtype,
346+ )
347+ 
348+ k_output = torch.empty(
349+ num_tokens,
350+ kv_size,
351+ device=qkv.device,
352+ dtype=qkv.dtype,
353+ )
354+ 
355+ v_output = torch.empty(
356+ num_tokens,
357+ kv_size,
358+ device=qkv.device,
359+ dtype=qkv.dtype,
360+ )
361+ 
362+ return q_output, k_output, v_output
363+