File size: 2,791 Bytes
677e207 | 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 | #
# Copyright (c) 2026 BEL ESPRIT D ACCORD TRUST HOLDINGS INC
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 X.AI Corp.
"""
Segment-Bound Masking for Ranker Attention.
Segment IDs: HISTORY=1, CANDIDATE=-1, PADDING=0
Ranker mask: history attends to all, candidates attend only to self.
"""
from dataclasses import dataclass
import jax
import jax.numpy as jnp
HISTORY_SEGMENT_ID = 1
CANDIDATE_SEGMENT_ID = -1
PADDING_SEGMENT_ID = 0
@dataclass(frozen=True)
class SegmentBounds:
history_lower: jax.Array
history_upper: jax.Array
candidate_lower: jax.Array
candidate_upper: jax.Array
@classmethod
def from_segment_ids(cls, segment_ids: jax.Array) -> "SegmentBounds":
segment_ids = jnp.asarray(segment_ids)
if segment_ids.ndim == 1:
segment_ids = segment_ids[None, :]
B, S = segment_ids.shape
history_mask = segment_ids == HISTORY_SEGMENT_ID
candidate_mask = segment_ids == CANDIDATE_SEGMENT_ID
def _bounds(mask):
any_mask = jnp.any(mask, axis=1)
first = jnp.argmax(mask, axis=1)
last = (S - 1) - jnp.argmax(mask[:, ::-1], axis=1)
lower = jnp.where(any_mask, first, 0)
upper = jnp.where(any_mask, last + 1, 0)
return lower.astype(jnp.int32), upper.astype(jnp.int32)
hl, hu = _bounds(history_mask)
cl, cu = _bounds(candidate_mask)
return cls(hl, hu, cl, cu)
def to_array(self) -> jax.Array:
return jnp.stack([self.history_lower, self.history_upper,
self.candidate_lower, self.candidate_upper], axis=1)
def ranker_mask(q_pos: jax.Array, kv_pos: jax.Array, bounds: SegmentBounds) -> jax.Array:
q_is_hist = (q_pos >= bounds.history_lower[..., None, None]) & \
(q_pos < bounds.history_upper[..., None, None])
q_is_cand = (q_pos >= bounds.candidate_lower[..., None, None]) & \
(q_pos < bounds.candidate_upper[..., None, None])
kv_is_hist = (kv_pos >= bounds.history_lower[..., None, None]) & \
(kv_pos < bounds.history_upper[..., None, None])
kv_is_cand = (kv_pos >= bounds.candidate_lower[..., None, None]) & \
(kv_pos < bounds.candidate_upper[..., None, None])
history_mask = kv_is_hist & (q_is_hist | q_is_cand)
candidate_self_mask = q_is_cand & kv_is_cand & (q_pos == kv_pos)
return history_mask | candidate_self_mask
def causal_ranker_mask(q_pos: jax.Array, kv_pos: jax.Array, bounds: SegmentBounds,
window_len: int = -1) -> jax.Array:
base = ranker_mask(q_pos, kv_pos, bounds)
causal = q_pos >= kv_pos
if window_len > 0:
causal = causal & (kv_pos > q_pos - window_len)
return base & causal
|