feat(MODEL-MM-indextts2): the DiT transformer stack, so the S2Mel estimator is COMPLETE front to tail (#634)

CMakeLists.txt

@@ -778,6 +778,7 @@ add_library(vllm STATIC
   src/vllm/model_executor/models/dit_tail.cpp
   src/vllm/model_executor/models/dit_skip.cpp
   src/vllm/model_executor/models/dit_front.cpp
+  src/vllm/model_executor/models/dit_stack.cpp
   src/vllm/model_executor/models/indextts2_s2mel_loader.cpp
   src/vllm/model_executor/models/vocos.cpp
   src/vllm/model_executor/models/lenreg.cpp

docs/FEATURES.md

@@ -196,7 +196,7 @@ on the committed fixture); reranking/classify models are not yet registered.
 | Video | ✅ correctness-gated | ✅ | ✅ | ☐ |
 | Audio | ✅ correctness-gated | ✅ | ◐ | ◐ |
 | Video+audio GENERATION (MiniMax-H3 DiT, LTX-2.5 DiT) | ◐ H3: all three modalities COHERENT on Q4_K_M (t2va, fl2va, ref2va; §8.20); the NVFP4 arm carries the patch grid; GGUF/NVFP4/bf16 loaders, pruned too (§8.21). LTX-2.5: a second lane, `SPIKE`, gated at reduced dims | ✅ H3 (vllm-omni, BF16-only, no quantized arm); LTX-2.5 only through the generic diffusers adapter, no native recipe ([vllm-omni#6066](https://github.com/vllm-project/vllm-omni/issues/6066)) | ☐ | ☐ |
-| Speech / audio GENERATION (TTS, vLLM-Omni lane) | ◐ IndexTTS-2.5 only: the DiT front end, blocks' skips and tail are ported, and the tail runs on REAL shipped weights; other stages gated at reduced dims. No render, no route (#634) | ✅ (vllm-omni: MOSS-TTS, Qwen3-TTS, Higgs Audio v3, Voxtral TTS, IndexTTS-2.5) | not assessed | not assessed |
+| Speech / audio GENERATION (TTS, vLLM-Omni lane) | ◐ IndexTTS-2.5 only: the S2Mel DiT is COMPLETE front to tail and gated against upstream; its tail runs on REAL shipped weights. Other stages gated at reduced dims. No render, no route (#634) | ✅ (vllm-omni: MOSS-TTS, Qwen3-TTS, Higgs Audio v3, Voxtral TTS, IndexTTS-2.5) | not assessed | not assessed |
 | MUSIC generation (MiniMax-Music3) | ☐ not generating. The W1 checkpoint LOADER has landed ([spec](../.agents/specs/minimax-music3.md), #672); no stage runs yet. Lyrics plus a structured description in, a multi-minute stereo song out | ☐ absent from the pin, from vLLM `main` and from `vllm-omni` alike | ◐ served by SGLang-Omni, a third repository, which loads the NATIVE checkpoint layout | ☐ |
 | Multimodal over the OpenAI server | ◐ image request path wired, forward pending | ✅ | ✅ | ◐ |
 

docs/USAGE.md

@@ -2342,6 +2342,9 @@ DIT_SRC=/path/to/index-tts/indextts/s2mel/modules \
 
 DIT_SRC=/path/to/index-tts/indextts/s2mel/modules \
   python3 scripts/gen-dit-front-goldens.py --out tests/vllm/models/dit_front_goldens.inc
+
+DIT_SRC=/path/to/index-tts/indextts/s2mel/modules \
+  python3 scripts/gen-dit-stack-goldens.py --out tests/vllm/models/dit_stack_goldens.inc
 ```
 
 The U-Net skip routing is recorded rather than generated into an `.inc`: this

include/vllm/model_executor/models/dit_stack.h

@@ -0,0 +1,60 @@
+// The S2Mel DiT transformer STACK: what sits between front end and tail (#634).
+//
+// Upstream `indextts/s2mel/modules/gpt_fast/model.py:161-191`
+// (Transformer.forward), index-tts @4f8792ff120cd3ea470dd511e997a17c86cddd10.
+// N blocks, the U-Net skip routing across them, and a final AdaptiveLayerNorm.
+//
+// This is composition only: the block is `dit::Block`, the routing is
+// `dit_skip::Plan`, and the per-layer skip merge is `dit_skip::ApplySkip`.
+// Nothing here reimplements any of them.
+//
+// The rotary table is an INPUT. Upstream precomputes `freqs_cis` once for the
+// whole model and indexes it by position, so passing it in keeps this a gate on
+// composition rather than on a second copy of that computation.
+//
+// Upstream builds a `skip_in_linear` on EVERY layer when uvit_skip_connection is
+// set, even the layers that never receive one, so the checkpoint carries
+// unused ones. They are loaded and left alone rather than treated as an error.
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "vllm/model_executor/models/dit.h"
+
+namespace vllm {
+namespace models {
+namespace dit_stack {
+
+struct LayerWeights {
+  dit::BlockWeights block;
+  // Present on every layer upstream; consulted only on receiving layers.
+  std::vector<float> skip_in_w;  // [dim, 2 * dim]
+  std::vector<float> skip_in_b;  // [dim]
+};
+
+struct Weights {
+  std::vector<LayerWeights> layers;
+  // transformer.norm, an AdaptiveLayerNorm like the per-block ones.
+  std::vector<float> norm_proj_w, norm_proj_b, norm_w;
+};
+
+struct Config {
+  int64_t dim = 0;
+  int64_t heads = 0;
+  int64_t head_dim = 0;
+  int64_t intermediate = 0;
+  int64_t frames = 0;
+  double eps = 1e-5;
+};
+
+// x is [frames, dim]; cond is [dim] (one conditioning vector, as upstream passes
+// t1 unsqueezed); freqs is the rotary table for these positions.
+// Returns [frames, dim].
+std::vector<float> Forward(const Config& cfg, const Weights& w,
+                           const std::vector<float>& x, const std::vector<float>& cond,
+                           const std::vector<float>& freqs);
+
+}  // namespace dit_stack
+}  // namespace models
+}  // namespace vllm

scripts/gen-dit-stack-goldens.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+"""Emit C++ goldens for the whole S2Mel DiT transformer STACK.
+
+Upstream `indextts/s2mel/modules/gpt_fast/model.py:161-191` (Transformer.forward)
+plus its TransformerBlock, index-tts @4f8792ff120cd3ea470dd511e997a17c86cddd10.
+This is what sits BETWEEN the ported front end and the ported tail: N blocks,
+the U-Net skip routing across them, and a final AdaptiveLayerNorm.
+
+The rotary table is emitted as an INPUT rather than recomputed on our side, so
+this gates the stack's composition and not a second implementation of
+`precompute_freqs_cis`.
+
+Usage: DIT_SRC=<path to indextts/s2mel/modules> python3 \
+           scripts/gen-dit-stack-goldens.py --out tests/vllm/models/dit_stack_goldens.inc
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import os
+import sys
+import types
+from pathlib import Path
+
+import torch
+
+
+def rnd(name: str, n: int, scale: float = 1.0) -> list:
+    h = 0xCBF29CE484222325
+    for ch in name.encode():
+        h = ((h ^ ch) * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF
+    out = []
+    for _ in range(n):
+        h = (h + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF
+        z = h
+        z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
+        z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF
+        z ^= z >> 31
+        out.append(((z >> 11) * (1.0 / 9007199254740992.0) * 2.0 - 1.0) * scale)
+    return out
+
+
+def tensor(name: str, shape, scale: float = 1.0) -> torch.Tensor:
+    n = 1
+    for d in shape:
+        n *= d
+    return torch.tensor(rnd(name, n, scale), dtype=torch.float64).reshape(shape).float()
+
+
+def fmt(values) -> str:
+    lines, row = [], []
+    for v in values:
+        row.append(f"{float(v):.9e}F")
+        if len(row) == 6:
+            lines.append("    " + ", ".join(row) + ",")
+            row = []
+    if row:
+        lines.append("    " + ", ".join(row) + ",")
+    return "\n".join(lines)
+
+
+DIM, HEADS, HEAD_DIM, DEPTH, FRAMES = 8, 2, 4, 5, 6
+INTERMEDIATE = None  # taken from the constructed model
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", required=True)
+    a = ap.parse_args()
+    src = Path(os.environ["DIT_SRC"])
+    sys.path.insert(0, str(src.parents[2]))
+    for name in ("munch",):
+        if name not in sys.modules:
+            stub = types.ModuleType(name)
+            stub.Munch = dict
+            sys.modules[name] = stub
+    spec = importlib.util.spec_from_file_location(
+        "indextts.s2mel.modules.gpt_fast.model", src / "gpt_fast" / "model.py"
+    )
+    gm = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(gm)
+
+    torch.manual_seed(0)
+    args = gm.ModelArgs(block_size=64, n_layer=DEPTH, n_head=HEADS, dim=DIM,
+                        head_dim=HEAD_DIM, vocab_size=16,
+                        uvit_skip_connection=True, time_as_token=False)
+    tr = gm.Transformer(args)
+    tr.setup_caches(1, 32, use_kv_cache=False)
+    tr.eval()
+
+    with torch.no_grad():
+        for pname, p in sorted(tr.named_parameters()):
+            p.copy_(tensor("stack." + pname, list(p.shape), 0.5))
+
+    x = tensor("stack.x", [1, FRAMES, DIM])
+    c = tensor("stack.c", [1, 1, DIM])
+    input_pos = torch.arange(FRAMES)
+    mask = torch.ones(1, 1, FRAMES, FRAMES, dtype=torch.bool)
+
+    with torch.no_grad():
+        out = tr(x, c, input_pos, mask)
+        freqs = tr.freqs_cis[input_pos]  # [FRAMES, head_dim/2, 2]
+
+    inter = tr.layers[0].feed_forward.w1.weight.shape[0]
+    names = sorted(n for n, _ in tr.named_parameters())
+
+    body = [
+        "// GENERATED by scripts/gen-dit-stack-goldens.py -- do not edit.",
+        "// Oracle: gpt_fast/model.py Transformer.forward + TransformerBlock,",
+        "// index-tts @4f8792ff120cd3ea470dd511e997a17c86cddd10, with",
+        "// uvit_skip_connection: true as the shipped config sets.",
+        "#pragma once",
+        "",
+        "#include <cstdint>",
+        "",
+        "namespace dit_stack_goldens {",
+        "",
+        f"inline constexpr int64_t kDim = {DIM};",
+        f"inline constexpr int64_t kHeads = {HEADS};",
+        f"inline constexpr int64_t kHeadDim = {HEAD_DIM};",
+        f"inline constexpr int64_t kDepth = {DEPTH};",
+        f"inline constexpr int64_t kFrames = {FRAMES};",
+        f"inline constexpr int64_t kIntermediate = {int(inter)};",
+        f"inline constexpr double kEps = {float(args.norm_eps):.9e};",
+        "",
+        "inline constexpr const char* kParamNames[] = {",
+    ]
+    body += [f'    "stack.{n}",' for n in names]
+    body += [
+        "};",
+        "",
+        "// The rotary table upstream used, [kFrames, kHeadDim/2, 2]. Emitted as",
+        "// an INPUT so this gates composition, not a second rotary implementation.",
+        "inline constexpr float kFreqs[] = {",
+        fmt(freqs.reshape(-1).tolist()),
+        "};",
+        "",
+        "// Transformer.forward output -- [kFrames, kDim].",
+        "inline constexpr float kOut[] = {",
+        fmt(out.reshape(-1).tolist()),
+        "};",
+        "",
+        "}  // namespace dit_stack_goldens",
+        "",
+    ]
+    Path(a.out).write_text("\n".join(body))
+    print(f"wrote {a.out}: depth {DEPTH}, intermediate {int(inter)}, "
+          f"{len(names)} params, out {tuple(out.shape)}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/vllm/model_executor/models/dit_stack.cpp

@@ -0,0 +1,67 @@
+// S2Mel DiT transformer stack. See dit_stack.h for the upstream anchors.
+#include "vllm/model_executor/models/dit_stack.h"
+
+#include <cstddef>
+#include <vector>
+
+#include "vllm/model_executor/models/dit_skip.h"
+#include "vt/dtype.h"
+
+namespace vllm {
+namespace models {
+namespace dit_stack {
+
+std::vector<float> Forward(const Config& cfg, const Weights& w,
+                           const std::vector<float>& x, const std::vector<float>& cond,
+                           const std::vector<float>& freqs) {
+  const int64_t layers = static_cast<int64_t>(w.layers.size());
+  VT_CHECK(layers > 0, "dit_stack: no layers");
+  VT_CHECK(cfg.dim > 0 && cfg.frames > 0, "dit_stack: dim and frames must be positive");
+  VT_CHECK(x.size() == static_cast<size_t>(cfg.frames * cfg.dim),
+           "dit_stack: x must be [frames, dim]");
+
+  const dit_skip::Schedule plan = dit_skip::Plan(layers);
+
+  std::vector<float> cur = x;
+  // The stack holds the OUTPUTS of emitting layers, in order, and receivers pop
+  // the most recent. `plan.source` already says which layer each receiver takes,
+  // so the stack here only has to carry the values.
+  std::vector<std::vector<float>> stack;
+
+  for (int64_t i = 0; i < layers; ++i) {
+    const LayerWeights& layer = w.layers[static_cast<size_t>(i)];
+
+    // Receive BEFORE the layer runs, merging with skip_in_linear.
+    if (plan.source[static_cast<size_t>(i)] >= 0) {
+      VT_CHECK(!stack.empty(), "dit_stack: a receiving layer found no skip");
+      VT_CHECK(!layer.skip_in_w.empty(),
+               "dit_stack: a receiving layer has no skip_in_linear");
+      const std::vector<float> skip = stack.back();
+      stack.pop_back();
+      cur = dit_skip::ApplySkip(cur, skip, cfg.frames, cfg.dim, layer.skip_in_w,
+                                layer.skip_in_b);
+    }
+
+    cur = dit::Block(cur, cond, cfg.frames, cfg.dim, cfg.heads, cfg.head_dim,
+                     cfg.intermediate, freqs, layer.block, cfg.eps);
+
+    // Emit AFTER, pushing this layer's own output.
+    bool emits = false;
+    for (const int64_t e : plan.emit) {
+      if (e == i) {
+        emits = true;
+      }
+    }
+    if (emits) {
+      stack.push_back(cur);
+    }
+  }
+
+  // transformer.norm: the same AdaptiveLayerNorm shape as the per-block norms.
+  return dit::AdaptiveLayerNorm(cur, cfg.frames, cfg.dim, cond, w.norm_proj_w,
+                                w.norm_proj_b, w.norm_w, cfg.eps);
+}
+
+}  // namespace dit_stack
+}  // namespace models
+}  // namespace vllm

tests/CMakeLists.txt

@@ -119,6 +119,7 @@ vllm_cpp_add_test(test_wavenet vllm/models/test_wavenet.cpp)
 vllm_cpp_add_test(test_dit_tail vllm/models/test_dit_tail.cpp)
 vllm_cpp_add_test(test_dit_skip vllm/models/test_dit_skip.cpp)
 vllm_cpp_add_test(test_dit_front vllm/models/test_dit_front.cpp)
+vllm_cpp_add_test(test_dit_stack vllm/models/test_dit_stack.cpp)
 vllm_cpp_add_test(test_indextts2_s2mel_loader vllm/models/test_indextts2_s2mel_loader.cpp)
 vllm_cpp_add_test(test_vocos vllm/models/test_vocos.cpp)
 vllm_cpp_add_test(test_lenreg vllm/models/test_lenreg.cpp)
@@ -136,6 +137,7 @@ target_include_directories(test_fvq PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/mod
 target_include_directories(test_wavenet PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models)
 target_include_directories(test_dit_tail PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models)
 target_include_directories(test_dit_front PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models)
+target_include_directories(test_dit_stack PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models)
 target_include_directories(test_w2vbert PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models)
 target_include_directories(test_campplus PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/vllm/models)
 vllm_cpp_add_test(test_gpt2 vllm/models/test_gpt2.cpp)