From 840cdee9f9bd5fa40f5ebbfd1d5207f976b89909 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 14 Aug 2026 12:12:33 +0000 Subject: [PATCH] feat(MODEL-MM-indextts2): the DiT transformer stack, so the S2Mel estimator is COMPLETE front to tail (#634) Ports `gpt_fast/model.py:161-191` (Transformer.forward): N blocks, the U-Net skip routing across them, and the final AdaptiveLayerNorm. With #738, #739, #741 and #745 the S2Mel DiT now exists end to end -- front end, stack, tail -- each gated against upstream's own modules. This is composition only. The block is `dit::Block`, the routing is `dit_skip::Plan`, the merge is `dit_skip::ApplySkip`, and nothing here reimplements any of them. The rotary table is passed IN, emitted by the generator from upstream's own `freqs_cis`, so the gate measures composition rather than a second copy of `precompute_freqs_cis`. One upstream quirk is recorded rather than tidied: a `skip_in_linear` is built on EVERY layer when `uvit_skip_connection` is set, including layers that never receive a skip, so the checkpoint carries unused ones. They load and are left alone. A case asserts that layer 0's skip_in_linear cannot affect the output at all, while a receiving layer's must. 3 cases / 52 assertions against upstream Transformer.forward. Five mutations, five caught, all compile_err=0, file byte-identical after: the skip merged against itself instead of the popped value, emitting a layer's INPUT instead of its output, FIFO instead of LIFO, the final norm dropped, and the receive moved after the block. A SIXTH mutation was withdrawn as badly designed rather than reported as a gate hole. Pushing before the block WITHOUT removing the push after it leaves the stack holding both, and because a layer's input IS the previous layer's output, every pop still returned the same tensor -- a genuine no-op, not a defect the gate missed. Rewritten to emit the input INSTEAD of the output, it fails. Still no render: the talker, the reference encoders and BigVGAN are unported, so nothing yet produces the conditioning this estimator consumes, and correctness against vLLM-Omni needs the oracle pin (#633). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- CMakeLists.txt | 1 + docs/FEATURES.md | 2 +- docs/USAGE.md | 3 + .../vllm/model_executor/models/dit_stack.h | 60 +++++++ scripts/gen-dit-stack-goldens.py | 154 ++++++++++++++++++ src/vllm/model_executor/models/dit_stack.cpp | 67 ++++++++ tests/CMakeLists.txt | 2 + tests/vllm/models/dit_stack_goldens.inc | 111 +++++++++++++ tests/vllm/models/test_dit_stack.cpp | 118 ++++++++++++++ 9 files changed, 517 insertions(+), 1 deletion(-) create mode 100644 include/vllm/model_executor/models/dit_stack.h create mode 100644 scripts/gen-dit-stack-goldens.py create mode 100644 src/vllm/model_executor/models/dit_stack.cpp create mode 100644 tests/vllm/models/dit_stack_goldens.inc create mode 100644 tests/vllm/models/test_dit_stack.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ab562d03..81add6235 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/docs/FEATURES.md b/docs/FEATURES.md index d2f6d3c03..634f8869b 100644 --- a/docs/FEATURES.md +++ b/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 | ✅ | ✅ | ◐ | diff --git a/docs/USAGE.md b/docs/USAGE.md index 0bf5c6aef..5621cac86 100644 --- a/docs/USAGE.md +++ b/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 diff --git a/include/vllm/model_executor/models/dit_stack.h b/include/vllm/model_executor/models/dit_stack.h new file mode 100644 index 000000000..4bc50d741 --- /dev/null +++ b/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 +#include + +#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 skip_in_w; // [dim, 2 * dim] + std::vector skip_in_b; // [dim] +}; + +struct Weights { + std::vector layers; + // transformer.norm, an AdaptiveLayerNorm like the per-block ones. + std::vector 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 Forward(const Config& cfg, const Weights& w, + const std::vector& x, const std::vector& cond, + const std::vector& freqs); + +} // namespace dit_stack +} // namespace models +} // namespace vllm diff --git a/scripts/gen-dit-stack-goldens.py b/scripts/gen-dit-stack-goldens.py new file mode 100644 index 000000000..447c50a3d --- /dev/null +++ b/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= 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 ", + "", + "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()) diff --git a/src/vllm/model_executor/models/dit_stack.cpp b/src/vllm/model_executor/models/dit_stack.cpp new file mode 100644 index 000000000..9bf970bd0 --- /dev/null +++ b/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 +#include + +#include "vllm/model_executor/models/dit_skip.h" +#include "vt/dtype.h" + +namespace vllm { +namespace models { +namespace dit_stack { + +std::vector Forward(const Config& cfg, const Weights& w, + const std::vector& x, const std::vector& cond, + const std::vector& freqs) { + const int64_t layers = static_cast(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(cfg.frames * cfg.dim), + "dit_stack: x must be [frames, dim]"); + + const dit_skip::Schedule plan = dit_skip::Plan(layers); + + std::vector 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> stack; + + for (int64_t i = 0; i < layers; ++i) { + const LayerWeights& layer = w.layers[static_cast(i)]; + + // Receive BEFORE the layer runs, merging with skip_in_linear. + if (plan.source[static_cast(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 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 diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3722500c2..68c3f3135 100644 --- a/tests/CMakeLists.txt +++ b/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) diff --git a/tests/vllm/models/dit_stack_goldens.inc b/tests/vllm/models/dit_stack_goldens.inc new file mode 100644 index 000000000..37e851bf2 --- /dev/null +++ b/tests/vllm/models/dit_stack_goldens.inc @@ -0,0 +1,111 @@ +// 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 + +namespace dit_stack_goldens { + +inline constexpr int64_t kDim = 8; +inline constexpr int64_t kHeads = 2; +inline constexpr int64_t kHeadDim = 4; +inline constexpr int64_t kDepth = 5; +inline constexpr int64_t kFrames = 6; +inline constexpr int64_t kIntermediate = 256; +inline constexpr double kEps = 1.000000000e-05; + +inline constexpr const char* kParamNames[] = { + "stack.layers.0.attention.wo.weight", + "stack.layers.0.attention.wqkv.weight", + "stack.layers.0.attention_norm.norm.weight", + "stack.layers.0.attention_norm.project_layer.bias", + "stack.layers.0.attention_norm.project_layer.weight", + "stack.layers.0.feed_forward.w1.weight", + "stack.layers.0.feed_forward.w2.weight", + "stack.layers.0.feed_forward.w3.weight", + "stack.layers.0.ffn_norm.norm.weight", + "stack.layers.0.ffn_norm.project_layer.bias", + "stack.layers.0.ffn_norm.project_layer.weight", + "stack.layers.0.skip_in_linear.bias", + "stack.layers.0.skip_in_linear.weight", + "stack.layers.1.attention.wo.weight", + "stack.layers.1.attention.wqkv.weight", + "stack.layers.1.attention_norm.norm.weight", + "stack.layers.1.attention_norm.project_layer.bias", + "stack.layers.1.attention_norm.project_layer.weight", + "stack.layers.1.feed_forward.w1.weight", + "stack.layers.1.feed_forward.w2.weight", + "stack.layers.1.feed_forward.w3.weight", + "stack.layers.1.ffn_norm.norm.weight", + "stack.layers.1.ffn_norm.project_layer.bias", + "stack.layers.1.ffn_norm.project_layer.weight", + "stack.layers.1.skip_in_linear.bias", + "stack.layers.1.skip_in_linear.weight", + "stack.layers.2.attention.wo.weight", + "stack.layers.2.attention.wqkv.weight", + "stack.layers.2.attention_norm.norm.weight", + "stack.layers.2.attention_norm.project_layer.bias", + "stack.layers.2.attention_norm.project_layer.weight", + "stack.layers.2.feed_forward.w1.weight", + "stack.layers.2.feed_forward.w2.weight", + "stack.layers.2.feed_forward.w3.weight", + "stack.layers.2.ffn_norm.norm.weight", + "stack.layers.2.ffn_norm.project_layer.bias", + "stack.layers.2.ffn_norm.project_layer.weight", + "stack.layers.2.skip_in_linear.bias", + "stack.layers.2.skip_in_linear.weight", + "stack.layers.3.attention.wo.weight", + "stack.layers.3.attention.wqkv.weight", + "stack.layers.3.attention_norm.norm.weight", + "stack.layers.3.attention_norm.project_layer.bias", + "stack.layers.3.attention_norm.project_layer.weight", + "stack.layers.3.feed_forward.w1.weight", + "stack.layers.3.feed_forward.w2.weight", + "stack.layers.3.feed_forward.w3.weight", + "stack.layers.3.ffn_norm.norm.weight", + "stack.layers.3.ffn_norm.project_layer.bias", + "stack.layers.3.ffn_norm.project_layer.weight", + "stack.layers.3.skip_in_linear.bias", + "stack.layers.3.skip_in_linear.weight", + "stack.layers.4.attention.wo.weight", + "stack.layers.4.attention.wqkv.weight", + "stack.layers.4.attention_norm.norm.weight", + "stack.layers.4.attention_norm.project_layer.bias", + "stack.layers.4.attention_norm.project_layer.weight", + "stack.layers.4.feed_forward.w1.weight", + "stack.layers.4.feed_forward.w2.weight", + "stack.layers.4.feed_forward.w3.weight", + "stack.layers.4.ffn_norm.norm.weight", + "stack.layers.4.ffn_norm.project_layer.bias", + "stack.layers.4.ffn_norm.project_layer.weight", + "stack.layers.4.skip_in_linear.bias", + "stack.layers.4.skip_in_linear.weight", + "stack.norm.norm.weight", + "stack.norm.project_layer.bias", + "stack.norm.project_layer.weight", +}; + +// 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[] = { + 1.000000000e+00F, 0.000000000e+00F, 1.000000000e+00F, 0.000000000e+00F, 5.403022766e-01F, 8.414709568e-01F, + 9.999499917e-01F, 9.999833070e-03F, -4.161468446e-01F, 9.092974067e-01F, 9.998000264e-01F, 1.999866590e-02F, + -9.899924994e-01F, 1.411200017e-01F, 9.995500445e-01F, 2.999549918e-02F, -6.536436081e-01F, -7.568024993e-01F, + 9.992001057e-01F, 3.998933360e-02F, 2.836622000e-01F, -9.589242935e-01F, 9.987502694e-01F, 4.997916520e-02F, +}; + +// Transformer.forward output -- [kFrames, kDim]. +inline constexpr float kOut[] = { + 9.827834964e-01F, 6.055158377e-01F, 1.263252974e+00F, -4.305707514e-01F, 5.831326842e-01F, -3.171992004e-01F, + -2.719482481e-01F, -1.203811839e-01F, 9.820024371e-01F, 6.012700200e-01F, 1.286735773e+00F, -3.967078030e-01F, + 5.839135647e-01F, -5.548285246e-01F, -2.741804123e-01F, -1.203172132e-01F, 9.755259752e-01F, 6.258884668e-01F, + 1.067016482e+00F, -4.042342901e-01F, 5.833311677e-01F, -5.459736586e-01F, -2.741442621e-01F, -1.203659624e-01F, + 9.724086523e-01F, 6.375790834e-01F, 1.222781301e+00F, -4.147224724e-01F, 5.838839412e-01F, -6.536991000e-01F, + -2.745609581e-01F, -1.203083768e-01F, 9.797583222e-01F, 6.385335326e-01F, 1.190497279e+00F, -3.871082962e-01F, + 5.839298368e-01F, -5.570108891e-01F, -2.747859061e-01F, -1.203286573e-01F, 9.674922824e-01F, 6.381499171e-01F, + 1.134287000e+00F, -3.850421607e-01F, 5.835525990e-01F, -5.281082392e-01F, -2.742534876e-01F, -1.203894019e-01F, +}; + +} // namespace dit_stack_goldens diff --git a/tests/vllm/models/test_dit_stack.cpp b/tests/vllm/models/test_dit_stack.cpp new file mode 100644 index 000000000..147f4fe3f --- /dev/null +++ b/tests/vllm/models/test_dit_stack.cpp @@ -0,0 +1,118 @@ +// The S2Mel DiT transformer stack against upstream goldens. See dit_stack.h. +#include +#include +#include + +#include "dit_stack_goldens.inc" +#include "doctest/doctest.h" +#include "vllm/model_executor/models/dit_stack.h" + +namespace { + +std::vector Rnd(const std::string& name, size_t n, double scale) { + uint64_t h = 0xCBF29CE484222325ULL; + for (unsigned char ch : name) { + h = (h ^ ch) * 0x100000001B3ULL; + } + std::vector out(n); + for (size_t i = 0; i < n; ++i) { + h += 0x9E3779B97F4A7C15ULL; + uint64_t z = h; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + z ^= z >> 31; + const double u = static_cast(z >> 11) * (1.0 / 9007199254740992.0); + out[i] = static_cast((u * 2.0 - 1.0) * scale); + } + return out; +} + +using namespace dit_stack_goldens; + +vllm::models::dit_stack::Config Cfg() { + vllm::models::dit_stack::Config c; + c.dim = kDim; + c.heads = kHeads; + c.head_dim = kHeadDim; + c.intermediate = kIntermediate; + c.frames = kFrames; + c.eps = kEps; + return c; +} + +vllm::models::dit_stack::Weights W() { + vllm::models::dit_stack::Weights w; + const size_t D = static_cast(kDim); + const size_t I = static_cast(kIntermediate); + const size_t QKV = static_cast((kHeads + 2 * kHeads) * kHeadDim); + for (int64_t i = 0; i < kDepth; ++i) { + const std::string p = "stack.layers." + std::to_string(i) + "."; + vllm::models::dit_stack::LayerWeights l; + l.block.wqkv = Rnd(p + "attention.wqkv.weight", QKV * D, 0.5); + l.block.wo = Rnd(p + "attention.wo.weight", D * D, 0.5); + l.block.w1 = Rnd(p + "feed_forward.w1.weight", I * D, 0.5); + l.block.w3 = Rnd(p + "feed_forward.w3.weight", I * D, 0.5); + l.block.w2 = Rnd(p + "feed_forward.w2.weight", D * I, 0.5); + l.block.attn_proj_w = Rnd(p + "attention_norm.project_layer.weight", 2 * D * D, 0.5); + l.block.attn_proj_b = Rnd(p + "attention_norm.project_layer.bias", 2 * D, 0.5); + l.block.attn_norm_w = Rnd(p + "attention_norm.norm.weight", D, 0.5); + l.block.ffn_proj_w = Rnd(p + "ffn_norm.project_layer.weight", 2 * D * D, 0.5); + l.block.ffn_proj_b = Rnd(p + "ffn_norm.project_layer.bias", 2 * D, 0.5); + l.block.ffn_norm_w = Rnd(p + "ffn_norm.norm.weight", D, 0.5); + l.skip_in_w = Rnd(p + "skip_in_linear.weight", D * 2 * D, 0.5); + l.skip_in_b = Rnd(p + "skip_in_linear.bias", D, 0.5); + w.layers.push_back(std::move(l)); + } + w.norm_proj_w = Rnd("stack.norm.project_layer.weight", 2 * D * D, 0.5); + w.norm_proj_b = Rnd("stack.norm.project_layer.bias", 2 * D, 0.5); + w.norm_w = Rnd("stack.norm.norm.weight", D, 0.5); + return w; +} + +} // namespace + +TEST_CASE("the whole DiT stack matches upstream Transformer.forward") { + const std::vector x = Rnd("stack.x", static_cast(kFrames * kDim), 1.0); + const std::vector c = Rnd("stack.c", static_cast(kDim), 1.0); + const std::vector freqs(kFreqs, + kFreqs + static_cast(kFrames * kHeadDim / 2 * 2)); + + const std::vector got = vllm::models::dit_stack::Forward(Cfg(), W(), x, c, freqs); + REQUIRE(got.size() == static_cast(kFrames * kDim)); + for (size_t i = 0; i < got.size(); ++i) { + CHECK(got[i] == doctest::Approx(kOut[i]).epsilon(3e-5)); + } +} + +TEST_CASE("the skip merge is actually consulted") { + // Perturbing ONLY a receiving layer's skip_in_linear must move the output. If + // the routing silently skipped the merge, the stack would still be a valid + // transformer and this is what separates the two. + const std::vector x = Rnd("stack.x", static_cast(kFrames * kDim), 1.0); + const std::vector c = Rnd("stack.c", static_cast(kDim), 1.0); + const std::vector freqs(kFreqs, + kFreqs + static_cast(kFrames * kHeadDim / 2 * 2)); + const auto base = vllm::models::dit_stack::Forward(Cfg(), W(), x, c, freqs); + + // Depth 5 routes 3 <- 1 and 4 <- 0, so layer 3 is a receiver. + auto w = W(); + w.layers[3].skip_in_w[0] += 1.0F; + const auto moved = vllm::models::dit_stack::Forward(Cfg(), w, x, c, freqs); + CHECK(moved != base); + + // Layer 0 never receives, so ITS skip_in_linear must not matter at all. + auto w0 = W(); + w0.layers[0].skip_in_w[0] += 1.0F; + CHECK(vllm::models::dit_stack::Forward(Cfg(), w0, x, c, freqs) == base); +} + +TEST_CASE("the final norm is applied") { + const std::vector x = Rnd("stack.x", static_cast(kFrames * kDim), 1.0); + const std::vector c = Rnd("stack.c", static_cast(kDim), 1.0); + const std::vector freqs(kFreqs, + kFreqs + static_cast(kFrames * kHeadDim / 2 * 2)); + const auto base = vllm::models::dit_stack::Forward(Cfg(), W(), x, c, freqs); + auto w = W(); + w.norm_w[0] += 1.0F; + CHECK(vllm::models::dit_stack::Forward(Cfg(), w, x, c, freqs) != base); +}