File size: 4,883 Bytes
e1acc15 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Direct register-fed SM86 MMA adapter, isolated from all engine dependencies."""
from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Protocol, cast
import torch
from . import _cuda
from .reference import admit
class Extension(Protocol):
def prepack(self, w: int, s: int, source: int, scales: int, n: int, k: int, stream: int) -> None: ...
def schedule(self, ids: int, schedule: int, slots: int, experts: int, status: int, stream: int) -> None: ...
def linear(self, x: int, ids: int, schedule: int, w: int, s: int, up: int, us: int, out: int,
uout: int, n: int, k: int, slots: int, experts: int, top: int, per_slot: bool,
fused: bool, stream: int) -> None: ...
def test_decode(self, codes: int, scales: int, out: int, count: int, stream: int) -> None: ...
@lru_cache(maxsize=1)
def extension() -> Extension:
from torch.utils.cpp_extension import load
return cast(Extension, load(
name='kx_mxfp4_regfed', sources=[str(Path(__file__).with_name('regfed.cu'))],
extra_cuda_cflags=['-O3','-gencode=arch=compute_86,code=sm_86','--fmad=false','-Xptxas=-v'],
extra_cflags=['-O2'], with_cuda=True, verbose=True))
@dataclass(frozen=True)
class RegisterExperts:
codes: torch.Tensor
scales: torch.Tensor
experts: int
outputs: int
inputs: int
ext: Extension
@classmethod
def from_checkpoint(cls, codes: torch.Tensor, scales: torch.Tensor,
device: torch.device) -> RegisterExperts:
admit(codes,scales)
e,n,half_k = codes.shape
if n%16 or (half_k*2)%64:
raise ValueError('register layout requires N%16=0 and K%64=0; use masked Triton for padding')
if int(scales.min()) == 0:
raise ValueError('register scale multiplier excludes E8M0 zero; use integer decoder')
if device.type != 'cuda' or device.index is None or torch.cuda.get_device_capability(device) != (8,6):
raise ValueError('explicit SM86 CUDA device required')
ext = extension() # prebuild before any capture; module has no static GPU state
with torch.cuda.device(device):
dst = torch.empty((e,n,half_k),dtype=torch.uint8,device=device)
ds = torch.empty(tuple(scales.shape),dtype=torch.uint8,device=device)
stream = torch.cuda.current_stream(device).cuda_stream
for i in range(e):
# One expert staging allocation only; stream directly into its final shard.
p,s = codes[i].contiguous().to(device),scales[i].contiguous().to(device)
ext.prepack(dst[i].data_ptr(),ds[i].data_ptr(),p.data_ptr(),s.data_ptr(),n,half_k*2,stream)
del p,s # same-stream allocator reuse follows the enqueued prepack
return cls(dst,ds,e,n,half_k*2,ext)
def linear_into(self, x: torch.Tensor, ids: torch.Tensor, out: torch.Tensor,
status: torch.Tensor, schedule: torch.Tensor, *, per_slot: bool=False,
up: RegisterExperts | None=None, up_out: torch.Tensor | None=None) -> None:
if ids.ndim != 2 or ids.shape[0]<1 or not 1<=ids.shape[1]<=3:
raise ValueError('expected nonempty [rows,top<=3] int32 IDs')
rows,top = ids.shape
if rows*top > 64:
raise ValueError('register kernel supports at most 64 routed slots per launch')
d = self.codes.device
_cuda(ids,d,torch.int32,(rows,top))
_cuda(x,d,torch.bfloat16,(rows*top if per_slot else rows,self.inputs))
_cuda(out,d,torch.float32,(rows,top,self.outputs))
_cuda(status,d,torch.int32,(1,))
_cuda(schedule,d,torch.int32,(rows*top,8))
if up is not None:
if (up.inputs,up.outputs,up.experts,up.codes.device)!=(self.inputs,self.outputs,self.experts,d):
raise ValueError('gate/up dimensions and devices must match')
if up_out is None:
raise ValueError('fused gate/up requires up_out')
_cuda(up_out,d,torch.float32,tuple(out.shape))
elif up_out is not None:
raise ValueError('up_out without up bank')
with torch.cuda.device(d):
stream = torch.cuda.current_stream(d).cuda_stream
self.ext.schedule(ids.data_ptr(),schedule.data_ptr(),rows*top,self.experts,status.data_ptr(),stream)
self.ext.linear(x.data_ptr(),ids.data_ptr(),schedule.data_ptr(),self.codes.data_ptr(),self.scales.data_ptr(),
0 if up is None else up.codes.data_ptr(),0 if up is None else up.scales.data_ptr(),
out.data_ptr(),0 if up_out is None else up_out.data_ptr(),self.outputs,self.inputs,
rows*top,self.experts,top,per_slot,up is not None,stream)
|