212 lines
7.5 KiB
Python
212 lines
7.5 KiB
Python
|
"""Task 38 · sigmoid_gate_topk_renorm
|
||
|
|
|
||
|
|
MoE 路由门算子 (DeepSeek-V2/V3 / Qwen3-MoE 风格) - 3-kernel Triton 实现:
|
||
|
|
|
||
|
|
K1 (_sigmoid_bias_kernel) : sel = sigmoid(routed) + bias
|
||
|
|
K2 (_topk_kernel) : indices = topk(sel, k) (自写, 迭代 tl.argmax)
|
||
|
|
K3 (_gate_finalize_kernel): gather 原始 routed → cat shared
|
||
|
|
→ sigmoid → renorm → scale → split 输出
|
||
|
|
|
||
|
|
国产 NPU 套路全套 (来自 D245 验证):
|
||
|
|
- input 指针 cast 到 int16/int32 (防 NaN 转换)
|
||
|
|
- stride / pid 用 int64 cast (防 overflow)
|
||
|
|
- with torch.get_device_module(x.device).device(x.device): 切设备
|
||
|
|
- enable_fp_fusion=False, num_warps=4
|
||
|
|
- 不用 tl.sigmoid, 用 1.0 / (1.0 + tl.exp(-x))
|
||
|
|
"""
|
||
|
|
import torch
|
||
|
|
import triton
|
||
|
|
import triton.language as tl
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# Kernel 1: sel = sigmoid(routed) + bias
|
||
|
|
# ============================================================
|
||
|
|
@triton.jit
|
||
|
|
def _sigmoid_bias_kernel(
|
||
|
|
routed_ptr, # [T, N] fp16/bf16
|
||
|
|
bias_ptr, # [N] fp32
|
||
|
|
sel_ptr, # [T, N] fp32 (中间 buffer)
|
||
|
|
N,
|
||
|
|
stride_t, # int (T 维度 stride)
|
||
|
|
BLOCK_N: tl.constexpr,
|
||
|
|
):
|
||
|
|
# input 指针 cast 防 NaN
|
||
|
|
if routed_ptr.dtype.element_ty.primitive_bitwidth == 16:
|
||
|
|
routed_ptr = routed_ptr.to(tl.pointer_type(tl.int16))
|
||
|
|
|
||
|
|
pid_t = tl.program_id(0).to(tl.int64)
|
||
|
|
pid_n = tl.program_id(1).to(tl.int64)
|
||
|
|
|
||
|
|
offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)
|
||
|
|
mask = offs < N
|
||
|
|
|
||
|
|
x = tl.load(routed_ptr + pid_t * stride_t + offs, mask=mask, other=0).to(tl.float32)
|
||
|
|
b = tl.load(bias_ptr + offs, mask=mask, other=0.0)
|
||
|
|
|
||
|
|
# 不用 tl.sigmoid (国产 NPU 不支持), 用 1/(1+exp(-x))
|
||
|
|
sig = 1.0 / (1.0 + tl.exp(-x))
|
||
|
|
sel = sig + b
|
||
|
|
|
||
|
|
tl.store(sel_ptr + pid_t * N + offs, sel, mask=mask)
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# Kernel 2: top-k 选 indices (自写, 迭代 argmax)
|
||
|
|
# ============================================================
|
||
|
|
@triton.jit
|
||
|
|
def _topk_kernel(
|
||
|
|
sel_ptr, # [T, N] fp32 (来自 K1)
|
||
|
|
idx_ptr, # [T, k] int32
|
||
|
|
N,
|
||
|
|
K: tl.constexpr,
|
||
|
|
BLOCK_N: tl.constexpr,
|
||
|
|
):
|
||
|
|
pid = tl.program_id(0).to(tl.int64)
|
||
|
|
|
||
|
|
offs = tl.arange(0, BLOCK_N)
|
||
|
|
mask = offs < N
|
||
|
|
|
||
|
|
sel = tl.load(sel_ptr + pid * N + offs, mask=mask, other=-float('inf'))
|
||
|
|
|
||
|
|
NEG_INF: tl.constexpr = float('-inf')
|
||
|
|
for i in tl.static_range(K):
|
||
|
|
# 找最大值的索引 (scalar)
|
||
|
|
idx = tl.argmax(sel, axis=0)
|
||
|
|
# 把 sel[idx] 设为 -inf (避免重复选)
|
||
|
|
is_max = (offs == idx)
|
||
|
|
sel = tl.where(is_max & mask, NEG_INF, sel)
|
||
|
|
# 写 indices
|
||
|
|
tl.store(idx_ptr + pid * K + i, idx.to(tl.int32))
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# Kernel 3: gather + cat + sigmoid + renorm + scale + split
|
||
|
|
# ============================================================
|
||
|
|
@triton.jit
|
||
|
|
def _gate_finalize_kernel(
|
||
|
|
routed_ptr, # [T, N] fp16/bf16
|
||
|
|
shared_ptr, # [T, S] fp16/bf16
|
||
|
|
idx_ptr, # [T, k] int32
|
||
|
|
route_scale, # float
|
||
|
|
global_scale_ptr, # [1] fp32
|
||
|
|
routed_w_ptr, # [T, k] output dtype (input dtype)
|
||
|
|
shared_w_ptr, # [T, S] output dtype (input dtype)
|
||
|
|
N, S,
|
||
|
|
stride_routed_t, stride_shared_t,
|
||
|
|
stride_routed_w_t, stride_shared_w_t,
|
||
|
|
K: tl.constexpr,
|
||
|
|
BLOCK_S: tl.constexpr, # next_pow2(S)
|
||
|
|
):
|
||
|
|
# input 指针 cast 防 NaN, output 不 cast (output 是计算结果, 不需保 NaN bits)
|
||
|
|
if routed_ptr.dtype.element_ty.primitive_bitwidth == 16:
|
||
|
|
routed_ptr = routed_ptr.to(tl.pointer_type(tl.int16))
|
||
|
|
shared_ptr = shared_ptr.to(tl.pointer_type(tl.int16))
|
||
|
|
|
||
|
|
pid = tl.program_id(0).to(tl.int64)
|
||
|
|
|
||
|
|
# ===== 加载 routed 的 K 个值 (按 indices gather) =====
|
||
|
|
idx_offs = tl.arange(0, K)
|
||
|
|
indices = tl.load(idx_ptr + pid * K + idx_offs).to(tl.int64)
|
||
|
|
routed_vals = tl.load(routed_ptr + pid * stride_routed_t + indices).to(tl.float32)
|
||
|
|
routed_sigmoid = 1.0 / (1.0 + tl.exp(-routed_vals))
|
||
|
|
|
||
|
|
# ===== 加载 shared 的 S 个值 =====
|
||
|
|
s_offs = tl.arange(0, BLOCK_S)
|
||
|
|
s_mask = s_offs < S
|
||
|
|
shared_vals = tl.load(
|
||
|
|
shared_ptr + pid * stride_shared_t + s_offs, mask=s_mask, other=0
|
||
|
|
).to(tl.float32)
|
||
|
|
shared_sigmoid = 1.0 / (1.0 + tl.exp(-shared_vals))
|
||
|
|
|
||
|
|
# ===== 计算 sum (routed + shared 一起归一化) =====
|
||
|
|
sum_routed = tl.sum(routed_sigmoid, axis=0)
|
||
|
|
sum_shared = tl.sum(shared_sigmoid, axis=0)
|
||
|
|
total_sum = sum_routed + sum_shared
|
||
|
|
|
||
|
|
# ===== scale =====
|
||
|
|
gs = tl.load(global_scale_ptr)
|
||
|
|
inv = (route_scale * gs) / total_sum
|
||
|
|
|
||
|
|
# ===== 写 routed_w =====
|
||
|
|
routed_w = (routed_sigmoid * inv).to(routed_w_ptr.dtype.element_ty)
|
||
|
|
tl.store(routed_w_ptr + pid * stride_routed_w_t + idx_offs, routed_w)
|
||
|
|
|
||
|
|
# ===== 写 shared_w =====
|
||
|
|
shared_w = (shared_sigmoid * inv).to(shared_w_ptr.dtype.element_ty)
|
||
|
|
tl.store(shared_w_ptr + pid * stride_shared_w_t + s_offs, shared_w, mask=s_mask)
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# Python wrapper
|
||
|
|
# ============================================================
|
||
|
|
def _next_pow2(x):
|
||
|
|
p = 1
|
||
|
|
while p < x:
|
||
|
|
p *= 2
|
||
|
|
return p
|
||
|
|
|
||
|
|
|
||
|
|
def sigmoid_gate_topk_renorm(logits, k, n_shared_experts, route_scale, global_scale, bias):
|
||
|
|
T, G = logits.shape
|
||
|
|
N = G - n_shared_experts
|
||
|
|
S = n_shared_experts
|
||
|
|
|
||
|
|
if logits.ndim != 2:
|
||
|
|
raise ValueError('logits must have shape [T, N+S]')
|
||
|
|
if not isinstance(k, int) or k <= 0 or k > N:
|
||
|
|
raise ValueError(f'k must be 0 < k <= N={N}, got {k}')
|
||
|
|
if not isinstance(n_shared_experts, int) or n_shared_experts < 0 or n_shared_experts > S:
|
||
|
|
raise ValueError(f'n_shared_experts must be 0 <= S={S}, got {n_shared_experts}')
|
||
|
|
if bias.shape != (N,):
|
||
|
|
raise ValueError(f'bias must have shape [{N}], got {tuple(bias.shape)}')
|
||
|
|
if global_scale.numel() != 1:
|
||
|
|
raise ValueError(f'global_scale must be a scalar tensor, got shape {tuple(global_scale.shape)}')
|
||
|
|
if logits.device.type in ('cpu', 'meta', 'mps'):
|
||
|
|
raise RuntimeError('a real Triton accelerator backend is required')
|
||
|
|
|
||
|
|
routed = logits[:, :N]
|
||
|
|
shared = logits[:, N:]
|
||
|
|
|
||
|
|
# ===== 中间 buffer =====
|
||
|
|
sel = torch.empty((T, N), dtype=torch.float32, device=logits.device)
|
||
|
|
indices = torch.empty((T, k), dtype=torch.int32, device=logits.device)
|
||
|
|
|
||
|
|
BLOCK_N = max(16, _next_pow2(N))
|
||
|
|
BLOCK_S = max(16, _next_pow2(S))
|
||
|
|
|
||
|
|
module = torch.get_device_module(logits.device)
|
||
|
|
with module.device(logits.device):
|
||
|
|
# ===== Kernel 1: sigmoid + bias =====
|
||
|
|
grid1 = (T, triton.cdiv(N, BLOCK_N))
|
||
|
|
_sigmoid_bias_kernel[grid1](
|
||
|
|
routed, bias, sel, N, routed.stride(0),
|
||
|
|
BLOCK_N=BLOCK_N,
|
||
|
|
enable_fp_fusion=False, num_warps=4,
|
||
|
|
)
|
||
|
|
|
||
|
|
# ===== Kernel 2: topk =====
|
||
|
|
grid2 = (T,)
|
||
|
|
_topk_kernel[grid2](
|
||
|
|
sel, indices, N, K=k, BLOCK_N=BLOCK_N,
|
||
|
|
enable_fp_fusion=False, num_warps=4,
|
||
|
|
)
|
||
|
|
|
||
|
|
# ===== Kernel 3: gate finalize =====
|
||
|
|
routed_w = torch.empty((T, k), dtype=logits.dtype, device=logits.device)
|
||
|
|
shared_w = torch.empty((T, S), dtype=logits.dtype, device=logits.device)
|
||
|
|
grid3 = (T,)
|
||
|
|
_gate_finalize_kernel[grid3](
|
||
|
|
routed, shared, indices, route_scale, global_scale,
|
||
|
|
routed_w, shared_w, N, S,
|
||
|
|
routed.stride(0), shared.stride(0),
|
||
|
|
routed_w.stride(0), shared_w.stride(0),
|
||
|
|
K=k, BLOCK_S=BLOCK_S,
|
||
|
|
enable_fp_fusion=False, num_warps=4,
|
||
|
|
)
|
||
|
|
|
||
|
|
return routed_w, indices, shared_w
|
||
|
|
|
||
|
|
|
||
|
|
reference = sigmoid_gate_topk_renorm
|