[feat] cosmos3 PR4: deepstack reasoner forward (image-conditioned reasoning)

Add the one new native piece for image-conditioned reasoning: a deepstack-capable
reasoner backbone forward. The other parts are already proven/reusable — the
Qwen3-VL vision_encoder (transformers, bit-exact vs framework) and the multimodal
prefill (masked-scatter of image features + get_rope_index, standard transformers
Qwen3-VL boilerplate, reusable like the tokenizer/encoder).

- `Cosmos3VFMTransformer.reason_forward(inputs_embeds, position_ids,
  deepstack_embeds, visual_pos_mask)`: runs the und (causal) layers over a single
  text(+image) sequence, injecting per-layer deepstack visual embeds at the visual
  positions for the first `len(deepstack_embeds)` layers (Qwen3-VL deepstack),
  then final `norm`. Mirrors the framework `reasoner_forward`.
- Parity (`test_cosmos3_reasoning_parity::test_deepstack_reasoner_forward_matches_framework`):
  native `reason_forward` vs framework `reasoner_forward` given identical
  inputs_embeds + positions + per-layer deepstack + visual mask — hidden states
  max abs diff = 0.0, mean = 0.0.

All five omni modalities are now bit-exact vs the framework (video/audio/action
gen, text reasoning, vision encoding, deepstack reasoner). Full suite: 150 passed.

fastvideo/models/dits/cosmos3.py

@@ -934,6 +934,44 @@ class Cosmos3VFMTransformer(BaseDiT):
         output["last_hidden_state"] = last_hidden_state
         return output
 
+    def reason_forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        deepstack_embeds: list[torch.Tensor] | None = None,
+        visual_pos_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Reasoner (causal) backbone forward over one text(+image) sequence.
+
+        All tokens are und (causal); the gen pathway is empty. Mirrors the
+        framework ``reasoner_forward``: run the und layers, optionally injecting
+        per-layer deepstack visual embeds at ``visual_pos_mask`` for the first
+        ``len(deepstack_embeds)`` layers (Qwen3-VL deepstack), then final
+        ``norm``. Returns post-norm hidden states ``[N, hidden]`` (apply
+        ``lm_head`` for logits). Used by image-conditioned reasoning, with the
+        text/vision embeds + ``position_ids`` produced by the multimodal prefill.
+
+        Args:
+            inputs_embeds: ``[N, hidden]`` text embeds with image features
+                already scattered in at the image-token positions.
+            position_ids: ``[3, N]`` (or ``[N]``) mRoPE positions.
+            deepstack_embeds: per-deepstack-layer ``[n_visual, hidden]`` embeds.
+            visual_pos_mask: ``[N]`` bool mask of visual token positions.
+        """
+        device = self.embed_tokens.weight.device
+        target_dtype = inputs_embeds.dtype
+        cos, sin = self.rotary_emb(position_ids, device=device, dtype=target_dtype)  # [N, head_dim]
+        und = inputs_embeds
+        gen = inputs_embeds.new_zeros((0, self.hidden_size))
+        cos_gen = cos.new_zeros((0, cos.shape[-1]))
+        sin_gen = sin.new_zeros((0, sin.shape[-1]))
+        for li, layer in enumerate(self.layers):
+            und, _ = layer(und, gen, cos, sin, cos_gen, sin_gen)
+            if deepstack_embeds is not None and visual_pos_mask is not None and li < len(deepstack_embeds):
+                und = und.clone()
+                und[visual_pos_mask] = und[visual_pos_mask] + deepstack_embeds[li].to(dtype=und.dtype)
+        return self.norm(und)  # [N, hidden]
+
     def _unpack_action(
         self,
         packed_preds: torch.Tensor,

tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py

@@ -121,3 +121,39 @@ class TestCosmos3ReasoningParity:
         fv = cosmos3_generate_reasoner_text(dit, input_ids.tolist(), max_new_tokens=8)
         print(f"\n[greedy_reason seed={seed}] framework={fw} native={fv}")
         assert fv == fw, f"token mismatch: native={fv} framework={fw}"
+
+    def test_deepstack_reasoner_forward_matches_framework(self):
+        """The deepstack reasoner backbone (image-conditioned reasoning's one new
+        native piece) matches the framework ``reasoner_forward`` given identical
+        prefill inputs (inputs_embeds + positions + per-layer deepstack embeds +
+        visual mask)."""
+        from cosmos_framework.model.vfm.mot.unified_mot import ReasonerKVCache
+
+        vfm, dit = self._build()
+        hidden = dit.hidden_size
+        n_layers = dit.num_hidden_layers
+        seq = 10
+        torch.manual_seed(5)
+        inputs_embeds = torch.randn(seq, hidden)
+        position_ids = torch.arange(seq).unsqueeze(0).expand(3, -1).contiguous()  # [3, seq]
+        visual_pos_mask = torch.zeros(seq, dtype=torch.bool)
+        visual_pos_mask[2:6] = True  # 4 "image" tokens
+        k = int(visual_pos_mask.sum())
+        deepstack = [torch.randn(k, hidden) for _ in range(n_layers)]  # one per layer
+
+        model = vfm.language_model.model
+        cache = ReasonerKVCache.empty(num_layers=len(model.layers))
+        with torch.no_grad():
+            hid_fw = model.reasoner_forward(
+                input_ids=None,
+                inputs_embeds=inputs_embeds.unsqueeze(0),
+                position_ids=position_ids.unsqueeze(1),  # [3, B=1, seq]
+                visual_pos_masks=visual_pos_mask.unsqueeze(0),
+                deepstack_visual_embeds=deepstack,
+                cache=cache,
+            )[0]  # [seq, hidden]
+            hid_fv = dit.reason_forward(inputs_embeds, position_ids, deepstack, visual_pos_mask)  # [seq, hidden]
+        assert hid_fw.shape == hid_fv.shape, f"shape fw={hid_fw.shape} fv={hid_fv.shape}"
+        mx, mn = _diffs(hid_fv, hid_fw)
+        print(f"\n[deepstack_reasoner] hidden max abs diff = {mx:.3e} mean abs diff = {mn:.3e}")
+        torch.testing.assert_close(hid_fv, hid_fw, atol=1e-4, rtol=1e-3)