已合并
dsv32 use triton rope #694
hw-zhoutianyang创建于 3月27日
dsv32 use triton rope #694
已合并
共 4 个文件变更+246-13
| @@ -42,6 +42,10 @@ from mindie_llm.runtime.utils.weight_prefetcher import weight_prefetcher | |||
| 42 | from mindie_llm.runtime.model_runner.forward_context_exp import ForwardContext, get_forward_context, AttentionMetadata | 42 | from mindie_llm.runtime.model_runner.forward_context_exp import ForwardContext, get_forward_context, AttentionMetadata |
| 43 | from mindie_llm.runtime.model_runner.input_buffer import input_buffer | 43 | from mindie_llm.runtime.model_runner.input_buffer import input_buffer |
| 44 | from .abstract import AttentionBackend, AttentionLayer, SelectAttentionImpl | 44 | from .abstract import AttentionBackend, AttentionLayer, SelectAttentionImpl |
| 45 | +from mindie_llm.runtime.ops.triton.triton_utils import is_triton_available | ||
| 46 | + | ||
| 47 | +if is_triton_available(): | ||
| 48 | + from mindie_llm.runtime.ops.triton.rope import rope_forward_triton_siso | ||
| 45 | 49 | ||
| 46 | 50 | ||
| 47 | torch.npu.config.allow_internal_format = True | 51 | torch.npu.config.allow_internal_format = True |
| @@ -633,22 +637,34 @@ class SfaBackendImpl(SelectAttentionImpl): | |||
| 633 | ): | 637 | ): |
| 634 | q = self.indexer.wq_b(q_c) | 638 | q = self.indexer.wq_b(q_c) |
| 635 | q = q.view(-1, self.indexer.n_heads, self.indexer.head_dim) | 639 | q = q.view(-1, self.indexer.n_heads, self.indexer.head_dim) |
| 636 | - q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.indexer.head_dim - self.qk_rope_head_dim], dim=-1) | 640 | + if is_triton_available() and not forward_context.is_prefill: |
| 637 | - | 641 | + cos = cos.view(-1, self.qk_rope_head_dim) |
| 638 | - q_pe = q_pe.unsqueeze(2) | 642 | + sin = sin.view(-1, self.qk_rope_head_dim) |
| 639 | - q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin) | 643 | + q = rope_forward_triton_siso(q, cos, sin, rope_dim=self.qk_rope_head_dim, is_neox_style=True) |
| 640 | - q_pe = q_pe.squeeze(2) | 644 | + else: |
| 641 | - q = torch.cat([q_pe, q_nope], dim=-1) | 645 | + q_pe, q_nope = torch.split( |
| 646 | + q, [self.qk_rope_head_dim, self.indexer.head_dim - self.qk_rope_head_dim], dim=-1 | ||
| 647 | + ) | ||
| 648 | + q_pe = q_pe.unsqueeze(2) | ||
| 649 | + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin) | ||
| 650 | + q_pe = q_pe.squeeze(2) | ||
| 651 | + q = torch.cat([q_pe, q_nope], dim=-1) | ||
| 642 | 652 | ||
| 643 | k_proj = self.indexer.wk(hidden_state) | 653 | k_proj = self.indexer.wk(hidden_state) |
| 644 | k = self.indexer.k_norm(k_proj).unsqueeze(1) | 654 | k = self.indexer.k_norm(k_proj).unsqueeze(1) |
| 645 | - k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.indexer.head_dim - self.qk_rope_head_dim], dim=-1) | 655 | + if is_triton_available() and not forward_context.is_prefill: |
| 646 | - k_pe = k_pe.unsqueeze(2) | 656 | + k = rope_forward_triton_siso(k, cos, sin, rope_dim=self.qk_rope_head_dim, is_neox_style=True) |
| 647 | - k_pe = torch_npu.npu_rotary_mul( | 657 | + else: |
| 648 | - k_pe, cos.view(-1, 1, 1, self.qk_rope_head_dim), sin.view(-1, 1, 1, self.qk_rope_head_dim) | 658 | + k_pe, k_nope = torch.split( |
| 649 | - ) | 659 | + k, [self.qk_rope_head_dim, self.indexer.head_dim - self.qk_rope_head_dim], dim=-1 |
| 650 | - k_pe = k_pe.squeeze(2) | 660 | + ) |
| 651 | - k = torch.cat([k_pe, k_nope], dim=-1) | 661 | + k_pe = k_pe.unsqueeze(2) |
| 662 | + k_pe = torch_npu.npu_rotary_mul( | ||
| 663 | + k_pe, cos.view(-1, 1, 1, self.qk_rope_head_dim), sin.view(-1, 1, 1, self.qk_rope_head_dim) | ||
| 664 | + ) | ||
| 665 | + k_pe = k_pe.squeeze(2) | ||
| 666 | + k = torch.cat([k_pe, k_nope], dim=-1) | ||
| 667 | + | ||
| 652 | # cp | 668 | # cp |
| 653 | cp_input_dict = attn_metadata.cp_input_dict | 669 | cp_input_dict = attn_metadata.cp_input_dict |
| 654 | if forward_context.is_prefill and self.cp_size > 1: | 670 | if forward_context.is_prefill and self.cp_size > 1: |
| @@ -45,6 +45,7 @@ from mindie_llm.runtime.config.mindie_llm_config import SpeculativeConfig | |||
| 45 | from mindie_llm.runtime.layers.sampling.sampler import Sampler | 45 | from mindie_llm.runtime.layers.sampling.sampler import Sampler |
| 46 | from mindie_llm.runtime.model_runner.spec_worker import auto_speculative_method_router, speculative_worker_selector | 46 | from mindie_llm.runtime.model_runner.spec_worker import auto_speculative_method_router, speculative_worker_selector |
| 47 | from mindie_llm.runtime.utils.npu.device_utils import get_npu_node_info | 47 | from mindie_llm.runtime.utils.npu.device_utils import get_npu_node_info |
| 48 | +from mindie_llm.runtime.ops.triton.triton_utils import init_device_properties_triton | ||
| 48 | 49 | ||
| 49 | # Allow tensor initialization and casting with internal format(e.g., NZ) | 50 | # Allow tensor initialization and casting with internal format(e.g., NZ) |
| 50 | torch.npu.config.allow_internal_format = True | 51 | torch.npu.config.allow_internal_format = True |
| @@ -215,6 +216,7 @@ class ModelRunnerExp: | |||
| 215 | if ENV_utils.async_inference: | 216 | if ENV_utils.async_inference: |
| 216 | self.sampler = Sampler(sampler_config) | 217 | self.sampler = Sampler(sampler_config) |
| 217 | set_mc2_token_capacity(self._max_batch_size, self.num_speculative_tokens + 1) | 218 | set_mc2_token_capacity(self._max_batch_size, self.num_speculative_tokens + 1) |
| 219 | + init_device_properties_triton() | ||
| 218 | 220 | ||
| 219 | def load_weights(self) -> None: | 221 | def load_weights(self) -> None: |
| 220 | """Load model weights and initialize rotary embeddings.""" | 222 | """Load model weights and initialize rotary embeddings.""" |
| @@ -15,6 +15,7 @@ | |||
| 15 | import torch | 15 | import torch |
| 16 | import triton | 16 | import triton |
| 17 | import triton.language as tl | 17 | import triton.language as tl |
| 18 | +from mindie_llm.runtime.ops.triton.triton_utils import get_vectorcore_num | ||
| 18 | 19 | ||
| 19 | 20 | ||
| 20 | 21 | ||
| @@ -206,3 +207,159 @@ def triton_apply_rope_partial_in_place(x, sin, cos): | |||
| 206 | head_num, | 207 | head_num, |
| 207 | ) | 208 | ) |
| 208 | return x.view(org_shape) | 209 | return x.view(org_shape) |
| 210 | + | ||
| 211 | + | ||
| 212 | + | ||
| 213 | +def _triton_rope_siso( | ||
| 214 | + qk_ptr, | ||
| 215 | + qk_row_stride, | ||
| 216 | + cos_ptr, | ||
| 217 | + cos_row_stride, | ||
| 218 | + sin_ptr, | ||
| 219 | + sin_row_stride, | ||
| 220 | + cos_sin_ptr, | ||
| 221 | + cos_sin_row_stride, | ||
| 222 | + pos_ptr, | ||
| 223 | + num_tokens, | ||
| 224 | + n_h: tl.constexpr, | ||
| 225 | + hd: tl.constexpr, | ||
| 226 | + rope_dim: tl.constexpr, | ||
| 227 | + pad_n_h: tl.constexpr, | ||
| 228 | + pad_rope_dim: tl.constexpr, | ||
| 229 | + BLOCK_SIZE: tl.constexpr, | ||
| 230 | + IS_NEOX_STYLE: tl.constexpr, | ||
| 231 | + USE_COS_SIN: tl.constexpr, | ||
| 232 | +): | ||
| 233 | + pid = tl.program_id(0).to(tl.int64) | ||
| 234 | + row_block_size = tl.num_programs(0) | ||
| 235 | + | ||
| 236 | + for row_idx in tl.range(pid, num_tokens, row_block_size): | ||
| 237 | + qk_start_ptr = qk_ptr + row_idx * qk_row_stride | ||
| 238 | + | ||
| 239 | + # #################################################################### | ||
| 240 | + # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position | ||
| 241 | + # m of this program instance | ||
| 242 | + # #################################################################### | ||
| 243 | + cos_offsets = tl.arange(0, pad_rope_dim // 2) | ||
| 244 | + sin_offsets = tl.arange(pad_rope_dim // 2, pad_rope_dim) | ||
| 245 | + cos_mask = cos_offsets < (rope_dim // 2) | ||
| 246 | + if USE_COS_SIN: | ||
| 247 | + pos_idx = tl.load(pos_ptr + row_idx).to(tl.int64) | ||
| 248 | + cos_start_ptr = cos_sin_ptr + pos_idx * cos_sin_row_stride | ||
| 249 | + cos_row = tl.load(cos_start_ptr + cos_offsets, mask=cos_mask, other=0).to(tl.float32) | ||
| 250 | + sin_row = tl.load(cos_start_ptr + sin_offsets, mask=cos_mask, other=0).to(tl.float32) | ||
| 251 | + else: | ||
| 252 | + cos_start_ptr = cos_ptr + row_idx * cos_row_stride | ||
| 253 | + sin_start_ptr = sin_ptr + row_idx * sin_row_stride | ||
| 254 | + cos_row = tl.load(cos_start_ptr + cos_offsets, mask=cos_mask, other=0).to(tl.float32) | ||
| 255 | + sin_row = tl.load(sin_start_ptr + cos_offsets, mask=cos_mask, other=0).to(tl.float32) | ||
| 256 | + | ||
| 257 | + # #################################################################### | ||
| 258 | + # Load the left and right half of q and k for the current | ||
| 259 | + # program instance (i.e. for the current token) separately | ||
| 260 | + # #################################################################### | ||
| 261 | + # left half of the head | ||
| 262 | + if IS_NEOX_STYLE: | ||
| 263 | + first_half_offsets = tl.arange(0, pad_n_h)[:, None] * hd + tl.arange(0, pad_rope_dim // 2)[None, :] | ||
| 264 | + else: | ||
| 265 | + first_half_offsets = tl.arange(0, pad_n_h)[:, None] * hd + (2 * tl.arange(0, pad_rope_dim // 2)[None, :]) | ||
| 266 | + | ||
| 267 | + first_mask = (tl.arange(0, pad_n_h)[:, None] < n_h) & ( | ||
| 268 | + tl.arange(0, pad_rope_dim // 2)[None, :] < (rope_dim // 2) | ||
| 269 | + ) | ||
| 270 | + qk_tile_1 = tl.load(qk_start_ptr + first_half_offsets, mask=first_mask, other=0).to(sin_row.dtype) | ||
| 271 | + | ||
| 272 | + # right half of the head | ||
| 273 | + if IS_NEOX_STYLE: | ||
| 274 | + second_half_offsets = first_half_offsets + (rope_dim // 2) | ||
| 275 | + else: | ||
| 276 | + second_half_offsets = first_half_offsets + 1 | ||
| 277 | + second_mask = first_mask | ||
| 278 | + qk_tile_2 = tl.load(qk_start_ptr + second_half_offsets, mask=second_mask, other=0).to(sin_row.dtype) | ||
| 279 | + | ||
| 280 | + # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin] | ||
| 281 | + new_qk_tile_1 = qk_tile_1 * cos_row - qk_tile_2 * sin_row | ||
| 282 | + tl.store(qk_start_ptr + first_half_offsets, new_qk_tile_1, mask=first_mask) | ||
| 283 | + | ||
| 284 | + new_qk_tile_2 = qk_tile_2 * cos_row + qk_tile_1 * sin_row | ||
| 285 | + tl.store(qk_start_ptr + second_half_offsets, new_qk_tile_2, mask=second_mask) | ||
| 286 | + | ||
| 287 | + | ||
| 288 | +def rope_forward_triton_siso( | ||
| 289 | + qk: torch.Tensor, | ||
| 290 | + cos: torch.Tensor = None, | ||
| 291 | + sin: torch.Tensor = None, | ||
| 292 | + cos_sin_cache: torch.Tensor = None, | ||
| 293 | + positions: torch.Tensor = None, | ||
| 294 | + rope_dim: int = -1, | ||
| 295 | + is_neox_style: bool = True, | ||
| 296 | +) -> torch.Tensor: | ||
| 297 | + if not qk.is_contiguous(): | ||
| 298 | + qk = qk.contiguous() | ||
| 299 | + | ||
| 300 | + num_tokens, n_head, head_dim = qk.shape | ||
| 301 | + assert rope_dim <= head_dim | ||
| 302 | + pad_rope_dim = triton.next_power_of_2(rope_dim) | ||
| 303 | + pad_n_head = triton.next_power_of_2(n_head) | ||
| 304 | + BLOCK_SIZE = pad_n_head | ||
| 305 | + num_vectorcore = get_vectorcore_num() | ||
| 306 | + n_row = min(num_tokens, num_vectorcore) | ||
| 307 | + | ||
| 308 | + if cos_sin_cache is not None and positions is not None: | ||
| 309 | + assert positions.shape[0] == num_tokens | ||
| 310 | + _triton_rope_siso[(n_row,)]( | ||
| 311 | + qk, | ||
| 312 | + qk.stride(0), | ||
| 313 | + None, | ||
| 314 | + None, | ||
| 315 | + None, | ||
| 316 | + None, | ||
| 317 | + cos_sin_cache, | ||
| 318 | + cos_sin_cache.stride(0), | ||
| 319 | + positions, | ||
| 320 | + num_tokens, | ||
| 321 | + n_head, | ||
| 322 | + head_dim, | ||
| 323 | + rope_dim, | ||
| 324 | + pad_n_head, | ||
| 325 | + pad_rope_dim, | ||
| 326 | + BLOCK_SIZE=BLOCK_SIZE, | ||
| 327 | + IS_NEOX_STYLE=is_neox_style, | ||
| 328 | + USE_COS_SIN=True, | ||
| 329 | + ) | ||
| 330 | + elif cos is not None and sin is not None: | ||
| 331 | + assert cos.shape[0] == num_tokens and sin.shape[0] == num_tokens | ||
| 332 | + cos = cos.view(num_tokens, -1) | ||
| 333 | + sin = sin.view(num_tokens, -1) | ||
| 334 | + if rope_dim == -1: | ||
| 335 | + # If rope_dim is not specified, we assume that input cos/sin is not | ||
| 336 | + # duplicated to rope_dim, which means rope_dim == cos.shape[-1] * 2 | ||
| 337 | + rope_dim = cos.shape[-1] * 2 | ||
| 338 | + _triton_rope_siso[(n_row,)]( | ||
| 339 | + qk, | ||
| 340 | + qk.stride(0), | ||
| 341 | + cos, | ||
| 342 | + cos.stride(0), | ||
| 343 | + sin, | ||
| 344 | + sin.stride(0), | ||
| 345 | + None, | ||
| 346 | + None, | ||
| 347 | + None, | ||
| 348 | + num_tokens, | ||
| 349 | + n_head, | ||
| 350 | + head_dim, | ||
| 351 | + rope_dim, | ||
| 352 | + pad_n_head, | ||
| 353 | + pad_rope_dim, | ||
| 354 | + BLOCK_SIZE=BLOCK_SIZE, | ||
| 355 | + IS_NEOX_STYLE=is_neox_style, | ||
| 356 | + USE_COS_SIN=False, | ||
| 357 | + ) | ||
| 358 | + else: | ||
| 359 | + raise ValueError( | ||
| 360 | + "Currently, rope_forward_triton supports passing:\n" | ||
| 361 | + "1. positions and original cos_sin_cache.\n" | ||
| 362 | + "2. cos and sin which are already selected by positions\n" | ||
| 363 | + "Please check whether you call rope_forward_triton correctly." | ||
| 364 | + ) | ||
| 365 | + return qk | ||
| @@ -0,0 +1,58 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved. | ||
| 2 | +# MindIE is licensed under Mulan PSL v2. | ||
| 3 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 4 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 5 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 7 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 8 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 9 | +# See the Mulan PSL v2 for more details. | ||
| 10 | + | ||
| 11 | +from typing import Any | ||
| 12 | + | ||
| 13 | +import torch | ||
| 14 | +from mindie_llm.utils.log.logging import logger | ||
| 15 | + | ||
| 16 | +_HAS_TRITON = None | ||
| 17 | +_NUM_AICORE = -1 | ||
| 18 | +_NUM_VECTORCORE = -1 | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +def is_triton_available() -> bool: | ||
| 22 | + global _HAS_TRITON | ||
| 23 | + if _HAS_TRITON is None: | ||
| 24 | + try: | ||
| 25 | + import triton # noqa | ||
| 26 | + import triton.language as tl # noqa | ||
| 27 | + | ||
| 28 | + _HAS_TRITON = True | ||
| 29 | + logger.info("Detected Triton installation. Triton-based operators are available.") | ||
| 30 | + except ImportError: | ||
| 31 | + _HAS_TRITON = False | ||
| 32 | + logger.info("Triton not found. Falling back to non-Triton implementations.") | ||
| 33 | + return _HAS_TRITON | ||
| 34 | + | ||
| 35 | + | ||
| 36 | +def init_device_properties_triton(): | ||
| 37 | + global _NUM_AICORE, _NUM_VECTORCORE | ||
| 38 | + if _NUM_AICORE == -1 and is_triton_available(): | ||
| 39 | + import triton | ||
| 40 | + | ||
| 41 | + device_properties: dict[str, Any] = triton.runtime.driver.active.utils.get_device_properties( | ||
| 42 | + torch.npu.current_device() | ||
| 43 | + ) | ||
| 44 | + _NUM_AICORE = device_properties.get("num_aicore", -1) | ||
| 45 | + _NUM_VECTORCORE = device_properties.get("num_vectorcore", -1) | ||
| 46 | + assert _NUM_AICORE > 0 and _NUM_VECTORCORE > 0, "Failed to detect device properties." | ||
| 47 | + | ||
| 48 | + | ||
| 49 | +def get_aicore_num(): | ||
| 50 | + global _NUM_AICORE | ||
| 51 | + assert _NUM_AICORE > 0, "Device properties not initialized. Please call init_device_properties_triton() first." | ||
| 52 | + return _NUM_AICORE | ||
| 53 | + | ||
| 54 | + | ||
| 55 | +def get_vectorcore_num(): | ||
| 56 | + global _NUM_VECTORCORE | ||
| 57 | + assert _NUM_VECTORCORE > 0, "Device properties not initialized. Please call init_device_properties_triton() first." | ||
| 58 | + return _NUM_VECTORCORE | ||