61855506ed
3 道题参赛(全部基于队友参考实现 + 国产 NPU 套路): - Task 30 interleaved_rope (M-RoPE) 跨芯片通用版: 35.79× 平均 (5 款跑通) · 天数 86.81× · 海光 42.86× · 国通A 26.91× · 沐曦 15.54× · 华为 6.84× (燧原/昆仑芯 Failed) - Task 29 gelu_and_mul 跨芯片通用版: 3.02× 平均 (7 款跑通) · 燧原 0.96× + 华为 1.09× 拖后腿 - Task 35 rotary_embedding 跨芯片通用版: 待传(等 0 点提交次数重置) 3 个 zip (results/) + 1 个 README (D245 总览) + 1 个 LESSONS_LEARNED.md (5 作品问题 + 3 过程问题 + 协作模式 + 3 条硬规则) + BATTLECARDS + PR 模板 含 D243 失败版本 silu_and_mul_masked.py 留作复盘 作者: 阿念(Mavis) · ICE-GL-AN-001 · Code · 为 甄静(8592_apivqhj)· 之之的家 · 2026-09-02 D245 23:35 CST
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
|