← Walkthroughmapa/models/varlen.py
Source snapshot · source SHA-256 1fea5b0cf2a1
1"""Block-diagonal attention over a flat, ragged token set.23Main path: _reference_block_diag() on CPU; gpu_block_diag() on CUDA.45Input: Q/K/V (M,H,hd) and sequence boundaries.6Output: context (M,H,hd), each token combining values only from its own block.7Next: attention.py joins the heads and projects back to the model width.89Tokens from every array in a session are laid end to end with no padding, and ``cu_seqlens``10marks where each array's block starts. A token attends inside its own block and nowhere11else. Two backends compute the same thing.1213On GPU the packed values are viewed as a jagged nested tensor, which routes scaled dot14product attention to its ragged kernel. Nothing is padded and no mask is built, so cost15follows the true block sizes. On CPU the block-diagonal mask is materialised and handed to16dense attention, which is only tractable at test sizes and is what the tests compare17against.18"""1920from __future__ import annotations2122import torch23import torch.nn.functional as F24from torch import Tensor2526# Finite rather than -inf: a row with no visible key softmaxes to uniform instead of NaN.27NEG_INF_MASK = -1e4282930def _segment_ids(cu_seqlens: Tensor, total: int) -> Tensor:31 """Cumulative block bounds ``(n_block+1,)`` to a block id per token ``(total,)``.3233 An empty block contributes nothing, which is exactly the number of tokens it holds.34 """35 lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long)36 seg = torch.arange(lengths.shape[0], device=cu_seqlens.device)37 return seg.repeat_interleave(lengths)383940def _reference_block_diag(q: Tensor, k: Tensor, v: Tensor, cu_seqlens: Tensor) -> Tensor:41 """Dense attention under a block-diagonal mask. CPU only: the mask is ``(total, total)``.4243 q, k, v are ``(total, n_heads, head_dim)``. Returns the same shape.44 """45 total = q.shape[0]46 seg = _segment_ids(cu_seqlens, total)47 same = seg[:, None] == seg[None, :]48 bias = torch.where(same, 0.0, NEG_INF_MASK).to(q.dtype)49 # Per head/block: softmax(Q @ K.T / sqrt(head_dim), dim=-1) @ V.50 # Rows are queries; softmax normalizes over key tokens.51 ctx = F.scaled_dot_product_attention(52 q.transpose(0, 1)[None], k.transpose(0, 1)[None], v.transpose(0, 1)[None],53 attn_mask=bias[None, None],54 )55 return ctx[0].transpose(0, 1).contiguous()565758def gpu_block_diag(59 q: Tensor, k: Tensor, v: Tensor, cu_seqlens_drop: Tensor, max_seqlen: int60) -> Tensor:61 """Block-diagonal attention through a jagged nested tensor. No pad and no mask.6263 ``nested_tensor_from_jagged`` views the packed ``(total, n_heads, head_dim)`` values as64 ``(n_block, len_i, n_heads, head_dim)``, so attention dispatches to the ragged kernel.65 The view preserves gradients, unlike the nested tensor constructor, which detaches.6667 ``cu_seqlens_drop`` is int64 and carries empty blocks already removed, because an empty68 jagged row makes the ragged backward produce NaN. Passing ``max_seqlen`` explicitly69 skips the kernel's own reduction over the offsets, which would sync the device on every70 call. The softmax scale is the default ``1/sqrt(head_dim)``, which is what the reference71 backend applies too.72 """7374 def _nt(x: Tensor) -> Tensor:75 return torch.nested.nested_tensor_from_jagged(76 x, offsets=cu_seqlens_drop, min_seqlen=1, max_seqlen=max_seqlen77 ).transpose(1, 2)7879 # Same within-block attention, without allocating a full total-by-total mask.80 ctx = F.scaled_dot_product_attention(_nt(q), _nt(k), _nt(v))81 return ctx.transpose(1, 2).values()