| """Inference batching with multiple answer slots per input. |
| |
| Right-padded, bidirectional batches; padding never becomes an attention key. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def attention(module, x, rope, valid_keys): |
| bsz, length, _ = x.shape |
| q, k, v = module.qkv(x).reshape(bsz, length, 3, module.num_heads, module.head_dim).unbind(2) |
| q, k = rope.apply_bshd(q, k) |
| q, k, v = [a.permute(0, 2, 1, 3).contiguous() for a in (q, k, v)] |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=valid_keys, |
| dropout_p=0.0, is_causal=False) |
| return module.proj_drop(module.proj(out.permute(0, 2, 1, 3).contiguous().reshape(bsz, length, module.attn_dim))) |
|
|
|
|
| @torch.inference_mode() |
| def batch_logits(model, sequences, positions, timestep=0.5, pad_id=1): |
| if model.training or not sequences or len(sequences) != len(positions): |
| raise ValueError("Batch requires an eval model and answer slots per input") |
| if any(not ps or any(not 0 <= p < len(s) for p in ps) for s, ps in zip(sequences, positions)): |
| raise ValueError("Answer slot outside input") |
| device = next(model.parameters()).device |
| width = max(map(len, sequences)) |
| x = torch.full((len(sequences), width), pad_id, dtype=torch.long, device=device) |
| for i, seq in enumerate(sequences): |
| x[i, :len(seq)] = torch.tensor(seq, dtype=torch.long, device=device) |
| lengths = torch.tensor([len(s) for s in sequences], device=device) |
| keys = (torch.arange(width, device=device)[None, :] < lengths[:, None])[:, None, None, :] |
| hidden = model.token_embed(x) |
| condition = model.time_embed(torch.full((len(sequences),), timestep, device=device, dtype=torch.float32)) |
| for block in model.blocks: |
| s1, c1, g1, s2, c2, g2 = block.adaLN(condition).chunk(6, -1) |
| normalized = block.norm1(hidden) * (1 + c1[:, None]) + s1[:, None] |
| hidden = hidden + g1[:, None] * attention(block.attn, normalized, model.rope, keys) |
| normalized = block.norm2(hidden) * (1 + c2[:, None]) + s2[:, None] |
| hidden = hidden + g2[:, None] * block.mlp(normalized) |
| hidden = model.final_ada(hidden, condition) |
| rows = [i for i, ps in enumerate(positions) for _ in ps] |
| columns = [p for ps in positions for p in ps] |
| slots = hidden[torch.tensor(rows, device=device), torch.tensor(columns, device=device)] |
| logits = model.logits_from_hidden(slots).float() |
| if not torch.isfinite(logits).all(): |
| raise RuntimeError("Non-finite batched answer logits") |
| return logits |
|
|
