File size: 3,027 Bytes
03c3dff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """BF16 inference fusion; FMA disabled, original cast points explicit.
Torch pow/mean is retained verbatim: no alternate RMS reduction is substituted.
The Q/K path materializes normalized pre-RoPE keys because meta pooling uses them.
"""
import torch
import triton as tr
import triton.language as tl
import triton.language.extra.cuda.libdevice as lib
@tr.jit
def norm_tail(X,V,W,Y,N:tl.constexpr,D:tl.constexpr,EPS:tl.constexpr,BLOCK:tl.constexpr):
i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
x=tl.load(X+i,i<N,0).to(tl.float32)
v=tl.load(V+i//D,i<N,0)
scale=lib.rsqrt(v+EPS)
n=(x*scale).to(tl.bfloat16).to(tl.float32)
w=tl.load(W+i%D,i<N,0).to(tl.float32)
tl.store(Y+i,n*w,i<N)
_norm_launchers={}
_rope_launchers={}
def tail(x,variance,weight,eps,fast=False):
y=torch.empty_like(x);n=x.numel()
key=(x.device.index,n,x.shape[-1],eps)
if fast and key in _norm_launchers:
_norm_launchers[key](x,variance,weight,y)
else:
compiled=norm_tail[(tr.cdiv(n,256),)](x,variance,weight,y,n,x.shape[-1],eps,256,enable_fp_fusion=False)
if fast:_norm_launchers[key]=compiled[(tr.cdiv(n,256),1,1)]
return y
def rms(x,weight,eps):
# EXACT original FP32 pow + torch.mean dispatch and shape.
f=x.to(torch.float32);variance=f.pow(2).mean(-1,keepdim=True)
return tail(x,variance,weight,eps,fast=True)
@tr.jit
def rope_kernel(X,C,S,Y,N:tl.constexpr,H:tl.constexpr,T:tl.constexpr,D:tl.constexpr,ST0:tl.constexpr,ST1:tl.constexpr,ST2:tl.constexpr,ST3:tl.constexpr,SY0:tl.constexpr,SY1:tl.constexpr,SY2:tl.constexpr,SY3:tl.constexpr,BLOCK:tl.constexpr):
i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK);d=i%D;t=i//D%T;h=i//(D*T)%H;b=i//(D*T*H)
ptr=b*ST0+h*ST1+t*ST2+d*ST3
rd=tl.where(d<D//2,d+D//2,d-D//2)
x=tl.load(X+ptr,i<N,0).to(tl.float32)
rot=tl.load(X+b*ST0+h*ST1+t*ST2+rd*ST3,i<N,0).to(tl.float32)
rot=tl.where(d<D//2,-rot,rot)
c=tl.load(C+(b*T+t)*D+d,i<N,0).to(tl.float32);s=tl.load(S+(b*T+t)*D+d,i<N,0).to(tl.float32)
a=(x*c).to(tl.bfloat16).to(tl.float32);z=(rot*s).to(tl.bfloat16).to(tl.float32)
tl.store(Y+b*SY0+h*SY1+t*SY2+d*SY3,a+z,i<N)
def rope_one(x,cos,sin):
b,h,t,d=x.shape
# TensorIterator canonicalizes dense degenerate dimensions (e.g. T=1),
# while preserving the physical order of non-contiguous dense inputs.
y=torch.empty(x.shape,device=x.device,dtype=x.dtype) if x.is_contiguous() else torch.empty_like(x)
n=x.numel()
key=(x.device.index,n,h,t,d,x.stride(),y.stride())
if key in _rope_launchers:
_rope_launchers[key](x,cos,sin,y)
else:
compiled=rope_kernel[(tr.cdiv(n,256),)](x,cos,sin,y,n,h,t,d,*x.stride(),*y.stride(),256,enable_fp_fusion=False)
_rope_launchers[key]=compiled[(tr.cdiv(n,256),1,1)]
return y
def rope(q,k,cos,sin,position_ids=None,unsqueeze_dim=1):
assert unsqueeze_dim==1 and cos.is_contiguous() and sin.is_contiguous()
return rope_one(q,cos,sin),rope_one(k,cos,sin)
def enabled(x):
return x.is_cuda and x.device.index == torch.cuda.current_device() and x.dtype == torch.bfloat16 and x.ndim >= 2 and x.shape[0] == 1 and not torch.is_grad_enabled()
|