[feat] cosmos3 PR4: text reasoning (und pathway + lm_head), bit-exact

Add the reasoning (VLM text-generation) capability. The framework reasoner uses
only the und (causal) pathway weights (no _moe_gen) + embed_tokens / norm /
lm_head — all already present in the native DiT — so a text-only forward + lm_head
reproduces it with no new model code.

- `cosmos3_generate_reasoner_text(transformer, input_ids, max_new_tokens)`:
  greedy text reasoning via the und causal backbone + lm_head (text-only prefill,
  re-prefill per step; KV-cache fast path is a later optimization). Mirrors the
  framework `generate_reasoner_text` (text-only path).
- Parity (`test_cosmos3_reasoning_parity.py`, framework = oracle):
  * prefill logits (native text-only forward + lm_head vs framework
    `reasoner_forward` + lm_head, per-position): max abs diff = 0.0, mean = 0.0;
  * greedy generation vs framework `generate_reasoner_text(do_sample=False)`:
    token-for-token identical across seeds.

Text reasoning is fully supported + bit-exact. Image-conditioned reasoning
additionally needs the Qwen3-VL `vision_encoder` (a transformers ViT, tracked
separately). Full suite: 146 passed, 0 skipped.

fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py

@@ -109,6 +109,60 @@ def cosmos3_tokenize_caption(
     return list(token_ids)
 
 
+# ===========================================================================
+# Reasoning (VLM text generation) — und (causal) pathway + lm_head
+# ===========================================================================
+def cosmos3_generate_reasoner_text(
+    transformer: Any,
+    input_ids: list[int],
+    max_new_tokens: int,
+    *,
+    eos_token_id: int | list[int] | None = None,
+) -> list[int]:
+    """Greedy text reasoning via the und (causal) backbone + ``lm_head``.
+
+    Mirrors the framework ``generate_reasoner_text`` (text-only prefill, greedy):
+    only the und-pathway weights (no ``_moe_gen``) + ``embed_tokens`` / ``norm`` /
+    ``lm_head`` participate; the generation pathway and the VFM multimodal
+    embedders are bypassed (no vision/sound/action tokens). Token-for-token
+    identical to the framework reasoner (``test_cosmos3_reasoning_parity``).
+
+    Re-prefills each step (no KV cache) — correctness-first; a KV-cache fast path
+    is a later optimization. Returns the newly generated token ids.
+    """
+    device = next(transformer.parameters()).device
+    ids = [int(x) for x in input_ids]
+    eos: set[int] = set()
+    if eos_token_id is not None:
+        eos = {int(eos_token_id)} if isinstance(eos_token_id, int) else {int(x) for x in eos_token_id}
+
+    new_tokens: list[int] = []
+    for _ in range(int(max_new_tokens)):
+        n = len(ids)
+        pos = torch.arange(n).unsqueeze(0).expand(3, -1).contiguous().to(device)
+        out = transformer(
+            text_ids=torch.tensor(ids, device=device, dtype=torch.long),
+            text_indexes=torch.arange(n, device=device),
+            position_ids=pos,
+            sequence_length=n,
+            split_lens=[n],
+            attn_modes=["causal"],
+            vision_tokens=[],
+            vision_token_shapes=[],
+            vision_sequence_indexes=torch.empty(0, dtype=torch.long, device=device),
+            vision_timesteps=torch.empty(0, device=device),
+            vision_mse_loss_indexes=torch.empty(0, dtype=torch.long, device=device),
+            vision_noisy_frame_indexes=[],
+        )
+        logits = transformer.lm_head(out["last_hidden_state"][n - 1])  # [vocab]
+        nxt = int(logits.argmax().item())
+        ids.append(nxt)
+        new_tokens.append(nxt)
+        if nxt in eos:
+            break
+    return new_tokens
+
+
 # ===========================================================================
 # VAE encode/decode bridge (normalize / denormalize, matching the framework)
 # ===========================================================================

tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical-parity test: FastVideo Cosmos3 text reasoning vs the framework.
+
+The omni model's reasoning (VLM text generation) uses ONLY the und (causal)
+pathway weights (no ``_moe_gen``) + ``embed_tokens`` / ``norm`` / ``lm_head``;
+the generation pathway and all multimodal embedders are bypassed. FastVideo's
+native ``Cosmos3VFMTransformer`` already contains exactly those (the und branch
+of the dual-pathway forward + ``lm_head``), so a text-only forward + ``lm_head``
+reproduces the framework reasoner.
+
+This pins:
+  * **prefill logits** — native text-only forward + ``lm_head`` vs the framework
+    ``language_model.model.reasoner_forward`` + ``lm_head`` (per-position); and
+  * **greedy generation** — ``cosmos3_generate_reasoner_text`` vs the framework
+    ``generate_reasoner_text(do_sample=False)`` (token-for-token).
+
+Framework is the parity ORACLE (CPU/float32 via SDPA monkey-patch). Text-only
+(image-conditioned reasoning additionally needs the Qwen3-VL ``vision_encoder``,
+tracked separately).
+
+Run:
+    cd <worktree> && <fv-cosmos3 python> -m pytest \
+        tests/local_tests/cosmos3/test_cosmos3_reasoning_parity.py -q -s
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+cosmos_framework = pytest.importorskip(
+    "cosmos_framework",
+    reason="cosmos_framework not installed; run in fv-cosmos3 env.",
+)
+
+from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # noqa: E402
+    cosmos3_generate_reasoner_text,
+)
+
+from .test_cosmos3_dit_parity import _copy_weights  # noqa: E402
+from .test_cosmos3_dit_parity_mrope import (  # noqa: E402
+    _build_tiny_cosmos3_mrope,
+    _build_tiny_fastvideo_dit_mrope,
+)
+from .test_cosmos3_reference_forward import _apply_sdpa_patches  # noqa: E402
+
+pytestmark = [pytest.mark.local]
+_apply_sdpa_patches()
+
+
+def _framework_prefill_logits(vfm, input_ids: torch.Tensor) -> torch.Tensor:
+    """Framework reasoner text-only prefill logits ``[1, T, vocab]`` (the oracle).
+
+    Mirrors ``_impl_generate_reasoner_text`` prefill: ``model.reasoner_forward``
+    then ``lm_head`` (here over ALL positions, not just the last)."""
+    from cosmos_framework.model.vfm.mot.unified_mot import ReasonerKVCache
+
+    causal_lm = vfm.language_model
+    model = causal_lm.model
+    cache = ReasonerKVCache.empty(num_layers=len(model.layers))
+    hidden = model.reasoner_forward(input_ids.unsqueeze(0), cache=cache)  # [1,T,hidden]
+    return causal_lm.lm_head(hidden)  # [1,T,vocab]
+
+
+def _native_prefill_logits(dit, input_ids: torch.Tensor) -> torch.Tensor:
+    """Native text-only forward + ``lm_head`` -> ``[T, vocab]``."""
+    n = input_ids.numel()
+    pos = torch.arange(n).unsqueeze(0).expand(3, -1).contiguous()
+    out = dit(
+        text_ids=input_ids.to(torch.long),
+        text_indexes=torch.arange(n),
+        position_ids=pos,
+        sequence_length=n,
+        split_lens=[n],
+        attn_modes=["causal"],
+        vision_tokens=[],
+        vision_token_shapes=[],
+        vision_sequence_indexes=torch.empty(0, dtype=torch.long),
+        vision_timesteps=torch.empty(0),
+        vision_mse_loss_indexes=torch.empty(0, dtype=torch.long),
+        vision_noisy_frame_indexes=[],
+    )
+    return dit.lm_head(out["last_hidden_state"])  # [T, vocab]
+
+
+def _diffs(a, b):
+    d = (a - b).abs()
+    return d.max().item(), d.mean().item()
+
+
+class TestCosmos3ReasoningParity:
+
+    def _build(self, seed_model=42, num_layers=2):
+        vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers)
+        dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers)
+        _copy_weights(vfm, dit)
+        return vfm, dit
+
+    @pytest.mark.parametrize("n_text", [4, 8, 12])
+    def test_reasoner_prefill_logits_match_framework(self, n_text):
+        vfm, dit = self._build()
+        torch.manual_seed(n_text)
+        input_ids = torch.randint(0, 60, (n_text,))
+        with torch.no_grad():
+            fw = _framework_prefill_logits(vfm, input_ids)[0]  # [T, vocab]
+            fv = _native_prefill_logits(dit, input_ids)  # [T, vocab]
+        assert fw.shape == fv.shape, f"shape fw={fw.shape} fv={fv.shape}"
+        mx, mn = _diffs(fv, fw)
+        print(f"\n[reasoner_prefill n={n_text}] logits max abs diff = {mx:.3e} mean abs diff = {mn:.3e}")
+        torch.testing.assert_close(fv, fw, atol=1e-4, rtol=1e-3)
+        # The greedy decisions (argmax per position) must agree exactly.
+        assert torch.equal(fv.argmax(-1), fw.argmax(-1))
+
+    @pytest.mark.parametrize("seed", [0, 1, 2, 3])
+    def test_greedy_reasoning_matches_framework(self, seed):
+        vfm, dit = self._build()
+        torch.manual_seed(seed)
+        input_ids = torch.randint(0, 60, (6,))
+        fw = vfm.generate_reasoner_text(
+            input_ids.unsqueeze(0), max_new_tokens=8, do_sample=False, return_only_new_tokens=True,
+        )[0].tolist()
+        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}"