feat(qwen4_exp): image input (vision tower + mRoPE) for Qwen3.8-Flash-Next

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),

python/freetoken/core.py

@@ -43,6 +43,14 @@ 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
+    # 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). ---
@@ -135,6 +143,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

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

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
 

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)

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

python/freetoken/message/backend.py

@@ -37,6 +37,15 @@ 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
+    # 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

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

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:

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__ = [