Hugging Face's logo

tiny-random
/
unlimited-ocr-2

Image-Text-to-Text
Transformers
Safetensors
unlimited-ocr
feature-extraction
custom_code
Model card Files
xet
 Community
unlimited-ocr-2 / modeling_deepseekv2.py
yujiepan's picture
yujiepan
Upload folder using huggingface_hub
facb9f1 verified
Raw History Blame Contribute Delete
25.6 kB
"""DASH overlay on top of the page-isolated (page2d) base modeling.
Thin monkey-patch: `_sparse_base_impl.py` carries the full base implementation; this file
only swaps the **prefill** path of SlidingWindowLlamaAttention for
dash_attention.DashAttention (entmax soft prior + log(prior) bias, gradients reaching the
router straight from the LM loss, no KL teacher) when
`sparse_attention_config.training_stage == "dash"`. Decode, the non-dash stages and all
other classes keep using the base implementation.
Three details the dash path relies on -- page isolation degrades silently if any of them
is dropped:
1. The main-branch rope goes through `self._main_rope(...)` rather than
`self.rotary_emb(...)`, otherwise sparse_main_rope_mode=page2d is bypassed and the main
branch falls back to 1D.
2. `self.indexer` is **not** set to None: `_main_rope_coords` returns None without an
indexer (main branch falls back to 1D), and every page/rope helper
(_build_page_aware_ids / _build_page_rope_pos / rotary_emb_2d /
_build_kl_candidate_mask) lives on the indexer.
3. Router scoring uses the indexer's own 2-axis rope (sparse_rope_*: one page pair at
theta=256 with the page index taken mod 2, plus 63 token pairs at theta=1e4, resampled)
together with a page-aware PER-QUERY candidate domain [B,S,nblk].
"""
import copy
import math
import os
import torch
import torch.nn.functional as F
if os.environ.get("DASH_DEBUG_ANOMALY") == "1":
torch.autograd.set_detect_anomaly(True)
print("[DASH] torch.autograd anomaly detection ENABLED", flush=True)
# Re-export the ENTIRE base namespace (including underscore helpers) so that
# modeling_unlimitedocr.py's `from .modeling_deepseekv2 import <name>` keeps working.
from . import _sparse_base_impl as _base # noqa: E402
globals().update({k: v for k, v in vars(_base).items() if not k.startswith("__")})
from ._sparse_base_impl import ( # noqa: E402
DeepseekV2PreTrainedModel as _BasePreTrained,
SlidingWindowLlamaAttention as _BaseAttn,
_llama_apply_rotary_pos_emb,
_llama_repeat_kv,
)
from .dash_attention import DashAttention # noqa: E402
from .dash_attention import entmax15 as _entmax15 # noqa: E402
from .dash_attention import _SCORE_NEG, _ACT_CLAMP, _MASK_NEG # noqa: E402
_orig_init = _BaseAttn.__init__
_orig_forward = _BaseAttn.forward
_DASH_LOGGED = False
_IDX_ROPE_WARNED = False
# id(sparse_attention_config) -> deep copy taken BEFORE DashAttention.__init__ rewrites it
# (_persist_router_cfg writes the live process state -- k_start=64, env DASH_STAGE0 default
# "indexer" -- over dash_topk_blocks / dash_stage0 / ...). Inference reads the snapshot.
_CKPT_SP = {}
# Inference switch, toggled by UnlimitedOCRForCausalLM.infer (see configure_dash_inference).
# When on, DASH layers route exactly like vLLM's `unlimited_page2d_dash` plugin
# (dash_router.DashRouterMath.route_and_select: top-M entmax window + hard top-K read set)
# in prefill AND decode; decode no longer falls back to dense full-cache attention.
_INFER = {"on": False}
_ENTMAX_NEG = -1.0e4
def _patched_init(self, config, layer_idx):
sp = getattr(config, "sparse_attention_config", None)
is_dash = isinstance(sp, dict) and sp.get("training_stage") == "dash"
if is_dash and id(sp) not in _CKPT_SP:
_CKPT_SP[id(sp)] = copy.deepcopy(sp)
if is_dash:
_saved_stage = sp.get("training_stage")
sp["training_stage"] = "warmup"
_orig_init(self, config, layer_idx)
if is_dash:
sp["training_stage"] = _saved_stage
self.dash_router = None
if is_dash and self.indexer is not None:
self.dash_router = DashAttention(config)
global _DASH_LOGGED
if not _DASH_LOGGED:
_DASH_LOGGED = True
print(f"[DASH] differentiable DASH attention active "
f"(entmax router, block_size={self.dash_router.block_size}, "
f"local_window={self.dash_router.local_window}, "
f"index_dim={self.dash_router.index_dim}, "
f"main_rope={self.main_rope_mode}, "
f"idx_rope={self.indexer.rope_mode}, "
f"page_bucket={self.indexer.image_bucket_constraint}); "
f"first dash layer_idx={layer_idx}", flush=True)
def _dash_page_ids(self, idx, src_ids, key_length):
"""Page-aware ids, built with the same calls as the base prefill path."""
if idx.image_bucket_constraint:
itm, ipi, qpi = idx._build_page_aware_ids(
src_ids, idx.image_token_id, idx.page_token_id, key_length)
else:
itm = idx._build_flat_image_token_mask(src_ids, idx.image_token_id, key_length)
ipi = qpi = None
idx._log_image_mask(itm, src_ids, key_length)
return itm, ipi, qpi
def _dash_block_cand(self, idx, router, pos_ids, itm, ipi, qpi, nblk, bsz, q_len):
"""PER-QUERY candidate domain [B,S,nblk], reusing the indexer's
_build_kl_candidate_mask (same-page image bucket + per-query fallback). Built in
q_chunk slices so the [B,S,nblk,block_size] intermediate is never materialised at
once (32k x 1024 x 32 bool is roughly 1 GiB)."""
if itm is None and (ipi is None or qpi is None):
return None
step = max(1, int(router.q_chunk))
out = []
for c0 in range(0, q_len, step):
c1 = min(c0 + step, q_len)
cm = idx._build_kl_candidate_mask(
pos_ids[:, c0:c1], nblk, itm, bsz, c1 - c0,
image_token_page_ids=ipi, query_token_page_ids=qpi) # [B,1,C,nblk]
out.append(cm[:, 0])
return torch.cat(out, dim=1) # [B,S,nblk]
def _dash_indexer_rope(self, idx, router, x, pos_ids, coords, main_pe):
"""cos/sin used for router scoring: the indexer's 2-axis rope (sparse_rope_*) when
the coordinates are available, otherwise the indexer's 1D rope. With
index_dim < head_dim the two axes cannot be sliced (cat(ang, ang) + rotate_half would
mix the phases of the two axes), so 1D is used instead."""
global _IDX_ROPE_WARNED
if coords is not None and getattr(idx, "rotary_emb_2d", None) is not None:
if router.index_dim == idx.head_dim:
return idx.rotary_emb_2d(x, coords[0], coords[1])
if not _IDX_ROPE_WARNED:
_IDX_ROPE_WARNED = True
print(f"[PageRope][WARN] sparse_index_dim={router.index_dim} != "
f"head_dim={idx.head_dim}: the 2-axis cos/sin cannot be sliced per "
f"frequency pair, router scoring falls back to 1D RoPE "
f"(the main branch is unaffected)", flush=True)
if getattr(idx, "rotary_emb_1d", None) is not None:
return idx.rotary_emb_1d(x, pos_ids)
return main_pe
def _dash_attn_forward(self, hidden_states, attention_mask, position_ids, past_kv):
idx = self.indexer
router = self.dash_router
num_heads = self.config.num_attention_heads
num_kv_heads = self.config.num_key_value_heads
head_dim = self.head_dim
num_kv_groups = self.num_key_value_groups
bsz, q_len, _ = hidden_states.size()
pos_ids = position_ids
if _base._pk.is_packed(pos_ids):
raise RuntimeError(
"[DASH] packing rows are not supported by the dash path "
"(fused e2e kernel requires attn_mask=None, and page rope would run the "
"<PAGE> cumsum across sample boundaries). Disable packing.")
q = self.q_proj(hidden_states).view(bsz, q_len, num_heads, head_dim).transpose(1, 2)
key_states = self.k_proj(hidden_states).view(bsz, q_len, num_kv_heads, head_dim).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(bsz, q_len, num_kv_heads, head_dim).transpose(1, 2)
src_ids = getattr(self, '_sparse_source_input_ids', None)
itm, ipi, qpi = _dash_page_ids(self, idx, src_ids, q_len)
coords = None
if idx.rope_mode == "page2d":
coords = self._main_rope_pos_2d
if coords is not None and coords[0].shape[-1] != q_len:
coords = None
if coords is None:
coords = idx._build_page_rope_pos(src_ids, q_len)
idx._rope_pos_2d = coords
if self.main_rope_mode == "page2d" and coords is not None:
self._main_rope_pos_2d = coords
cos, sin = self._main_rope(value_states, pos_ids, past_kv, q_len)
else:
self._main_rope_pos_2d = None
cos, sin = self.rotary_emb(value_states, pos_ids)
q, key_states = _llama_apply_rotary_pos_emb(q, key_states, cos, sin)
if past_kv is not None:
key_states, value_states = past_kv.update(key_states, value_states, self.layer_idx)
k = _llama_repeat_kv(key_states, num_kv_groups)
v = _llama_repeat_kv(value_states, num_kv_groups)
idx_pe = _dash_indexer_rope(self, idx, router, value_states, pos_ids, coords, (cos, sin))
nblk = -(-k.shape[2] // router.block_size)
block_cand = _dash_block_cand(self, idx, router, pos_ids, itm, ipi, qpi,
nblk, bsz, q_len)
router_mask = attention_mask
if router_mask is not None and bool((router_mask[:, :, -1, :] > -1.0).all()):
router_mask = None
if _INFER["on"]:
attn_output = _infer_prefill_attend(self, router, hidden_states, q, k, v, idx_pe,
pos_ids, block_cand, router_mask, past_kv)
else:
attn_output = router(hidden_states, q, k, v,
position_embeddings=idx_pe,
position_ids=pos_ids,
attention_mask=router_mask,
image_block_mask=None,
block_cand=block_cand,
page_ids=(ipi, qpi))
attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, q_len, -1)
attn_output = self.o_proj(attn_output)
return attn_output, None, past_kv, None
def _patched_forward(self, *args, **kwargs):
if self.dash_router is not None:
def _get(name, idx_, default=None):
if name in kwargs:
return kwargs[name]
if len(args) > idx_:
return args[idx_]
return default
hidden_states = _get('hidden_states', 0)
attention_mask = _get('attention_mask', 1)
position_ids = _get('position_ids', 2)
past_kv = _get('past_key_value', 3)
if past_kv is None:
past_kv = kwargs.get('past_key_values', None)
bsz, q_len, _ = hidden_states.size()
if q_len > 1:
if position_ids is None:
position_ids = torch.arange(
q_len, device=hidden_states.device).unsqueeze(0).expand(bsz, -1)
return _dash_attn_forward(self, hidden_states, attention_mask,
position_ids, past_kv)
if (_INFER["on"] and past_kv is not None
and self.layer_idx in getattr(past_kv, "_dash_idx_k", {})):
return _dash_decode_forward(self, hidden_states, position_ids, past_kv)
return _orig_forward(self, *args, **kwargs)
_BaseAttn.__init__ = _patched_init
_BaseAttn.forward = _patched_forward
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
_orig_init_weights = _BasePreTrained._init_weights
def _patched_init_weights(self, module):
_orig_init_weights(self, module)
if isinstance(module, DashAttention):
module.head_cls.data.zero_()
_BasePreTrained._init_weights = _patched_init_weights
# ===========================================================================
# Inference path -- parity with vLLM's `unlimited_page2d_dash` plugin
# (plugin_page2d_dash/unlimited_page2d_dash: dash_router.py + dash_attention_layer.py).
# Only active between configure_dash_inference(model, True) / (model, False), i.e. inside
# UnlimitedOCRForCausalLM.infer; training never enters these functions.
#
# vs. the training forward (DashAttention.forward) the inference router differs in:
# * routing params come from the checkpoint snapshot (_CKPT_SP): dash_topk_blocks (hard
# top-K read budget), dash_stage0, per-layer dash_scale_per_layer / dash_floor_per_layer
# -- the live config dict has been overwritten by _persist_router_cfg at construction;
# * entmax runs on the top-M (DASH_INFER_TOPM, deploy default 16) scores, then the read
# set is the first K candidate slots (DashRouterMath.route_and_select, budget=topk);
# * decode (q_len == 1) is ALSO block-sparse (vLLM DASH_DECODE_DENSE=0): head_cls block
# summaries over the cached router keys, page-aware candidate mask, bias-aware attend
# over the selected blocks + the token-level local window.
# ===========================================================================
def _sched_slot(sp, layer_idx):
freq = sp.get("sparse_attention_freq")
if not isinstance(freq, (list, tuple)) or layer_idx >= len(freq):
return int(layer_idx)
return sum(1 for f in freq[:layer_idx] if f)
def _per_layer_val(sp, key, layer_idx):
arr = sp.get(key)
if not isinstance(arr, (list, tuple)) or not arr:
return None
slot = _sched_slot(sp, int(layer_idx))
return float(arr[slot]) if 0 <= slot < len(arr) else None
def configure_dash_inference(decoder, on, topm=16, page_scope_prefill=True, state=None):
"""Switch the vLLM-parity DASH inference path on/off for ``decoder`` (DeepseekV2Model).
on=True returns a state dict; pass it back with on=False to restore.
"""
if not on:
if state is not None:
_base._PAGE_SCOPE_PREFILL = state["page_scope_prefill"]
for attn, stage0 in state["stage0"]:
attn.dash_router.stage0 = stage0
_INFER["on"] = False
return None
state = {"page_scope_prefill": _base._PAGE_SCOPE_PREFILL, "stage0": []}
cfg = decoder.config
live = getattr(cfg, "sparse_attention_config", None) or {}
sp = _CKPT_SP.get(id(live), live)
stage0 = str(os.environ.get("DASH_STAGE0") or sp.get("dash_stage0") or "indexer").lower()
budget = sp.get("dash_topk_blocks") or sp.get("dash_image_fallback_remote") or 12
for layer in decoder.layers:
attn = getattr(layer, "self_attn", None)
router = getattr(attn, "dash_router", None)
if router is None:
continue
li = attn.layer_idx
scale = _per_layer_val(sp, "dash_scale_per_layer", li)
floor = _per_layer_val(sp, "dash_floor_per_layer", li)
if scale is None:
scale = float(sp.get("dash_scaling_factor", router.scaling_factor))
if floor is None:
floor = float(sp.get("dash_prior_floor", router._prior_floor))
router._infer_params = dict(
idx_scale=float(scale) / math.sqrt(router.index_dim), floor=float(floor),
sigma=float(sp.get("dash_sigma", router.sigma)), K=int(budget), topm=int(topm))
state["stage0"].append((attn, router.stage0))
router.stage0 = stage0
# vLLM page_state.py isolates PREFILL queries too (image query -> its own page, prompt
# marker -> raw page, trailing instruction -> unconstrained) == training PAGE_SCOPE_PREFILL=1.
_base._PAGE_SCOPE_PREFILL = bool(page_scope_prefill)
_INFER["on"] = True
return state
@torch.no_grad()
def _infer_route_select(block_scores, block_cand, p):
"""vLLM DashRouterMath.route_and_select (budget_topk=True) with training's candidacy.
block_scores [B,Lq,nblk] fp32 (_SCORE_NEG = not a remote block); block_cand
[B,Lq,nblk] bool page-aware candidate domain or None.
Returns (topi [B,Lq,K] int64, bias [B,Lq,K] fp32, valid [B,Lq,K] bool).
"""
remote = block_scores > (_SCORE_NEG * 0.5)
cand = remote if block_cand is None else (remote & block_cand)
cand = torch.where(~cand.any(dim=-1, keepdim=True), remote, cand)
x = torch.where(cand, block_scores * p["idx_scale"], block_scores.new_full((), _ENTMAX_NEG))
x = torch.nan_to_num(x, nan=_ENTMAX_NEG, posinf=_ENTMAX_NEG, neginf=_ENTMAX_NEG)
nblk = x.shape[-1]
M = max(1, min(int(p["topm"]), nblk)) if int(p["topm"]) > 0 else nblk
xv, xi = x.topk(M, dim=-1) # descending
prior = _entmax15(xv, dim=-1) # tau* over the top-M window
is_cand = cand.gather(-1, xi)
K = max(1, min(int(p["K"]), M))
xi, prior, valid = xi[..., :K], prior[..., :K], is_cand[..., :K]
logp = torch.log(prior.clamp_min(p["floor"]))
cnt = valid.sum(dim=-1, keepdim=True).clamp(min=1)
mu = torch.where(valid, logp, torch.zeros_like(logp)).sum(dim=-1, keepdim=True) / cnt
bias = torch.where(valid, (logp - mu) * (1.0 / p["sigma"]), torch.zeros_like(logp))
return xi, bias.float(), valid
def _router_qk(router, hidden_states, pe):
"""Router index q/k exactly as DashAttention._block_scores_headcls builds them."""
B, S, _ = hidden_states.shape
d = router.index_dim
hs = torch.nan_to_num(hidden_states, nan=0.0, posinf=0.0, neginf=0.0)
idx_q = router.q_norm(router.q_proj(hs).view(B, S, 1, d)).transpose(1, 2)
idx_k = router.k_norm(router.k_proj(hs).view(B, S, 1, d)).transpose(1, 2)
if pe is not None:
cos, sin = pe
idx_q, idx_k = _llama_apply_rotary_pos_emb(idx_q, idx_k, cos[..., :d], sin[..., :d])
idx_q = torch.nan_to_num(idx_q[:, 0].float(), nan=0.0, posinf=0.0, neginf=0.0)
idx_k = torch.nan_to_num(idx_k[:, 0].float(), nan=0.0, posinf=0.0, neginf=0.0)
return idx_q, idx_k
@torch.no_grad()
def _block_scores_infer(router, idx_q, idx_k, q_pos, nblk):
"""[B,Lq,nblk] fp32 remote block scores for queries ``idx_q`` at positions ``q_pos``
against ALL keys ``idx_k`` [B,S,d] (prefill: S == Lq; decode: the cached keys).
Same math as DashAttention._block_scores_headcls / _block_scores."""
B, S, d = idx_k.shape
bs = router.block_size
inv_sqrt_d = 1.0 / math.sqrt(d)
Sp = nblk * bs
blk_ids = torch.arange(nblk, device=idx_k.device)
out = []
if router.stage0 == "head_cls":
if Sp > S:
kk = F.pad(idx_k, (0, 0, 0, Sp - S))
real = F.pad(idx_k.new_ones(B, S, dtype=torch.bool), (0, Sp - S), value=False)
else:
kk = idx_k[:, :Sp]
real = idx_k.new_ones(B, Sp, dtype=torch.bool)
kk = kk.view(B, nblk, bs, d)
vmask = real.view(B, nblk, bs)
hc = router.head_cls.to(kk.dtype).view(1, 1, 1, d)
aw = ((kk * hc).sum(-1) * inv_sqrt_d).clamp(-_ACT_CLAMP, _ACT_CLAMP)
neg = torch.finfo(aw.dtype).min
row_max = torch.where(vmask, aw, aw.new_full((), neg)).amax(dim=-1, keepdim=True)
row_max = torch.where(vmask.any(dim=-1, keepdim=True), row_max, torch.zeros_like(row_max))
e = torch.exp(aw - row_max) * vmask.to(aw.dtype)
w = e / e.sum(dim=-1, keepdim=True).clamp_min(1e-6)
k_bar = torch.nan_to_num((w.unsqueeze(-1) * kk).sum(dim=2), nan=0.0, posinf=0.0, neginf=0.0)
chunk_valid = vmask.any(dim=-1) # [B,nblk]
for c0 in range(0, idx_q.shape[1], router.q_chunk):
c1 = min(c0 + router.q_chunk, idx_q.shape[1])
z = torch.matmul(idx_q[:, c0:c1], k_bar.transpose(-1, -2)) * inv_sqrt_d
z = z.clamp(-_ACT_CLAMP, _ACT_CLAMP)
qblk = q_pos[:, c0:c1] // bs
keep = (blk_ids[None, None, :] <= (qblk[:, :, None] - router.local_blocks)) \
& chunk_valid[:, None, :]
out.append(z.masked_fill(~keep, _SCORE_NEG))
return torch.cat(out, dim=1)
# stage0 == "indexer": block-max of token-level scores, remote-only (causal minus local)
kpos = torch.arange(S, device=idx_k.device)
for c0 in range(0, idx_q.shape[1], router.q_chunk):
c1 = min(c0 + router.q_chunk, idx_q.shape[1])
sc = torch.matmul(idx_q[:, c0:c1], idx_k.transpose(-1, -2)).clamp(-_ACT_CLAMP, _ACT_CLAMP)
pos = q_pos[:, c0:c1]
future = kpos[None, None, :] > pos[:, :, None]
local = (kpos[None, None, :] > (pos[:, :, None] - router.local_window)) & ~future
sc = sc.masked_fill(future | local, _SCORE_NEG)
if Sp > S:
sc = F.pad(sc, (0, Sp - S), value=_SCORE_NEG)
out.append(sc.view(B, c1 - c0, nblk, bs).amax(dim=-1))
return torch.cat(out, dim=1)
def _infer_prefill_attend(self, router, hidden_states, q, k, v, idx_pe, pos_ids,
block_cand, attn_mask, past_kv):
"""Prefill (q_len > 1) under _INFER: vLLM-parity routing, training attend kernels."""
B, H, S, D = q.shape
nblk = (S + router.block_size - 1) // router.block_size
p = router._infer_params
with torch.autocast(device_type=q.device.type, enabled=False):
idx_q, idx_k = _router_qk(router, hidden_states, idx_pe)
if past_kv is not None:
if not hasattr(past_kv, "_dash_idx_k"):
past_kv._dash_idx_k = {}
past_kv._dash_idx_k[self.layer_idx] = idx_k
sc = _block_scores_infer(router, idx_q, idx_k, pos_ids, nblk)
topi, bias, valid = _infer_route_select(sc, block_cand, p)
if router.use_triton and q.is_cuda and attn_mask is None:
out = router._forward_triton(q, k, v, topi, bias, valid, S, nblk, pos_ids)
if out is not None:
return out
return router._forward_dense(q, k, v, topi, bias, valid, S, nblk, attn_mask)
@torch.no_grad()
def _decode_attend(q, k_all, v_all, topi, bias, valid, q_pos, block_size, local_window, scale):
"""Bias-aware block-sparse attention for decode queries (fp32 math).
Selected remote blocks add their ``bias``; the token-level local window
(q_pos - local_window, q_pos] is attended without bias and wins on overlap --
same read set / weights as vLLM dash_sparse_attn decode and training _forward_dense.
"""
B, H, Lq, D = q.shape
S = k_all.shape[2]
nblk = (S + block_size - 1) // block_size
kpos = torch.arange(S, device=q.device)
src = torch.where(valid, bias, bias.new_full((), _MASK_NEG))
blk = torch.full((B, Lq, nblk), _MASK_NEG, device=q.device, dtype=torch.float32)
blk = blk.scatter(-1, topi.long(), src.float())
base = blk.index_select(-1, kpos // block_size) # [B,Lq,S]
qp = q_pos[:, :, None]
causal = kpos[None, None, :] <= qp
local = causal & (kpos[None, None, :] > (qp - local_window))
keep = causal & (local | (base > _MASK_NEG * 0.5))
m = torch.where(local, torch.zeros_like(base), base).masked_fill(~keep, float("-inf"))
kk = k_all.float()
vv = v_all.float()
if kk.shape[1] != H:
kk = _llama_repeat_kv(kk, H // kk.shape[1])
vv = _llama_repeat_kv(vv, H // vv.shape[1])
s = torch.matmul(q.float(), kk.transpose(-1, -2)) * scale + m[:, None]
return torch.matmul(torch.softmax(s, dim=-1), vv) # [B,H,Lq,D]
def _dash_decode_forward(self, hidden_states, position_ids, past_kv):
"""Decode step (q_len == 1) of a DASH layer under _INFER (vLLM DASH_DECODE_DENSE=0)."""
idx = self.indexer
router = self.dash_router
nh = self.config.num_attention_heads
nkv = self.config.num_key_value_heads
hd = self.head_dim
bsz, q_len, _ = hidden_states.size()
if position_ids is None:
past = past_kv.get_seq_length(self.layer_idx)
position_ids = torch.arange(past, past + q_len, device=hidden_states.device)[None].expand(bsz, -1)
q = self.q_proj(hidden_states).view(bsz, q_len, nh, hd).transpose(1, 2)
k = self.k_proj(hidden_states).view(bsz, q_len, nkv, hd).transpose(1, 2)
v = self.v_proj(hidden_states).view(bsz, q_len, nkv, hd).transpose(1, 2)
# (tok_pos, page_pos) from the incremental decode state machine (advances it once).
coords = None
if self.main_rope_mode == "page2d" and self.rotary_emb_2d_main is not None:
coords = self._main_rope_coords(q_len, past_kv)
if coords is not None:
cos, sin = self.rotary_emb_2d_main(v, coords[0], coords[1])
else:
cos, sin = self.rotary_emb(v, position_ids)
q, k = _llama_apply_rotary_pos_emb(q, k, cos, sin)
k_all, v_all = past_kv.update(k, v, self.layer_idx)
idx_pe = _dash_indexer_rope(self, idx, router, v, position_ids,
coords if idx.rope_mode == "page2d" else None, (cos, sin))
p = router._infer_params
with torch.autocast(device_type=hidden_states.device.type, enabled=False):
idx_q, idx_k_new = _router_qk(router, hidden_states, idx_pe)
idx_k = torch.cat([past_kv._dash_idx_k[self.layer_idx], idx_k_new], dim=1)
past_kv._dash_idx_k[self.layer_idx] = idx_k
S = k_all.shape[2]
nblk = (S + router.block_size - 1) // router.block_size
sc = _block_scores_infer(router, idx_q, idx_k, position_ids, nblk) # [B,Lq,nblk]
allow = getattr(self, "_dash_dec_allow", None) # [B,nblk] or None
cand = None
if allow is not None and router.image_only:
if allow.shape[-1] < nblk:
allow = F.pad(allow, (0, nblk - allow.shape[-1]), value=False)
cand = allow[:, None, :nblk].expand(bsz, q_len, nblk)
topi, bias, valid = _infer_route_select(sc, cand, p)
out = _decode_attend(q, k_all, v_all, topi, bias, valid, position_ids,
router.block_size, router.local_window, router._scale)
out = out.to(hidden_states.dtype).transpose(1, 2).contiguous().reshape(bsz, q_len, -1)
return self.o_proj(out), None, past_kv, None