File size: 7,850 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 | # SPDX-License-Identifier: Apache-2.0
"""One layer's EXL3 expert bank on one TP rank (torch only: no vLLM, no exllamav3 import).
Shared by the vLLM method (moe_method.py, where the tensors are registered Parameters) and the
kernel tests (tests/test_kernels_gpu.py, where the bank owns its tensors).
Layout: experts sit in K-sorted physical slots (plan.LayerPlan.perm). Trellis words of all experts
share two flat int16 banks (w13 = [w1 | w3] per slot, w2), sign/scale vectors are [E, ...] fp16:
w13_suh [E, 2, H] (w1, w3: full input width, not cut)
w13_svh [E, 2, I_loc] (cut along the intermediate dim)
w2_suh [E, I_loc] (cut)
w2_svh [E, H] (not cut)
Each rank keeps intermediate channels [rank * I_loc, (rank + 1) * I_loc) of every expert, cut from
the full exllamav3 tensors at load time (plan.shard_slices; exact because the cut is 128-aligned).
"""
from __future__ import annotations
import torch
from .hostcopy import copy_from_host
from .kernels import GroupTables, ProjTables
from .plan import (
PROJ_DOWN,
PROJ_GATE,
PROJ_UP,
BankLayout,
LayerPlan,
shard_slices,
tp_shard,
trellis_shape,
)
TENSOR_NAMES = ("w13_trellis", "w2_trellis", "w13_suh", "w13_svh", "w2_suh", "w2_svh")
class Exl3Bank:
def __init__(self, lp: LayerPlan, hidden: int, inter_full: int, tp_size: int, tp_rank: int):
self.lp = lp
self.hidden = hidden
self.inter_full = inter_full
self.tp_size = tp_size
self.tp_rank = tp_rank
i0, i1 = tp_shard(inter_full, tp_size, tp_rank)
self.inter_local = i1 - i0
self.bl = BankLayout.build(lp, hidden, self.inter_local)
self.t: dict[str, torch.Tensor] = {}
@property
def num_experts(self) -> int:
return self.lp.num_experts
def allocate(self, device: torch.device | str | None = None) -> dict[str, torch.Tensor]:
"""Empty tensors (device None = torch's current default device, which is what vLLM's
create_weights relies on)."""
E, H, I = self.num_experts, self.hidden, self.inter_local
kw = {} if device is None else {"device": device}
f16 = torch.float16
self.t = {
"w13_trellis": torch.empty(self.bl.total13, dtype=torch.int16, **kw),
"w2_trellis": torch.empty(self.bl.total2, dtype=torch.int16, **kw),
"w13_suh": torch.empty(E, 2, H, dtype=f16, **kw),
"w13_svh": torch.empty(E, 2, I, dtype=f16, **kw),
"w2_suh": torch.empty(E, I, dtype=f16, **kw),
"w2_svh": torch.empty(E, H, dtype=f16, **kw),
}
return self.t
def bind(self, tensors: dict[str, torch.Tensor]) -> None:
"""Point the bank at externally owned tensors (vLLM Parameters' .data after loading)."""
for n in TENSOR_NAMES:
self.t[n] = tensors[n]
@staticmethod
def dest_name(kind: str, proj: str) -> str:
side = "w2" if proj == PROJ_DOWN else "w13"
return f"{side}_{kind}"
def load_piece(self, kind: str, proj: str, expert_id: int, w: torch.Tensor,
dst: torch.Tensor | None = None) -> None:
"""Copy this rank's part of one full exllamav3 tensor (trellis / suh / svh) of logical expert
`expert_id` into its slot. Raises on any shape or dtype the plan does not predict."""
if proj not in (PROJ_GATE, PROJ_UP, PROJ_DOWN):
raise ValueError(f"bad projection {proj!r}")
e = int(expert_id)
p = self.lp.perm[e]
b = self.lp.bits[e]
bl = self.bl
gate_up = proj in (PROJ_GATE, PROJ_UP)
sl = shard_slices(proj, self.inter_full, self.tp_size, self.tp_rank)
if dst is None:
dst = self.t[self.dest_name(kind, proj)]
if kind == "trellis":
k = b.k_gu if gate_up else b.k_d
rows, cols = (self.hidden, self.inter_full) if gate_up else (self.inter_full, self.hidden)
want = trellis_shape(rows, cols, k)
if tuple(w.shape) != want or w.dtype != torch.int16:
raise ValueError(f"expert {e} {proj}.trellis: got {tuple(w.shape)} {w.dtype}; the plan's "
f"K={k} means {want} int16 (plan and checkpoint disagree)")
a, z = sl.trellis_range
src = w.narrow(sl.trellis_dim, a, z - a)
if proj == PROJ_GATE:
off, n = bl.off_w1[p], bl.words_w1[p]
elif proj == PROJ_UP:
off, n = bl.off_w3[p], bl.words_w1[p]
else:
off, n = bl.off_w2[p], bl.words_w2[p]
copy_from_host(dst[off:off + n].view(src.shape), src) # w1/w3 source is a strided column cut
elif kind in ("suh", "svh"):
if w.dtype != torch.float16:
raise ValueError(f"expert {e} {proj}.{kind}: dtype {w.dtype}, want float16")
rng = sl.suh_range if kind == "suh" else sl.svh_range
src = w if rng is None else w[rng[0]:rng[1]]
d = dst[p, 0 if proj == PROJ_GATE else 1] if gate_up else dst[p]
if tuple(d.shape) != tuple(src.shape):
raise ValueError(f"expert {e} {proj}.{kind}: {tuple(src.shape)} vs slot {tuple(d.shape)}")
copy_from_host(d, src)
else:
raise ValueError(f"bad tensor kind {kind!r}")
def views(self, p: int):
"""((trellis view, suh, svh, K) for w1, w3, w2) of physical slot p."""
bl, t = self.bl, self.t
b = self.lp.bits[self.lp.inv_perm[p]]
s1 = trellis_shape(self.hidden, self.inter_local, b.k_gu)
s2 = trellis_shape(self.inter_local, self.hidden, b.k_d)
w13, w2 = t["w13_trellis"], t["w2_trellis"]
return (
(w13[bl.off_w1[p]:bl.off_w1[p] + bl.words_w1[p]].view(s1), t["w13_suh"][p, 0], t["w13_svh"][p, 0], b.k_gu),
(w13[bl.off_w3[p]:bl.off_w3[p] + bl.words_w1[p]].view(s1), t["w13_suh"][p, 1], t["w13_svh"][p, 1], b.k_gu),
(w2[bl.off_w2[p]:bl.off_w2[p] + bl.words_w2[p]].view(s2), t["w2_suh"][p], t["w2_svh"][p], b.k_d),
)
def tables(self) -> tuple[list[GroupTables], tuple[ProjTables, ProjTables, ProjTables]]:
"""Device pointer tables (int64, physical slot order) and one GroupTables per K group.
Pointers are raw addresses: rebuild them if the tensors ever move."""
bl, t = self.bl, self.t
E, H, I = self.num_experts, self.hidden, self.inter_local
dev = t["w13_trellis"].device
b13, b2 = t["w13_trellis"].data_ptr(), t["w2_trellis"].data_ptr()
s13, v13 = t["w13_suh"].data_ptr(), t["w13_svh"].data_ptr()
s2, v2 = t["w2_suh"].data_ptr(), t["w2_svh"].data_ptr()
fb = 2 # int16 / fp16 bytes
for n in TENSOR_NAMES:
if not t[n].is_contiguous():
raise ValueError(f"{n} must be contiguous for raw pointer tables")
def tab(vals):
return torch.tensor(vals, dtype=torch.int64, device=dev)
g = ProjTables(tab([b13 + fb * o for o in bl.off_w1]),
tab([s13 + fb * (p * 2 + 0) * H for p in range(E)]),
tab([v13 + fb * (p * 2 + 0) * I for p in range(E)]))
u = ProjTables(tab([b13 + fb * o for o in bl.off_w3]),
tab([s13 + fb * (p * 2 + 1) * H for p in range(E)]),
tab([v13 + fb * (p * 2 + 1) * I for p in range(E)]))
d = ProjTables(tab([b2 + fb * o for o in bl.off_w2]),
tab([s2 + fb * p * I for p in range(E)]),
tab([v2 + fb * p * H for p in range(E)]))
groups = [GroupTables(grp.k_gu, grp.k_d, grp.lo, grp.hi,
g.slice(grp.lo, grp.hi), u.slice(grp.lo, grp.hi), d.slice(grp.lo, grp.hi))
for grp in self.lp.groups]
return groups, (g, u, d)
|