Spaces:
Sleeping
Sleeping
| """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 | |
| 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] | |