File size: 11,451 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | """Load-time prepack/admission and allocation-free, current-stream dense calls."""
from __future__ import annotations
from dataclasses import dataclass
import fcntl
import os
from functools import lru_cache
from pathlib import Path
from typing import Protocol, cast
import numpy as np
from numpy.typing import NDArray
import torch
def require_lock() -> None:
"""Qualification-only guard: verify the inherited exclusive flock, not a flag."""
locks = {'GPU-7744b44b-8a13-b06d-a67e-f958fd74b8fc': 'gpu-7744b44b.lock',
'GPU-bc62df2a-59cf-9828-7d43-558d1355bab7': 'gpu-bc62df2a.lock'}
visible = os.environ.get('CUDA_VISIBLE_DEVICES', '')
if visible not in locks:
raise RuntimeError('qualification must pin the development UUID')
lock=Path.home()/'dsv41-offload/locks'/locks[visible]
fd=int(os.environ.get('KX_FP8DENSE_LOCK_FD','9'))
held=os.fstat(fd)
actual=lock.stat()
if (held.st_ino,held.st_dev)!=(actual.st_ino,actual.st_dev):
raise RuntimeError('inherited descriptor is not the GPU lock')
fcntl.flock(fd,fcntl.LOCK_EX|fcntl.LOCK_NB)
with lock.open('a') as probe:
try:fcntl.flock(probe,fcntl.LOCK_EX|fcntl.LOCK_NB)
except BlockingIOError:return
raise RuntimeError('exclusive GPU lock was not held')
class Extension(Protocol):
def simt(self, x: int, w: int, s: int, y: int, m: int, n: int, k: int,
xs: int, ys: int, groups: int, f32: bool, stream: int) -> None: ...
def prepack(self, dst: int, src: int, n: int, k: int, stream: int) -> None: ...
def linear(self, x: int, w: int, s: int, y: int, scratch: int, m: int, n: int, k: int,
xs: int, ys: int, groups: int, splits: int, rows: int, f32: bool, fast: 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_fp8dense', sources=[str(Path(__file__).with_name('dense.cu'))],
extra_cuda_cflags=['-O3', '-gencode=arch=compute_86,code=sm_86', '--fmad=false', '-Xptxas=-v'],
extra_cflags=['-O2'], with_cuda=True, verbose=True))
def domain_histogram(codes: torch.Tensor, scales: torch.Tensor) -> NDArray[np.int64]:
"""CPU load admission; return exact present (scale,code) counts, <=32-row temporary."""
if codes.device.type != 'cpu' or scales.device.type != 'cpu' or codes.dtype != torch.uint8 or scales.dtype != torch.uint8:
raise ValueError('load admission requires CPU uint8 codes/scales')
if codes.ndim != 3 or scales.shape != (codes.shape[0], codes.shape[1] // 32, codes.shape[2] // 32):
raise ValueError('codes [G,N,K] and scales [G,N/32,K/32] required')
c, s = codes.numpy(), scales.numpy()
hist = np.zeros(65536, dtype=np.int64)
for group in range(c.shape[0]):
for row in range(0, c.shape[1], 32):
block = c[group, row:row+32].reshape(32, -1, 32).astype(np.int32)
pairs = (s[group, row//32].astype(np.int32)[None, :, None] << 8) | block
hist += np.bincount(pairs.ravel(), minlength=65536)
used = np.flatnonzero(hist)
cc, ss = used & 255, used >> 8
mag = cc & 127
if np.any(mag == 127) or np.any(ss == 255):
raise ValueError('NaN E4M3/E8M0 value in checkpoint')
exponent, mantissa = (mag >> 3), mag & 7
values = np.where(exponent == 0, mantissa/512.0, np.ldexp(1.0+mantissa/8.0, exponent-7))
values = np.ldexp(values, ss-127)
nonzero = values != 0
# Register exponent addition is exact for every normal BF16 weight. Refuse
# the subnormal/overflow domains rather than silently flushing or wrapping.
if np.any(nonzero & ((values < 2.0**-126) | (values > float(torch.finfo(torch.bfloat16).max)))):
raise ValueError('decoded weight outside finite normal BF16 domain')
return hist
def _matrix(tensor: torch.Tensor, device: torch.device, dtype: torch.dtype, shape: tuple[int, ...]) -> None:
if tensor.device != device or tensor.dtype != dtype or tuple(tensor.shape) != shape or not tensor.is_contiguous():
raise ValueError(f'expected contiguous {dtype} {shape} on {device}, got {tensor.dtype} {tuple(tensor.shape)} {tensor.device}')
@dataclass(frozen=True)
class DenseFP8:
codes: torch.Tensor
scales: torch.Tensor
groups: int
outputs: int
inputs: int
ext: Extension
fast_scales: bool
layout: str
@classmethod
def from_checkpoint(cls, codes: torch.Tensor, scales: torch.Tensor, device: torch.device, *,
groups: int = 1, row_offset: int = 0, col_offset: int = 0, layout: str = "register") -> DenseFP8:
"""codes are a [G*N,K] rank shard; scales are the FULL checkpoint grid.
row_offset/col_offset locate the shard in that full grid. Load once,
outside capture. G=4 implements wo_a's four block-diagonal groups.
Source bytes need not be contiguous. Only one temporary device shard
is retained during prepack, never a decoded-weight bank.
"""
if codes.ndim != 2 or groups not in (1, 4) or codes.shape[0] % (32*groups) or codes.shape[1] % 32:
raise ValueError('FP8 codes [G*N,K], G=1/4, N/K positive multiples of 32 required')
if min(codes.shape) <= 0 or row_offset < 0 or col_offset < 0 or row_offset % 32 or col_offset % 32:
raise ValueError('positive shapes and nonnegative block-aligned shard offsets required')
n, k = codes.shape[0]//groups, codes.shape[1]
if scales.ndim != 2 or scales.shape[0] < (row_offset+groups*n)//32 or scales.shape[1] < (col_offset+k)//32:
raise ValueError('full tensor scale grid does not cover the shard')
sc = scales[row_offset//32:(row_offset+groups*n)//32, col_offset//32:(col_offset+k)//32].contiguous()
domain_histogram(codes.reshape(groups, n, k), sc.view(groups, n//32, k//32))
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')
if layout not in ('register', 'simt'):
raise ValueError('layout must be register or simt')
fast_scales=bool(int(sc.max())<=134)
if layout=='simt' and (not fast_scales or k%128):
raise ValueError('SIMT requires scales<=134 and K%128=0')
ext = extension() # compile before capture; no device-global caches/state
with torch.cuda.device(device):
ds = sc.to(device)
source = codes.contiguous().to(device)
dst = torch.empty((groups, n, k), dtype=torch.uint8, device=device) if layout=='register' else source.view(groups,n,k)
st = torch.cuda.current_stream(device).cuda_stream
if layout=='register':
for group in range(groups):
ext.prepack(dst[group].data_ptr(), source[group*n:(group+1)*n].data_ptr(), n, k, st)
else:
dst=source.view(groups,n,k)
return cls(dst, ds, groups, n, k, ext, fast_scales, layout)
def configuration(self, rows: int) -> tuple[int, int]:
"""Offline-measured (split count, M/8 register tiles), no autotuning.
Callers allocate [splits,M,G,N] FP32 scratch when splits>1 and pass
both values explicitly to linear_into. Unlisted small test shapes
use the no-scratch generic kernel.
"""
if not 1<=rows<=128:
raise ValueError('rows must be 1..128')
if self.layout=='simt':
return 1,1
n,k=self.outputs,self.inputs
if k==15360 and n==5120:
return (8,2) if rows<=32 else (2,16)
if k==5120 and n in (512,1152,1280,2304):
return 16,1
if (n,k)==(16384,1280):
return 1,1
if k in (4096,1152) and n in (1024,5120):
return (8,1) if rows<=8 else (4,1)
return 1,min(16,1<<((rows-1)//8).bit_length())
def linear_into(self, x: torch.Tensor, out: torch.Tensor, *, scratch: torch.Tensor | None = None,
splits: int = 1, row_tiles: int | None = None) -> None:
"""x already QDQ'd, BF16 [M,G*K]; out BF16/FP32 [M,G*N].
No allocation, sync, host tensor read, or lazy build. FP32 split-K
scratch is [splits,M,G,N]; reduction is fixed-order, with one final RNE.
Rows may be strided, columns must be contiguous. For the measured
dispatch, call configuration(M) and pass both values explicitly;
split=1 is the generic control requiring no scratch.
"""
if x.ndim != 2 or not 1 <= x.shape[0] <= 128 or x.shape[1] != self.groups*self.inputs:
raise ValueError('x must be [1..128,G*K]')
m, d = x.shape[0], self.codes.device
if x.device != d or x.dtype != torch.bfloat16 or x.stride(1) != 1 or x.stride(0) < x.shape[1] or x.stride(0) % 2 or x.data_ptr() % 4:
raise ValueError('BF16 input on weight device with contiguous columns required')
if out.shape != (m, self.groups*self.outputs) or out.device != d or out.dtype not in (torch.bfloat16, torch.float32):
raise ValueError('output shape/device/dtype mismatch')
if out.stride(1) != 1 or out.stride(0) < out.shape[1]:
raise ValueError('output columns must be contiguous and rows nonoverlapping')
if splits not in range(1, 33) or splits > self.inputs//16:
raise ValueError('splits must be 1..32 and no greater than K/16')
rows = row_tiles if row_tiles is not None else min(16, 1 << ((m-1)//8).bit_length())
if rows not in (1, 2, 4, 8, 16):
raise ValueError('row_tiles must be 1,2,4,8,16')
if splits > 1:
if scratch is None:
raise ValueError('split-K requires caller-owned scratch')
_matrix(scratch, d, torch.float32, (splits, m, self.groups, self.outputs))
for buf in (out, scratch):
if buf is not None and buf.untyped_storage().data_ptr() in (x.untyped_storage().data_ptr(), self.codes.untyped_storage().data_ptr(), self.scales.untyped_storage().data_ptr()):
raise ValueError('output/scratch must not alias inputs or weights')
if scratch is not None and scratch.untyped_storage().data_ptr() == out.untyped_storage().data_ptr():
raise ValueError('output/scratch must not alias')
with torch.cuda.device(d):
if self.layout=='simt':
if splits!=1 or m>20 or x.data_ptr()%8 or x.stride(0)%4:
raise ValueError('SIMT needs M<=20, splits=1 and 8-byte aligned rows')
self.ext.simt(x.data_ptr(),self.codes.data_ptr(),self.scales.data_ptr(),out.data_ptr(),m,
self.outputs,self.inputs,x.stride(0),out.stride(0),self.groups,
out.dtype==torch.float32,torch.cuda.current_stream(d).cuda_stream)
return
self.ext.linear(x.data_ptr(), self.codes.data_ptr(), self.scales.data_ptr(), out.data_ptr(),
0 if scratch is None else scratch.data_ptr(), m, self.outputs, self.inputs,
x.stride(0), out.stride(0), self.groups, splits, rows,
out.dtype == torch.float32, self.fast_scales, torch.cuda.current_stream(d).cuda_stream)
|