File size: 7,361 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 | """Allocation-free native packed expert pipeline; integration remains with N1."""
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 offload.dspark.kernels.mxfp4.regfed import RegisterExperts
class Extension(Protocol):
def schedule(self, ids: int, dst: int, status: int, slots: int, experts: int, i64: bool, group: int, stream: int) -> None: ...
def launch(self, x: int, rw: int, sched: int, w: int, s: int, up: int, us: int, w2: int, s2: int,
h: int, out: int, gout: int, uout: int, n: int, k: int, slots: int, out_stride: int,
warps: int, stream: int) -> None: ...
@lru_cache(maxsize=1)
def extension() -> Extension:
from torch.utils.cpp_extension import load
return cast(Extension, load(name='kx_fusedmoe', sources=[str(Path(__file__).with_name('fusedmoe.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 _tensor(t: torch.Tensor, d: torch.device, dtype: torch.dtype, shape: tuple[int, ...], *,
contiguous: bool = True) -> None:
if t.device != d or t.dtype != dtype or tuple(t.shape) != shape or (contiguous and not t.is_contiguous()):
raise ValueError(f'expected {dtype} {shape} on {d}, contiguous={contiguous}; got {t.dtype} {t.shape} {t.device}')
@dataclass(frozen=True)
class Workspace:
schedule: torch.Tensor
hq: torch.Tensor
@classmethod
def allocate(cls, rows: int, intermediate: int, device: torch.device) -> Workspace:
"""Allocate before capture. The workspace is exclusively owned by its stream."""
if not 1 <= rows <= 20 or intermediate % 32:
raise ValueError('workspace requires 1..20 rows and 32-aligned intermediate')
return cls(torch.empty(1 + rows*3*17, dtype=torch.int32, device=device),
torch.empty(rows*3, intermediate, dtype=torch.bfloat16, device=device))
VARIANTS: dict[int, str] = {
4: '32 channels/warp, 4 warps/CTA', 8: '32 channels/warp, 8 warps/CTA',
16: '16 channels/warp, shared amax, 4 warps/CTA',
32: '16 channels/warp, shared amax, 8 warps/CTA',
64: '16 channels/warp, compensated FP32, 4 warps/CTA',
128: '16 channels/warp, compensated FP32, 8 warps/CTA',
256: 'CTA split-K2, ordinary FP32', 512: 'CTA split-K4, ordinary FP32',
1024: 'CTA split-K4, compensated FP32', 2048: 'CTA split-K2, compensated FP32',
4096: 'CTA split-K2, K16 compensated FP32', 8192: 'CTA split-K4, K16 compensated FP32',
16384: 'CTA split-K8, two output warps, K16 compensated FP32',
65536: 'CTA split-K2, group16 reuse, ordinary FP32',
131072: 'CTA split-K2, K64 compensated FP32 with high/low join',
}
@dataclass(frozen=True)
class FusedMoE:
gate: RegisterExperts
up: RegisterExperts
down: RegisterExperts
ext: Extension
variant: int = 131072
@classmethod
def from_register_banks(cls, gate: RegisterExperts, up: RegisterExperts, down: RegisterExperts,
*, variant: int = 131072) -> FusedMoE:
"""Reuse existing prepacked KX banks without copies or changed weight values."""
if (gate.experts, gate.outputs, gate.inputs, gate.codes.device) != (
up.experts, up.outputs, up.inputs, up.codes.device):
raise ValueError('gate/up dimensions and devices must match')
if (down.experts, down.inputs, down.outputs, down.codes.device) != (
gate.experts, gate.outputs, gate.inputs, gate.codes.device):
raise ValueError('down must be the gate/up transposed shape on the same device')
if gate.outputs % 64 or gate.inputs % 64 or variant not in (4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 65536, 131072):
raise ValueError('64-aligned dimensions and a supported variant required')
if gate.codes.device.type != 'cuda' or torch.cuda.get_device_capability(gate.codes.device) != (8, 6):
raise ValueError('SM86 CUDA weights required')
for bank in (gate, up, down):
_tensor(bank.codes, gate.codes.device, torch.uint8, (bank.experts, bank.outputs, bank.inputs//2))
_tensor(bank.scales, gate.codes.device, torch.uint8, (bank.experts, bank.outputs, bank.inputs//32))
return cls(gate, up, down, extension(), variant)
def forward_into(self, xq: torch.Tensor, ids: torch.Tensor, rw: torch.Tensor, out: torch.Tensor,
status: torch.Tensor, workspace: Workspace, *,
projections: tuple[torch.Tensor, torch.Tensor] | None = None) -> None:
"""xq -> FP32 SLOT partials; no allocation, host synchronization, or CPU ID read.
Optional FP32 projection buffers expose the exact MMA sums for qualification.
They add stores and must be omitted from performance runs.
"""
d = self.gate.codes.device
rows = xq.shape[0]
n, k = self.gate.outputs, self.gate.inputs
if not 1 <= rows <= 20 or ids.dtype not in (torch.int32, torch.int64):
raise ValueError('1..20 rows and int32/int64 route IDs required')
_tensor(xq, d, torch.bfloat16, (rows, k))
_tensor(ids, d, ids.dtype, (rows, 3))
_tensor(rw, d, torch.float32, (rows, 3))
_tensor(out, d, torch.float32, (rows, 3, k), contiguous=False)
if out.stride(1) != k or out.stride(2) != 1 or out.stride(0) < 3*k:
raise ValueError('SLOT output needs contiguous channels/slots and nonoverlapping rows')
_tensor(status, d, torch.int32, (1,))
_tensor(workspace.schedule, d, torch.int32, (1+rows*3*17,))
_tensor(workspace.hq, d, torch.bfloat16, (rows*3, n))
if projections is not None:
for p in projections:
_tensor(p, d, torch.float32, (rows, 3, n))
with torch.cuda.device(d):
st = torch.cuda.current_stream(d).cuda_stream
self.ext.schedule(ids.data_ptr(), workspace.schedule.data_ptr(), status.data_ptr(), rows*3,
self.gate.experts, ids.dtype == torch.int64, 16 if self.variant == 65536 else 8, st)
self.ext.launch(xq.data_ptr(), rw.data_ptr(), workspace.schedule.data_ptr(), self.gate.codes.data_ptr(),
self.gate.scales.data_ptr(), self.up.codes.data_ptr(), self.up.scales.data_ptr(),
self.down.codes.data_ptr(), self.down.scales.data_ptr(), workspace.hq.data_ptr(), out.data_ptr(),
0 if projections is None else projections[0].data_ptr(),
0 if projections is None else projections[1].data_ptr(), n, k, rows*3, out.stride(0), self.variant, st)
def select_variant(stage: int, rows: int) -> int:
"""Static measured profile; choose before capture, without reading device IDs.
CARRY is numerically qualified on every stage/rank. Plain K2 is qualified
on stage1 and is faster at the measured ten-row shape. Other shapes retain
CARRY. This profile does not meet all six latency targets; see RESULTS.md.
"""
if stage not in (0, 1, 2) or not 1 <= rows <= 20:
raise ValueError('expected stage 0..2 and rows 1..20')
return 256 if stage == 1 and rows == 10 else 131072
|