File size: 4,940 Bytes
073de96 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Row lengths and positions: the bucket of a record, its three M-RoPE position planes and the cos / sin tables the row
graph takes as inputs (the graph holds no COS / SIN / position table of its own).
Buckets: one row graph per length L in BUCKETS; a record of T tokens runs on the smallest L >= T, right-padded. A
record longer than the largest graph is refused (`RequestTooLong`), never truncated.
Position planes [3, T] (t, h, w) of the author's model (transformers Qwen3_5Model.get_rope_index; closed form recorded
with the reference, P = token_offset = the <|vision_start|> index, one picture of H x W merged tokens, N = H * W):
text-only record: (i, i, i)
i <= P -> (i, i, i); picture token k at P+1+k -> (P+1, P+1 + k // W, P+1 + k % W);
i >= P+N+1 -> i - (N - max(H, W)) on all three planes
Pad positions (right of the record) continue the planes' maximum + 1, + 2, ... on all three planes; their outputs are
never read and the graph's `valid` mask hides them as keys.
cos / sin (the author's Qwen3_5TextRotaryEmbedding: default rope, theta 1e7, rotary dim 64 of the 256-wide heads,
attention scaling 1, interleaved mrope_section [11, 11, 10]):
inv_freq[j] = 1 / theta^(2j / 64) j = 0..31, every step in float32 as the author's code
angle[p][i, j] = float32(inv_freq[j] * pos[p][i]) p = t, h, w
interleave: j % 3 == 1 and j < 33 -> h plane, j % 3 == 2 and j < 30 -> w plane, every other j -> t plane
cos = cos(cat(angle, angle)), sin = sin(cat(angle, angle)) float32 [1, L, 64]
The cosine and sine are taken in float64 on the float32 angle and rounded to float32 (torch's float32 cos / sin sit
within 1 ulp of these)."""
from __future__ import annotations
import numpy as np
from .errors import RequestTooLong
BUCKETS = (256, 512, 1024, 2048, 4096)
ROTARY_DIM = 64
ROPE_THETA = 10000000.0
MROPE_SECTION = (11, 11, 10)
def pick_length(tokens: int, buckets=BUCKETS) -> int:
"""The smallest row graph length that holds `tokens`; RequestTooLong past the largest."""
for length in sorted(buckets):
if tokens <= length:
return length
raise RequestTooLong(f"the record has {tokens} tokens; the largest row graph holds {max(buckets)}")
def text_positions(tokens: int) -> np.ndarray:
"""[3, T] int64 planes of a record without a picture: (i, i, i)."""
return np.broadcast_to(np.arange(tokens, dtype=np.int64), (3, tokens)).copy()
def image_positions(tokens: int, token_offset: int, h: int, w: int) -> np.ndarray:
"""[3, T] int64 planes of a record with one picture of h x w merged tokens right after the prefix."""
p, n = int(token_offset), int(h) * int(w)
if not p + 1 + n < tokens:
raise ValueError(f"a picture of {n} tokens at {p} does not fit a record of {tokens} tokens")
pos = text_positions(tokens)
k = np.arange(n, dtype=np.int64)
pos[0, p + 1:p + 1 + n] = p + 1
pos[1, p + 1:p + 1 + n] = p + 1 + k // w
pos[2, p + 1:p + 1 + n] = p + 1 + k % w
pos[:, p + n + 1:] = np.arange(p + n + 1, tokens, dtype=np.int64) - (n - max(h, w))
return pos
def record_positions(tokens: int, token_offset: int | None = None, image_grid=None) -> np.ndarray:
"""[3, T] planes of a record (text-only when `image_grid` is None)."""
if image_grid is None:
return text_positions(tokens)
return image_positions(tokens, token_offset, *image_grid)
def padded_positions(pos3: np.ndarray, length: int) -> np.ndarray:
"""[3, L]: the record's planes, then maximum + 1, + 2, ... on all three planes for the pad positions."""
n = pos3.shape[1]
if n > length:
raise ValueError(f"{n} positions do not fit length {length}")
out = np.empty((3, length), dtype=np.int64)
out[:, :n] = pos3
out[:, n:] = int(pos3.max()) + 1 + np.arange(length - n, dtype=np.int64)[None, :]
return out
def inv_freq(theta: float = ROPE_THETA, rotary_dim: int = ROTARY_DIM) -> np.ndarray:
"""float32 [rotary_dim / 2] = 1 / theta^(arange(0, dim, 2) / dim), every step in float32."""
expo = np.arange(0, rotary_dim, 2, dtype=np.float32) / np.float32(rotary_dim)
return (np.float32(1.0) / np.power(np.float32(theta), expo, dtype=np.float32)).astype(np.float32)
def mrope_tables(pos3: np.ndarray, freq: np.ndarray | None = None,
section: tuple[int, int, int] = MROPE_SECTION) -> tuple[np.ndarray, np.ndarray]:
"""[3, L] position planes -> cos, sin float32 [1, L, 2 * len(freq)]."""
freq = inv_freq() if freq is None else freq
ang = (pos3.astype(np.float32)[:, :, None] * freq[None, None, :]).astype(np.float32) # [3, L, d/2]
a = ang[0].copy()
for plane, offset in ((1, 1), (2, 2)):
stop = section[plane] * 3
a[:, offset:stop:3] = ang[plane][:, offset:stop:3]
emb = np.concatenate([a, a], axis=-1).astype(np.float64)
return np.cos(emb).astype(np.float32)[None], np.sin(emb).astype(np.float32)[None]
|