Spaces:
Paused
Paused
| """Strict update-from-IO kernel — the counter update without a materialized weight-gradient. | |
| `fused_update.triton_counter_update` consumes a dense grad_w (formed by cuBLAS). This module forms | |
| grad_w[o,i] = sum_m grad_out[m,o] * x[m,i] INSIDE the kernel, per packed lane, so the dense [out,in] | |
| gradient is never materialized -- the strict memory-native update (the analogue of the engine's | |
| OpenCL apply_update_fused). One program per output row streams the M=batch*seq dimension, | |
| accumulates the four packed-lane grad_w vectors, then does exact RMS + scale + deterministic | |
| stochastic-rounding tick + re-pack in one pass. | |
| Determinism (hash-SR) makes it reproducible, so counter_update_from_io_hashsr is the CPU reference | |
| and the kernel is validated bit-quantified against it on a GPU (the FP reduction order differs from | |
| cuBLAS, so a handful of weights round the other way -- the same SR-boundary noise as fused_update). | |
| Honest scope: this is a MEMORY play, not a speed one -- forming grad_w in-kernel is a hand-written | |
| GEMM and loses to cuBLAS. Use it when not materializing grad_w matters; otherwise the cuBLAS-grad_w | |
| + fused_update path is faster. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| from .fused_update import HAS_TRITON, counter_update_hashsr | |
| __all__ = ["counter_update_from_io_hashsr", "HAS_TRITON", "triton_counter_update_from_io"] | |
| def counter_update_from_io_hashsr(codes: torch.Tensor, scale: torch.Tensor, v: torch.Tensor, | |
| x: torch.Tensor, grad_out: torch.Tensor, **kw) -> torch.Tensor: | |
| """CPU reference: form grad_w = grad_out^T x, then the deterministic-SR update. Identical math | |
| to the kernel (which forms grad_w in registers instead of as a tensor).""" | |
| grad_w = grad_out.transpose(0, 1).to(torch.float32) @ x.to(torch.float32) | |
| return counter_update_hashsr(codes, scale, v, grad_w, **kw) | |
| if HAS_TRITON: | |
| import triton | |
| import triton.language as tl | |
| from .fused_update import _tick | |
| def _counter_update_from_io_kernel(state_ptr, scale_ptr, v_ptr, x_ptr, go_ptr, seed_ptr, | |
| C, in_features, M, lr, lr_scale, rms_beta, rms_eps, | |
| BLOCK_G: tl.constexpr, BLOCK_M: tl.constexpr, | |
| LAGGED: tl.constexpr = False): | |
| # One program per output row. Forms grad_w[row, :] per packed lane (no dense grad_w), then | |
| # exact RMS + scale + SR tick + re-pack. state packed [out, (in/4)*3]; x [M,in]; go [M,out]. | |
| row = tl.program_id(0) | |
| seed = tl.load(seed_ptr).to(tl.uint32) | |
| gpr = in_features // 4 | |
| lv = 2 * C - 1 | |
| Cf = C * 1.0 | |
| gids = tl.arange(0, BLOCK_G) | |
| gmask = gids < gpr | |
| col0 = gids * 4 | |
| # accumulate the four lane grad_w vectors over the M dimension | |
| gw0 = tl.zeros((BLOCK_G,), tl.float32); gw1 = tl.zeros((BLOCK_G,), tl.float32) | |
| gw2 = tl.zeros((BLOCK_G,), tl.float32); gw3 = tl.zeros((BLOCK_G,), tl.float32) | |
| for m0 in range(0, M, BLOCK_M): | |
| moff = m0 + tl.arange(0, BLOCK_M) | |
| mmask = moff < M | |
| gocol = tl.load(go_ptr + moff * tl.num_programs(0) + row, mask=mmask, other=0.0) # [BLOCK_M] | |
| xb = x_ptr + moff[:, None] * in_features | |
| x0 = tl.load(xb + (col0 + 0)[None, :], mask=mmask[:, None] & gmask[None, :], other=0.0) | |
| x1 = tl.load(xb + (col0 + 1)[None, :], mask=mmask[:, None] & gmask[None, :], other=0.0) | |
| x2 = tl.load(xb + (col0 + 2)[None, :], mask=mmask[:, None] & gmask[None, :], other=0.0) | |
| x3 = tl.load(xb + (col0 + 3)[None, :], mask=mmask[:, None] & gmask[None, :], other=0.0) | |
| gw0 += tl.sum(gocol[:, None] * x0, axis=0) | |
| gw1 += tl.sum(gocol[:, None] * x1, axis=0) | |
| gw2 += tl.sum(gocol[:, None] * x2, axis=0) | |
| gw3 += tl.sum(gocol[:, None] * x3, axis=0) | |
| # decode the four lane ternary weights from the packed state | |
| base = row * gpr * 3 + gids * 3 | |
| b0 = tl.load(state_ptr + base + 0, mask=gmask, other=0).to(tl.int32) | |
| b1 = tl.load(state_ptr + base + 1, mask=gmask, other=0).to(tl.int32) | |
| b2 = tl.load(state_ptr + base + 2, mask=gmask, other=0).to(tl.int32) | |
| c0 = b0 & 0x3F | |
| c1 = ((b0 >> 6) | (b1 << 2)) & 0x3F | |
| c2 = ((b1 >> 4) | (b2 << 4)) & 0x3F | |
| c3 = (b2 >> 2) & 0x3F | |
| t0 = (c0 // lv - 1).to(tl.float32); t1 = (c1 // lv - 1).to(tl.float32) | |
| t2 = (c2 // lv - 1).to(tl.float32); t3 = (c3 // lv - 1).to(tl.float32) | |
| # exact row stats over the whole row (masked lanes contributed 0) | |
| g_sq = (tl.sum(gw0 * gw0) + tl.sum(gw1 * gw1) + tl.sum(gw2 * gw2) + tl.sum(gw3 * gw3)) / in_features | |
| grad_s = (tl.sum(gw0 * t0) + tl.sum(gw1 * t1) + tl.sum(gw2 * t2) + tl.sum(gw3 * t3)) / tl.sqrt(in_features.to(tl.float32)) | |
| s_old = tl.load(scale_ptr + row) | |
| v_old = tl.load(v_ptr + row) | |
| vv = rms_beta * v_old + (1.0 - rms_beta) * g_sq | |
| tl.store(v_ptr + row, vv) | |
| denom = tl.maximum(tl.sqrt(v_old if LAGGED else vv), rms_eps) | |
| s_new = tl.minimum(tl.maximum(s_old - lr_scale * grad_s, 1e-5), 10.0) | |
| # SR tick per lane, then re-pack 4 codes / 3 bytes | |
| nc0 = _tick(c0, gw0, row, col0 + 0, in_features, lv, C, Cf, lr, denom, s_old, s_new, seed) | |
| nc1 = _tick(c1, gw1, row, col0 + 1, in_features, lv, C, Cf, lr, denom, s_old, s_new, seed) | |
| nc2 = _tick(c2, gw2, row, col0 + 2, in_features, lv, C, Cf, lr, denom, s_old, s_new, seed) | |
| nc3 = _tick(c3, gw3, row, col0 + 3, in_features, lv, C, Cf, lr, denom, s_old, s_new, seed) | |
| p0 = (nc0 | (nc1 << 6)) & 0xFF | |
| p1 = ((nc1 >> 2) | (nc2 << 4)) & 0xFF | |
| p2 = ((nc2 >> 4) | (nc3 << 2)) & 0xFF | |
| tl.store(state_ptr + base + 0, p0.to(tl.uint8), mask=gmask) | |
| tl.store(state_ptr + base + 1, p1.to(tl.uint8), mask=gmask) | |
| tl.store(state_ptr + base + 2, p2.to(tl.uint8), mask=gmask) | |
| tl.store(scale_ptr + row, s_new) | |
| def triton_counter_update_from_io(state_packed: torch.Tensor, scale: torch.Tensor, v: torch.Tensor, | |
| x: torch.Tensor, grad_out: torch.Tensor, *, C: int, lr: float, | |
| lr_scale: float, rms_beta: float, rms_eps: float, seed: int, | |
| lagged: bool = False) -> None: | |
| """One-launch strict update: forms grad_w in-kernel from (x, grad_out) -- no dense gradient -- | |
| and applies the RMS + SR counter update. Mutates state_packed, scale, v in place. | |
| Bit-quantified-equal to counter_update_from_io_hashsr (verify on GPU). lagged=True uses the | |
| previous-step v for the denominator (per-element tick).""" | |
| if not HAS_TRITON: | |
| raise RuntimeError("triton not available") | |
| out = scale.numel() | |
| in_features = x.shape[1] | |
| M = x.shape[0] | |
| assert in_features % 4 == 0 | |
| gpr = in_features // 4 | |
| BLOCK_G = triton.next_power_of_2(gpr) | |
| seed_t = torch.tensor([int(seed) & 0xFFFFFFFF], dtype=torch.int64, device=x.device) | |
| _counter_update_from_io_kernel[(out,)]( | |
| state_packed, scale.reshape(out), v.reshape(out), x.contiguous(), grad_out.contiguous(), | |
| seed_t, C, in_features, M, lr, lr_scale, rms_beta, rms_eps, | |
| BLOCK_G=BLOCK_G, BLOCK_M=128, LAGGED=lagged, | |
| ) | |