| """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 |
|
|
