CLEF-Screen: JointSchemaHead redo FT (seed1, 3 epochs) from jev-retrain MI355X run 2026-11-04
2b605e8 verified | """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] | |