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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

Author: gdevenyi

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

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

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/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(

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,

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,14 +99,21 @@ 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]:
         """The PLE layers in decoder order -- the seam the loader attaches table backends to."""
         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