| """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() |
| 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): |
| |
| 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 |
| 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) |
|
|
