From 9e73f2cf126843de238aff50538d9bf575805262 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 14:19:37 -0400 Subject: [PATCH 1/5] feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) Shard the dense weights per rank at load (attention qkv by head, GDN in_proj as its six parts with the matching conv1d channels and A_log/dt_bias, shared-expert gate_up per part; o_proj/out_proj/down_proj row-parallel; embed/lm_head by vocab rows) and the NVFP4 expert banks along the intermediate axis, so every rank holds half the experts and each MoE layer needs one all-reduce (routed + gate * shared are combined before the reduce). Router, QSA indexer, norms, hyper-connections and PLE stay replicated so all ranks select the same blocks and n-gram rows. Also: LinearColParallelMerged(local_output_sizes=) for the kv-replicated case and distributed_timeout 60 -> 1800 s (ranks reach their first collective minutes apart behind a 100+ GiB load). Limits: offload backend with bf16 dense projections; fp8_block / nvfp4 dense checkpoints raise under TP. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/engine/config.py | 2 +- python/freetoken/layers/linear.py | 8 +- python/freetoken/models/nvfp4_banks.py | 123 ++++++++---------- .../freetoken/models/qwen4_exp/attention.py | 39 ++++-- python/freetoken/models/qwen4_exp/gdn.py | 66 +++++++--- python/freetoken/models/qwen4_exp/moe.py | 29 ++++- python/freetoken/models/qwen4_exp/weight.py | 69 +++++++++- tests/models/qwen4_exp/test_tp_shard.py | 92 +++++++++++++ tests/models/test_nvfp4_banks_tp.py | 76 +++++++++++ 9 files changed, 388 insertions(+), 116 deletions(-) create mode 100644 tests/models/qwen4_exp/test_tp_shard.py create mode 100644 tests/models/test_nvfp4_banks_tp.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index bcbe6bcf2..e190cac18 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -74,7 +74,7 @@ class EngineConfig: # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None - distributed_timeout: float = 60.0 + distributed_timeout: float = 1800.0 # ranks reach the first collective minutes apart on a 100+ GiB offload load use_dummy_weight: bool = False use_pynccl: bool = True max_seq_len_override: int | None = None diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py index f707e0429..15c973a98 100644 --- a/python/freetoken/layers/linear.py +++ b/python/freetoken/layers/linear.py @@ -59,10 +59,14 @@ def __init__( input_size: int, output_sizes: List[int], has_bias: bool, + local_output_sizes: List[int] | None = None, ): - # check that all output sizes are divisible by tp_size + # check that all output sizes are divisible by tp_size (a caller that replicates + # GQA kv heads across ranks passes the per-rank sizes explicitly) tp_info = get_tp_info() - tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes] + if local_output_sizes is None: + local_output_sizes = [div_even(size, tp_info.size) for size in output_sizes] + tp_output_sizes = local_output_sizes output_size = sum(output_sizes) tp_output_size = sum(tp_output_sizes) super().__init__(input_size, output_size, input_size, tp_output_size, has_bias) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 6b933ff1d..6980c9ed8 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -9,7 +9,8 @@ import safetensors import torch -from freetoken.utils import download_hf_weight +from freetoken.distributed import get_tp_info +from freetoken.utils import div_even, download_hf_weight from tqdm import tqdm LayerToBank = Callable[[int, object], int | None] @@ -78,6 +79,41 @@ def _alloc_nvfp4_host_banks(num_layers: int, E: int, H: int, I: int): }, num_layers) +def _tp_slice(inter: int) -> tuple[int, int]: + """``(i_local, i_lo)``: this rank's slice of the intermediate axis. TP shards every expert + along I (gate/up rows, down columns), the ``stream_moe_expert_sources`` convention, so the + routed output is a partial sum the MoE layer all-reduces.""" + tp = get_tp_info() + i_local = div_even(inter, tp.size) + assert i_local % 16 == 0, f"NVFP4 TP shard {i_local} must cover whole 16-wide scale blocks" + return i_local, tp.rank * i_local + + +class _Placer: + """Writes one checkpoint expert tensor into its bank slot (this rank's I slice only).""" + + def __init__(self, banks: dict, inter: int): + self.b = banks + self.i_local, self.i_lo = _tp_slice(inter) + + def put(self, layer: int, expert: int, role: str, kind: str, tensor, global_scale=None): + n, lo = self.i_local, self.i_lo + b = self.b + if role == "down": # [H, I/2] codes, [H, I/16] scales, [H] global + if kind == "weight": + b["down_packed"][layer][expert] = tensor[:, lo // 2 : (lo + n) // 2] + else: + b["down_scale"][layer][expert] = tensor[:, lo // 16 : (lo + n) // 16] + b["down_global"][layer][expert] = global_scale + return + rows = slice(0, n) if role == "gate" else slice(n, 2 * n) # gate | up on the row axis + if kind == "weight": + b["gate_up_packed"][layer][expert, rows] = tensor[lo : lo + n] + else: + b["gate_up_scale"][layer][expert, rows] = tensor[lo : lo + n] + b["gate_up_global"][layer][expert, rows] = global_scale # per-tensor scalar + + def load_nvfp4_expert_source_banks( model_path: str, config, @@ -149,13 +185,9 @@ def load_nvfp4_expert_source_banks( globals_map[key] = _ingest_global(spec, f.get_tensor(name)) drop_page_cache(path) - _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill - gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] - gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] - gate_up_global = [b.tensor for b in _hb["gate_up_global"]] - down_packed = [b.tensor for b in _hb["down_packed"]] - down_scale = [b.tensor for b in _hb["down_scale"]] - down_global = [b.tensor for b in _hb["down_global"]] + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, _tp_slice(I)[0]) # unpinned; pinned after fill + banks = {name: [b.tensor for b in layers] for name, layers in _hb.items()} + place = _Placer(banks, I) from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline @@ -170,30 +202,11 @@ def _load(sink) -> int: expert = int(match.group("expert")) proj = match.group("proj") role = spec.proj_to_role[proj] + if role not in ("gate", "up", "down"): + raise ValueError(f"{spec.desc}: unknown projection role {role!r}") kind = _canon_kind(spec, match.group("kind")) - tensor = f.get_tensor(name) - if kind == "weight": - if role == "gate": - gate_up_packed[bank_layer_id][expert, :I] = tensor - elif role == "up": - gate_up_packed[bank_layer_id][expert, I:] = tensor - elif role == "down": - down_packed[bank_layer_id][expert] = tensor - else: - raise ValueError(f"{spec.desc}: unknown projection role {role!r}") - else: - global_scale = globals_map[(layer, expert, proj)] - if role == "gate": - gate_up_scale[bank_layer_id][expert, :I] = tensor - gate_up_global[bank_layer_id][expert, :I] = global_scale - elif role == "up": - gate_up_scale[bank_layer_id][expert, I:] = tensor - gate_up_global[bank_layer_id][expert, I:] = global_scale - elif role == "down": - down_scale[bank_layer_id][expert] = tensor - down_global[bank_layer_id][expert] = global_scale - else: - raise ValueError(f"{spec.desc}: unknown projection role {role!r}") + place.put(bank_layer_id, expert, role, kind, f.get_tensor(name), + None if kind == "weight" else globals_map[(layer, expert, proj)]) tracker.note(bank_layer_id) placed += 1 drop_page_cache(path) @@ -207,14 +220,7 @@ def _load(sink) -> int: expected = num_layers * E * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" - return { - "gate_up_packed": gate_up_packed, - "gate_up_scale": gate_up_scale, - "gate_up_global": gate_up_global, - "down_packed": down_packed, - "down_scale": down_scale, - "down_global": down_global, - } + return banks def load_nvfp4_expert_source_banks_parallel( @@ -273,13 +279,9 @@ def load_nvfp4_expert_source_banks_parallel( ) drop_page_cache(path) - _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill - gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] - gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] - gate_up_global = [b.tensor for b in _hb["gate_up_global"]] - down_packed = [b.tensor for b in _hb["down_packed"]] - down_scale = [b.tensor for b in _hb["down_scale"]] - down_global = [b.tensor for b in _hb["down_global"]] + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, _tp_slice(I)[0]) # unpinned; pinned after fill + banks = {name: [b.tensor for b in layers] for name, layers in _hb.items()} + place = _Placer(banks, I) from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline @@ -296,24 +298,8 @@ def _load(sink) -> int: proj = match.group("proj") role = spec.proj_to_role[proj] kind = _canon_kind(spec, match.group("kind")) - if kind == "weight": - if role == "gate": - gate_up_packed[bank_layer_id][expert, :I] = tensor - elif role == "up": - gate_up_packed[bank_layer_id][expert, I:] = tensor - else: - down_packed[bank_layer_id][expert] = tensor - else: - g = globals_map[(layer, expert, proj)] - if role == "gate": - gate_up_scale[bank_layer_id][expert, :I] = tensor - gate_up_global[bank_layer_id][expert, :I] = g - elif role == "up": - gate_up_scale[bank_layer_id][expert, I:] = tensor - gate_up_global[bank_layer_id][expert, I:] = g - else: - down_scale[bank_layer_id][expert] = tensor - down_global[bank_layer_id][expert] = g + place.put(bank_layer_id, expert, role, kind, tensor, + None if kind == "weight" else globals_map[(layer, expert, proj)]) tracker.note(bank_layer_id) placed += 1 return placed @@ -326,14 +312,7 @@ def _load(sink) -> int: expected = num_layers * E * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" - return { - "gate_up_packed": gate_up_packed, - "gate_up_scale": gate_up_scale, - "gate_up_global": gate_up_global, - "down_packed": down_packed, - "down_scale": down_scale, - "down_global": down_global, - } + return banks __all__ = [ diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index d6ab2867a..db77d2b6e 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -19,9 +19,16 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.distributed import get_tp_info +from freetoken.layers import ( + BaseOP, + GemmaPlusOneRMSNorm, + LinearColParallelMerged, + LinearOProj, + LinearReplicated, +) from freetoken.layers.rotary import get_rope -from freetoken.utils import nvtx_annotate +from freetoken.utils import div_even, nvtx_annotate if TYPE_CHECKING: from freetoken.core import Batch @@ -120,11 +127,21 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: self.head_dim = config.head_dim self.qo_attn_dim = self.num_q * self.head_dim self.kv_attn_dim = self.num_kv * self.head_dim - self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + # TP: q heads split across ranks, kv heads split or (num_kv < tp) replicated; the + # indexer stays replicated so every rank selects the same blocks. + tp = get_tp_info() + self._local_num_q = div_even(self.num_q, tp.size) + self._local_num_kv = div_even(self.num_kv, tp.size, allow_replicate=True) + self._local_qo_dim = self._local_num_q * self.head_dim + self._local_kv_dim = self._local_num_kv * self.head_dim + self._qkv_split = [self._local_qo_dim * 2, self._local_kv_dim, self._local_kv_dim] self.qkv_proj = LinearColParallelMerged( - config.hidden_size, self._qkv_split, has_bias=False + config.hidden_size, + [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim], + has_bias=False, + local_output_sizes=self._qkv_split, ) - self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False) + self.o_proj = LinearOProj(self.qo_attn_dim, config.hidden_size, has_bias=False) self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) rotary = config.rotary_config @@ -140,21 +157,21 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: @nvtx_annotate("QSA") def forward(self, x: torch.Tensor, batch: Batch) -> torch.Tensor: qg, k, v = self.qkv_proj.forward(x).split(self._qkv_split, dim=-1) - qg = qg.view(-1, self.num_q, self.head_dim * 2) + qg = qg.view(-1, self._local_num_q, self.head_dim * 2) q = qg[..., : self.head_dim].contiguous() - gate = qg[..., self.head_dim :].reshape(-1, self.qo_attn_dim) - k = k.contiguous().view(-1, self.num_kv, self.head_dim) + gate = qg[..., self.head_dim :].reshape(-1, self._local_qo_dim) + k = k.contiguous().view(-1, self._local_num_kv, self.head_dim) v = v.contiguous() self.q_norm.forward_inplace(q) self.k_norm.forward_inplace(k) q, k = self.rotary.forward( - batch.positions, q.view(-1, self.qo_attn_dim), k.view(-1, self.kv_attn_dim) + batch.positions, q.view(-1, self._local_qo_dim), k.view(-1, self._local_kv_dim) ) index = self.indexer.forward(x) o = get_global_ctx().attn_backend.qsa_forward( - q.view(-1, self.num_q, self.head_dim), k, v, index, self.layer_id, batch + q.view(-1, self._local_num_q, self.head_dim), k, v, index, self.layer_id, batch ) - gated = o.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate) + gated = o.reshape(-1, self._local_qo_dim) * torch.sigmoid(gate) return self.o_proj.forward(gated) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index 69838153f..fc562e76c 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -3,8 +3,10 @@ import torch import torch.nn.functional as F from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen -from freetoken.layers import BaseOP, LinearColParallelMerged +from freetoken.layers import BaseOP, LinearColParallelMerged, LinearOProj +from freetoken.utils import div_even from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged @@ -80,6 +82,15 @@ def __init__( self.value_dim = num_v_heads * head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim self.conv_kernel_size = conv_kernel_size + # TP-local head counts: k heads and their v-head groups split evenly across ranks + # (the fla kernels take the GQA ratio from the shapes); the state pool is sharded the + # same way (kvcache.linear_state_pool._linear_local_dims). + tp = get_tp_info() + self._local_num_k_heads = div_even(num_k_heads, tp.size, allow_replicate=True) + self._local_num_v_heads = div_even(num_v_heads, tp.size, allow_replicate=True) + self._local_key_dim = self._local_num_k_heads * head_k_dim + self._local_value_dim = self._local_num_v_heads * head_v_dim + self._local_conv_dim = 2 * self._local_key_dim + self._local_value_dim # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). @@ -87,7 +98,14 @@ def __init__( self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 - self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] + self._in_proj_split = [ + self._local_conv_dim, self._local_value_dim, + self._local_num_v_heads, self._local_num_v_heads, + ] + if tp.size > 1: + assert not self._fp8 and attn_quant == "none", ( + "qwen4_exp TP shards bf16 GDN projections only" + ) if self._fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( @@ -98,21 +116,27 @@ def __init__( ) else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. - self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False) - self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size) + self.in_proj = LinearColParallelMerged( + hidden_size, [self.conv_dim, self.value_dim, num_v_heads, num_v_heads], + has_bias=False, local_output_sizes=self._in_proj_split, + ) + self.conv1d = _DepthwiseConv1d(self._local_conv_dim, conv_kernel_size) # Recurrence-gating params kept in fp32 (exp/softplus is precision-sensitive, # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a # per-call .float() upcast in the decode wrapper. The weight loader exempts # *.A_log / *.dt_bias from the model-dtype downcast. - self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32) - self.A_log = torch.empty(num_v_heads, dtype=torch.float32) + self.dt_bias = torch.empty(self._local_num_v_heads, dtype=torch.float32) + self.A_log = torch.empty(self._local_num_v_heads, dtype=torch.float32) self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps, activation=output_gate) # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. - self.out_proj = make_replicated_quant( - expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False - ) + if tp.size > 1: # row-parallel over the local v heads, all-reduce inside + self.out_proj = LinearOProj(self.value_dim, hidden_size, has_bias=False) + else: + self.out_proj = make_replicated_quant( + expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False + ) def _gate_params(self, a: torch.Tensor, b: torch.Tensor): beta = b.sigmoid() @@ -178,7 +202,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: else: proj = self.in_proj.forward(hidden_states) conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1) - z = z.reshape(total, self.num_v_heads, self.head_v_dim) + nk, nv = self._local_num_k_heads, self._local_num_v_heads + kd, vd = self._local_key_dim, self._local_value_dim + z = z.reshape(total, nv, self.head_v_dim) li = pool.local_index(self.layer_id) if batch.is_decode: @@ -187,10 +213,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # no clone, no external l2norm). q/k stay at num_k_heads (kernel handles GQA). mixed = self._conv_decode(conv_in, fla.cache_indices, pool) # [B, conv_dim] B = mixed.shape[0] - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, B, self.num_v_heads, self.head_v_dim).to(dtype) + qf, kf, vf = torch.split(mixed, [kd, kd, vd], dim=-1) + q = qf.reshape(1, B, nk, self.head_k_dim).to(dtype) + k = kf.reshape(1, B, nk, self.head_k_dim).to(dtype) + v = vf.reshape(1, B, nv, self.head_v_dim).to(dtype) core_out = gdn_decode_fla( q, k, v, a, b, A_log=self.A_log, dt_bias=self.dt_bias, state_source=pool.recurrent_states[li], indices=fla.cache_indices, @@ -200,13 +226,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: mixed = self._conv_prefill( conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state) # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads. - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, total, self.num_v_heads, self.head_v_dim).to(dtype) + qf, kf, vf = torch.split(mixed, [kd, kd, vd], dim=-1) + q = qf.reshape(1, total, nk, self.head_k_dim).to(dtype) + k = kf.reshape(1, total, nk, self.head_k_dim).to(dtype) + v = vf.reshape(1, total, nv, self.head_v_dim).to(dtype) g, beta = self._gate_params(a, b) - g = g.reshape(1, total, self.num_v_heads) - beta = beta.float().reshape(1, total, self.num_v_heads) + g = g.reshape(1, total, nv) + beta = beta.float().reshape(1, total, nv) # The chunk kernel reads + writes back initial_state[cache_indices] in place; # fresh sequences (cached_len==0) must start from a zeroed slot. if fla.fresh_state_indices is not None: diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 9ef65c0a8..e743a3af0 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -4,8 +4,12 @@ from typing import TYPE_CHECKING import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator, get_tp_info from freetoken.kernel.triton.moe_shared_gate import shared_gate_mul_add, shared_gate_sigmoid -from freetoken.layers.moe import make_moe_layer +from freetoken.layers import LinearRowParallel, silu_and_mul +from freetoken.layers.moe import OffloadMoELayer, make_moe_layer from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE if TYPE_CHECKING: @@ -16,9 +20,15 @@ class Qwen4ExpMoE(Qwen3_5MoE): """Qwen3_5MoE with the shared-expert gate on triton instead of gemv + sigmoid + mul + add. Same weights, same state dict. The gate reduction stays ahead of the routed experts, which may write into ``hidden_states`` in place. + + TP: the offload experts are sharded along the intermediate axis and the bf16 shared expert is + row-parallel, so both produce partial sums; ``routed + gate * shared`` is linear in them and + is reduced once (one all-reduce per MoE layer instead of two). """ def __init__(self, config: ModelConfig, layer_id: int | None = None) -> None: + self._comm = DistributedCommunicator() + self._tp_size = get_tp_info().size if getattr(config, "expert_quant", "none") != "fp8_block": super().__init__(config, layer_id=layer_id) return @@ -37,9 +47,22 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) router_logits = self.gate.forward(hidden_states) - shared = self.shared_expert.forward(hidden_states) gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) - routed = self.experts.forward(hidden_states=hidden_states, router_logits=router_logits) + se, ex = self.shared_expert, self.experts + if ( + self._tp_size > 1 + and isinstance(se.down_proj, LinearRowParallel) + and isinstance(ex, OffloadMoELayer) + ): + shared = F.linear(silu_and_mul(se.gate_up_proj.forward(hidden_states)), se.down_proj.weight) + if get_global_ctx().batch.is_prefill: + routed = ex.prefill_forward(hidden_states, router_logits) + else: + routed = ex.decode_forward(hidden_states, router_logits) + out = self._comm.all_reduce(shared_gate_mul_add(routed, shared, gate)) + return out.view(num_tokens, hidden_dim) + shared = se.forward(hidden_states) + routed = ex.forward(hidden_states=hidden_states, router_logits=router_logits) return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index f8d2a7494..b6868f6ba 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -21,13 +21,13 @@ import safetensors import torch from freetoken.distributed import get_tp_info -from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.loader import drop_page_cache, iter_weight_files, shard_tensor from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, load_nvfp4_expert_source_banks, ) from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import download_hf_weight +from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight from freetoken.utils.progress import byte_bar from tqdm import tqdm @@ -137,6 +137,58 @@ def _try_fuse( return None +def _shard_rows( + t: torch.Tensor, parts: list[tuple[int, int]], rank: int, world: int +) -> torch.Tensor: + """Column-parallel slice (dim 0) of a ``[part0 | part1 | ...]`` fusion; ``parts`` gives each + part as ``(heads, rows_per_head)``. Heads split evenly across ranks; a part with fewer heads + than ranks (GQA kv) replicates head ``rank * heads // world``, the ``div_even(..., + allow_replicate=True)`` convention of the TP-aware layers.""" + out, off = [], 0 + for heads, rows in parts: + local = div_even(heads, world, allow_replicate=True) + first = rank * heads // world + out.append(t[off + first * rows : off + (first + local) * rows]) + off += heads * rows + assert off == t.shape[0], f"fusion parts {parts} cover {off} rows, tensor has {t.shape[0]}" + return torch.cat(out, dim=0) + + +def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.Tensor: + """TP shard of one state-dict tensor (fused projections included); identity at TP=1. + + Column-parallel (dim 0, by head): attention ``qkv_proj`` [q|gate per head | k | v], GDN + ``in_proj`` [q | k | v | z | b | a] and the matching ``conv1d`` channels, ``A_log`` / + ``dt_bias``, shared-expert ``gate_up_proj``. Row-parallel (dim 1): ``o_proj``, + ``out_proj``, shared-expert ``down_proj``. Vocab rows: ``embed_tokens`` / ``lm_head``. + Everything else (router, indexer, norms, HC, PLE, shared-expert gate) is replicated. + """ + if world == 1: + return t + if name.endswith(".self_attn.qkv_proj.weight"): + q = (config.num_qo_heads, 2 * config.head_dim) + kv = (config.num_kv_heads, config.head_dim) + return _shard_rows(t, [q, kv, kv], rank, world) + if ".linear_attn." in name: + g = config.linear_attention_group() + k = (g.num_key_heads, g.key_head_dim) + v = (g.num_value_heads, g.value_head_dim) + if name.endswith(".in_proj.weight"): + return _shard_rows(t, [k, k, v, v, (v[0], 1), (v[0], 1)], rank, world) + if name.endswith(".conv1d.weight"): + return _shard_rows(t, [k, k, v], rank, world) + if name.endswith((".A_log", ".dt_bias")): + return _shard_rows(t, [(v[0], 1)], rank, world) + if name.endswith(".out_proj.weight"): + return t.chunk(world, dim=1)[rank].clone() + return t + if name.endswith(".shared_expert.gate_up_proj.weight"): + half = t.shape[0] // 2 + return _shard_rows(t, [(half, 1), (half, 1)], rank, world) + # o_proj / down_proj: dim 1; embed_tokens / lm_head: vocab rows; others unchanged. + return shard_tensor(name, t, rank=rank, world_size=world, num_kv_heads=None) + + def iter_weights( model_path: str, device: torch.device, @@ -157,16 +209,18 @@ def iter_weights( ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from :func:`load_nvfp4_expert_sources`. """ - if get_tp_info().size > 1: - raise NotImplementedError("qwen4_exp weight loading supports TP=1 only") if not include_non_moe: return + from .config import parse_config + + tp = get_tp_info() + config = parse_config(cached_load_hf_config(model_path)) fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), desc="Loading weights", - disable=not get_tp_info().is_primary(), + disable=not tp.is_primary(), ): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: for raw_name in f.keys(): @@ -177,9 +231,10 @@ def iter_weights( fused = _try_fuse(name, tensor, fuse_buf) if fused is not None: if fused != (): # () means buffered, not yet complete - yield fused + name, tensor = fused + yield name, _shard(name, tensor, config, tp.rank, tp.size) continue - yield name, tensor + yield name, _shard(name, tensor, config, tp.rank, tp.size) assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" diff --git a/tests/models/qwen4_exp/test_tp_shard.py b/tests/models/qwen4_exp/test_tp_shard.py new file mode 100644 index 000000000..87a305ef0 --- /dev/null +++ b/tests/models/qwen4_exp/test_tp_shard.py @@ -0,0 +1,92 @@ +"""TP sharding of the qwen4_exp dense weights (pure tensor math, no TP runtime needed).""" + +from types import SimpleNamespace + +import torch +from freetoken.models.qwen4_exp.weight import _shard, _shard_rows + + +def _cfg(): + g = SimpleNamespace( + num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8 + ) + return SimpleNamespace( + num_qo_heads=6, num_kv_heads=2, head_dim=8, linear_attention_group=lambda: g + ) + + +def _gather(name, t, cfg, world, dim=0): + return torch.cat([_shard(name, t, cfg, r, world) for r in range(world)], dim=dim) + + +def test_shard_rows_splits_each_part_by_head(): + t = torch.arange(4 * 4 + 2 * 4).reshape( + -1, 1 + ) # part A: 4 heads x 4 rows, part B: 2 heads x 4 + s = [_shard_rows(t, [(4, 4), (2, 4)], r, 4) for r in range(4)] + assert all(x.shape[0] == 8 for x in s) + assert s[0][:4].flatten().tolist() == [0, 1, 2, 3] + assert s[3][:4].flatten().tolist() == [12, 13, 14, 15] + # 2 kv heads over 4 ranks replicate: rank * heads // world -> heads 0, 0, 1, 1 + assert s[0][4:].equal(s[1][4:]) and s[2][4:].flatten().tolist() == [20, 21, 22, 23] + + +def test_shard_round_trips_the_fused_projections(): + cfg, world = _cfg(), 2 + qkv = torch.randn(6 * 16 + 2 * 8 + 2 * 8, 5) + assert ( + _gather("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, world).shape + == qkv.shape + ) + q0 = _shard("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, 0, world) + assert q0.shape[0] == 3 * 16 + 8 + 8 + torch.testing.assert_close(q0[:48], qkv[:48]) # q heads 0-2 (16 rows each: q|gate) + torch.testing.assert_close(q0[48:56], qkv[96:104]) # kv head 0 of k + torch.testing.assert_close(q0[56:64], qkv[112:120]) # kv head 0 of v + kd, vd, nv = 4 * 8, 12 * 8, 12 + in_proj = torch.randn(kd + kd + vd + vd + nv + nv, 5) + s = _shard("model.layers.1.linear_attn.in_proj.weight", in_proj, cfg, 1, world) + assert s.shape[0] == (kd + kd + vd + vd + nv + nv) // 2 + torch.testing.assert_close(s[:16], in_proj[16:32]) # q: k heads 2,3 + torch.testing.assert_close(s[-6:], in_proj[-6:]) # a: v heads 6-11 + conv = torch.randn(kd + kd + vd, 1, 4) + assert _shard( + "model.layers.1.linear_attn.conv1d.weight", conv, cfg, 0, world + ).shape == ((kd + kd + vd) // 2, 1, 4) + a_log = torch.randn(nv) + torch.testing.assert_close( + _shard("model.layers.1.linear_attn.A_log", a_log, cfg, 1, world), a_log[6:] + ) + + +def test_shard_row_parallel_and_vocab_and_replicated(): + cfg, world = _cfg(), 2 + for name in ( + "model.layers.0.self_attn.o_proj.weight", + "model.layers.1.linear_attn.out_proj.weight", + "model.layers.0.mlp.shared_expert.down_proj.weight", + ): + t = torch.randn(3, 8) + torch.testing.assert_close(_gather(name, t, cfg, world, dim=1), t) + gate_up = torch.randn(2 * 6, 3) + s1 = _shard( + "model.layers.0.mlp.shared_expert.gate_up_proj.weight", gate_up, cfg, 1, world + ) + torch.testing.assert_close(s1, torch.cat([gate_up[3:6], gate_up[9:12]])) + emb = torch.randn(10, 3) + torch.testing.assert_close( + _gather("model.embed_tokens.weight", emb, cfg, world), emb + ) + torch.testing.assert_close(_gather("lm_head.weight", emb, cfg, world), emb) + for name in ( + "model.layers.0.mlp.gate.weight", + "model.layers.0.self_attn.indexer.index_qk_proj.weight", + "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + "model.layers.1.ple.value_proj.weight", + "model.layers.0.mlp.shared_expert_gate.weight", + ): + t = torch.randn(4, 6) + assert _shard(name, t, cfg, 1, world).equal(t) + assert _shard("model.layers.0.self_attn.qkv_proj.weight", gate_up, cfg, 0, 1).equal( + gate_up + ) diff --git a/tests/models/test_nvfp4_banks_tp.py b/tests/models/test_nvfp4_banks_tp.py new file mode 100644 index 000000000..13bf2a2c7 --- /dev/null +++ b/tests/models/test_nvfp4_banks_tp.py @@ -0,0 +1,76 @@ +"""TP sharding of the NVFP4 expert source banks: the two ranks' banks concatenated along the +intermediate axis must equal the unsharded placement.""" + +import torch +from freetoken.distributed.info import DistributedInfo +from freetoken.models import nvfp4_banks +from freetoken.models.nvfp4_banks import _alloc_nvfp4_host_banks, _Placer + +E, H, INTER = 2, 64, 32 + + +def _place(monkeypatch, rank, world): + monkeypatch.setattr( + nvfp4_banks, "get_tp_info", lambda: DistributedInfo(rank, world) + ) + hb = _alloc_nvfp4_host_banks(1, E, H, INTER // world) + banks = {name: [b.tensor for b in layers] for name, layers in hb.items()} + placer = _Placer(banks, INTER) + torch.manual_seed(0) + for e in range(E): + for role in ("gate", "up"): + placer.put( + 0, + e, + role, + "weight", + torch.randint(0, 255, (INTER, H // 2), dtype=torch.uint8), + ) + placer.put( + 0, + e, + role, + "weight_scale", + torch.randn(INTER, H // 16).to(torch.float8_e4m3fn), + torch.tensor(0.5 + e, dtype=torch.float16), + ) + placer.put( + 0, + e, + "down", + "weight", + torch.randint(0, 255, (H, INTER // 2), dtype=torch.uint8), + ) + placer.put( + 0, + e, + "down", + "weight_scale", + torch.randn(H, INTER // 16).to(torch.float8_e4m3fn), + torch.tensor(2.0 + e, dtype=torch.float16), + ) + return {k: v[0] for k, v in banks.items()} + + +def test_rank_banks_concatenate_to_the_full_placement(monkeypatch): + full = _place(monkeypatch, 0, 1) + r0, r1 = _place(monkeypatch, 0, 2), _place(monkeypatch, 1, 2) + n = INTER // 2 + for name in ( + "gate_up_packed", + "gate_up_scale", + "gate_up_global", + ): # rows: [gate I | up I] + gate = torch.cat([r0[name][:, :n], r1[name][:, :n]], dim=1) + up = torch.cat([r0[name][:, n:], r1[name][:, n:]], dim=1) + assert torch.equal( + torch.cat([gate, up], dim=1).view(torch.uint8), full[name].view(torch.uint8) + ), name + for name in ("down_packed", "down_scale"): # columns + assert torch.equal( + torch.cat([r0[name], r1[name]], dim=2).view(torch.uint8), + full[name].view(torch.uint8), + ), name + assert torch.equal(r0["down_global"], full["down_global"]) and torch.equal( + r1["down_global"], full["down_global"] + ) From d16d63492e969f9ae7f5efb497ea71f8a9505c39 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 15:09:33 -0400 Subject: [PATCH 2/5] fix(qwen4_exp): load the HF config for sharding only when TP > 1 tests/models/qwen4_exp/test_weight.py feeds iter_weights a synthetic checkpoint whose config.json has no model_type; at TP=1 nothing is sharded, so do not touch the config. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/weight.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index b6868f6ba..5a0013bfb 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -215,7 +215,9 @@ def iter_weights( from .config import parse_config tp = get_tp_info() - config = parse_config(cached_load_hf_config(model_path)) + # The sharding geometry needs the HF config; TP=1 never shards, so keep the plain path free + # of a config load (synthetic test checkpoints carry no model_type). + config = parse_config(cached_load_hf_config(model_path)) if tp.size > 1 else None fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), From 70d9e8dfbabddd675ed9af9d8b69b985b3aa3baa Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 14:21:15 -0400 Subject: [PATCH 3/5] feat(qwen4_exp): image input (vision tower + mRoPE) for Qwen3.8-Flash-Next Serve images through the OpenAI chat endpoint. Opt-in with FREETOKEN_LOAD_VISION=1. Model side: the HF Qwen4ExpVisionModel runs inside a BaseOP whose tensors load as ``visual.*`` with the dense weights (meta build, assign-on-load, rotary buffer rebuilt on the device), so the expert-cache planner counts them and --dummy-weight works. Soft tokens replace the image placeholders before the hyper-connection repeat. mRoPE: ``mrope.py`` ports HF get_rope_index (3-D T/H/W positions, decode delta) and the interleaved cos|sin rows. A prefill batch with image tokens gets a per-token cos|sin table that the existing rope kernels index by row (attention and the QSA indexer); decode reads the normal cache at position + delta. The table carries index_ratio - 1 lead rows per request so a straddling indexer group can be roped at its first token. Text-only batches alias positions and run the same kernels as before. Request path: image_url parts (inline data: URLs only, 16 MiB cap) are decoded in the API server; the tokenizer worker runs the checkpoint's image processor (FREETOKEN_IMAGE_MAX_PIXELS, default 1280*28*28) and expands each <|image_pad|>; the scheduler encodes the images on every TP rank and computes the rope positions before admission. Image prompts must fit one prefill chunk (rejected with an error otherwise). The wire encoder now carries N-D tensors. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/attention/qsa_sparse.py | 20 ++- python/freetoken/core.py | 11 ++ python/freetoken/engine/graph.py | 5 + python/freetoken/layers/rotary.py | 5 +- python/freetoken/message/backend.py | 6 + python/freetoken/message/tokenizer.py | 2 + python/freetoken/message/utils.py | 8 +- .../freetoken/models/qwen4_exp/attention.py | 8 +- python/freetoken/models/qwen4_exp/config.py | 11 +- python/freetoken/models/qwen4_exp/model.py | 76 ++++++++- python/freetoken/models/qwen4_exp/mrope.py | 112 ++++++++++++ python/freetoken/models/qwen4_exp/vision.py | 80 +++++++++ python/freetoken/models/qwen4_exp/weight.py | 12 +- python/freetoken/scheduler/prefill.py | 19 ++- python/freetoken/scheduler/scheduler.py | 52 ++++++ python/freetoken/scheduler/utils.py | 2 + python/freetoken/server/generation.py | 63 ++++++- python/freetoken/server/openai_api.py | 4 +- python/freetoken/tokenizer/server.py | 16 +- python/freetoken/tokenizer/tokenize.py | 70 +++++++- tests/models/qwen4_exp/common.py | 4 + tests/models/qwen4_exp/test_mrope.py | 161 ++++++++++++++++++ tests/models/qwen4_exp/test_mrope_gpu.py | 79 +++++++++ tests/models/qwen4_exp/test_qsa_backend.py | 1 + tests/models/qwen4_exp/test_skeleton.py | 6 +- tests/scheduler/test_cost_accounting_core.py | 11 +- 26 files changed, 805 insertions(+), 39 deletions(-) create mode 100644 python/freetoken/models/qwen4_exp/mrope.py create mode 100644 python/freetoken/models/qwen4_exp/vision.py create mode 100644 tests/models/qwen4_exp/test_mrope.py create mode 100644 tests/models/qwen4_exp/test_mrope_gpu.py diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 4a28dc852..85c901211 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -88,6 +88,10 @@ class QSASparseMetadata(BaseAttnMetadata): cmp_rows: torch.Tensor | None = None # [T] int32, compressed slab destination ring_rows: torch.Tensor | None = None # [T] int32, flat ring row or -1 positions: torch.Tensor | None = None # [T] int32, logical query positions + # mRoPE: rope position per token (positions + Req.mrope_delta), or -- when mrope_cos_sin is + # set (a prefill batch with image tokens) -- the token's row in that per-token cos|sin table. + rope_positions: torch.Tensor | None = None # [T] int32 + mrope_cos_sin: torch.Tensor | None = None # [T, rotary_dim] fp32 or None # fmt: on def get_last_indices(self, bs: int) -> torch.Tensor: @@ -296,6 +300,10 @@ def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None: """Per-token slab row and ring row for this forward; the other QSA layers reuse it (it is layer-invariant). Pure device arithmetic: no host sync, graph-capturable.""" md.positions = batch.positions + md.rope_positions = ( + batch.positions if batch.rope_positions is None else batch.rope_positions + ) + md.mrope_cos_sin = batch.mrope_cos_sin out_loc = batch.out_loc.to(torch.int64) positions = batch.positions.to(torch.int64) rows = torch.arange(out_loc.numel(), device=self.device) @@ -338,10 +346,16 @@ def _update_index_cache(self, index, md: QSASparseMetadata, slot: int) -> None: pooled, first, ) + if md.rope_positions is not md.positions: + # Rope the pooled key at its group's FIRST token (HF: cos/sin indexed by group start). + # That token's rope position / table row sits at the same offset from this token's as + # the token indices do. Decode groups that close after an image are text (the chat + # template puts >= 4 tokens after <|vision_end|>), so ``first + delta`` is exact. + first = first + (md.rope_positions - md.positions) qsa_index_norm_rope( pooled, first, - self._index_rope_cache(), + self._index_rope_cache() if md.mrope_cos_sin is None else md.mrope_cos_sin, index.k_norm_weight, index.eps, self.kvcache.cmp_k_cache(slot), @@ -366,8 +380,8 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor: ) qsa_index_norm_rope( index.q.view(-1, self.index_head_dim), - positions, - self._index_rope_cache(), + md.rope_positions, + self._index_rope_cache() if md.mrope_cos_sin is None else md.mrope_cos_sin, index.q_norm_weight, index.eps, q_index.view(-1, self.index_head_dim), diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539cb..22fcf56ab 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -43,6 +43,11 @@ class Req: # Optional precomputed multimodal soft-token embeddings (GPU, [num_image_tokens, # hidden]) scattered at image-token positions during this request's prefill. mm_embeds: torch.Tensor | None = None + # mRoPE (qwen4_exp image prompts): ``[3, prompt_len]`` T/H/W rope positions of the prompt + # tokens (CPU) and the offset decode adds to a token index (``max_pos + 1 - prompt_len``, + # <= 0). None / 0 for text-only prompts, where rope position == token index. + mrope_positions: torch.Tensor | None = None + mrope_delta: int = 0 # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or # until allocated from LinearStatePool. Set by the scheduler (P2). --- @@ -135,6 +140,12 @@ class Batch: attn_metadata: BaseAttnMetadata = field(init=False) # concatenated multimodal soft-token embeddings for a prefill batch (or None) mm_embeds: torch.Tensor | None = field(default=None, init=False) + # Rope positions per token (``positions + Req.mrope_delta``); the same tensor as ``positions`` + # when no request in the batch carries an image. Set by the scheduler / graph buffer. + rope_positions: torch.Tensor | None = field(default=None, init=False) + # Prefill batches with image tokens: per-token mRoPE cos|sin rows ``[T, rotary_dim]`` (fp32) + # that the attention layers use as the rope cache with ``positions = arange(T)``. + mrope_cos_sin: torch.Tensor | None = field(default=None, init=False) # Prefill log stats snapshotted at schedule time (before forward's complete_one() # advances cached_len), so the prefill log reports the tokens actually forwarded and # the prefix-cache hit -- matching SGLang's #new-token / #cached-token. Set by the diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f2025025..38944a0b4 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -24,6 +24,7 @@ class GraphCaptureBuffer: input_ids: torch.Tensor out_loc: torch.Tensor positions: torch.Tensor + rope_positions: torch.Tensor # positions + per-request mRoPE delta (== positions, text-only) logits: torch.Tensor table_idx: torch.Tensor # per-request slot id for GatedDeltaNet state gather/scatter # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once. @@ -35,6 +36,7 @@ def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuf input_ids=torch.zeros(bs, dtype=torch.int32, device=device), out_loc=torch.zeros(bs, dtype=torch.int32, device=device), positions=torch.zeros(bs, dtype=torch.int32, device=device), + rope_positions=torch.zeros(bs, dtype=torch.int32, device=device), logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), table_idx=torch.zeros(bs, dtype=torch.int32, device=device), fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device), @@ -48,6 +50,7 @@ def set_batch(self, batch: Batch) -> None: batch.input_ids = self.input_ids[_slice] batch.out_loc = self.out_loc[_slice] batch.positions = self.positions[_slice] + batch.rope_positions = self.rope_positions[_slice] batch.linear_table_idx = self.table_idx[_slice] # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the # persistent table_idx slot map, so the captured kernels see stable addresses. @@ -61,6 +64,8 @@ def copy_from(self, batch: Batch) -> None: if batch.out_loc is not None: self.out_loc[_slice] = batch.out_loc self.positions[_slice] = batch.positions + rope = batch.positions if batch.rope_positions is None else batch.rope_positions + self.rope_positions[_slice] = rope if batch.linear_table_idx is not None: self.table_idx[_slice] = batch.linear_table_idx diff --git a/python/freetoken/layers/rotary.py b/python/freetoken/layers/rotary.py index 3756a299b..e55aabe4f 100644 --- a/python/freetoken/layers/rotary.py +++ b/python/freetoken/layers/rotary.py @@ -73,13 +73,16 @@ def forward( positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor, + cos_sin_cache: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor]: + # ``cos_sin_cache`` overrides the position-indexed cache with a per-row table (mRoPE: + # row i holds token i's cos|sin, positions = arange), same [rows, rotary_dim] layout. self.apply_rope_with_cos_sin_cache_inplace( positions=positions, query=query, key=key, head_size=self.head_size, - cos_sin_cache=self._cos_sin_cache, + cos_sin_cache=self._cos_sin_cache if cos_sin_cache is None else cos_sin_cache, is_neox=self.is_neox, ) return query, key diff --git a/python/freetoken/message/backend.py b/python/freetoken/message/backend.py index c42ecc5a6..d42a83c06 100644 --- a/python/freetoken/message/backend.py +++ b/python/freetoken/message/backend.py @@ -37,6 +37,12 @@ class UserMsg(BaseBackendMsg): # Optional precomputed multimodal soft-token embeddings (GPU tensor). Only used by # the in-process offline path; remains None for the (serialized) online path. mm_embeds: torch.Tensor | None = None + # Online image path: processor outputs from the tokenizer worker (CPU fp32 ``pixel_values`` + # [patches, C*T*P*P] + ``image_grid_thw`` [N, 3]); the scheduler encodes them on its rank + # into mm_embeds / mrope_positions / mrope_delta (see Req) before admission. + mm_inputs: dict | None = None + mrope_positions: torch.Tensor | None = None + mrope_delta: int = 0 @dataclass diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 33b75c785..2fb9000ae 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -72,6 +72,8 @@ class TokenizeMsg(BaseTokenizerMsg): sampling_params: SamplingParams chat_template_kwargs: Dict[str, Any] | None = None tools: List[Dict[str, Any]] | None = None + # Encoded image files, in the order their ``{"type": "image"}`` parts appear in ``text``. + images: List[bytes] | None = None @dataclass diff --git a/python/freetoken/message/utils.py b/python/freetoken/message/utils.py index ee92adf5d..e17be3596 100644 --- a/python/freetoken/message/utils.py +++ b/python/freetoken/message/utils.py @@ -32,10 +32,11 @@ def serialize_type(self) -> Dict: serialized = {} if isinstance(self, torch.Tensor): - assert self.dim() == 1, "we can only serialize 1D tensor for now" serialized["__type__"] = "Tensor" - serialized["buffer"] = self.numpy().tobytes() + serialized["buffer"] = self.contiguous().numpy().tobytes() serialized["dtype"] = str(self.dtype) + if self.dim() != 1: # 1-D stays shape-less (wire compatible); N-D carries its shape + serialized["shape"] = list(self.shape) return serialized # normal type @@ -64,14 +65,13 @@ def _deserialize_any(cls_map: Dict[str, Type], data: Any) -> Any: def deserialize_type(cls_map: Dict[str, Type], data: Dict) -> Any: type_name = data["__type__"] - # we can only serialize 1D tensor for now if type_name == "Tensor": buffer = data["buffer"] dtype_str = data["dtype"].replace("torch.", "") np_dtype = getattr(np, dtype_str) assert isinstance(buffer, bytes) np_tensor = np.frombuffer(buffer, dtype=np_dtype) - return torch.from_numpy(np_tensor.copy()) + return torch.from_numpy(np_tensor.copy()).reshape(data.get("shape", (-1,))) cls = cls_map.get(type_name) if cls is None: diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index db77d2b6e..2698a8970 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -164,8 +164,14 @@ def forward(self, x: torch.Tensor, batch: Batch) -> torch.Tensor: v = v.contiguous() self.q_norm.forward_inplace(q) self.k_norm.forward_inplace(k) + # mRoPE: rope_positions is positions + the request's delta (text after an image), or + # row indices into batch.mrope_cos_sin for a prefill batch that carries image tokens. + rope_pos = batch.positions if batch.rope_positions is None else batch.rope_positions q, k = self.rotary.forward( - batch.positions, q.view(-1, self._local_qo_dim), k.view(-1, self._local_kv_dim) + rope_pos, + q.view(-1, self._local_qo_dim), + k.view(-1, self._local_kv_dim), + cos_sin_cache=batch.mrope_cos_sin, ) index = self.indexer.forward(x) o = get_global_ctx().attn_backend.qsa_forward( diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index bb5d1dff5..6db46bcb4 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -12,6 +12,7 @@ ModelConfig, RotaryConfig, SlotStateSpec, + vision_load_enabled, ) @@ -40,6 +41,10 @@ class Qwen4ExpArgs: index_head_dim: int index_budget: int index_ratio: int + # mRoPE: which rotary frequencies follow the image H / W axes (HF mrope_section, interleaved). + mrope_section: Tuple[int, ...] = (11, 11, 10) + # Vision patch merge (2 -> one soft token per 2x2 patches); only read when vision is loaded. + spatial_merge_size: int = 2 @property def index_topk_blocks(self) -> int: @@ -233,6 +238,8 @@ def _quant(probe: str) -> str: if isinstance(eos_token_id, (list, tuple)): eos_token_id = eos_token_id[0] + # Vision is opt-in (FREETOKEN_LOAD_VISION=1): the tower is ~0.9 GiB bf16 per rank. + vision_config = getattr(hf_config, "vision_config", None) if vision_load_enabled() else None qwen4_args = Qwen4ExpArgs( hidden_size=text.hidden_size, hc_count=int(text.hc_count), @@ -251,6 +258,8 @@ def _quant(probe: str) -> str: index_head_dim=int(text.indexer_head_dim), index_budget=int(text.indexer_budget), index_ratio=int(text.indexer_compress_ratio), + mrope_section=tuple(int(v) for v in rope_params.get("mrope_section", (11, 11, 10))), + spatial_merge_size=int(getattr(vision_config, "spatial_merge_size", 2)), ) return ModelConfig( @@ -278,7 +287,7 @@ def _quant(probe: str) -> str: use_qk_norm=True, model_type=getattr(hf_config, "model_type", "qwen4_exp"), architectures=getattr(hf_config, "architectures", ["Qwen4ExpForConditionalGeneration"]), - vision_config=None, # served text-only + vision_config=vision_config, image_token_id=getattr(hf_config, "image_token_id", None), attention_groups=groups, expert_quant=expert_quant, diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 19c31e2fc..7d02d6d15 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -1,4 +1,4 @@ -"""Qwen3.8-Flash-Next decoder stack (text-only). +"""Qwen3.8-Flash-Next decoder stack (text, plus images when the vision tower is loaded). The residual state is ``R [T, hc_count*hidden]`` end to end: the embedding is repeated over the ``hc_count`` streams, every layer mixes them down to one ``[T, hidden]`` block input and injects @@ -99,6 +99,7 @@ def __init__(self, config: ModelConfig) -> None: self.hyper_connection_mixer = GatedResidual(config, use_combine=False) # plain tuple (not an OP child), so it never shows up in the state dict self._ple = tuple(layer.ple for layer in self.layers.op_list if layer.ple is not None) + self._image_token_id = config.image_token_id @property def ple_layers(self) -> List[PLELayer]: @@ -106,7 +107,13 @@ def ple_layers(self) -> List[PLELayer]: return list(self._ple) def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: - hidden = self.embed_tokens.forward(input_ids).repeat(1, self.hc_count) + hidden = self.embed_tokens.forward(input_ids) + if batch.mm_embeds is not None: + # image soft tokens replace the placeholder embeddings (HF order: before the + # hc_count repeat); the whole image run sits in this prefill chunk (prefill.py) + mask = input_ids == self._image_token_id + hidden = hidden.masked_scatter(mask.unsqueeze(-1), batch.mm_embeds.to(hidden.dtype)) + hidden = hidden.repeat(1, self.hc_count) meta = None if self._ple: from .ple import build_ple_metadata, commit_ngram_context @@ -126,6 +133,7 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: class Qwen4ExpForCausalLM(BaseLLMModel): def __init__(self, config: ModelConfig) -> None: self._config = config + self._inv_freq: torch.Tensor | None = None self.model = Qwen4ExpModel(config) if getattr(config, "lm_head_quant", "none") == "nvfp4": from freetoken.kernel.triton.nvfp4_linear import Nvfp4LMHead @@ -141,8 +149,72 @@ def __init__(self, config: ModelConfig) -> None: tie_word_embeddings=config.tie_word_embeddings, tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, ) + if config.vision_config is not None: # FREETOKEN_LOAD_VISION=1 + from .vision import Qwen4ExpVisionTower + + self.visual = Qwen4ExpVisionTower(config.vision_config) super().__init__() + @property + def has_vision(self) -> bool: + return hasattr(self, "visual") + + @torch.inference_mode() + def encode_images(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + """Vision tower + merger on processor outputs: ``[num_image_tokens, hidden]`` (device).""" + if not self.has_vision: + raise RuntimeError("image inputs need the vision tower: start with FREETOKEN_LOAD_VISION=1") + return self.visual.forward(pixel_values, grid_thw) + + def prepare_mm_inputs( + self, input_ids: torch.Tensor, mm_inputs: dict + ) -> tuple[torch.Tensor, torch.Tensor, int]: + """One prompt's ``(mm_embeds [n, hidden] on device, mrope_positions [3, L] CPU, mrope_delta)``.""" + from .mrope import rope_index + + grid = mm_inputs["image_grid_thw"] + embeds = self.encode_images(mm_inputs["pixel_values"], grid) + image_token_id = self._config.image_token_id + n = int((input_ids == image_token_id).sum()) + if n != embeds.shape[0]: + raise ValueError(f"{n} image placeholder tokens for {embeds.shape[0]} image features") + pos, delta = rope_index( + input_ids, grid, image_token_id, self._config.qwen4_args.spatial_merge_size + ) + return embeds, pos, delta + + def mrope_table(self, reqs, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """``(rope_positions, mrope_cos_sin)`` for a prefill batch with image tokens; see + :func:`mrope.mrope_table`.""" + from .mrope import mrope_table + + return mrope_table( + reqs, + self._config.qwen4_args.index_ratio, + self._inv_freq_on(device), + self._config.qwen4_args.mrope_section, + device, + ) + + def _inv_freq_on(self, device: torch.device) -> torch.Tensor: + """Attention rope frequencies (fp32, the RotaryEmbedding formula) on ``device``.""" + if self._inv_freq is None or self._inv_freq.device != device: + rc = self._config.rotary_config + self._inv_freq = 1.0 / ( + rc.base + ** (torch.arange(0, rc.rotary_dim, 2, dtype=torch.float, device=device) / rc.rotary_dim) + ) + return self._inv_freq + + def mrope_cos_sin(self, positions: torch.Tensor) -> torch.Tensor: + """``[T, rotary_dim]`` fp32 cos|sin rows for 3-D ``positions [3, T]`` (same frequencies as + the attention rope cache, so text rows equal the cache rows bit for bit).""" + from .mrope import mrope_cos_sin + + return mrope_cos_sin( + positions, self._inv_freq_on(positions.device), self._config.qwen4_args.mrope_section + ) + def load_host_tables(self, engine_config) -> int: """Attach the PLE n-gram table (pinned checkpoint bank, or zeros for dummy weights); returns the pinned host bytes the engine reserves from its pin budget.""" ple_layers = self.model.ple_layers diff --git a/python/freetoken/models/qwen4_exp/mrope.py b/python/freetoken/models/qwen4_exp/mrope.py new file mode 100644 index 000000000..9aa9157b0 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/mrope.py @@ -0,0 +1,112 @@ +"""mRoPE for Qwen3.8-Flash-Next image prompts. + +Ports two pieces of HF ``modeling_qwen4_exp``: ``get_rope_index`` (the 3-D T/H/W positions of a +prompt and the delta the text after the last image continues from) and +``Qwen4ExpTextRotaryEmbedding.apply_interleaved_mrope`` (the per-token cos|sin row). Text-only +prompts reduce to ``positions`` on all three axes and a zero delta, so the serving path only +builds a table for prefill batches that carry images; decode reads the normal cache at +``position + delta``. + +The cos|sin table has the same ``[rows, rotary_dim]`` layout as ``RotaryEmbedding._cos_sin_cache`` +(cos on the first half, sin on the second), so the flashinfer/triton rope kernels and the QSA +indexer's ``qsa_index_norm_rope`` take it as the cache with ``positions = arange(rows)``. +""" + +from __future__ import annotations + +import torch + + +def rope_index( + input_ids: torch.Tensor, + grid_thw: torch.Tensor | None, + image_token_id: int, + merge: int, +) -> tuple[torch.Tensor, int]: + """``([3, L] int64 positions, delta)`` for one prompt (HF ``get_rope_index``, images only). + + Text runs advance all three axes by their length. An image of ``t x h x w`` patches takes + ``t * (h // merge) * (w // merge)`` placeholder tokens at T/H/W offsets from the current + position and advances it by ``max(h, w) // merge``. ``delta`` is what decode adds to a + token index to get its rope position (``max_pos + 1 - L``, <= 0). + """ + ids = input_ids.tolist() + length = len(ids) + out = torch.empty(3, length, dtype=torch.int64) + cur = i = img = 0 + while i < length: + if ids[i] == image_token_id: + t, h, w = (int(v) for v in grid_thw[img]) + img += 1 + hh, ww = h // merge, w // merge + n = t * hh * ww + assert ids[i : i + n] == [image_token_id] * n, ( + "image placeholder run is too short" + ) + tt, hp, wp = torch.meshgrid( + torch.arange(t), torch.arange(hh), torch.arange(ww), indexing="ij" + ) + out[0, i : i + n] = tt.reshape(-1) + cur + out[1, i : i + n] = hp.reshape(-1) + cur + out[2, i : i + n] = wp.reshape(-1) + cur + cur += max(hh, ww) + i += n + continue + j = i + while j < length and ids[j] != image_token_id: + j += 1 + out[:, i:j] = torch.arange(cur, cur + (j - i)) + cur += j - i + i = j + assert grid_thw is None or img == len(grid_thw), "more images than placeholder runs" + return out, int(out.max()) + 1 - length + + +def mrope_cos_sin( + positions: torch.Tensor, inv_freq: torch.Tensor, section: tuple[int, ...] +) -> torch.Tensor: + """``[T, 2 * len(inv_freq)]`` fp32 cos|sin rows for 3-D ``positions [3, T]``. + + Interleaved layout (HF ``apply_interleaved_mrope``): frequency ``k`` rotates by the T axis, + except ``k = 1 + 3m`` (m < section[1]) which use H and ``k = 2 + 3m`` (m < section[2]) W. + """ + freqs = positions.to(inv_freq.dtype).unsqueeze(-1) * inv_freq # [3, T, n] + f = freqs[0].clone() + for axis in (1, 2): + idx = slice(axis, section[axis] * 3, 3) + f[:, idx] = freqs[axis][:, idx] + return torch.cat((f.cos(), f.sin()), dim=-1) + + +def mrope_table( + reqs, + ratio: int, + inv_freq: torch.Tensor, + section: tuple[int, ...], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Rope rows for a prefill batch that carries image tokens: ``(rope_positions [T] int32 = each + token's row, table [rows, rotary_dim] fp32 cos|sin)``. + + Every request also gets rows for up to ``ratio - 1`` positions before its chunk: the QSA + indexer ropes a pooled key at its group's FIRST token, which for a chunked text prompt sharing + the batch may come from the previous chunk (image prompts are never chunked). Requests without + ``mrope_positions`` use their token index on all three axes, so their rows equal the rope + cache rows exactly. + """ + segments, rows, base = [], [], 0 + for r in reqs: + pre = min(r.cached_len, ratio - 1) + start = r.cached_len - pre + if getattr(r, "mrope_positions", None) is not None: + seg = r.mrope_positions[:, start : r.device_len] + else: + seg = torch.arange(start, r.device_len).expand(3, -1) + segments.append(seg) + rows.append(torch.arange(base + pre, base + seg.shape[1], dtype=torch.int32)) + base += seg.shape[1] + table = mrope_cos_sin(torch.cat(segments, dim=1).to(device), inv_freq, section) + return torch.cat(rows).to(device), table + + +__all__ = ["mrope_cos_sin", "mrope_table", "rope_index"] diff --git a/python/freetoken/models/qwen4_exp/vision.py b/python/freetoken/models/qwen4_exp/vision.py new file mode 100644 index 000000000..23c1c47a7 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/vision.py @@ -0,0 +1,80 @@ +"""Qwen3.8-Flash-Next vision tower: HF ``Qwen4ExpVisionModel`` behind the BaseOP state-dict contract. + +The tower is bf16, ~0.9 GiB, never quantized and identical on every TP rank (each rank encodes +its own copy of the request's images; the soft tokens are tiny next to the tower). Its tensors +travel under the ``visual.*`` keys of the model state dict so the engine loads them with the +dense weights: counted in the weight bytes before the expert-cache planner sizes the cache, +dummy-weight builds fill them like any other tensor, and TP ranks all see the same keys. + +Meta-device trap: ``create_model`` runs under ``torch.device("meta")``. The HF module built +there is shapes only; ``load_state_dict`` rebuilds a meta module, assigns the loaded tensors +(no second copy on the GPU) and re-creates the non-persistent rotary ``inv_freq`` buffer on the +device, which ``assign=True`` would otherwise leave on meta. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from freetoken.layers import BaseOP +from freetoken.layers.base import _concat_prefix + + +def _vision_model_cls(): + from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpVisionModel + + return Qwen4ExpVisionModel + + +class Qwen4ExpVisionTower(BaseOP): + def __init__(self, vision_config: Any) -> None: + # The standalone vision config carries no attention backend; sdpa splits packed images + # by cu_seqlens, so several images per call are fine. + vision_config._attn_implementation = "sdpa" + self._config = vision_config + self._hf = _vision_model_cls()(vision_config) + self.spatial_merge_size = int(vision_config.spatial_merge_size) + + def state_dict(self, *, prefix: str = "", result=None): + result = result if result is not None else {} + for name, tensor in self._hf.state_dict().items(): + result[_concat_prefix(prefix, name)] = tensor + return result + + def load_state_dict( + self, state_dict, *, prefix: str = "", _internal: bool = False + ) -> None: + own = { + k: state_dict.pop(_concat_prefix(prefix, k)) for k in self._hf.state_dict() + } + device = next(iter(own.values())).device + with torch.device("meta"): + hf = _vision_model_cls()(self._config) + hf.load_state_dict(own, strict=True, assign=True) + head_dim = self._config.hidden_size // self._config.num_heads + with torch.device(device): + hf.rotary_pos_emb = type(hf.rotary_pos_emb)(head_dim // 2) + stuck = [n for n, b in hf.named_buffers() if b.is_meta] + assert not stuck, f"vision buffers left on meta: {stuck}" + self._hf = hf.eval().requires_grad_(False) + if not _internal and state_dict: + raise RuntimeError( + f"Unexpected keys in state_dict: {list(state_dict.keys())}" + ) + + @torch.inference_mode() + def forward( + self, pixel_values: torch.Tensor, grid_thw: torch.Tensor + ) -> torch.Tensor: + """``[sum(t*h*w) / merge**2, text_hidden]`` soft tokens for the packed images.""" + param = next(self._hf.parameters()) + out = self._hf( + pixel_values.to(param.device, param.dtype), + grid_thw=grid_thw.to(param.device), + return_dict=True, + ) + return out.pooler_output + + +__all__ = ["Qwen4ExpVisionTower"] diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 5a0013bfb..bfca63a3c 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -6,7 +6,7 @@ * :func:`load_ple_table` -- the 47.7 GiB FP8 n-gram table, 128 checkpoint shards concatenated into one pinned :class:`HostBank`. * :func:`load_nvfp4_expert_sources` -- the routed NVFP4 experts, into the offload cache's source banks. -Dropped: ``mtp.*`` (speculative head, including its stacked ``mtp.layers.0.mlp.experts.*``) and ``model.visual.*`` (served text-only). +Dropped: ``mtp.*`` (speculative head, including its stacked ``mtp.layers.0.mlp.experts.*``); ``model.visual.*`` is dropped unless vision is opted in (``FREETOKEN_LOAD_VISION=1``), then it loads as ``visual.*``. """ from __future__ import annotations @@ -20,6 +20,7 @@ import safetensors import torch +from freetoken.models.config import vision_load_enabled from freetoken.distributed import get_tp_info from freetoken.models.loader import drop_page_cache, iter_weight_files, shard_tensor from freetoken.models.nvfp4_banks import ( @@ -100,8 +101,13 @@ def _rename(raw_name: str) -> str | None: """Checkpoint key -> FreeToken state-dict key, or None to skip.""" - if raw_name.startswith(("mtp.", "model.visual.", "visual.")): + if raw_name.startswith("mtp."): return None + if raw_name.startswith(("model.visual.", "visual.")): + # the tower's HF module keys, under the model's ``visual`` op (opt-in, else dropped) + if not vision_load_enabled(): + return None + return "visual." + raw_name.split("visual.", 1)[1] if _PLE_TABLE_INFIX in raw_name: return None # n-gram table + its scale: load_ple_table if _EXPERT_RE.search(raw_name): @@ -163,7 +169,7 @@ def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.T ``out_proj``, shared-expert ``down_proj``. Vocab rows: ``embed_tokens`` / ``lm_head``. Everything else (router, indexer, norms, HC, PLE, shared-expert gate) is replicated. """ - if world == 1: + if world == 1 or name.startswith("visual."): return t if name.endswith(".self_attn.qkv_proj.weight"): q = (config.num_qo_heads, 2 * config.head_dim) diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index f5bc8f7a3..511fdb286 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -172,10 +172,10 @@ def _add_one_req( device_ids = self.table_manager.token_pool[table_idx, _slice] device_ids.copy_(_maybe_pinned(pending_req.input_ids[_slice]), non_blocking=True) if is_chunked and pending_req.mm_embeds is not None: - raise NotImplementedError( - "Multimodal prompts must fit in a single prefill chunk; increase " - "--max-extend-tokens or shrink the prompt." - ) + # An image prompt prefills in one chunk (its soft tokens and mRoPE table cover the + # whole run): wait for a pass with the budget free. Admission already refused + # prompts longer than the budget, so this cannot starve. + return None req = CLS( input_ids=pending_req.input_ids[: cached_len + chunk_size], table_idx=table_idx, @@ -185,6 +185,8 @@ def _add_one_req( cache_handle=cache_handle, sampling_params=pending_req.sampling_params, mm_embeds=pending_req.mm_embeds, + mrope_positions=pending_req.mrope_positions, + mrope_delta=pending_req.mrope_delta, ) # Hybrid GDN per-request state slots (None for non-hybrid). On a fresh admit these are # freshly allocated; on a chunked continuation they are inherited from the prior chunk. @@ -245,7 +247,14 @@ class PrefillManager: def add_one_req(self, req: UserMsg) -> None: self.pending_list.append( - PendingReq(req.uid, req.input_ids, req.sampling_params, mm_embeds=req.mm_embeds) + PendingReq( + req.uid, + req.input_ids, + req.sampling_params, + mm_embeds=req.mm_embeds, + mrope_positions=req.mrope_positions, + mrope_delta=req.mrope_delta, + ) ) def schedule_next_batch(self, prefill_budget: int) -> Batch | None: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b0..1cd26cb45 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -520,6 +520,12 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: logger.warning_rank0( f"Adjust max_tokens to {max_output_len} for request {msg.uid}." ) + if msg.mm_inputs is not None: + error = self._prepare_multimodal(msg) + if error is not None: + logger.warning_rank0(f"Rejecting request {msg.uid}: {error}") + self.send_result([ErrorReplyMsg(uid=msg.uid, error=error)]) + return self.prefill_manager.add_one_req(msg) elif isinstance(msg, AbortBackendMsg): logger.debug_rank0("Aborting request %d", msg.uid) @@ -759,6 +765,29 @@ def _log_cache_geometry(self, event: str) -> None: except Exception as e: # noqa: BLE001 logger.warning(f"could not log cache geometry: {e!r}") + def _prepare_multimodal(self, msg: UserMsg) -> str | None: + """Encode the request's images on this rank (every TP rank sees the same UserMsg) into + ``mm_embeds`` + mRoPE positions; returns the client-facing error instead of raising.""" + model = self.engine.model + if not getattr(model, "has_vision", False): + return "image inputs need the vision tower: start the server with FREETOKEN_LOAD_VISION=1" + if len(msg.input_ids) > self.prefill_budget: + # image prompts must prefill in one chunk (prefill.py never splits them) + return ( + f"image prompts must fit in one prefill chunk: {len(msg.input_ids)} tokens > " + f"{self.prefill_budget} (--max-extend-tokens); use a smaller image or prompt" + ) + try: + with self.engine_stream_ctx: # ordered before the prefill that reads mm_embeds + msg.mm_embeds, msg.mrope_positions, msg.mrope_delta = model.prepare_mm_inputs( + msg.input_ids, msg.mm_inputs + ) + except Exception as exc: # noqa: BLE001 -- a bad image must not take the scheduler down + logger.warning_rank0(f"image encoding failed for request {msg.uid}: {exc!r}") + return f"could not encode images: {exc}" + msg.mm_inputs = None + return None + def _prepare_batch(self, batch: Batch) -> ForwardInput: self.engine.graph_runner.pad_batch(batch) self._forward_iter += 1 @@ -781,6 +810,9 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: if batch.is_prefill: self._gather_multimodal(batch) batch.positions = _make_positions(batch, self.device) + batch.rope_positions, batch.mrope_cos_sin = _make_rope_positions( + batch, self.device, self.engine.model + ) input_mapping = _make_input_tuple(batch, self.device) write_mapping = _make_write_tuple(batch, self.device) batch.out_loc = self.engine.page_table[input_mapping] @@ -890,6 +922,26 @@ def _make_positions(batch: Batch, device: torch.device) -> torch.Tensor: return indices_host.to(device, non_blocking=True) +def _make_rope_positions( + batch: Batch, device: torch.device, model +) -> tuple[torch.Tensor, torch.Tensor | None]: + """``(rope_positions, mrope_cos_sin)`` for the batch (see ``Batch``). Text-only batches + alias ``positions`` (no extra work); requests decoding after an image add their + ``mrope_delta``; a prefill batch holding image tokens gets a per-token cos|sin table from + the model (``mrope_table``) and each token's table row as its rope position.""" + reqs = batch.padded_reqs + if batch.is_prefill and any(r.mrope_positions is not None for r in reqs): + return model.mrope_table(reqs, device) + if all(r.mrope_delta == 0 for r in reqs): + return batch.positions, None + delta_host = torch.empty(len(batch.positions), dtype=torch.int32, pin_memory=True) + offset = 0 + for req in reqs: + delta_host[offset : offset + req.extend_len].fill_(req.mrope_delta) + offset += req.extend_len + return batch.positions + delta_host.to(device, non_blocking=True), None + + def _make_input_tuple(batch: Batch, device: torch.device) -> Indice2D: mapping_host = torch.empty(len(batch.positions), dtype=torch.int64, pin_memory=True) offset = 0 diff --git a/python/freetoken/scheduler/utils.py b/python/freetoken/scheduler/utils.py index 0d6dd5129..7dcc21066 100644 --- a/python/freetoken/scheduler/utils.py +++ b/python/freetoken/scheduler/utils.py @@ -18,6 +18,8 @@ class PendingReq: sampling_params: SamplingParams chunked_req: ChunkedReq | None = None mm_embeds: torch.Tensor | None = None + mrope_positions: torch.Tensor | None = None + mrope_delta: int = 0 @property def input_len(self) -> int: diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..9e14e980a 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -14,6 +14,8 @@ from __future__ import annotations import asyncio +import base64 +import binascii import json import time from collections.abc import AsyncIterator @@ -139,6 +141,7 @@ class GenSpec: chat_template_kwargs: dict[str, Any] = field(default_factory=dict) template_tools: list[dict[str, Any]] | None = None # tools the model sees (TokenizeMsg.tools) parser_tools: list[dict[str, Any]] | None = None # tools for FunctionCallParser; None disables parsing + images: list[bytes] = field(default_factory=list) # encoded image files, in message order @property def parse_tools(self) -> bool: @@ -187,18 +190,23 @@ def pick(value, key, framework): ) -def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: +def render_messages( + messages: list[dict[str, Any]], images: list[bytes] | None = None +) -> list[dict[str, Any]]: """Normalize OpenAI-shaped message dicts for the chat template: flatten text - content parts to a string and decode tool-call arguments from JSON. Raises - ValueError on a non-text content part (text-only server). Shared by all adapters.""" - return [_render_message(m) for m in messages] + content parts to a string and decode tool-call arguments from JSON. Shared by all adapters. + ``images`` (OpenAI chat only): accept ``image_url`` parts, decode them (``data:`` URLs) into + the list in message order and keep ``{"type": "image"}`` parts in the content for the + template. Without it a non-text part raises ValueError (text-only adapters).""" + return [_render_message(m, images) for m in messages] -def _render_message(message: dict[str, Any]) -> dict[str, Any]: + +def _render_message(message: dict[str, Any], images: list[bytes] | None = None) -> dict[str, Any]: m = dict(message) content = m.get("content") if isinstance(content, list): - m["content"] = _flatten_text_parts(content) + m["content"] = _render_parts(content, images) # Templates read different reasoning keys (reasoning_content: most; reasoning: # gemma4; thinking: gpt-oss) — accept any, emit both. reasoning = m.get("reasoning_content") or m.get("reasoning") or m.get("thinking") @@ -230,6 +238,48 @@ def _render_message(message: dict[str, Any]) -> dict[str, Any]: return m +# Decoded image bytes above this are refused (a 4000x3000 JPEG is ~3 MiB; PNGs a few times that). +IMAGE_MAX_BYTES = 16 << 20 + + +def _render_parts(parts: list[Any], images: list[bytes] | None) -> str | list[dict[str, Any]]: + """Text-only part lists flatten to a string (unchanged behaviour); a list carrying images keeps + ``{"type": "text"}`` / ``{"type": "image"}`` parts, which the chat template renders as text and + one ``<|image_pad|>`` placeholder per image (expanded by the tokenizer worker).""" + if images is None or not any( + isinstance(p, dict) and p.get("type") == "image_url" for p in parts + ): + return _flatten_text_parts(parts) + out: list[dict[str, Any]] = [] + for part in parts: + ptype = part.get("type") if isinstance(part, dict) else None + if ptype == "text": + out.append({"type": "text", "text": part.get("text") or ""}) + elif ptype == "image_url": + images.append(_decode_image_url(part.get("image_url"))) + out.append({"type": "image"}) + else: + raise ValueError(f"Unsupported content part type: {ptype}") + return out + + +def _decode_image_url(image_url: Any) -> bytes: + """OpenAI ``image_url`` part -> image file bytes. Only inline ``data:`` URLs: fetching a remote + URL from the server would be an outbound request on the client's behalf.""" + url = image_url.get("url") if isinstance(image_url, dict) else image_url + if not isinstance(url, str) or not url.startswith("data:"): + raise ValueError("image_url must be a data: URL (base64-encoded image); remote URLs are not fetched") + header, sep, payload = url.partition(",") + if not sep or ";base64" not in header: + raise ValueError("image_url data: URL must be base64-encoded") + if len(payload) > IMAGE_MAX_BYTES * 4 // 3 + 4: + raise ValueError(f"image exceeds {IMAGE_MAX_BYTES >> 20} MiB") + try: + return base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError(f"image_url data: URL is not valid base64: {exc}") from exc + + def _flatten_text_parts(parts: list[Any]) -> str: texts: list[str] = [] for part in parts: @@ -269,6 +319,7 @@ async def submit_generation(spec: GenSpec, state: Any) -> int: sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, tools=spec.template_tools, + images=spec.images or None, ) ) return uid diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd263..89c6582ab 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -66,8 +66,9 @@ def chat_request_to_genspec( thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + images: list[bytes] = [] return GenSpec( - messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]), + messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages], images), sampling_params=resolve_sampling( temperature=req.temperature, top_k=req.top_k, @@ -80,6 +81,7 @@ def chat_request_to_genspec( chat_template_kwargs=ctk, template_tools=_tools_for_template(req), parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None), + images=images, ) diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862d0..6ef0d9143 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -85,7 +85,9 @@ def _tokenize_requests( messages: List[TokenizeMsg], logger: Any, ) -> tuple[List[TokenizeMsg], List[torch.Tensor], List[UserReply]]: - """Tokenize independently, returning backend work plus terminal frontend errors. + """Tokenize independently, returning backend work plus terminal frontend errors. The + returned tensors carry the request's vision-processor outputs as ``.mm_inputs`` (None for + text-only requests). Successful tokenization deliberately emits no prompt-token reply: accounting starts only when the scheduler later confirms first-prefill admission. @@ -95,7 +97,8 @@ def _tokenize_requests( errors: List[UserReply] = [] for msg in messages: try: - tokens = tokenize_manager.tokenize([msg])[0] + tokens, mm_inputs = tokenize_manager.tokenize_one(msg) + tokens.mm_inputs = mm_inputs except Exception as exc: # noqa: BLE001 — isolate, never crash the worker logger.warning(f"tokenization failed for request {msg.uid}: {exc!r}") errors.append( @@ -147,7 +150,7 @@ def tokenize_worker( from .detokenize import DetokenizeManager from .tokenize import TokenizeManager - tokenize_manager = TokenizeManager(tokenizer) + tokenize_manager = TokenizeManager(tokenizer, model_path=tokenizer_path) detokenize_manager = DetokenizeManager( tokenizer, load_eos_token_ids(tokenizer_path, tokenizer) ) @@ -254,7 +257,12 @@ def tokenize_worker( ) if ok_msgs: backend = [ - UserMsg(uid=msg.uid, input_ids=t, sampling_params=msg.sampling_params) + UserMsg( + uid=msg.uid, + input_ids=t, + sampling_params=msg.sampling_params, + mm_inputs=getattr(t, "mm_inputs", None), + ) for msg, t in zip(ok_msgs, ok_tensors, strict=True) ] send_backend.put(backend[0] if len(backend) == 1 else BatchBackendMsg(data=backend)) diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 0636b3428..3a00941f2 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -45,10 +45,39 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An _EFFORT_PROBE_MESSAGES = [{"role": "user", "content": "ping"}] +# Image resize bounds in pixels (Qwen2-VL processor ``size``). The cap bounds the prompt: at +# patch 16 / merge 2 one soft token covers 32x32 pixels, so the default is ~1000 tokens per +# image and a 4000x3000 photo is downscaled instead of costing 11k tokens. +IMAGE_MIN_PIXELS = 4 * 28 * 28 +IMAGE_MAX_PIXELS = int(os.getenv("FREETOKEN_IMAGE_MAX_PIXELS", str(1280 * 28 * 28))) + + +def _expand_image_tokens( + input_ids: torch.Tensor, image_token_id: int, counts: list[int] +) -> torch.Tensor: + """Replace the i-th ``image_token_id`` placeholder with ``counts[i]`` copies.""" + hits = (input_ids == image_token_id).nonzero().flatten().tolist() + if len(hits) != len(counts): + raise ValueError( + f"{len(hits)} image placeholders in the rendered prompt for {len(counts)} images" + ) + pieces, prev = [], 0 + for hit, n in zip(hits, counts, strict=True): + pieces.append(input_ids[prev:hit]) + pieces.append(torch.full((n,), image_token_id, dtype=input_ids.dtype)) + prev = hit + 1 + pieces.append(input_ids[prev:]) + return torch.cat(pieces) + class TokenizeManager: - def __init__(self, tokenizer: PreTrainedTokenizerBase) -> None: + def __init__( + self, tokenizer: PreTrainedTokenizerBase, model_path: str | None = None + ) -> None: self.tokenizer = tokenizer + # image requests need the checkpoint's image processor + placeholder id (worker only) + self._model_path = model_path + self._vision: tuple[Any, int] | None = None self._dsv4_encoder = _load_dsv4_encoder_if_needed(tokenizer) self._effort_profile: EffortProfile | None = None self._thinking_profile: ThinkingProfile | None = None @@ -74,6 +103,45 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[torch.Tensor]: results.append(input_ids.view(-1).to(torch.int32)) return results + def tokenize_one(self, msg: TokenizeMsg) -> tuple[torch.Tensor, dict[str, Any] | None]: + """``tokenize`` for one message plus, when it carries images, the vision processor's + outputs (``pixel_values`` fp32 [patches, C*T*P*P], ``image_grid_thw`` [N, 3]) with every + ``<|image_pad|>`` placeholder expanded to its image's soft-token count.""" + input_ids = self.tokenize([msg])[0] + if not msg.images: + return input_ids, None + import io + + from PIL import Image + + processor, image_token_id = self._vision_processor() + images = [Image.open(io.BytesIO(data)).convert("RGB") for data in msg.images] + out = processor(images=images, return_tensors="pt") + grid = out["image_grid_thw"] + counts = (grid.prod(-1) // processor.merge_size**2).tolist() + # pixel_values stay fp32 on the wire (numpy cannot carry bf16); the tower casts them + return _expand_image_tokens(input_ids, image_token_id, counts), { + "pixel_values": out["pixel_values"], + "image_grid_thw": grid, + } + + def _vision_processor(self) -> tuple[Any, int]: + if self._vision is None: + if self._model_path is None: + raise ValueError("image inputs are not supported here") + from freetoken.utils.hf import cached_load_hf_config + from transformers import AutoImageProcessor + + image_token_id = getattr(cached_load_hf_config(self._model_path), "image_token_id", None) + if image_token_id is None: + raise ValueError("this model has no image placeholder token") + processor = AutoImageProcessor.from_pretrained( + self._model_path, + size={"shortest_edge": IMAGE_MIN_PIXELS, "longest_edge": IMAGE_MAX_PIXELS}, + ) + self._vision = (processor, int(image_token_id)) + return self._vision + def render_prompt(self, msg: TokenizeMsg) -> str: """The template/encoder half of ``tokenize``, exposed so the frontend can validate a request before committing an SSE stream. Sanitizes diff --git a/tests/models/qwen4_exp/common.py b/tests/models/qwen4_exp/common.py index 1f9c117bd..03d8f3534 100644 --- a/tests/models/qwen4_exp/common.py +++ b/tests/models/qwen4_exp/common.py @@ -206,6 +206,8 @@ def req(self, table_idx: int, cached_len: int, device_len: int) -> SimpleNamespa cached_len=cached_len, device_len=device_len, extend_len=device_len - cached_len, + mrope_positions=None, + mrope_delta=0, ) def step(self, req: SimpleNamespace) -> None: @@ -231,6 +233,8 @@ def batch(self, reqs, phase: str) -> SimpleNamespace: is_prefill=phase == "prefill", is_decode=phase == "decode", positions=positions, + rope_positions=None, # text-only: the layer falls back to positions + mrope_cos_sin=None, out_loc=out_loc, attn_metadata=None, active_table_idx=torch.tensor( diff --git a/tests/models/qwen4_exp/test_mrope.py b/tests/models/qwen4_exp/test_mrope.py new file mode 100644 index 000000000..3c48acb23 --- /dev/null +++ b/tests/models/qwen4_exp/test_mrope.py @@ -0,0 +1,161 @@ +"""mRoPE for qwen4_exp image prompts: FreeToken's port against the HF reference (CPU). + +The HF checks need a transformers with ``qwen4_exp``; the pure-FreeToken checks always run. +""" + +from __future__ import annotations + +import base64 +from types import MethodType, SimpleNamespace + +import pytest +import torch +from freetoken.models.qwen4_exp.mrope import mrope_cos_sin, rope_index +from freetoken.server.generation import render_messages +from freetoken.tokenizer.tokenize import _expand_image_tokens + +IMAGE = 248056 +MERGE = 2 +SECTION = (11, 11, 10) + + +def _prompt( + grids: list[tuple[int, int, int]], text_len: int = 5, tail: int = 7 +) -> torch.Tensor: + ids = list(range(1, text_len + 1)) + for t, h, w in grids: + ids += [IMAGE] * (t * (h // MERGE) * (w // MERGE)) + [50, 51] + ids += list(range(100, 100 + tail)) + return torch.tensor(ids, dtype=torch.int32) + + +def test_rope_index_text_only(): + ids = _prompt([]) + pos, delta = rope_index(ids, None, IMAGE, MERGE) + assert delta == 0 + assert torch.equal(pos, torch.arange(len(ids)).expand(3, -1)) + + +def test_rope_index_image_layout(): + grid = torch.tensor([[1, 4, 6]]) + ids = _prompt([(1, 4, 6)], text_len=3, tail=2) + pos, delta = rope_index(ids, grid, IMAGE, MERGE) + # text 0..2, image at start 3: T=3, H in 3..4, W in 3..5, then text resumes at 3 + max(2, 3) + img = pos[:, 3:9] + assert img[0].tolist() == [3] * 6 + assert img[1].tolist() == [3, 3, 3, 4, 4, 4] + assert img[2].tolist() == [3, 4, 5, 3, 4, 5] + assert pos[:, 9].tolist() == [6, 6, 6] + assert delta == int(pos.max()) + 1 - len(ids) + assert delta <= 0 + + +def test_mrope_cos_sin_text_rows_equal_cache(): + inv_freq = 1.0 / (1e7 ** (torch.arange(0, 64, 2, dtype=torch.float) / 64)) + pos = torch.arange(0, 37).expand(3, -1) + table = mrope_cos_sin(pos, inv_freq, SECTION) + freqs = torch.einsum("i,j -> ij", torch.arange(37, dtype=torch.float), inv_freq) + assert torch.equal(table, torch.cat((freqs.cos(), freqs.sin()), dim=-1)) + + +def test_expand_image_tokens(): + ids = torch.tensor([1, IMAGE, 2, IMAGE, 3], dtype=torch.int32) + out = _expand_image_tokens(ids, IMAGE, [3, 2]) + assert out.tolist() == [1, IMAGE, IMAGE, IMAGE, 2, IMAGE, IMAGE, 3] + with pytest.raises(ValueError): + _expand_image_tokens(ids, IMAGE, [3]) + + +def test_render_messages_images(): + png = b"\x89PNG fake" + url = "data:image/png;base64," + base64.b64encode(png).decode() + messages = [ + {"role": "system", "content": [{"type": "text", "text": "sys"}]}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": url}}, + {"type": "text", "text": "hi"}, + ], + }, + ] + images: list[bytes] = [] + out = render_messages(messages, images) + assert out[0]["content"] == "sys" # text-only lists still flatten + assert out[1]["content"] == [{"type": "image"}, {"type": "text", "text": "hi"}] + assert images == [png] + with pytest.raises(ValueError): # text-only adapters keep refusing image parts + render_messages(messages) + with pytest.raises(ValueError): # remote URLs are never fetched + render_messages( + [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://x/y.png"}} + ], + } + ], + [], + ) + + +# ----------------------------------------------------------------------------- HF reference +hf = pytest.importorskip("transformers.models.qwen4_exp.modeling_qwen4_exp") + + +def _hf_rope_index(ids: torch.Tensor, grid: torch.Tensor | None): + cfg = SimpleNamespace( + image_token_id=IMAGE, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + vision_config=SimpleNamespace(spatial_merge_size=MERGE), + ) + stub = SimpleNamespace(config=cfg) + stub.get_vision_position_ids = MethodType( + hf.Qwen4ExpModel.get_vision_position_ids, stub + ) + ids = ids.view(1, -1).long() + pos, delta = hf.Qwen4ExpModel.get_rope_index( + stub, ids, mm_token_type_ids=(ids == IMAGE).int(), image_grid_thw=grid + ) + return pos[:, 0], int(delta.reshape(-1)[0]) + + +@pytest.mark.parametrize( + "grids", + [[(1, 4, 6)], [(1, 8, 8), (1, 2, 12)], [(1, 6, 2)]], +) +def test_rope_index_matches_hf(grids): + ids = _prompt(grids) + grid = torch.tensor(grids) + pos, delta = rope_index(ids, grid, IMAGE, MERGE) + ref_pos, ref_delta = _hf_rope_index(ids, grid) + assert torch.equal(pos, ref_pos), (pos, ref_pos) + assert delta == ref_delta + + +def test_mrope_cos_sin_matches_hf(): + text_cfg = SimpleNamespace( + head_dim=256, + hidden_size=2560, + num_attention_heads=24, + max_position_embeddings=262144, + rope_parameters={ + "rope_type": "default", + "rope_theta": 1e7, + "partial_rotary_factor": 0.25, + "mrope_section": list(SECTION), + "mrope_interleaved": True, + }, + ) + rotary = hf.Qwen4ExpTextRotaryEmbedding(text_cfg) + pos = torch.stack( + [torch.randint(0, 5000, (3, 29)) for _ in range(1)], dim=1 + ) # [3, 1, 29] + x = torch.zeros(1, 29, 256) + cos, sin = rotary(x, pos) + table = mrope_cos_sin(pos[:, 0], rotary.inv_freq.float(), SECTION) + torch.testing.assert_close(table[:, :32], cos[0, :, :32], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(table[:, 32:], sin[0, :, :32], atol=1e-5, rtol=1e-5) diff --git a/tests/models/qwen4_exp/test_mrope_gpu.py b/tests/models/qwen4_exp/test_mrope_gpu.py new file mode 100644 index 000000000..33f847023 --- /dev/null +++ b/tests/models/qwen4_exp/test_mrope_gpu.py @@ -0,0 +1,79 @@ +"""mRoPE plumbing through the QSA layer (GPU): the per-token table path (what an image-carrying +prefill batch uses) must reproduce the position-indexed cache path exactly for text, including +a chunked continuation sharing the batch (the indexer ropes a straddling group at its first +token, which lives in the table's lead rows) and the decode `positions + delta` path. +""" + +from __future__ import annotations + +import torch +from freetoken.models.qwen4_exp.mrope import mrope_table + +from .common import Fixture, parsed_config, requires_cuda + +QSA_LAYER = 3 +SECTION = (11, 11, 10) + + +def _inv_freq(config, device): + rc = config.rotary_config + return 1.0 / ( + rc.base + ** ( + torch.arange(0, rc.rotary_dim, 2, dtype=torch.float, device=device) + / rc.rotary_dim + ) + ) + + +def _run(config, table_path: bool, seed: int = 5): + """Prefill B (chunk 1) -> prefill [A, B chunk 2] -> one decode step; return the two outputs.""" + fixture = Fixture(config, num_pages=64) + attn = fixture.layer(QSA_LAYER) + ratio = config.qwen4_args.index_ratio + gen = torch.Generator(device=fixture.device).manual_seed(seed) + hidden = config.hidden_size + x_b1 = torch.randn( + 41, hidden, device=fixture.device, dtype=fixture.dtype, generator=gen + ) + x_ab = torch.randn( + 37 + 30, hidden, device=fixture.device, dtype=fixture.dtype, generator=gen + ) + x_dec = torch.randn( + 2, hidden, device=fixture.device, dtype=fixture.dtype, generator=gen + ) + assert 41 % ratio, ( + "chunk 1 must end mid-group so a group straddles the chunk boundary" + ) + + b = fixture.req(1, 0, 41) + # chunk 1 always takes the cache path + attn.forward(x_b1, fixture.batch([b], "prefill")) + b.cached_len, b.device_len, b.extend_len = 41, 71, 30 + fixture.allocate(1, 41, 71) + a = fixture.req(0, 0, 37) + batch = fixture.batch([a, b], "prefill") + if table_path: + batch.rope_positions, batch.mrope_cos_sin = mrope_table( + [a, b], ratio, _inv_freq(config, fixture.device), SECTION, fixture.device + ) + assert batch.mrope_cos_sin.shape[0] == 37 + 30 + (ratio - 1) + out_prefill = attn.forward(x_ab, batch).clone() + + fixture.step(a) + fixture.step(b) + batch = fixture.batch([a, b], "decode") + if table_path: + # decode never has a table; exercise the explicit rope_positions path (delta 0) + batch.rope_positions = batch.positions + 0 + out_decode = attn.forward(x_dec, batch).clone() + return out_prefill, out_decode + + +@requires_cuda +def test_table_path_matches_cache_path(): + config = parsed_config() + ref_prefill, ref_decode = _run(config, table_path=False) + tab_prefill, tab_decode = _run(config, table_path=True) + assert torch.equal(ref_prefill, tab_prefill) + assert torch.equal(ref_decode, tab_decode) diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 1d3b944ce..da2c8027c 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -174,6 +174,7 @@ def test_decode_graph_replay_matches_eager(): capture_batch = SimpleNamespace( padded_reqs=[dummy] * bs, reqs=[dummy] * bs, phase="decode", size=bs, padded_size=bs, is_prefill=False, is_decode=True, positions=static["positions"], + rope_positions=None, mrope_cos_sin=None, out_loc=static["out_loc"], attn_metadata=None, active_table_idx=None, ) fixture.backend.prepare_for_capture(capture_batch) diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 00f7c43f3..6e18f73de 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -346,7 +346,9 @@ def test_qsa_layer_matches_hf_dense(): x = (torch.randn(seq_len, config.hidden_size, device=device, dtype=dtype) * 0.5) positions = torch.arange(seq_len, device=device, dtype=torch.int64) req = SimpleNamespace(extend_len=seq_len, cached_len=0, table_idx=1) - batch = SimpleNamespace(padded_reqs=[req], reqs=[req], positions=positions) + batch = SimpleNamespace( + padded_reqs=[req], reqs=[req], positions=positions, rope_positions=None, mrope_cos_sin=None + ) backend = TorchDenseQSAReference(config, num_slots=4, max_len=64, device=device, dtype=dtype) _fresh_ctx(attn_backend=backend) @@ -480,6 +482,7 @@ def test_decoder_stack_prefill_and_decode(monkeypatch): padded_reqs=reqs, reqs=reqs, size=len(reqs), is_prefill=True, is_decode=False, input_ids=torch.tensor(flat, dtype=torch.int64, device=device), positions=torch.cat([torch.arange(len(p)) for p in prompts]).to(device), + rope_positions=None, mrope_cos_sin=None, mm_embeds=None, attn_metadata=SimpleNamespace(get_last_indices=lambda bs: last[:bs]), ) with ctx.forward_batch(batch): @@ -495,6 +498,7 @@ def test_decoder_stack_prefill_and_decode(monkeypatch): padded_reqs=reqs, reqs=reqs, size=len(reqs), is_prefill=False, is_decode=True, input_ids=torch.tensor([14] * len(reqs), dtype=torch.int64, device=device), positions=torch.tensor([len(p) for p in prompts], dtype=torch.int64, device=device), + rope_positions=None, mrope_cos_sin=None, mm_embeds=None, attn_metadata=None, ) with ctx.forward_batch(decode): diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index 8942c4cac..9ea850a11 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -50,8 +50,8 @@ def _tokenize_msg(uid: int) -> TokenizeMsg: def test_successful_tokenization_does_not_account_prompt_before_admission(): class Tokenizer: - def tokenize(self, messages): - return [torch.tensor([10, 11, 12], dtype=torch.int32)] + def tokenize_one(self, msg): + return torch.tensor([10, 11, 12], dtype=torch.int32), None ok, tensors, errors = _tokenize_requests(Tokenizer(), [_tokenize_msg(1)], _Logger()) assert [msg.uid for msg in ok] == [1] @@ -61,11 +61,10 @@ def tokenize(self, messages): def test_tokenization_failure_and_empty_prompt_are_terminal_without_usage(): class Tokenizer: - def tokenize(self, messages): - uid = messages[0].uid - if uid == 2: + def tokenize_one(self, msg): + if msg.uid == 2: raise ValueError("bad template") - return [torch.empty(0, dtype=torch.int32)] + return torch.empty(0, dtype=torch.int32), None logger = _Logger() ok, tensors, errors = _tokenize_requests( From 28fd56de58d0ea0284f87bfdc177cc013d5d0831 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 18:40:14 -0400 Subject: [PATCH 4/5] feat(qwen4_exp): chunked prefill for image prompts Image prompts no longer have to fit in one prefill chunk (they were refused above --max-extend-tokens, 8192 by default, which a long agent context hits at once). prefill.py chunks them like text; the scheduler scatters, per chunk, only the soft-token rows whose placeholders fall inside that chunk, skipping the rows earlier chunks consumed. The mRoPE table already windows per chunk. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/model.py | 7 +++++- python/freetoken/models/qwen4_exp/mrope.py | 4 +-- python/freetoken/scheduler/prefill.py | 5 ---- python/freetoken/scheduler/scheduler.py | 26 +++++++++++++------ tests/scheduler/test_mm_window.py | 29 ++++++++++++++++++++++ 5 files changed, 55 insertions(+), 16 deletions(-) create mode 100644 tests/scheduler/test_mm_window.py diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 7d02d6d15..30b53bf30 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -110,7 +110,8 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: hidden = self.embed_tokens.forward(input_ids) if batch.mm_embeds is not None: # image soft tokens replace the placeholder embeddings (HF order: before the - # hc_count repeat); the whole image run sits in this prefill chunk (prefill.py) + # hc_count repeat); a chunked prompt gets the rows of the placeholders inside + # this chunk (scheduler._gather_multimodal) mask = input_ids == self._image_token_id hidden = hidden.masked_scatter(mask.unsqueeze(-1), batch.mm_embeds.to(hidden.dtype)) hidden = hidden.repeat(1, self.hc_count) @@ -159,6 +160,10 @@ def __init__(self, config: ModelConfig) -> None: def has_vision(self) -> bool: return hasattr(self, "visual") + @property + def image_token_id(self) -> int: + return self._config.image_token_id + @torch.inference_mode() def encode_images(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: """Vision tower + merger on processor outputs: ``[num_image_tokens, hidden]`` (device).""" diff --git a/python/freetoken/models/qwen4_exp/mrope.py b/python/freetoken/models/qwen4_exp/mrope.py index 9aa9157b0..a9b9167f7 100644 --- a/python/freetoken/models/qwen4_exp/mrope.py +++ b/python/freetoken/models/qwen4_exp/mrope.py @@ -89,8 +89,8 @@ def mrope_table( token's row, table [rows, rotary_dim] fp32 cos|sin)``. Every request also gets rows for up to ``ratio - 1`` positions before its chunk: the QSA - indexer ropes a pooled key at its group's FIRST token, which for a chunked text prompt sharing - the batch may come from the previous chunk (image prompts are never chunked). Requests without + indexer ropes a pooled key at its group's FIRST token, which for a chunked prompt may come + from the previous chunk. Requests without ``mrope_positions`` use their token index on all three axes, so their rows equal the rope cache rows exactly. """ diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index 511fdb286..ea874272d 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -171,11 +171,6 @@ def _add_one_req( _slice = slice(cached_len, cached_len + chunk_size) device_ids = self.table_manager.token_pool[table_idx, _slice] device_ids.copy_(_maybe_pinned(pending_req.input_ids[_slice]), non_blocking=True) - if is_chunked and pending_req.mm_embeds is not None: - # An image prompt prefills in one chunk (its soft tokens and mRoPE table cover the - # whole run): wait for a pass with the budget free. Admission already refused - # prompts longer than the budget, so this cannot starve. - return None req = CLS( input_ids=pending_req.input_ids[: cached_len + chunk_size], table_idx=table_idx, diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 1cd26cb45..0ec233f55 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -771,12 +771,6 @@ def _prepare_multimodal(self, msg: UserMsg) -> str | None: model = self.engine.model if not getattr(model, "has_vision", False): return "image inputs need the vision tower: start the server with FREETOKEN_LOAD_VISION=1" - if len(msg.input_ids) > self.prefill_budget: - # image prompts must prefill in one chunk (prefill.py never splits them) - return ( - f"image prompts must fit in one prefill chunk: {len(msg.input_ids)} tokens > " - f"{self.prefill_budget} (--max-extend-tokens); use a smaller image or prompt" - ) try: with self.engine_stream_ctx: # ordered before the prefill that reads mm_embeds msg.mm_embeds, msg.mrope_positions, msg.mrope_delta = model.prepare_mm_inputs( @@ -852,11 +846,18 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: def _gather_multimodal(self, batch: Batch) -> None: """Concatenate per-request vision soft tokens (in request order) for a prefill - batch so the model can scatter them at image-token positions. ``req.mm_embeds`` + batch so the model can scatter them at image-token positions. A chunked prompt + contributes only the placeholders inside its current chunk. ``req.mm_embeds`` is kept (not cleared) so the cache manager can recognize multimodal requests and keep them out of the shared prefix cache (image placeholders share a token id but carry per-image content).""" - parts = [req.mm_embeds for req in batch.reqs if req.mm_embeds is not None] + image_token_id = getattr(self.engine.model, "image_token_id", None) + parts = [ + _mm_embeds_window(req, image_token_id) + for req in batch.reqs + if req.mm_embeds is not None + ] + parts = [p for p in parts if p.shape[0]] if parts: batch.mm_embeds = torch.cat(parts, dim=0) @@ -922,6 +923,15 @@ def _make_positions(batch: Batch, device: torch.device) -> torch.Tensor: return indices_host.to(device, non_blocking=True) +def _mm_embeds_window(req, image_token_id: int) -> torch.Tensor: + """The rows of ``req.mm_embeds`` for this prefill window: the image placeholders in + ``input_ids[cached_len:]``, skipping those an earlier chunk already scattered.""" + is_image = req.input_ids == image_token_id + before = int(is_image[: req.cached_len].sum()) + within = int(is_image[req.cached_len :].sum()) + return req.mm_embeds[before : before + within] + + def _make_rope_positions( batch: Batch, device: torch.device, model ) -> tuple[torch.Tensor, torch.Tensor | None]: diff --git a/tests/scheduler/test_mm_window.py b/tests/scheduler/test_mm_window.py new file mode 100644 index 000000000..c81f6d3e9 --- /dev/null +++ b/tests/scheduler/test_mm_window.py @@ -0,0 +1,29 @@ +"""Chunked prefill of an image prompt: each chunk scatters only the placeholder rows it holds.""" + +from types import SimpleNamespace + +import torch + +from freetoken.scheduler.scheduler import _mm_embeds_window + +IMG = 7 + + +def test_window_rows_follow_the_placeholders_across_chunks(): + ids = torch.tensor([1, 2, IMG, IMG, IMG, 3, IMG, IMG, 4, 5]) + embeds = ( + torch.arange(5).unsqueeze(1).float() + ) # one row per placeholder, prompt order + chunks = [ + (0, 4), + (4, 8), + (8, 10), + ] # a cut inside the first image run, one after the second + got = [] + for start, end in chunks: + req = SimpleNamespace(input_ids=ids[:end], cached_len=start, mm_embeds=embeds) + got.append(_mm_embeds_window(req, IMG)) + assert [g.shape[0] for g in got] == [2, 3, 0] + assert torch.equal(torch.cat(got), embeds) + whole = SimpleNamespace(input_ids=ids, cached_len=0, mm_embeds=embeds) + assert torch.equal(_mm_embeds_window(whole, IMG), embeds) From ed0982ae006fa4f541cd5ecc0b5c565fa775ce2d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 20:10:14 -0400 Subject: [PATCH 5/5] feat(qwen4_exp): key image prompts in the prefix cache by image content hash Image placeholders share one token id, so multimodal requests were kept out of the shared prefix cache and every turn of a conversation holding an image re-prefilled the whole context (measured: 20 turns of 109k-122k tokens, ~35 s each at TP=2). The tokenizer worker now emits cache_ids next to the expanded input_ids: the same tokens, with each image's placeholder run replaced by ids derived from a blake2b hash of the image bytes (>= 2**30, above any vocabulary, hash + offset within the run). The cache manager keys match/insert on cache_ids when present, else on input_ids; the model still reads input_ids, and the per-chunk placeholder window already skips cached placeholders. The multimodal exclusions in match_req and the three cache_req paths are gone. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/core.py | 3 ++ python/freetoken/message/backend.py | 3 ++ python/freetoken/scheduler/cache.py | 61 +++++++++---------------- python/freetoken/scheduler/prefill.py | 6 +++ python/freetoken/scheduler/scheduler.py | 8 ++-- python/freetoken/scheduler/utils.py | 3 ++ python/freetoken/tokenizer/server.py | 20 +++++--- python/freetoken/tokenizer/tokenize.py | 31 ++++++++++++- tests/models/qwen4_exp/test_mrope.py | 22 +++++++++ tests/scheduler/test_mm_cache_key.py | 60 ++++++++++++++++++++++++ 10 files changed, 165 insertions(+), 52 deletions(-) create mode 100644 tests/scheduler/test_mm_cache_key.py diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 22fcf56ab..e1aeed4bb 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -48,6 +48,9 @@ class Req: # <= 0). None / 0 for text-only prompts, where rope position == token index. mrope_positions: torch.Tensor | None = None mrope_delta: int = 0 + # Prefix-cache key ids (image prompts: placeholder runs replaced per image content hash); + # None when input_ids are the key. Same length as input_ids, sliced with it. + cache_ids: torch.Tensor | None = None # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or # until allocated from LinearStatePool. Set by the scheduler (P2). --- diff --git a/python/freetoken/message/backend.py b/python/freetoken/message/backend.py index d42a83c06..3ef3be115 100644 --- a/python/freetoken/message/backend.py +++ b/python/freetoken/message/backend.py @@ -43,6 +43,9 @@ class UserMsg(BaseBackendMsg): mm_inputs: dict | None = None mrope_positions: torch.Tensor | None = None mrope_delta: int = 0 + # Prefix-cache key ids for image prompts: placeholder runs replaced per image content + # hash (tokenizer.tokenize._image_cache_ids); None for text prompts (input_ids are the key). + cache_ids: torch.Tensor | None = None @dataclass diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 44adde42f..91db763ce 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -29,6 +29,13 @@ def _swa_eviction_interval() -> int: _SWA_RETAIN_GAP = 16 +def _key_ids(req) -> torch.Tensor: + """The ids the prefix cache keys a request on: ``cache_ids`` when set (image prompts, the + placeholder runs replaced per image content hash), else ``input_ids``.""" + ids = getattr(req, "cache_ids", None) + return req.input_ids if ids is None else ids + + class CacheManager: def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str, linear_state_pool=None, swa_pool=None, sliding_window_size=None): @@ -93,10 +100,10 @@ def _make_prefix_cache(self, device, page_size, type): def match_req(self, req: PendingReq) -> MatchResult: input_len = req.input_len assert input_len > 0, "Input length must be greater than 0." - # Multimodal requests must not reuse a shared prefix: image-placeholder tokens - # have identical ids across images but carry different content (and KV), so a - # match would serve the wrong image's KV. Match against the empty prefix. - ids = req.input_ids[:0] if req.mm_embeds is not None else req.input_ids[: input_len - 1] + # Image prompts key on cache_ids (placeholder runs per image content hash): the raw + # placeholder ids are shared across images, so keyed on them a hit could serve another + # image's KV. + ids = _key_ids(req)[: input_len - 1] if self.is_swa: from freetoken.kvcache.swa_radix_cache import SWACacheHandle m = self.prefix_cache.match_prefix(ids) @@ -299,18 +306,7 @@ def cache_req(self, req: Req, *, finished: bool) -> None: # We should free it if the request has finished. page_indices = self.page_table[req.table_idx, : req.cached_len] old_handle = req.cache_handle - # Multimodal requests are never inserted into the shared prefix cache (see - # ``match_req``). Their KV pages stay owned by the active request and are freed - # on completion; nothing is exposed for cross-request reuse. - if req.mm_embeds is not None: - self.unlock(old_handle) - if finished: - tail = self._padded_tail(req, old_handle.cached_len) - if self.swa_paged: - self._free_swa(tail) - self._free(tail) - return - insert_ids = req.input_ids[: req.cached_len] + insert_ids = _key_ids(req)[: req.cached_len] cached_len, new_handle = self.prefix_cache.insert_prefix(insert_ids, page_indices) # unlock until all operations on handle is done self.unlock(old_handle) @@ -350,13 +346,6 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: old_handle = req.cache_handle page_indices = self.page_table[req.table_idx, : req.cached_len] - if req.mm_embeds is not None: - self.unlock(old_handle) - if finished: - self._free(page_indices[old_handle.cached_len :]) - self._free_req_slots(req) - return - if finished: # A pending freeze (the tool-call anchor, or a prefill ×64 track the request # finished too early to chunk-commit) is a strictly shorter prefix than the live @@ -376,7 +365,7 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: frozen_idx = 1 - req.mamba_next_track_idx frozen = req.mamba_ping_pong[frozen_idx] prefix_len, mamba_exist = self.prefix_cache.insert( - req.input_ids[:L], page_indices[:L], frozen) + _key_ids(req)[:L], page_indices[:L], frozen) pool.free([s for s in req.mamba_ping_pong if mamba_exist or s != frozen]) req.mamba_ping_pong = None self._free(page_indices[free_upto : max(free_upto, prefix_len)]) @@ -390,7 +379,7 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: keep_live = False if insert_len == req.cached_len and insert_len > 0: prefix_len, mamba_exist = self.prefix_cache.insert( - req.input_ids[:insert_len], page_indices[:insert_len], req.linear_slot_idx) + _key_ids(req)[:insert_len], page_indices[:insert_len], req.linear_slot_idx) self.unlock(old_handle) self._free(page_indices[free_upto : max(free_upto, prefix_len)]) keep_live = not mamba_exist # tree now owns linear_slot_idx @@ -413,13 +402,13 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None: frozen_idx = 1 - req.mamba_next_track_idx # the slot the forward just wrote frozen = req.mamba_ping_pong[frozen_idx] prefix_len, mamba_exist = self.prefix_cache.insert( - req.input_ids[:L], page_indices[:L], frozen) + _key_ids(req)[:L], page_indices[:L], frozen) self.unlock(old_handle) self._free(page_indices[old_handle.cached_len : prefix_len]) # Lock the committed snapshot node FIRST: the replacement-slot alloc below can trigger # evict_mamba (via ensure_mamba_slots), which would otherwise reclaim this still-unlocked # just-donated node -- freeing its KV pages under the still-decoding request. - m = self.prefix_cache.match_prefix(req.input_ids[:L]) + m = self.prefix_cache.match_prefix(_key_ids(req)[:L]) # Same re-point as the generic path: the dedup free above returned this request's own # pages for [old_handle.cached_len, prefix_len) while its row still named them. if prefix_len > old_handle.cached_len: @@ -444,14 +433,6 @@ def _cache_req_swa(self, req: Req, *, finished: bool) -> None: old_handle = req.cache_handle page_indices = self.page_table[req.table_idx, : req.cached_len] - if req.mm_embeds is not None: - self.unlock(old_handle) - if finished: - tail = self._padded_tail(req, old_handle.cached_len) - self._free_swa(tail) - self._free(tail) - return - insert_len = align_down(req.cached_len, self.page_size) freed = page_indices[:0] if insert_len > 0: @@ -464,7 +445,7 @@ def _cache_req_swa(self, req: Req, *, finished: bool) -> None: # unfinished chunk's frontier is already > 0 and must be honored (else insert adopts # sentinel slots -> the request's later SWA gathers read slot 0 -> corruption). _, freed = self.prefix_cache.insert( - req.input_ids[:insert_len], page_indices[:insert_len], + _key_ids(req)[:insert_len], page_indices[:insert_len], swa_evicted_seqlen=req.swa_evicted_seqlen, update_kv_after_len=old_handle.cached_len) self.unlock(old_handle) @@ -492,8 +473,8 @@ def _cache_req_swa(self, req: Req, *, finished: bool) -> None: ) if keep_from > 0: self._free_swa( - self.prefix_cache.trim_head_swa(req.input_ids[:prompt_len], keep_from)) - self.prefix_cache.match_prefix(req.input_ids[:prompt_len]) + self.prefix_cache.trim_head_swa(_key_ids(req)[:prompt_len], keep_from)) + self.prefix_cache.match_prefix(_key_ids(req)[:prompt_len]) else: # inc_lock is node-granular, and the suffix insert just made this chunk's whole # extend one node: locking it would pin the entire chunk's swa for all of decode, @@ -504,8 +485,8 @@ def _cache_req_swa(self, req: Req, *, finished: bool) -> None: keep_from = align_down( max(insert_len - self.sliding_window_size - _SWA_RETAIN_GAP, 0), self.page_size) if keep_from > 0: - self.prefix_cache.match_prefix(req.input_ids[:keep_from]) - m = self.prefix_cache.match_prefix(req.input_ids[:insert_len]) + self.prefix_cache.match_prefix(_key_ids(req)[:keep_from]) + m = self.prefix_cache.match_prefix(_key_ids(req)[:insert_len]) # Re-point the page table to the tree's live slots for the committed region. Any dup # slots insert reclaimed had their full->swa mapping reset to the 0 sentinel; unlike the # full pool (KV survives in place until realloc), a stale swa mapping would make the diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index ea874272d..0bf5bf6db 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -182,6 +182,11 @@ def _add_one_req( mm_embeds=pending_req.mm_embeds, mrope_positions=pending_req.mrope_positions, mrope_delta=pending_req.mrope_delta, + cache_ids=( + None + if pending_req.cache_ids is None + else pending_req.cache_ids[: cached_len + chunk_size] + ), ) # Hybrid GDN per-request state slots (None for non-hybrid). On a fresh admit these are # freshly allocated; on a chunked continuation they are inherited from the prior chunk. @@ -249,6 +254,7 @@ def add_one_req(self, req: UserMsg) -> None: mm_embeds=req.mm_embeds, mrope_positions=req.mrope_positions, mrope_delta=req.mrope_delta, + cache_ids=req.cache_ids, ) ) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 0ec233f55..e1a695624 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -847,10 +847,10 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: def _gather_multimodal(self, batch: Batch) -> None: """Concatenate per-request vision soft tokens (in request order) for a prefill batch so the model can scatter them at image-token positions. A chunked prompt - contributes only the placeholders inside its current chunk. ``req.mm_embeds`` - is kept (not cleared) so the cache manager can recognize multimodal requests and - keep them out of the shared prefix cache (image placeholders share a token id but - carry per-image content).""" + contributes only the placeholders inside its current chunk; ``req.mm_embeds`` is + kept (not cleared) so later chunks can slice theirs. Image prompts take part in the + prefix cache through ``req.cache_ids`` (per-image content hash, see + ``tokenizer.tokenize._image_cache_ids``): a hit skips the cached placeholders too.""" image_token_id = getattr(self.engine.model, "image_token_id", None) parts = [ _mm_embeds_window(req, image_token_id) diff --git a/python/freetoken/scheduler/utils.py b/python/freetoken/scheduler/utils.py index 7dcc21066..4e0128d8e 100644 --- a/python/freetoken/scheduler/utils.py +++ b/python/freetoken/scheduler/utils.py @@ -20,6 +20,9 @@ class PendingReq: mm_embeds: torch.Tensor | None = None mrope_positions: torch.Tensor | None = None mrope_delta: int = 0 + cache_ids: torch.Tensor | None = ( + None # prefix-cache key ids (image prompts), else input_ids + ) @property def input_len(self) -> int: diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 6ef0d9143..2b4446676 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -80,6 +80,18 @@ def _send_generation_replies( _put_user_replies(send_frontend, terminal_errors) +def _backend_msg(msg, t) -> UserMsg: + mm_inputs = getattr(t, "mm_inputs", None) + cache_ids = mm_inputs.pop("cache_ids", None) if mm_inputs else None + return UserMsg( + uid=msg.uid, + input_ids=t, + sampling_params=msg.sampling_params, + mm_inputs=mm_inputs, + cache_ids=cache_ids, + ) + + def _tokenize_requests( tokenize_manager: Any, messages: List[TokenizeMsg], @@ -257,13 +269,7 @@ def tokenize_worker( ) if ok_msgs: backend = [ - UserMsg( - uid=msg.uid, - input_ids=t, - sampling_params=msg.sampling_params, - mm_inputs=getattr(t, "mm_inputs", None), - ) - for msg, t in zip(ok_msgs, ok_tensors, strict=True) + _backend_msg(msg, t) for msg, t in zip(ok_msgs, ok_tensors, strict=True) ] send_backend.put(backend[0] if len(backend) == 1 else BatchBackendMsg(data=backend)) if len(abort_msg) > 0: diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 3a00941f2..af724a341 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -2,6 +2,7 @@ import importlib.util import json +import hashlib import os import threading from types import ModuleType @@ -52,6 +53,28 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An IMAGE_MAX_PIXELS = int(os.getenv("FREETOKEN_IMAGE_MAX_PIXELS", str(1280 * 28 * 28))) +_IMAGE_KEY_BASE = 1 << 30 # image cache-key ids start above every vocabulary + + +def _image_cache_ids( + input_ids: torch.Tensor, image_token_id: int, counts: list[int], hashes: list[int] +) -> torch.Tensor: + """Prefix-cache key ids for an image prompt: ``input_ids`` with the i-th placeholder run + replaced by ids derived from the i-th image's content hash. Placeholders share one token id + across images, so keyed on raw ids a prefix hit could serve another image's KV; the derived + ids (``>= 2**30``, above any vocabulary, ``hash + offset`` within the run) make each run + unique to its image and the prompt safe to share. The model still reads ``input_ids``.""" + ids = input_ids.clone() + hits = (ids == image_token_id).nonzero().flatten() + mask = _IMAGE_KEY_BASE - 1 + pos = 0 + for n, h in zip(counts, hashes, strict=True): + run = (torch.arange(n, dtype=torch.int64) + (h & mask)) & mask + ids[hits[pos : pos + n]] = (run + _IMAGE_KEY_BASE).to(ids.dtype) + pos += n + return ids + + def _expand_image_tokens( input_ids: torch.Tensor, image_token_id: int, counts: list[int] ) -> torch.Tensor: @@ -119,10 +142,16 @@ def tokenize_one(self, msg: TokenizeMsg) -> tuple[torch.Tensor, dict[str, Any] | out = processor(images=images, return_tensors="pt") grid = out["image_grid_thw"] counts = (grid.prod(-1) // processor.merge_size**2).tolist() + hashes = [ + int.from_bytes(hashlib.blake2b(data, digest_size=8).digest(), "big") + for data in msg.images + ] + input_ids = _expand_image_tokens(input_ids, image_token_id, counts) # pixel_values stay fp32 on the wire (numpy cannot carry bf16); the tower casts them - return _expand_image_tokens(input_ids, image_token_id, counts), { + return input_ids, { "pixel_values": out["pixel_values"], "image_grid_thw": grid, + "cache_ids": _image_cache_ids(input_ids, image_token_id, counts, hashes), } def _vision_processor(self) -> tuple[Any, int]: diff --git a/tests/models/qwen4_exp/test_mrope.py b/tests/models/qwen4_exp/test_mrope.py index 3c48acb23..333ba3cbf 100644 --- a/tests/models/qwen4_exp/test_mrope.py +++ b/tests/models/qwen4_exp/test_mrope.py @@ -159,3 +159,25 @@ def test_mrope_cos_sin_matches_hf(): table = mrope_cos_sin(pos[:, 0], rotary.inv_freq.float(), SECTION) torch.testing.assert_close(table[:, :32], cos[0, :, :32], atol=1e-5, rtol=1e-5) torch.testing.assert_close(table[:, 32:], sin[0, :, :32], atol=1e-5, rtol=1e-5) + + +def test_image_cache_ids(): + from freetoken.tokenizer.tokenize import _image_cache_ids + + ids = _expand_image_tokens( + torch.tensor([1, IMAGE, 2, IMAGE, 3], dtype=torch.int32), IMAGE, [3, 2] + ) + a = _image_cache_ids(ids, IMAGE, [3, 2], [0xAAAA, 0xBBBB]) + b = _image_cache_ids(ids, IMAGE, [3, 2], [0xAAAA, 0xCCCC]) + text = ids != IMAGE + assert a.dtype == ids.dtype and torch.equal( + a[text], ids[text] + ) # text ids untouched + assert bool((a[~text] >= 1 << 30).all()) # never a vocabulary id + assert torch.equal(a[:5], b[:5]) and not torch.equal( + a[5:7], b[5:7] + ) # keyed per image + assert torch.equal( + a, _image_cache_ids(ids, IMAGE, [3, 2], [0xAAAA, 0xBBBB]) + ) # deterministic + assert len(set(a[1:4].tolist())) == 3 # one run, distinct ids per position diff --git a/tests/scheduler/test_mm_cache_key.py b/tests/scheduler/test_mm_cache_key.py new file mode 100644 index 000000000..dc7415a60 --- /dev/null +++ b/tests/scheduler/test_mm_cache_key.py @@ -0,0 +1,60 @@ +"""Image prompts take part in the prefix cache through ``cache_ids``: the same image hits, a +different image (same placeholder ids) matches only the text prefix, and a text prompt never +matches a placeholder run.""" + +import torch + +from freetoken.core import Req, SamplingParams +from freetoken.scheduler.cache import CacheManager +from freetoken.scheduler.utils import PendingReq + +IMAGE = 9 +KEY = 1 << 30 + + +def _ids(*toks): + return torch.tensor(toks, dtype=torch.int32) + + +def _finish(cm, uid, input_ids, cache_ids, table_idx): + """Admit, 'prefill', and commit one request into the tree.""" + sp = SamplingParams(max_tokens=1) + mr = cm.match_req(PendingReq(uid, input_ids, sp, cache_ids=cache_ids)) + req = Req( + input_ids=input_ids, + table_idx=table_idx, + cached_len=mr.cuda_handle.cached_len, + output_len=1, + uid=uid, + sampling_params=sp, + cache_handle=mr.cuda_handle, + mm_embeds=None if cache_ids is None else torch.zeros(3, 2), + cache_ids=cache_ids, + ) + cm.lock(mr.cuda_handle) + cm.allocate_paged([req]) + req.complete_one() + cm.cache_req(req, finished=True) + return req + + +def _cached(cm, input_ids, cache_ids=None): + return cm.match_req( + PendingReq(99, input_ids, SamplingParams(max_tokens=1), cache_ids=cache_ids) + ).cuda_handle.cached_len + + +def test_image_prompt_hits_only_for_the_same_image(): + cm = CacheManager(64, 1, torch.zeros(4, 64, dtype=torch.int32), "radix") + prompt = _ids(1, 2, IMAGE, IMAGE, IMAGE, 3) + same = _ids(1, 2, KEY + 7, KEY + 8, KEY + 9, 3) + other = _ids(1, 2, KEY + 70, KEY + 71, KEY + 72, 3) + _finish(cm, 1, prompt, same, table_idx=0) + # match_req keys on all but the last token, so a follow-up turn hits the whole prompt + follow = torch.cat([prompt, _ids(4, 5)]) + assert _cached(cm, follow, torch.cat([same, _ids(4, 5)])) == len(prompt) + # a different image behind the same placeholder ids only matches the text before it + assert _cached(cm, follow, torch.cat([other, _ids(4, 5)])) == 2 + # a text prompt with the raw placeholder ids never matches a keyed run + assert _cached(cm, follow) == 2 + cm.check_integrity()