Spaces:
Sleeping
Sleeping
File size: 3,081 Bytes
3afc977 | 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 | """Block diffusion sampling with a KV cache (Nemotron-style; Arriola Block
Diffusion).
The masked region is generated block by block, left to right. The clean context
(prefix+suffix) is encoded once into a per-layer K/V cache; each committed block
appends its K/V. A block is refined for `n_inner` ReMDM steps against the cache —
bidirectional within the block, revising and re-masking its weakest tiles — then
committed. Per inner step we only recompute the current block's queries against
the cache, not the whole sequence: that is the latency win AR-style decoding gets
from its cache and plain diffusion does not.
Latency knobs: block_len (tokens per block) and n_inner (refinement steps).
"""
from __future__ import annotations
import torch
import torch.nn.functional as F
@torch.no_grad()
def sample(model, ids_row, region_row, attn_row, tok, cfg, n_inner: int):
"""ids_row,region_row,attn_row: 1-D tensors of length T. Returns the predicted
region tokens in order (list of ids)."""
device = ids_row.device
pos_all = torch.arange(ids_row.size(0), device=device)
# Clean context = real, non-region positions (bos, prefix, suffix, eos).
ctx_sel = attn_row & (~region_row)
ctx_idx = pos_all[ctx_sel]
caches = model.encode_context(ids_row[ctx_idx].unsqueeze(0), ctx_idx.unsqueeze(0))
region_idx = pos_all[region_row]
R = region_idx.numel()
tile = cfg.tile_size
result = {}
for b0 in range(0, R, cfg.block_len):
blk_pos = region_idx[b0 : b0 + cfg.block_len]
Lb = blk_pos.numel()
blk = torch.full((Lb,), tok.mask_id, dtype=torch.long, device=device)
masked = torch.ones(Lb, dtype=torch.bool, device=device)
for inner in range(n_inner):
logits, _ = model.decode_block(blk.unsqueeze(0), blk_pos.unsqueeze(0), caches)
probs = F.softmax(logits[0], dim=-1)
conf, pred = probs.max(dim=-1) # (Lb,)
# Fill ONLY the currently masked positions; keep committed tokens.
blk = torch.where(masked, pred, blk)
masked = torch.zeros(Lb, dtype=torch.bool, device=device)
if inner == n_inner - 1:
break
# Re-mask the lowest-confidence tiles (revision) for the next step.
frac_masked = 1.0 - (inner + 1) / n_inner
n_tiles = (Lb + tile - 1) // tile
keep_masked = round(n_tiles * frac_masked)
if keep_masked <= 0:
continue
tconf = torch.stack([conf[t * tile : (t + 1) * tile].mean() for t in range(n_tiles)])
for t in torch.argsort(tconf)[:keep_masked].tolist():
lo, hi = t * tile, min((t + 1) * tile, Lb)
blk[lo:hi] = tok.mask_id
masked[lo:hi] = True
# Commit: append this block's (final) K/V to the cache for later blocks.
_, caches = model.decode_block(blk.unsqueeze(0), blk_pos.unsqueeze(0), caches)
for k in range(Lb):
result[int(blk_pos[k])] = int(blk[k])
return [result[int(p)] for p in region_idx]
|