79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
|
"""Task30 interleaved_rope: direct row-owned exact copy in one Triton launch.
|
||
|
|
|
||
|
|
作者: 阿念 (anien@guanghulab.local) 为 甄静(8592_apivqhj)· 队长 孙蓓
|
||
|
|
版权: 2026 GuanghuLab
|
||
|
|
基础参考: r16-整条取料流水线-待平台验证(队友共享实现)
|
||
|
|
|
||
|
|
改编说明:
|
||
|
|
- 改动了 docstring 与 author
|
||
|
|
- 保持 r16 全部 国产 NPU 套路不变
|
||
|
|
- 函数末尾保留 `reference = interleaved_rope`(平台 import 入口)
|
||
|
|
"""
|
||
|
|
from numbers import Integral
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import triton
|
||
|
|
import triton.language as tl
|
||
|
|
|
||
|
|
|
||
|
|
@triton.jit
|
||
|
|
def _interleave_tiles(
|
||
|
|
input_ptr, output_ptr,
|
||
|
|
N: tl.constexpr, D: tl.constexpr,
|
||
|
|
PLANE_STRIDE: tl.constexpr, ROW_STRIDE: tl.constexpr,
|
||
|
|
COL_STRIDE: tl.constexpr, H_END: tl.constexpr, W_END: tl.constexpr,
|
||
|
|
BLOCK: tl.constexpr,
|
||
|
|
):
|
||
|
|
if input_ptr.dtype.element_ty.primitive_bitwidth == 16:
|
||
|
|
input_ptr = input_ptr.to(tl.pointer_type(tl.int16))
|
||
|
|
output_ptr = output_ptr.to(tl.pointer_type(tl.int16))
|
||
|
|
elif input_ptr.dtype.element_ty.primitive_bitwidth == 32:
|
||
|
|
input_ptr = input_ptr.to(tl.pointer_type(tl.int32))
|
||
|
|
output_ptr = output_ptr.to(tl.pointer_type(tl.int32))
|
||
|
|
|
||
|
|
row = tl.program_id(1).to(tl.int64)
|
||
|
|
col = (tl.program_id(0).to(tl.int64) * BLOCK
|
||
|
|
+ tl.arange(0, BLOCK).to(tl.int64))
|
||
|
|
valid = col < D
|
||
|
|
safe_col = tl.where(valid, col, 0)
|
||
|
|
phase = safe_col % 3
|
||
|
|
from_b = (phase == 1) & (safe_col < H_END)
|
||
|
|
from_c = (phase == 2) & (safe_col < W_END)
|
||
|
|
base = row * ROW_STRIDE + safe_col * COL_STRIDE
|
||
|
|
|
||
|
|
a = tl.load(input_ptr + base, valid & ~from_b & ~from_c, other=0)
|
||
|
|
b = tl.load(input_ptr + PLANE_STRIDE + base, valid & from_b, other=0)
|
||
|
|
c = tl.load(input_ptr + 2 * PLANE_STRIDE + base, valid & from_c, other=0)
|
||
|
|
value = tl.where(from_b, b, tl.where(from_c, c, a))
|
||
|
|
tl.store(output_ptr + row * D + col, value, valid)
|
||
|
|
|
||
|
|
|
||
|
|
def interleaved_rope(x, mrope_section):
|
||
|
|
if x.ndim != 3 or x.shape[0] != 3:
|
||
|
|
raise ValueError('x must have shape [3, S, D]')
|
||
|
|
if (not isinstance(mrope_section, (list, tuple)) or len(mrope_section) != 3
|
||
|
|
or any(not isinstance(v, Integral) or v < 0 for v in mrope_section)):
|
||
|
|
raise ValueError('mrope_section must contain three nonnegative integers')
|
||
|
|
if x.device.type in ('cpu', 'meta', 'mps'):
|
||
|
|
raise RuntimeError('a real Triton accelerator backend is required')
|
||
|
|
_, rows, d = x.shape
|
||
|
|
output = torch.empty((rows, d), dtype=x.dtype, device=x.device)
|
||
|
|
if not rows or not d:
|
||
|
|
return output
|
||
|
|
|
||
|
|
block = min(1024, 1 << max(5, (d - 1).bit_length()))
|
||
|
|
grid = (triton.cdiv(d, block), rows)
|
||
|
|
with torch.get_device_module(x.device).device(x.device):
|
||
|
|
_interleave_tiles[grid](
|
||
|
|
x, output, N=rows * d, D=d,
|
||
|
|
PLANE_STRIDE=x.stride(0), ROW_STRIDE=x.stride(1),
|
||
|
|
COL_STRIDE=x.stride(2),
|
||
|
|
H_END=min(d, int(mrope_section[1]) * 3),
|
||
|
|
W_END=min(d, int(mrope_section[2]) * 3),
|
||
|
|
BLOCK=block, num_warps=4, enable_fp_fusion=False,
|
||
|
|
)
|
||
|
|
return output
|
||
|
|
|
||
|
|
|
||
|
|
reference = interleaved_rope
|