"""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: # Preserve the production per-column FP32 tree reduction and BF16 bias. # A small sub-tile bounds register use even on the uncached path. 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) # Every tile's payload precedes its acq_rel retirement. Only the final # retiring CTA continues: no grid barrier or assumption about scheduling. 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) # Release flag follows all payload lanes' system fences. 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) # Invalid rows retain an error flag and are discarded by the host. # Use an in-range internal continuation so a later sequential Markov # position cannot dereference token -1 before that host check. 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)) # A first rank must never wait on the GPU while the other rank JITs. # Compile both specializations before launching either. Warmups before # _capture_one prime every position; capture itself cannot compile. 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) # HYPOTHESIS (PP4 Round8): compile-only warmup leaves # Triton's CUDA handles lazy. Resolve both devices' # launchers before either peer can wait on the GPU. # Reading CompiledKernel.run loads, but never launches. 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"]