File size: 2,422 Bytes
e65937c | 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 | """Utility functions for DFlash draft model."""
import torch
def get_base_indices_for_anchored_blocks(
anchor_positions: torch.Tensor, # shape: [1, num_anchors]
block_size: int,
) -> torch.Tensor: # shape: [num_anchors*block_size]
anchor_positions = anchor_positions.to(dtype=torch.long).view(-1)
# dtype: long, shape: [num_anchors]
offsets = torch.arange(block_size, device=anchor_positions.device, dtype=torch.long)
idx = (
anchor_positions[:, None] + offsets[None, :]
) # shape: [num_anchors, block_size]
return idx.reshape(-1)
def select_anchors(
loss_mask: torch.Tensor, # shape: [1, total_seq_len]
num_anchors: int,
block_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Randomly select anchor positions from valid tokens in sequence.
Args:
loss_mask: Binary mask indicating valid positions [1, total_seq_len]
n: Number of anchors to select per batch item
block_size: Block size (last block_size positions excluded)
Returns:
tuple: (anchors, anchor_valid)
- anchors: Selected anchor indices [num_anchors]
- anchor_valid: Boolean mask for valid anchors [num_anchors]
"""
if loss_mask.ndim != 2: # noqa: PLR2004
raise ValueError(f"Expected [B, T], got {loss_mask.shape}")
if block_size <= 0:
raise ValueError(f"Expected block size > 0, got {block_size}")
valid_mask = loss_mask.bool().clone()
valid_mask[:, -block_size:] = False
valid_indices = torch.nonzero(valid_mask.squeeze(0), as_tuple=False).squeeze(
-1
) # shape: [num_non_zero]
device = loss_mask.device
anchors = torch.zeros(num_anchors, dtype=torch.long, device=device)
anchor_valid = torch.zeros(num_anchors, dtype=torch.bool, device=device)
k = min(num_anchors, valid_indices.numel())
# Constrain value of k for torch dynamo
torch._check(k <= valid_indices.numel()) # noqa: SLF001
torch._check(k >= 0) # noqa: SLF001
perm = torch.randperm(valid_indices.numel(), device=loss_mask.device)
# Contiguous anchors let flex attention use dense (fast) blocks instead of
# scattered all-partial (slow) ones; the order never affects the loss.
anchors[:k] = torch.sort(torch.gather(valid_indices, 0, perm[:k])).values
anchor_valid[:k] = True
return anchors, anchor_valid
# shape: [num_anchors], [num_anchors]
|