| """Packed fast path: whole requests concatenated into one padding-free sequence per forward pass. |
| |
| Each packed unit is a whole request (or, for over-length requests, a chunk of its questions) with one readout |
| per question marker. Requests stay independent: |
| |
| - full-attention layers: varlen FlashAttention (``torch.nn.attention.varlen``) with window (-1, -1), i.e. |
| noncausal within each request, which is what the checkpoint's padding-masked noncausal SDPA computes; |
| - linear-attention layers: ``fla``'s chunked gated delta rule with ``cu_seqlens`` (state reset per request) and a |
| depthwise causal convolution that does not read across request boundaries; |
| - positions restart at 0 for every request; image requests get their own M-RoPE positions (``get_rope_index``); |
| - readouts are gathered at the marker positions only. |
| |
| Optional speedups (``FastOptions``): torch.compile of sub-modules or whole layers, static pack sizes with CUDA |
| graphs (filler sequences reach the bucket size; request boundaries reach the eager kernels through a side |
| channel), last-layer pruning at marker rows, and FP8 matmuls. |
| |
| The decoder layers are re-expressed here with the checkpoint's own modules and weights; the Transformers model |
| is not patched, so the reference path of the same ``Decider`` stays exact. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import queue |
| import threading |
| from collections.abc import Iterator |
| from typing import TYPE_CHECKING, Any |
|
|
| import torch |
| import torch.nn.functional as F |
| from transformers.activations import ACT2FN |
| from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb |
|
|
| from .protocol import Encoded |
| from .types import Request |
|
|
| if TYPE_CHECKING: |
| from .model import Decider, FastOptions |
|
|
| FILLER_MAX = 1024 |
| _PACK: dict[str, Any] = {} |
|
|
|
|
| def _require_kernels() -> tuple[Any, Any]: |
| try: |
| from fla.ops.gated_delta_rule import chunk_gated_delta_rule |
| except ImportError as error: |
| raise ImportError( |
| "The fast path needs flash-linear-attention (pip install flash-linear-attention); " |
| "use fast=False for the reference path" |
| ) from error |
| try: |
| from torch.nn.attention.varlen import varlen_attn |
| except ImportError as error: |
| raise ImportError("The fast path needs torch>=2.10 with torch.nn.attention.varlen") from error |
| return chunk_gated_delta_rule, varlen_attn |
|
|
|
|
| _chunk_rule, _varlen = None, None |
|
|
|
|
| @torch.compiler.disable |
| def _gated_delta(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor) -> Any: |
| out, _ = _chunk_rule(q, k, v, g=g, beta=beta, initial_state=None, output_final_state=False, |
| use_qk_l2norm_in_kernel=True, cu_seqlens=_PACK["cu_long"], cu_seqlens_cpu=_PACK["cu_cpu"]) |
| return out |
|
|
|
|
| @torch.compiler.disable |
| def _attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: float) -> torch.Tensor: |
| |
| q, k, v = (x[0].transpose(0, 1) for x in (q, k, v)) |
| cu, mx = _PACK["cu"], _PACK["max"] |
| out = _varlen(q, k, v, cu, cu, mx, mx, scale=scale, window_size=(-1, -1), enable_gqa=k.shape[1] != q.shape[1]) |
| return out.unsqueeze(0) |
|
|
|
|
| def packed_conv(x: torch.Tensor, weight: torch.Tensor, seqpos: torch.Tensor) -> torch.Tensor: |
| """Depthwise causal conv over (1, C, T) that never reads taps before the start of a request (fp32 sums).""" |
| taps = weight.shape[-1] |
| xf = x.float() |
| out = xf * weight[:, taps - 1, None].float() |
| for shift in range(1, taps): |
| shifted = F.pad(xf, (shift, 0))[..., : x.shape[-1]] |
| out = out + shifted * (seqpos >= shift).float() * weight[:, taps - 1 - shift, None].float() |
| return out.to(x.dtype) |
|
|
|
|
| def linear_attention(mod: torch.nn.Module, x: torch.Tensor, seqpos: torch.Tensor, conv: Any) -> torch.Tensor: |
| """Qwen3_5GatedDeltaNet.forward without cache or padding, on a packed (1, T, D) sequence.""" |
| _, length, _ = x.shape |
| mixed = mod.in_proj_qkv(x).transpose(1, 2) |
| z = mod.in_proj_z(x).reshape(1, length, -1, mod.head_v_dim) |
| b = mod.in_proj_b(x) |
| a = mod.in_proj_a(x) |
| if mod.conv1d.bias is not None: |
| raise ValueError("Unexpected convolution bias") |
| mixed = ACT2FN[mod.activation](conv(mixed, mod.conv1d.weight.squeeze(1), seqpos)).transpose(1, 2) |
| query, key, value = torch.split(mixed, [mod.key_dim, mod.key_dim, mod.value_dim], dim=-1) |
| query = query.reshape(1, length, -1, mod.head_k_dim) |
| key = key.reshape(1, length, -1, mod.head_k_dim) |
| value = value.reshape(1, length, -1, mod.head_v_dim) |
| beta = b.sigmoid() |
| g = -mod.A_log.float().exp() * F.softplus(a.float() + mod.dt_bias) |
| if mod.num_v_heads // mod.num_k_heads > 1: |
| query = query.repeat_interleave(mod.num_v_heads // mod.num_k_heads, dim=2) |
| key = key.repeat_interleave(mod.num_v_heads // mod.num_k_heads, dim=2) |
| core = _gated_delta(query, key, value, g, beta) |
| core = mod.norm(core.reshape(-1, mod.head_v_dim), z.reshape(-1, mod.head_v_dim)) |
| return mod.out_proj(core.reshape(1, length, -1)) |
|
|
|
|
| def full_attention(attn: torch.nn.Module, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| _, length, _ = x.shape |
| shape = (1, length, -1, attn.head_dim) |
| query, gate = torch.chunk(attn.q_proj(x).view(1, length, -1, attn.head_dim * 2), 2, dim=-1) |
| gate = gate.reshape(1, length, -1) |
| query = attn.q_norm(query.view(shape)).transpose(1, 2) |
| key = attn.k_norm(attn.k_proj(x).view(shape)).transpose(1, 2) |
| value = attn.v_proj(x).view(shape).transpose(1, 2) |
| query, key = apply_rotary_pos_emb(query, key, cos, sin) |
| out = _attention(query, key, value, attn.scaling) |
| return attn.o_proj(out.reshape(1, length, -1).contiguous() * torch.sigmoid(gate)) |
|
|
|
|
| def decoder_layer( |
| layer: torch.nn.Module, h: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, seqpos: torch.Tensor, |
| conv: Any = packed_conv, |
| ) -> torch.Tensor: |
| x = layer.input_layernorm(h) |
| if layer.block_type == "linear_attention": |
| x = linear_attention(layer.linear_attn, x, seqpos, conv) |
| else: |
| x = full_attention(layer.self_attn, x, cos, sin) |
| h = h + x |
| return h + layer.mlp(layer.post_attention_layernorm(h)) |
|
|
|
|
| def last_layer_at( |
| layer: torch.nn.Module, h: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, sel: torch.Tensor, |
| cu_q: torch.Tensor, max_q: int, cu_k: torch.Tensor, max_k: int, |
| ) -> torch.Tensor: |
| """Last (full-attention) layer with K/V for every token but queries, output and MLP only at rows ``sel``.""" |
| attn = layer.self_attn |
| x = layer.input_layernorm(h) |
| length, hd, n = x.shape[1], attn.head_dim, sel.shape[0] |
| key = attn.k_norm(attn.k_proj(x).view(1, length, -1, hd)).transpose(1, 2) |
| value = attn.v_proj(x).view(1, length, -1, hd).transpose(1, 2) |
| query, gate = torch.chunk(attn.q_proj(x[:, sel]).view(1, n, -1, hd * 2), 2, dim=-1) |
| gate = gate.reshape(1, n, -1) |
| query = attn.q_norm(query).transpose(1, 2) |
| query, _ = apply_rotary_pos_emb(query, query, cos[:, sel], sin[:, sel]) |
| _, key = apply_rotary_pos_emb(key, key, cos, sin) |
| q, k, v = (t[0].transpose(0, 1).contiguous() for t in (query, key, value)) |
| out = _varlen(q, k, v, cu_q, cu_k, max_q, max_k, scale=attn.scaling, window_size=(-1, -1), |
| enable_gqa=k.shape[1] != q.shape[1]) |
| out = attn.o_proj(out.reshape(1, n, -1) * torch.sigmoid(gate)) |
| hs = h[:, sel] + out |
| return hs + layer.mlp(layer.post_attention_layernorm(hs)) |
|
|
|
|
| class FP8Linear(torch.nn.Module): |
| """Rowwise-scaled float8_e4m3fn matmul: per-output-channel weight scales, per-token activation scales.""" |
|
|
| def __init__(self, linear: torch.nn.Linear) -> None: |
| super().__init__() |
| weight = linear.weight.data.float() |
| scale = weight.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) / 448.0 |
| self.register_buffer("wq", (weight / scale).to(torch.float8_e4m3fn)) |
| self.register_buffer("w_scale", scale.t().contiguous()) |
| self.bias = linear.bias |
| self.out_features = linear.out_features |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| shape = x.shape |
| flat = x.reshape(-1, shape[-1]) |
| a_scale = flat.abs().amax(dim=-1, keepdim=True).float().clamp(min=1e-12) / 448.0 |
| xq = (flat.float() / a_scale).to(torch.float8_e4m3fn) |
| out = torch._scaled_mm(xq, self.wq.t(), scale_a=a_scale, scale_b=self.w_scale, bias=self.bias, |
| out_dtype=torch.bfloat16) |
| return out.reshape(*shape[:-1], self.out_features) |
|
|
|
|
| def to_fp8(module: torch.nn.Module) -> None: |
| for name, child in list(module.named_children()): |
| if isinstance(child, torch.nn.Linear) and child.in_features % 16 == 0 and child.out_features % 16 == 0 \ |
| and child.out_features >= 256: |
| setattr(module, name, FP8Linear(child)) |
| else: |
| to_fp8(child) |
|
|
|
|
| class FastPath: |
| def __init__(self, decider: Decider, options: FastOptions) -> None: |
| global _chunk_rule, _varlen |
| if decider.device.type != "cuda": |
| raise ValueError("The fast path needs a CUDA device") |
| _chunk_rule, _varlen = _require_kernels() |
| if options.cuda_graphs and options.compile != "layer": |
| raise ValueError("cuda_graphs=True needs compile='layer'") |
| if options.budget() < 1: |
| raise ValueError("The packing budget must be positive") |
| self.decider = decider |
| self.options = options |
| self.lm = decider.backbone.language_model |
| self.layers = list(self.lm.layers) |
| if self.layers[-1].block_type != "full_attention": |
| raise ValueError("Last-layer pruning expects a full-attention last layer") |
| self.image_token_id = decider.backbone.config.image_token_id |
| if options.fp8: |
| to_fp8(self.lm.layers) |
| torch.cuda.empty_cache() |
| torch._dynamo.config.cache_size_limit = max(torch._dynamo.config.cache_size_limit, 256) |
| torch._dynamo.config.accumulated_cache_size_limit = max( |
| torch._dynamo.config.accumulated_cache_size_limit, 4096) |
| self.conv = packed_conv |
| self.layer_fn: Any = decoder_layer |
| self.graph_fn: Any = None |
| if options.compile == "sub": |
| self.conv = torch.compile(packed_conv, dynamic=True) |
| for layer in self.layers: |
| mods = [layer.mlp, layer.input_layernorm, layer.post_attention_layernorm] |
| mods += [layer.linear_attn.norm] if layer.block_type == "linear_attention" else \ |
| [layer.self_attn.q_norm, layer.self_attn.k_norm] |
| for mod in mods: |
| mod.forward = torch.compile(mod.forward, dynamic=True) |
| elif options.compile == "layer": |
| if options.cuda_graphs: |
| self.graph_fn = torch.compile(decoder_layer, mode="reduce-overhead", dynamic=False) |
| self.layer_fn = torch.compile(decoder_layer, dynamic=True) |
| self.last_fn: Any = last_layer_at if options.compile is None else torch.compile(last_layer_at, dynamic=True) |
| self.buckets = tuple(sorted(options.buckets)) if options.cuda_graphs else () |
|
|
| |
|
|
| def run(self, requests: list[Request]) -> Iterator[tuple[Request, dict[str, list[float]], list[Encoded]]]: |
| """Yield (request, logits per question, chunks) in input order.""" |
| encoder, max_length = self.decider.encoder, self.decider.max_length |
| for request in requests: |
| encoder.validate(request) |
| budget = self.options.budget() |
|
|
| def plans() -> Iterator[tuple[int, list[Encoded]]]: |
| for index, request in enumerate(requests): |
| yield index, encoder.plan(request, max_length) |
|
|
| source = _prefetch(plans()) if self.options.prefetch and len(requests) > 1 else plans() |
| pending: dict[int, list[Encoded]] = {} |
| results: dict[int, dict[str, list[float]]] = {} |
| remaining: dict[int, int] = {} |
| queue_: list[tuple[int, Encoded]] = [] |
| queued_tokens = 0 |
| next_out = 0 |
|
|
| def flush(force: bool) -> None: |
| nonlocal queue_, queued_tokens |
| while queue_ and (force or queued_tokens >= budget): |
| pack, used = [], 0 |
| for item in queue_: |
| if pack and used + item[1].length > budget: |
| break |
| pack.append(item) |
| used += item[1].length |
| for (index, chunk), rows in zip(pack, self.forward([c for _, c in pack]), strict=True): |
| results.setdefault(index, {}).update(zip(chunk.keys, rows, strict=True)) |
| remaining[index] -= 1 |
| queue_ = queue_[len(pack):] |
| queued_tokens -= used |
| if not force and queued_tokens < budget: |
| break |
|
|
| for index, plan in source: |
| pending[index] = plan |
| remaining[index] = len(plan) |
| for chunk in plan: |
| queue_.append((index, chunk)) |
| queued_tokens += chunk.length |
| flush(force=False) |
| while next_out in pending and remaining[next_out] == 0: |
| yield requests[next_out], results.pop(next_out), pending.pop(next_out) |
| next_out += 1 |
| flush(force=True) |
| while next_out in pending: |
| yield requests[next_out], results.pop(next_out), pending.pop(next_out) |
| next_out += 1 |
|
|
| |
|
|
| def _positions(self, chunk: Encoded) -> torch.Tensor: |
| if chunk.image_grid_thw is None: |
| return torch.arange(chunk.length).expand(3, -1) |
| positions, _ = self.decider.backbone.get_rope_index( |
| chunk.input_ids[None], chunk.mm_token_type_ids[None], image_grid_thw=chunk.image_grid_thw) |
| return positions[:, 0] |
|
|
| @torch.inference_mode() |
| def forward(self, chunks: list[Encoded]) -> list[list[list[float]]]: |
| device = self.decider.device |
| lengths = [chunk.length for chunk in chunks] |
| total = sum(lengths) |
| bucket = next((b for b in self.buckets if b >= total), None) |
| ids = [chunk.input_ids for chunk in chunks] |
| positions = [self._positions(chunk) for chunk in chunks] |
| if bucket is not None: |
| fill = bucket - total |
| while fill > 0: |
| size = min(fill, FILLER_MAX) |
| lengths.append(size) |
| ids.append(torch.full((size,), self.decider.encoder.pad_id, dtype=torch.long)) |
| positions.append(torch.arange(size).expand(3, -1)) |
| fill -= size |
| cu_cpu = F.pad(torch.tensor(lengths).cumsum(0), (1, 0)) |
| seqpos = torch.cat([torch.arange(n) for n in lengths]).to(device, non_blocking=True) |
| input_ids = torch.cat(ids).to(device, non_blocking=True) |
| rope = torch.cat(positions, dim=1)[:, None, :].to(device, non_blocking=True) |
| cu = cu_cpu.to(device, torch.int32, non_blocking=True) |
| _PACK.update(cu=cu, cu_long=cu.long(), cu_cpu=cu_cpu, max=max(lengths)) |
|
|
| h = self.lm.embed_tokens(input_ids)[None] |
| pixels = [chunk.pixel_values for chunk in chunks if chunk.pixel_values is not None] |
| if pixels: |
| grids = torch.cat([c.image_grid_thw for c in chunks if c.image_grid_thw is not None]).to(device) |
| features = self.decider.backbone.get_image_features(torch.cat(pixels).to(device), grids, |
| return_dict=True).pooler_output |
| features = torch.cat(features).to(h.dtype) |
| mask = (input_ids == self.image_token_id)[None, :, None].expand_as(h) |
| h = h.masked_scatter(mask, features) |
| cos, sin = self.lm.rotary_emb(h, rope) |
|
|
| starts = cu_cpu[: len(chunks)].tolist() |
| sel = torch.cat([chunk.markers + start for chunk, start in zip(chunks, starts, strict=True)]).to(device) |
| static = bucket is not None and self.graph_fn is not None |
| |
| fn = self.graph_fn if static else self.layer_fn if total <= self.options.budget() else decoder_layer |
| layers = self.layers[:-1] if self.options.prune_last else self.layers |
| if static: |
| torch.compiler.cudagraph_mark_step_begin() |
| for layer in layers: |
| h = fn(layer, h, cos, sin, seqpos, self.conv) |
| if static: |
| h = h.clone() |
| if self.options.prune_last: |
| per_seq = torch.tensor([0] + [len(c.keys) for c in chunks]).cumsum(0) |
| cu_q = per_seq.to(device, torch.int32) |
| h = self.last_fn(self.layers[-1], h, cos, sin, sel, cu_q, max(len(c.keys) for c in chunks), |
| cu[: len(chunks) + 1], max(c.length for c in chunks)) |
| else: |
| h = h[:, sel] |
| logits = self.decider.readout(self.lm.norm(h)[0]).float().cpu() |
| out, cursor = [], 0 |
| for chunk in chunks: |
| out.append([logits[cursor + j, :count].tolist() for j, count in enumerate(chunk.counts)]) |
| cursor += len(chunk.keys) |
| return out |
|
|
| def warmup(self) -> None: |
| """Run one dummy pack per static bucket (or two sizes when dynamic) so compilation happens up front.""" |
| request = {"state": "warmup", "questions": {"q": {"type": "noul", "instructions": "Warm up?"}}} |
| chunk = self.decider.encoder.plan(request, None)[0] |
| sizes = self.buckets or (512, 4096) |
| for size in sizes: |
| count = max(1, min(size // chunk.length, 64)) if not self.buckets else 1 |
| for _ in range(3 if self.buckets else 2): |
| self.forward([chunk] * count) |
| torch.cuda.synchronize() |
|
|
|
|
| def _prefetch(source: Iterator[Any], depth: int = 64) -> Iterator[Any]: |
| """Run ``source`` (tokenization and image preprocessing) in a background thread; re-raise its errors.""" |
| items: queue.Queue[Any] = queue.Queue(maxsize=depth) |
| done = object() |
|
|
| def produce() -> None: |
| try: |
| for item in source: |
| items.put(item) |
| except BaseException as error: |
| items.put(error) |
| items.put(done) |
|
|
| threading.Thread(target=produce, daemon=True).start() |
| while (item := items.get()) is not done: |
| if isinstance(item, BaseException): |
| raise item |
| yield item |
|
|
