| """Single-launch Markov/top32/dyadic24 selector for the explicit experiment. |
| |
| Each tile publishes 32 candidates. The last retiring tile reduces them and |
| draws, so no block waits for unscheduled blocks. On TP2 each last tile exchanges |
| only 32 candidates with its peer, with bounded release/acquire and consumption |
| acknowledgment. Emulation reduces both shards in the same launch. No conditioned |
| vocabulary tensor is allocated. The original head logits remain BF16. |
| """ |
| from __future__ import annotations |
|
|
| from collections.abc import Iterator |
| from contextlib import contextmanager |
| from typing import Any |
|
|
| import torch |
| import triton |
| import triton.language as tl |
|
|
| from .sampled import _pack, _unpack |
|
|
|
|
| @triton.jit |
| def _clock() -> tl.tensor: |
| return tl.inline_asm_elementwise("mov.u64 $0, %globaltimer;", constraints="=l", args=[], |
| dtype=tl.uint64, is_pure=False, pack=1) |
|
|
|
|
| @triton.jit |
| def _acquire(p: tl.tensor) -> tl.tensor: |
| return tl.inline_asm_elementwise("ld.acquire.sys.global.u32 $0, [$1];", constraints="=r,l", args=[p], |
| dtype=tl.uint32, is_pure=False, pack=1) |
|
|
|
|
| @triton.jit |
| def _release(p: tl.tensor, v: tl.tensor) -> None: |
| tl.inline_asm_elementwise("st.release.sys.global.u32 [$1], $2; mov.u32 $0, 0;", |
| constraints="=r,l,r", args=[p, v.to(tl.uint32)], dtype=tl.uint32, |
| is_pure=False, pack=1) |
|
|
|
|
| @triton.jit |
| def _wait(p: tl.tensor, epoch: tl.tensor, error: tl.tensor, timeout: tl.constexpr) -> tl.tensor: |
| start = _clock() |
| value = _acquire(p) |
| err = _acquire(error) |
| while (value != epoch) & (err == 0) & (_clock() - start < timeout): |
| value = _acquire(p) |
| err = _acquire(error) |
| return (value == epoch) & (err == 0) |
|
|
|
|
| @triton.jit |
| def _kernel(L0: tl.tensor, L1: tl.tensor, P0: tl.tensor, P1: tl.tensor, |
| S0: tl.tensor, S1: tl.tensor, C0: tl.tensor, C1: tl.tensor, |
| W10: tl.tensor, W11: tl.tensor, W20: tl.tensor, W21: tl.tensor, |
| KEYS: tl.tensor, MASK: tl.tensor, CAND: tl.tensor, BAD: tl.tensor, |
| DONE: tl.tensor, EPOCH: tl.tensor, FLAGS: tl.tensor, ACK: tl.tensor, STICKY: tl.tensor, |
| INBOX: tl.tensor, PEER_INBOX: tl.tensor, PEER_FLAGS: tl.tensor, |
| PEER_ACK: tl.tensor, PEER_ERROR: tl.tensor, |
| OUT0: tl.tensor, OUT1: tl.tensor, IDS: tl.tensor, Q: tl.tensor, ERROR: tl.tensor, |
| V: tl.constexpr, V_SECOND: tl.constexpr, V0: tl.constexpr, V_SECOND_OFFSET: tl.constexpr, |
| LS: tl.constexpr, LS_SECOND: tl.constexpr, PS: tl.constexpr, |
| CS: tl.constexpr, CS_SECOND: tl.constexpr, WS: tl.constexpr, RANK_DIM: tl.constexpr, |
| DEPTH: tl.constexpr, POS: tl.constexpr, NT: tl.constexpr, |
| PAD: tl.constexpr, EMULATED: tl.constexpr, TILE: tl.constexpr, |
| TIMEOUT: tl.constexpr) -> None: |
| b, tile = tl.program_id(0), tl.program_id(1) |
| shard = tile // NT |
| t = tile % NT |
| if EMULATED and shard == 1: |
| base, prev, slots, cache, w1, w2 = L1, P1, S1, C1, W11, W21 |
| vocabulary, base_stride, cache_stride, offset = V_SECOND, LS_SECOND, CS_SECOND, V_SECOND_OFFSET |
| else: |
| base, prev, slots, cache, w1, w2 = L0, P0, S0, C0, W10, W20 |
| vocabulary, base_stride, cache_stride, offset = V, LS, CS, V0 |
| tok = tl.load(prev + b * PS).to(tl.int64) |
| slot = tl.load(slots + tok).to(tl.int64) |
| j = tl.arange(0, TILE) |
| cols = t * TILE + j |
| if slot >= 0: |
| lg = tl.load(base + b * base_stride + cols, cols < vocabulary, 0.).to(tl.float32) |
| bias = tl.load(cache + slot * cache_stride + cols, cols < vocabulary, 0.).to(tl.float32) |
| cached_x = (lg + bias).to(tl.bfloat16).to(tl.float32) |
| bad = tl.sum(((tl.extra.cuda.libdevice.isnan(cached_x) | (cached_x == float("inf"))) & (cols < vocabulary)).to(tl.int32), 0) |
| cached_x = tl.where(cols < vocabulary, cached_x, float("-inf")) |
| cached_x = tl.where(cached_x == 0, 0., cached_x) |
| top = tl.topk(_pack(cached_x, cols + offset), 32) |
| else: |
| |
| |
| rr = tl.arange(0, RANK_DIM) |
| emb = tl.load(w1 + tok * RANK_DIM + rr).to(tl.float32) |
| top = tl.full((32,), 0, tl.uint64) |
| bad = tl.full((), 0, tl.int32) |
| sj = tl.arange(0, 64) |
| for sub in range(TILE // 64): |
| scols = t * TILE + sub * 64 + sj |
| weight = tl.load(w2 + scols[:, None].to(tl.int64) * WS + rr[None, :], |
| scols[:, None] < vocabulary, 0.).to(tl.float32) |
| ubias = tl.sum(weight * emb[None, :], 1).to(tl.bfloat16).to(tl.float32) |
| ul = tl.load(base + b * base_stride + scols, scols < vocabulary, 0.).to(tl.float32) |
| ux = (ul + ubias).to(tl.bfloat16).to(tl.float32) |
| bad += tl.sum(((tl.extra.cuda.libdevice.isnan(ux) | (ux == float("inf"))) & (scols < vocabulary)).to(tl.int32), 0) |
| ux = tl.where(scols < vocabulary, ux, float("-inf")) |
| ux = tl.where(ux == 0, 0., ux) |
| partial = tl.topk(_pack(ux, scols + offset), 32) |
| top = tl.topk(tl.reshape(tl.join(top, partial), (64,)), 32) |
| r = tl.arange(0, 32) |
| count: tl.constexpr = NT * (2 if EMULATED else 1) |
| tl.store(CAND + (b * count + tile) * 32 + r, top) |
| tl.store(BAD + b * count + tile, bad) |
| |
| |
| tl.inline_asm_elementwise("fence.sc.gpu; mov.u64 $0, $1;", constraints="=l,l", args=[top], |
| dtype=tl.uint64, is_pure=False, pack=1) |
| tl.debug_barrier() |
| retired = tl.atomic_add(DONE + b, 1, sem="acq_rel", scope="gpu") |
| if retired == count - 1: |
| jj = tl.arange(0, PAD) |
| keys = tl.load(CAND + b * count * 32 + jj, jj < count * 32, 0) |
| local = tl.topk(keys, 32) |
| errs = tl.load(BAD + b * count + jj, jj < count, 0) |
| valid = tl.sum(errs, 0) == 0 |
| if not EMULATED: |
| epoch = tl.load(EPOCH + b).to(tl.uint32) |
| peer_inbox = tl.cast(PEER_INBOX, tl.pointer_type(tl.uint64)) |
| peer_flags = tl.cast(PEER_FLAGS, tl.pointer_type(tl.uint32)) |
| peer_ack = tl.cast(PEER_ACK, tl.pointer_type(tl.uint32)) |
| peer_error = tl.cast(PEER_ERROR, tl.pointer_type(tl.uint32)) |
| if not valid: |
| tl.atomic_or(STICKY + b, 2, sem="relaxed", scope="sys") |
| tl.atomic_or(peer_error + b, 2, sem="relaxed", scope="sys") |
| healthy = (_acquire(STICKY + b) == 0) & (epoch != 0xffffffff) |
| if healthy: |
| tl.store(peer_inbox + b * 32 + r, local) |
| |
| tl.inline_asm_elementwise("fence.sc.sys; mov.u64 $0, $1;", constraints="=l,l", args=[local], |
| dtype=tl.uint64, is_pure=False, pack=1) |
| tl.debug_barrier() |
| _release(peer_flags + b, epoch) |
| healthy = _wait(FLAGS + b, epoch, STICKY + b, TIMEOUT) |
| if healthy: |
| merge_j = tl.arange(0, 64) |
| own = tl.gather(local, merge_j % 32, 0) |
| peer = tl.load(INBOX + b * 32 + merge_j % 32, volatile=True) |
| local = tl.topk(tl.where(merge_j < 32, own, peer), 32) |
| else: |
| tl.atomic_or(STICKY + b, 1, sem="relaxed", scope="sys") |
| tl.atomic_or(peer_error + b, 1, sem="relaxed", scope="sys") |
| valid = valid & healthy |
| selected_logits, support = _unpack(local) |
| positive = tl.load(MASK + b) |
| valid = valid & (tl.max(selected_logits, 0) > float("-inf")) |
| mx = tl.max(selected_logits, 0) |
| w = tl.exp((selected_logits - mx) / .7).to(tl.float64) |
| tickets = tl.floor(w / tl.sum(w, 0) * 16777216.).to(tl.int32) |
| left = 16777216 - tl.sum(tickets, 0) |
| valid = valid & (left >= 0) & (left < 32) |
| tickets += tl.where(r == 0, left, 0) |
| token = tl.sum(tl.where(r == 0, support, 0), 0) |
| packed = tl.sort((support.to(tl.uint64) << 32) | tickets.to(tl.uint32).to(tl.uint64), descending=False) |
| support = (packed >> 32).to(tl.int32) |
| tickets = packed.to(tl.int32) |
| if positive & valid: |
| seed = tl.load(KEYS + b * 3).to(tl.uint64) |
| rnd = tl.load(KEYS + b * 3 + 1).to(tl.uint64) |
| pos = tl.load(KEYS + b * 3 + 2).to(tl.uint64) + POS |
| u, _, _, _ = tl.philox(seed, rnd.to(tl.uint32), (rnd >> 32).to(tl.uint32), |
| pos.to(tl.uint32), (pos >> 32).to(tl.uint32) ^ 0x50503301, 10) |
| ticket = (u & 0xffffff).to(tl.int32) |
| chosen = tl.min(tl.where(tl.cumsum(tickets, 0) > ticket, r, 32), 0) |
| token = tl.sum(tl.where(r == chosen, support, 0), 0) |
| |
| |
| |
| tl.store(OUT0 + b * DEPTH + POS, tl.where(positive & ~valid, 0, token)) |
| if EMULATED: |
| tl.store(OUT1 + b * DEPTH + POS, tl.where(positive & ~valid, 0, token)) |
| tl.store(IDS + (b * DEPTH + POS) * 32 + r, support) |
| tl.store(Q + (b * DEPTH + POS) * 32 + r, |
| tl.where(positive & valid, tickets.to(tl.float32) / 16777216., 0.)) |
| tl.store(ERROR + b * DEPTH + POS, (positive & ~valid).to(tl.int32)) |
| if not EMULATED: |
| if healthy: |
| tl.debug_barrier() |
| _release(peer_ack + b, epoch) |
| healthy = _wait(ACK + b, epoch, STICKY + b, TIMEOUT) |
| if not healthy: |
| tl.atomic_or(STICKY + b, 1, sem="relaxed", scope="sys") |
| tl.atomic_or(peer_error + b, 1, sem="relaxed", scope="sys") |
| tl.store(ERROR + b * DEPTH + POS, 1) |
| tl.store(EPOCH + b, epoch + 1) |
| tl.store(DONE + b, 0) |
|
|
|
|
| class FusedContext: |
| """Fixed-address per-shape state. Own it for one worker, never concurrently. |
| |
| Peer access comes from the worker's existing qualified P2P communicator. |
| Sticky peer errors require worker teardown; they are never cleared on retry. |
| """ |
|
|
| def __init__(self, model: Any, batch: int, depth: int, tile: int = 1024): |
| if tuple(model.ALL) != (0, 1) or not 1 <= batch <= 4 or depth not in (4, 5): |
| raise ValueError("fused sampled selection requires full K2 TP2, B1..4, depth4/5") |
| if tile not in (512, 1024, 2048): |
| raise ValueError("sampled tile must be 512, 1024 or 2048") |
| self.model, self.batch, self.depth, self.tile = model, batch, depth, tile |
| self.state: dict[int, dict[str, torch.Tensor]] = {} |
| self.primed: set[tuple[int, int, int, int]] = set() |
| self.capture_expected: frozenset[int] = frozenset() |
| v = max(rk.v1 - rk.v0 for rk in model.ranks) |
| self.nt = triton.cdiv(v, tile) |
| for rank in model.ALL: |
| dev = model.devs[rank] |
| count = self.nt * (2 if model.emulated else 1) |
| self.state[rank] = { |
| "keys": torch.zeros(batch, 3, dtype=torch.int64, device=dev), |
| "mask": torch.ones(batch, dtype=torch.bool, device=dev), |
| "cand": torch.empty(batch, count * 32, dtype=torch.uint64, device=dev), |
| "bad": torch.empty(batch, count, dtype=torch.int32, device=dev), |
| "done": torch.zeros(batch, dtype=torch.int32, device=dev), |
| "epoch": torch.ones(batch, dtype=torch.int32, device=dev), |
| "flags": torch.zeros(batch, dtype=torch.int32, device=dev), |
| "ack": torch.zeros(batch, dtype=torch.int32, device=dev), |
| "sticky": torch.zeros(batch, dtype=torch.int32, device=dev), |
| "inbox": torch.empty(batch, 32, dtype=torch.uint64, device=dev), |
| "ids": torch.zeros(batch, depth, 32, dtype=torch.int32, device=dev), |
| "q": torch.zeros(batch, depth, 32, dtype=torch.float32, device=dev), |
| "error": torch.zeros(batch, depth, dtype=torch.int32, device=dev), |
| } |
|
|
| @contextmanager |
| def capturing(self, ranks: tuple[int, ...]) -> Iterator[None]: |
| """Declare capture intent independently of CUDA's current-stream state. |
| |
| A graph accidentally capturing another device must never launch a peer |
| kernel eagerly: its absent partner would poison both selector states. |
| """ |
| if self.capture_expected: |
| raise RuntimeError("nested sampled capture is not supported") |
| self.capture_expected = frozenset((0,) if self.model.emulated else ranks) |
| try: |
| yield |
| finally: |
| self.capture_expected = frozenset() |
|
|
| def update(self, keys: torch.Tensor, mask: torch.Tensor) -> None: |
| if keys.shape != (self.batch, 3) or keys.dtype != torch.int64 or mask.shape != (self.batch,) or mask.dtype != torch.bool: |
| raise ValueError("fused sampled RNG/mask ABI mismatch") |
| for r, state in self.state.items(): |
| with torch.cuda.device(self.model.devs[r]): |
| state["keys"].copy_(keys, non_blocking=True) |
| state["mask"].copy_(mask, non_blocking=True) |
|
|
| def select(self, ranks: Any, logits: dict[int, torch.Tensor], previous: dict[int, torch.Tensor], |
| outputs: dict[int, torch.Tensor], position: int) -> None: |
| model = self.model |
| active = (0,) if model.emulated else ranks |
| if self.capture_expected: |
| if frozenset(active) != self.capture_expected: |
| raise RuntimeError("sampled capture rank ownership differs") |
| for r in active: |
| with torch.cuda.device(model.devs[r]): |
| if not torch.cuda.is_current_stream_capturing(): |
| raise RuntimeError(f"sampled rank {r}: expected capture on its current device stream; refusing eager peer launch") |
| prepared: list[tuple[int, tuple[Any, ...], tuple[int, int, int, int]]] = [] |
| for r in active: |
| rk, other = model.ranks[r], model.ranks[1 if model.emulated else r] |
| state, peer = self.state[r], self.state[1 - r] |
| v = rk.v1 - rk.v0 |
| second = 1 if model.emulated else r |
| base = logits[r][:, position] |
| args = ( |
| base, logits[second][:, position], previous[r], previous[second], |
| rk.markov.slot_of_token, other.markov.slot_of_token, rk.markov.cache, other.markov.cache, |
| rk.markov.w1, other.markov.w1, rk.markov.w2, other.markov.w2, |
| state["keys"], state["mask"], state["cand"], state["bad"], state["done"], state["epoch"], |
| state["flags"], state["ack"], state["sticky"], state["inbox"], |
| peer["inbox"].data_ptr(), peer["flags"].data_ptr(), peer["ack"].data_ptr(), peer["sticky"].data_ptr(), |
| outputs[r], outputs[second], state["ids"], state["q"], state["error"], |
| v, other.v1 - other.v0, rk.v0, other.v0, base.stride(0), logits[second][:, position].stride(0), |
| previous[r].stride(0), rk.markov.cache.stride(0), other.markov.cache.stride(0), rk.markov.w2.stride(0), |
| rk.markov.w1.shape[1], self.depth, position, self.nt, |
| triton.next_power_of_2(self.nt * (2 if model.emulated else 1) * 32), |
| model.emulated, self.tile, 50_000_000) |
| prepared.append((r, args, (r, position, base.stride(0), previous[r].stride(0)))) |
| grid = (self.batch, self.nt * (2 if model.emulated else 1)) |
| |
| |
| |
| kernel: Any = _kernel |
| if not model.emulated: |
| for r, args, key in prepared: |
| if key not in self.primed: |
| with torch.cuda.device(model.devs[r]): |
| if torch.cuda.is_current_stream_capturing() or tuple(ranks) != tuple(model.ALL): |
| raise RuntimeError("paired sampled kernels must be primed together before capture") |
| compiled = kernel.warmup(*args, grid=grid, num_warps=4, enable_fp_fusion=False) |
| |
| |
| |
| |
| if compiled is None or not callable(compiled.run): |
| raise RuntimeError("paired sampled warmup did not load CUDA launch handles") |
| self.primed.add(key) |
| for r, args, _ in prepared: |
| with torch.cuda.device(model.devs[r]): |
| kernel[grid](*args, num_warps=4, enable_fp_fusion=False) |
|
|
| def result(self) -> tuple[torch.Tensor, torch.Tensor]: |
| for state in self.state.values() if not self.model.emulated else (self.state[0],): |
| if torch.any(state["error"]).item() or torch.any(state["sticky"]).item(): |
| raise RuntimeError("fused sampled selector failed; discard reply and tear down worker") |
| return self.state[0]["ids"], self.state[0]["q"] |
|
|
