Spaces:
Running on Zero
Running on Zero
| """Triton kernels for FastDepth's fused frame body (fused=True). | |
| _add_ln x' = x + a [+ c[b]] with the same rounding order as torch (each add rounded to the depth dtype), stored; | |
| y = LayerNorm(x') (fp32 statistics, affine), stored. Replaces 2-3 elementwise kernels + a LayerNorm. | |
| _topk_sample per row: logits/temp -> top-k mask -> exponential-race sample (== argmax(p/E), E ~ Exp(1)), i.e. the same | |
| distribution as the torch path (topk + masked_fill + softmax + exponential_ + div + argmax), one kernel. | |
| RNG = Philox (tl.rand) keyed by a device-side seed that the captured graph increments every replay. | |
| """ | |
| import torch | |
| import triton | |
| import triton.language as tl | |
| def _add_ln(X, A, C, XO, Y, W, Bb, eps, T, | |
| sxb, sxt, sab, sat, scb, sob, sot, syb, syt, | |
| N: tl.constexpr, HAS_C: tl.constexpr): | |
| row = tl.program_id(0) | |
| b = row // T | |
| t = row % T | |
| c = tl.arange(0, N) | |
| dt = XO.dtype.element_ty | |
| x = (tl.load(X + b * sxb + t * sxt + c).to(tl.float32) + tl.load(A + b * sab + t * sat + c).to(tl.float32)).to(dt) | |
| if HAS_C: | |
| x = (x.to(tl.float32) + tl.load(C + b * scb + c).to(tl.float32)).to(dt) | |
| tl.store(XO + b * sob + t * sot + c, x) | |
| xf = x.to(tl.float32) | |
| mean = tl.sum(xf, 0) / N | |
| d = xf - mean | |
| var = tl.sum(d * d, 0) / N | |
| y = d / tl.sqrt(var + eps) * tl.load(W + c).to(tl.float32) + tl.load(Bb + c).to(tl.float32) | |
| tl.store(Y + b * syb + t * syt + c, y.to(Y.dtype.element_ty)) | |
| def add_ln(x, a, ln, c=None): | |
| """x, a: [B, T, D] (any strides, last dim contiguous); c: [B, 1, D] or None; ln: nn.LayerNorm(D). | |
| Returns (x + a [+ c], LayerNorm(...)) as new contiguous tensors.""" | |
| B, T, D = x.shape | |
| xo = torch.empty(B, T, D, device=x.device, dtype=x.dtype) | |
| y = torch.empty_like(xo) | |
| cc = c if c is not None else x | |
| _add_ln[(B * T,)](x, a, cc, xo, y, ln.weight, ln.bias, ln.eps, T, | |
| x.stride(0), x.stride(1), a.stride(0), a.stride(1), cc.stride(0), xo.stride(0), xo.stride(1), | |
| y.stride(0), y.stride(1), N=D, HAS_C=c is not None) | |
| return xo, y | |
| def _topk_sample(L, Out, Seed, Ctr, level, inv_temp, s_l, V: tl.constexpr, K: tl.constexpr, GREEDY: tl.constexpr, | |
| PER_ROW: tl.constexpr): | |
| row = tl.program_id(0) | |
| c = tl.arange(0, V) | |
| x = tl.load(L + row * s_l + c).to(tl.float32) * inv_temp | |
| xs = tl.sort(x, descending=True) | |
| kth = tl.sum(tl.where(c == K - 1, xs, 0.0), 0) | |
| xm = tl.where(x >= kth, x, -float("inf")) | |
| if GREEDY: | |
| score = xm | |
| else: | |
| if PER_ROW: # per-request stream: seed of this row's request, counter = frames generated by it | |
| seed = tl.load(Seed + row) | |
| ctr = tl.load(Ctr + row) | |
| u = tl.rand(seed, (ctr * 32 + level) * V + c) | |
| else: | |
| seed = tl.load(Seed) | |
| u = tl.rand(seed, (row * 64 + level) * V + c) | |
| u = tl.maximum(u, 1e-12) | |
| score = xm - tl.log(-tl.log(u)) | |
| tl.store(Out + row, tl.argmax(score, 0).to(tl.int64)) | |
| def topk_sample(logits, seed_t, level, temp, top_k, greedy=False, ctr_t=None): | |
| """seed_t [1] (engine-wide stream) or, with ctr_t, per-row seeds [B] and counters [B] (reproducible per request).""" | |
| B, V = logits.shape | |
| out = torch.empty(B, device=logits.device, dtype=torch.long) | |
| per_row = ctr_t is not None | |
| _topk_sample[(B,)](logits, out, seed_t, ctr_t if per_row else seed_t, level, 1.0 / temp, logits.stride(0), V=V, K=top_k, | |
| GREEDY=greedy, PER_ROW=per_row) | |
| return out | |
| def _row_exponential(Out, Seed, Ctr, level, s_o, V: tl.constexpr, BLOCK: tl.constexpr): | |
| """Out[row, :V] = Exp(1) noise from the row's own Philox stream (seed, counter): per-request reproducible sampling.""" | |
| row = tl.program_id(0) | |
| seed = tl.load(Seed + row) | |
| ctr = tl.load(Ctr + row) | |
| for start in range(0, V, BLOCK): | |
| c = start + tl.arange(0, BLOCK) | |
| u = tl.rand(seed, (ctr * 32 + level) * 4096 + c) | |
| u = tl.maximum(u, 1e-12) | |
| tl.store(Out + row * s_o + c, -tl.log(u), mask=c < V) | |
| def row_exponential(B, V, seeds, ctrs, level=0): | |
| out = torch.empty(B, V, device=seeds.device, dtype=torch.float32) | |
| _row_exponential[(B,)](out, seeds, ctrs, level, out.stride(0), V=V, BLOCK=1024) | |
| return out | |