{g.items.map((item) => renderButton(item, "hw", sel.hw))}
{Array.from({ length: maxHwCols - g.items.length }).map((_, i) => (
@@ -1431,13 +1458,13 @@ export const Deployment = ({ config, benchmarks }) => {
key={mode}
style={{
...(index === runModes.length - 1
- ? s.runModeChipLast(runMode === mode)
- : s.runModeChip(runMode === mode)),
+ ? s.runModeChipLast(activeRunMode === mode)
+ : s.runModeChip(activeRunMode === mode)),
...(runModes.length === 1 ? { borderRadius: 7 } : {}),
}}
onClick={() => setRunMode(mode)}
role="tab"
- aria-selected={runMode === mode}
+ aria-selected={activeRunMode === mode}
>
{mode === "docker" ? "Docker" : "Python"}
@@ -1473,38 +1500,40 @@ export const Deployment = ({ config, benchmarks }) => {
{/* Playground link — scrollIntoView, not an href, so the hash (which
carries the selection) isn't overwritten. */}
-
-
Need to go beyond the verified matrix?
-
{
- const el = document.getElementById("playground");
- if (el) el.scrollIntoView({ behavior: "smooth", block: "start" });
- }}
+ {config.showPlaygroundLink !== false && (
+
- Open the Playground →
-
-
+ Need to go beyond the verified matrix?
+ {
+ const el = document.getElementById("playground");
+ if (el) el.scrollIntoView({ behavior: "smooth", block: "start" });
+ }}
+ style={{
+ background: "transparent",
+ border: "none",
+ padding: 0,
+ color: isDark ? "#FDBA74" : "#C2410C",
+ cursor: "pointer",
+ fontSize: "12px",
+ fontWeight: 600,
+ textDecoration: "underline",
+ textUnderlineOffset: "2px",
+ }}
+ >
+ Open the Playground →
+
+
+ )}
{/* cURL modal */}
{modal === "curl" && (
diff --git a/docs_new/src/snippets/_playground.jsx b/docs_new/src/snippets/_playground.jsx
index 2b994998a..9cd398e9e 100644
--- a/docs_new/src/snippets/_playground.jsx
+++ b/docs_new/src/snippets/_playground.jsx
@@ -1390,6 +1390,9 @@ export const Playground = ({ config }) => {
const di = config.dockerImages || {};
const image = di[`${sel.hw}|${sel.quant}|${sel.strategy}`]
|| di[`${sel.hw}|${sel.quant}`] || di[sel.hw] || "lmsysorg/sglang:dev";
+ const dockerRunCommand = typeof config.dockerRunCommand === "function"
+ ? config.dockerRunCommand(sel)
+ : (config.dockerRunCommand || "sglang serve");
const portFlag = f.find((x) => x.split(/[\s=]/)[0] === "--port");
const servePort = portFlag ? portFlag.slice("--port".length).trim() : "{{PORT}}";
// Mirrors `multiNodeDockerFlags` on the _deployment.jsx HARDWARE_CATALOG
@@ -1406,11 +1409,12 @@ export const Playground = ({ config }) => {
(multinode || pdMode) ? " --network host" : ` -p ${servePort}:${servePort}`,
...(multinode ? fabricFlags.map((x) => " " + x) : []),
" -v ~/.cache/huggingface:/root/.cache/huggingface",
+ ...(config.dockerMounts || []).map((mount) => ` -v ${mount}`),
` --env "HF_TOKEN={{HF_TOKEN}}"`,
...cellEnv.map((e) => ` --env ${e}`),
" --ipc=host",
` ${image}`,
- " sglang serve",
+ ` ${dockerRunCommand}`,
...f.map((x) => " " + x),
];
cmd = dockerLines.join(" \\\n");
diff --git a/docs_new/src/snippets/configs/MiniMaxAI/minimax-h3.jsx b/docs_new/src/snippets/configs/MiniMaxAI/minimax-h3.jsx
new file mode 100644
index 000000000..fc2579ecd
--- /dev/null
+++ b/docs_new/src/snippets/configs/MiniMaxAI/minimax-h3.jsx
@@ -0,0 +1,618 @@
+// MiniMax-H3 diffusion deployment matrix. Consumed by _deployment.jsx.
+//
+// The mode, quantization, and encoder choices are deployment overlays because
+// they do not change which base hardware topology fits. Request sampling
+// controls remain in the generated cURL instead of being mixed into this
+// deployment matrix.
+// Hardware/profile cells remain deliberately small and carry an honest
+// verification state for the exact platform, rather than inheriting a result
+// measured on a different GPU.
+
+
+export const config = {
+ modelName: "MiniMax-H3",
+
+ supportedHardware: [
+ "b200",
+ "b300",
+ "h200",
+ "h100",
+ "mi300x",
+ "mi355x",
+ "rtx5090",
+ ],
+ hardware: [
+ { id: "rtx5090", label: "RTX 5090", vram: "32GB", vendor: "consumer" },
+ ],
+ groupHardware: false,
+
+ matchDims: [
+ {
+ id: "profile",
+ title: "Deployment Profile",
+ showWhen: (s) => ["b200", "b300", "h200", "h100"].includes(s.hw),
+ options: [
+ { id: "resident", label: "Resident" },
+ {
+ id: "fsdp",
+ label: "FSDP sharded",
+ showWhen: (s) =>
+ ["b200", "b300", "h200", "h100"].includes(s.hw),
+ },
+ {
+ id: "offload",
+ label: "Layerwise offload",
+ showWhen: (s) => s.hw === "rtx5090",
+ },
+ ],
+ },
+ ],
+
+ overlayDims: [
+ {
+ id: "weights",
+ title: "Checkpoint Weights",
+ default: "fl2va",
+ options: [
+ {
+ id: "fl2va",
+ label: "FL2VA (First-and-Last-Frame-to-Video-and-Audio)",
+ flags: ["--model-variant fl2va"],
+ },
+ {
+ id: "ref2va",
+ label: "Ref2VA (Reference-to-Video-and-Audio)",
+ flags: ["--model-variant ref2va"],
+ },
+ ],
+ },
+ {
+ id: "mode",
+ title: "Request Mode",
+ default: "t2va",
+ options: [
+ {
+ id: "t2va",
+ label: "Text only",
+ showWhen: (s) => s.weights === "fl2va",
+ },
+ {
+ id: "i2va",
+ label: "First frame",
+ showWhen: (s) => s.weights === "fl2va",
+ },
+ {
+ id: "l2va",
+ label: "Last frame",
+ showWhen: (s) => s.weights === "fl2va",
+ },
+ {
+ id: "fl2va",
+ label: "First + last frames",
+ showWhen: (s) => s.weights === "fl2va",
+ },
+ {
+ id: "ref_image",
+ label: "Image reference",
+ showWhen: (s) => s.weights === "ref2va",
+ },
+ {
+ id: "ref_image_audio",
+ label: "Image + audio",
+ showWhen: (s) => s.weights === "ref2va",
+ },
+ {
+ id: "v2v",
+ label: "Video reference",
+ showWhen: (s) => s.weights === "ref2va",
+ },
+ {
+ id: "video_audio",
+ label: "Video + soundtrack",
+ showWhen: (s) => s.weights === "ref2va",
+ },
+ {
+ id: "audio_only",
+ label: "Audio reference",
+ showWhen: (s) => s.weights === "ref2va",
+ },
+ {
+ id: "mixed_ref",
+ label: "Mixed references",
+ showWhen: (s) => s.weights === "ref2va",
+ },
+ ],
+ },
+ {
+ id: "quant",
+ title: "Online Quantization",
+ default: "bf16",
+ showWhen: (s) => ["b200", "b300"].includes(s.hw),
+ options: [
+ { id: "bf16", label: "Off — Native BF16/FP32" },
+ {
+ id: "fp8",
+ label: "FP8 — Approximate",
+ showWhen: (s) => ["b200", "b300"].includes(s.hw),
+ disabled: (s) => s.profile !== "resident",
+ disableReason:
+ "The documented FP8 operating point keeps the transformer resident; FSDP combinations have not been validated.",
+ flags: ["--quantization fp8"],
+ hints: [
+ "Online FP8 is approximate. Validate both video and audio quality;",
+ "verified B200 and B300 runs reduced memory; re-benchmark latency on the target workload.",
+ ],
+ },
+ ],
+ },
+ {
+ id: "encoder",
+ title: "Text Encoder Parallel",
+ default: "auto",
+ options: [
+ {
+ id: "auto",
+ label: "Auto (recommended)",
+ hints: [
+ "Auto uses folding for the single-request recipes below and can",
+ "select data parallel encoding for a compatible TP1 request batch.",
+ ],
+ },
+ {
+ id: "fold",
+ label: "Fold (single-request)",
+ flags: ["--encoder-parallel fold"],
+ hints: [
+ "Fold shards the resident Qwen3-VL encoder across the replica and is",
+ "best suited to single-node GPUs with fast peer-to-peer links.",
+ ],
+ },
+ {
+ id: "dp",
+ label: "DP (batched throughput)",
+ disabled: (s) =>
+ s.hw === "rtx5090" ||
+ (s.hw === "h100" && s.profile === "resident"),
+ disableReason:
+ "Encoder DP requires TP1 and DiT DP1; this verified recipe uses TP2.",
+ flags: [
+ "--encoder-parallel dp",
+ "--batching-max-size {{BATCHING_MAX_SIZE}}",
+ ],
+ hints: [
+ "DP distributes a compatible multi-request text batch across ranks;",
+ "it does not improve a batch of one and replicates encoder weights.",
+ ],
+ },
+ {
+ id: "replicate",
+ label: "Replicate (compatibility)",
+ flags: ["--encoder-parallel replicate"],
+ },
+ ],
+ },
+ ],
+
+ modelNames: {
+ default: "MiniMaxAI/MiniMax-H3",
+ },
+
+ placeholders: {
+ HOST_IP: {
+ target: "command",
+ label: "Bind host",
+ default: "0.0.0.0",
+ },
+ PORT: {
+ target: "command",
+ label: "Bind port",
+ default: "30010",
+ },
+ HF_TOKEN: {
+ target: "command",
+ label: "HF token (Docker)",
+ default: "
",
+ },
+ MEDIA_DIR: {
+ target: "command",
+ label: "Host media directory (Docker)",
+ default: "/data/minimax-h3",
+ },
+ CURL_HOST: {
+ target: "curl",
+ label: "Server host",
+ default: "localhost",
+ },
+ CURL_PORT: {
+ target: "curl",
+ label: "Server port",
+ default: "30010",
+ },
+ NUM_OUTPUTS: {
+ target: "curl",
+ label: "Outputs per prompt (1-10)",
+ default: "1",
+ },
+ BATCHING_MAX_SIZE: {
+ target: "command",
+ label: "Maximum request batch size",
+ default: "2",
+ },
+ DURATION_SECONDS: {
+ target: "curl",
+ label: "Duration (seconds, 4-15)",
+ default: "5",
+ },
+ FIRST_FRAME: {
+ target: "curl",
+ label: "FL2VA first frame URI",
+ default: "file:///data/minimax-h3/first-frame.png",
+ },
+ LAST_FRAME: {
+ target: "curl",
+ label: "FL2VA last frame URI",
+ default: "file:///data/minimax-h3/last-frame.png",
+ },
+ INPUT_VIDEO: {
+ target: "curl",
+ label: "First video URI",
+ default: "file:///data/minimax-h3/video-1.mp4",
+ },
+ INPUT_VIDEO_START_SECONDS: {
+ target: "curl",
+ label: "First video start (seconds)",
+ default: "0",
+ },
+ SECOND_INPUT_VIDEO: {
+ target: "curl",
+ label: "Second video URI (mixed ref)",
+ default: "file:///data/minimax-h3/video-2.mp4",
+ },
+ SECOND_INPUT_VIDEO_START_SECONDS: {
+ target: "curl",
+ label: "Second video start (seconds)",
+ default: "0",
+ },
+ REFERENCE_IMAGE: {
+ target: "curl",
+ label: "First reference image URI",
+ default: "file:///data/minimax-h3/reference-1.png",
+ },
+ SECOND_REFERENCE_IMAGE: {
+ target: "curl",
+ label: "Second reference image URI",
+ default: "file:///data/minimax-h3/reference-2.png",
+ },
+ REFERENCE_AUDIO: {
+ target: "curl",
+ label: "First reference audio URI",
+ default: "file:///data/minimax-h3/reference-1.mp3",
+ },
+ SECOND_REFERENCE_AUDIO: {
+ target: "curl",
+ label: "Second reference audio URI",
+ default: "file:///data/minimax-h3/reference-2.mp3",
+ },
+ },
+
+ curl: (s) => {
+ const request = {
+ model: "{{MODEL_NAME}}",
+ prompt:
+ "Night-vision bedroom footage: while the owner sleeps, three cats burst in playing tiny brass instruments at full volume, freeze, then march out as if nothing happened.",
+ seconds: "{{DURATION_SECONDS}}",
+ task: "t2va",
+ conditions: [],
+ target: {
+ short_edge: 768,
+ aspect_ratio: "16:9",
+ duration_seconds: "{{DURATION_SECONDS}}",
+ },
+ num_outputs_per_prompt: "{{NUM_OUTPUTS}}",
+ num_inference_steps: 50,
+ flow_shift: 12.0,
+ audio_flow_shift: 3.0,
+ seed: 1101,
+ };
+ const imageReference = (uri) => ({
+ type: "image",
+ uri,
+ role: "reference",
+ });
+ const audioReference = (uri) => ({
+ type: "audio",
+ uri,
+ role: "reference",
+ });
+ const videoReference = (uri, start, type = "video") => ({
+ type,
+ uri,
+ role: "reference",
+ start_time_seconds: start,
+ });
+
+ if (["i2va", "l2va", "fl2va"].includes(s.mode)) {
+ request.task = "fl2va";
+ request.prompt =
+ "Continue naturally between the supplied endpoint frame or frames, with synchronized ambient sound.";
+ request.target.aspect_ratio = "auto";
+ request.seed = 2101;
+ request.conditions = [];
+ if (s.mode !== "l2va") {
+ request.conditions.push({
+ type: "image",
+ uri: "{{FIRST_FRAME}}",
+ role: "keyframe",
+ frame_index: 0,
+ });
+ }
+ if (s.mode !== "i2va") {
+ request.conditions.push({
+ type: "image",
+ uri: "{{LAST_FRAME}}",
+ role: "keyframe",
+ frame_index: -1,
+ });
+ }
+ } else if (s.mode === "ref_image") {
+ request.task = "ref2va";
+ request.prompt = "Use as the visual subject and style reference.";
+ request.target.aspect_ratio = "auto";
+ request.conditions = [imageReference("{{REFERENCE_IMAGE}}")];
+ request.seed = 3101;
+ } else if (s.mode === "ref_image_audio") {
+ request.task = "ref2va";
+ request.prompt =
+ "Use as the visual subject and as the sound reference.";
+ request.target.aspect_ratio = "auto";
+ request.conditions = [
+ imageReference("{{REFERENCE_IMAGE}}"),
+ audioReference("{{REFERENCE_AUDIO}}"),
+ ];
+ request.seed = 3102;
+ } else if (s.mode === "v2v" || s.mode === "video_audio") {
+ request.task = "ref2va";
+ request.prompt =
+ s.mode === "video_audio"
+ ? "Follow and its required soundtrack with coherent synchronized motion."
+ : "Follow the appearance and motion of ; use its soundtrack when present.";
+ request.conditions = [
+ videoReference(
+ "{{INPUT_VIDEO}}",
+ "{{INPUT_VIDEO_START_SECONDS}}",
+ s.mode === "video_audio" ? "video_audio" : "video",
+ ),
+ ];
+ request.seed = s.mode === "video_audio" ? 4102 : 4101;
+ } else if (s.mode === "audio_only") {
+ request.task = "ref2va";
+ request.prompt = "Build a coherent visual scene around .";
+ request.conditions = [audioReference("{{REFERENCE_AUDIO}}")];
+ request.seed = 3103;
+ } else if (s.mode === "mixed_ref") {
+ request.task = "ref2va";
+ request.prompt =
+ "Combine , , , , , and in their one-based modality order.";
+ request.conditions = [
+ imageReference("{{REFERENCE_IMAGE}}"),
+ imageReference("{{SECOND_REFERENCE_IMAGE}}"),
+ audioReference("{{REFERENCE_AUDIO}}"),
+ audioReference("{{SECOND_REFERENCE_AUDIO}}"),
+ videoReference("{{INPUT_VIDEO}}", "{{INPUT_VIDEO_START_SECONDS}}"),
+ videoReference(
+ "{{SECOND_INPUT_VIDEO}}",
+ "{{SECOND_INPUT_VIDEO_START_SECONDS}}",
+ ),
+ ];
+ request.seed = 3104;
+ }
+
+ const body = JSON.stringify(request, null, 2).replace(
+ /"{{(NUM_OUTPUTS|DURATION_SECONDS|INPUT_VIDEO_START_SECONDS|SECOND_INPUT_VIDEO_START_SECONDS)}}"/g,
+ "{{$1}}",
+ );
+ return `curl -sS -X POST http://{{CURL_HOST}}:{{CURL_PORT}}/v1/videos \\
+ -H 'Content-Type: application/json' \\
+ -d '${body}'`;
+ },
+
+ dockerMounts: ["{{MEDIA_DIR}}:/data/minimax-h3:ro"],
+
+ dockerRunCommand: (s) =>
+ ["mi300x", "mi355x"].includes(s.hw)
+ ? `bash -lc 'python -m pip install -e "/sgl-workspace/sglang/python[diffusion_hip]" && exec sglang serve "$@"' --`
+ : `bash -lc 'python -m pip install -e "/sgl-workspace/sglang/python[diffusion]" && exec sglang serve "$@"' --`,
+
+ // Publish AMD Docker only after an H3-capable ROCm image has been validated.
+ runModes: (s) =>
+ ["mi300x", "mi355x"].includes(s.hw)
+ ? ["python"]
+ : ["python", "docker"],
+
+ dockerImages: {
+ b200: "lmsysorg/sglang:dev",
+ b300: "lmsysorg/sglang:dev",
+ h200: "lmsysorg/sglang:dev",
+ h100: "lmsysorg/sglang:dev",
+ },
+
+ showPlaygroundLink: false,
+
+ cells: [
+ {
+ match: { hw: "b200", profile: "resident" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 8",
+ "--ulysses-degree 8",
+ "--performance-mode speed",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ },
+ {
+ match: { hw: "b300", profile: "resident" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 8",
+ "--ulysses-degree 8",
+ "--performance-mode speed",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "This is the B300 topology used for the documented benchmark sweep, not a claimed minimum GPU count.",
+ },
+ {
+ match: { hw: "h200", profile: "resident" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 4",
+ "--ulysses-degree 4",
+ "--performance-mode speed",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ },
+ {
+ match: { hw: "b300", profile: "fsdp" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 8",
+ "--ulysses-degree 8",
+ "--performance-mode speed",
+ "--use-fsdp-inference true",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "FSDP reduces resident DiT memory but adds per-block parameter collectives. Prefer Resident when the full pipeline fits.",
+ },
+ {
+ match: { hw: "h200", profile: "fsdp" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 4",
+ "--ulysses-degree 4",
+ "--performance-mode speed",
+ "--use-fsdp-inference true",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "FSDP reduces resident DiT memory but adds per-block parameter collectives. Prefer Resident when the full pipeline fits.",
+ },
+ {
+ match: { hw: "b200", profile: "fsdp" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 4",
+ "--ulysses-degree 4",
+ "--performance-mode speed",
+ "--use-fsdp-inference true",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "The 4-GPU FSDP path is lossless but slower than the 8-GPU resident recipe.",
+ },
+ {
+ match: { hw: "h100", profile: "resident" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 4",
+ "--tp-size 2",
+ "--ulysses-degree 2",
+ "--performance-mode speed",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "Fastest measured 4× H100 80 GB topology. TP4 + Ulysses1 lowers peak memory at a small latency cost.",
+ },
+ {
+ match: { hw: "h100", profile: "fsdp" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 4",
+ "--ulysses-degree 4",
+ "--performance-mode speed",
+ "--use-fsdp-inference true",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "Capacity path on 4× H100 80 GB. Prefer the resident TP2 + Ulysses2 profile for latency.",
+ },
+ {
+ match: { hw: "mi300x", profile: "resident" },
+ nnodes: 1,
+ verified: true,
+ env: ["SGLANG_USE_AITER=1"],
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 8",
+ "--ulysses-degree 8",
+ "--performance-mode speed",
+ "--attention-backend aiter",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "Validated on 1×, 2×, 4×, and 8× MI300X with BF16 and AITER packed attention. The picker emits the fastest measured 8-GPU topology; set --num-gpus and --ulysses-degree to the same lower count for a measured capacity recipe.",
+ },
+ {
+ match: { hw: "mi355x", profile: "resident" },
+ nnodes: 1,
+ verified: true,
+ env: ["SGLANG_USE_AITER=1"],
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 8",
+ "--ulysses-degree 8",
+ "--performance-mode speed",
+ "--attention-backend aiter",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "Validated on 1×, 2×, 4×, and 8× MI355X with BF16 and AITER packed attention. The picker emits the fastest measured 8-GPU topology; set --num-gpus and --ulysses-degree to the same lower count for a measured capacity recipe.",
+ },
+ {
+ match: { hw: "rtx5090", profile: "offload" },
+ nnodes: 1,
+ verified: true,
+ flags: [
+ "--model-path {{MODEL_NAME}}",
+ "--num-gpus 2",
+ "--tp-size 2",
+ "--ulysses-degree 1",
+ "--performance-mode memory",
+ "--layerwise-offload-components dit,text_encoder,vae",
+ "--dit-offload-prefetch-size 1",
+ "--dit-layerwise-resident-layers 20",
+ "--enable-torch-compile false",
+ "--host {{HOST_IP}}",
+ "--port {{PORT}}",
+ ],
+ warn:
+ "Validated lossless BF16/FP32 recipe on 2× RTX 5090 (32 GB each) with a 384 GiB-class host. TP2 avoids the full per-rank DiT replication observed with Ulysses2 on PCIe.",
+ },
+ ],
+};
diff --git a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
index ab4452945..9a107f301 100644
--- a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
+++ b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
@@ -6,6 +6,8 @@
#include
#include
+#include
+
#include
#include
@@ -42,7 +44,8 @@ constexpr uint32_t active_mask() {
}
}
-SGL_DEVICE float load_cache_value(const float* ptr, int64_t idx) {
+template
+SGL_DEVICE CacheDType load_cache_value(const CacheDType* ptr, int64_t idx) {
#ifdef USE_ROCM
return ptr[idx];
#else
@@ -50,7 +53,15 @@ SGL_DEVICE float load_cache_value(const float* ptr, int64_t idx) {
#endif
}
-template
+template <
+ int64_t kHeadDim,
+ int64_t kRopeDim,
+ bool kIsNeox,
+ bool kUsePDL,
+ typename DType,
+ typename CacheDType,
+ bool kRoundNormBeforeRope,
+ typename IdType>
__global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ params) {
using namespace device;
@@ -63,14 +74,17 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread;
constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2;
constexpr uint32_t kActiveMask = active_mask();
- constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
+ constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(CacheDType);
static_assert(kElemsPerThread % 2 == 0, "Each lane must own an even number of elements");
static_assert(kRopeDim > 0 && kRopeDim <= kHeadDim, "Invalid rope dimension");
static_assert(kRopeDim % kElemsPerThread == 0, "rope_dim must align with per-lane vector width");
static_assert(
- !kIsNeox || (kRotaryLanes >= 2 && ((kRotaryLanes & (kRotaryLanes - 1)) == 0)),
- "NeoX fused qknorm+rope requires rotary lane count to be a power of 2");
+ !kIsNeox || (kRotaryLanes >= 2 && kRotaryLanes % 2 == 0),
+ "NeoX fused qknorm+rope requires an even rotary lane count");
+ static_assert(
+ !kRoundNormBeforeRope || std::is_same_v,
+ "Rounded QKNorm+RoPE requires cache and activation dtypes to match");
using Packed = packed_t;
using Storage = AlignedVector;
@@ -98,6 +112,53 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
auto input_vec = load_as(input, lane_id);
const auto weight_vec = load_as(weight_ptr, lane_id);
+ if constexpr (kRoundNormBeforeRope) {
+ auto output_vec = norm::apply_norm_warp(input_vec, weight_vec, eps);
+ const auto pos = static_cast(static_cast(positions)[token_id]);
+ const auto cos_ptr = static_cast(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
+ const auto sin_ptr = cos_ptr + kRopeDim / 2;
+
+ if constexpr (kIsNeox) {
+ if (lane_id < kRotaryLanes) {
+ const auto partner_lane =
+ lane_id < kHalfRotaryLanes ? lane_id + kHalfRotaryLanes : lane_id - kHalfRotaryLanes;
+#pragma unroll
+ for (uint32_t j = 0; j < kVecSize; ++j) {
+ auto partner_vec = output_vec[j];
+ auto partner_bits = reinterpret_cast(partner_vec);
+ partner_bits = __shfl_sync(kActiveMask, partner_bits, partner_lane);
+ reinterpret_cast(partner_vec) = partner_bits;
+ auto& values = unpack(output_vec[j]);
+ const auto& partner_values = unpack(partner_vec);
+#pragma unroll
+ for (uint32_t i = 0; i < 2; ++i) {
+ const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i;
+ const auto cos = load_cache_value(cos_ptr, half_idx);
+ const auto sin = load_cache_value(sin_ptr, half_idx);
+ values[i] = lane_id < kHalfRotaryLanes ? values[i] * cos - partner_values[i] * sin
+ : values[i] * cos + partner_values[i] * sin;
+ }
+ }
+ }
+ } else {
+ if (lane_id < kRotaryLanes) {
+#pragma unroll
+ for (uint32_t j = 0; j < kVecSize; ++j) {
+ auto& values = unpack(output_vec[j]);
+ const auto half_idx = lane_id * kElemsPerThread / 2 + j;
+ const auto cos = load_cache_value(cos_ptr, half_idx);
+ const auto sin = load_cache_value(sin_ptr, half_idx);
+ const auto x = values[0];
+ const auto y = values[1];
+ values[0] = x * cos - y * sin;
+ values[1] = y * cos + x * sin;
+ }
+ }
+ }
+ store_as(const_cast(input), output_vec, lane_id);
+ continue;
+ }
+
float elems[kElemsPerThread];
float sum_of_squares = 0.0f;
@@ -122,27 +183,28 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
if constexpr (kIsNeox) {
if (lane_id < kRotaryLanes) {
const auto pos = static_cast(static_cast(positions)[token_id]);
- const auto cos_ptr = static_cast(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
+ const auto cos_ptr =
+ static_cast(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
+ const auto partner_lane = lane_id < kHalfRotaryLanes ? lane_id + kHalfRotaryLanes : lane_id - kHalfRotaryLanes;
#pragma unroll
for (uint32_t i = 0; i < kElemsPerThread; ++i) {
- float swapped = __shfl_xor_sync(kActiveMask, elems[i], kHalfRotaryLanes);
+ float swapped = __shfl_sync(kActiveMask, elems[i], partner_lane);
if (lane_id < kHalfRotaryLanes) {
swapped = -swapped;
}
- int dim_idx = static_cast(lane_id * kElemsPerThread + i);
- dim_idx = (dim_idx * 2) % kRopeDim;
- const int half_idx = dim_idx / 2;
- const float cos = load_cache_value(cos_ptr, half_idx);
- const float sin = load_cache_value(sin_ptr, half_idx);
+ const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + i;
+ const float cos = cast(load_cache_value(cos_ptr, half_idx));
+ const float sin = cast(load_cache_value(sin_ptr, half_idx));
elems[i] = elems[i] * cos + swapped * sin;
}
}
} else {
if (lane_id < kRotaryLanes) {
const auto pos = static_cast(static_cast(positions)[token_id]);
- const auto cos_ptr = static_cast(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
+ const auto cos_ptr =
+ static_cast(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
const auto sin_ptr = cos_ptr + kRopeDim / 2;
#pragma unroll
@@ -150,8 +212,8 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
const float x = elems[i];
const float y = elems[i + 1];
const int half_idx = static_cast(lane_id * kElemsPerThread + i) / 2;
- const float cos = load_cache_value(cos_ptr, half_idx);
- const float sin = load_cache_value(sin_ptr, half_idx);
+ const float cos = cast(load_cache_value(cos_ptr, half_idx));
+ const float sin = cast(load_cache_value(sin_ptr, half_idx));
elems[i] = x * cos - y * sin;
elems[i + 1] = y * cos + x * sin;
}
@@ -168,11 +230,19 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
PDLTriggerSecondary();
}
-template
+template <
+ int64_t kHeadDim,
+ int64_t kRopeDim,
+ bool kIsNeox,
+ bool kUsePDL,
+ typename DType,
+ typename CacheDType,
+ bool kRoundNormBeforeRope>
struct QKNormRopeKernel {
static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
template
- static constexpr auto kernel = fused_qknorm_rope_warp;
+ static constexpr auto kernel =
+ fused_qknorm_rope_warp;
static void
run(const tvm::ffi::TensorView q,
@@ -201,7 +271,7 @@ struct QKNormRopeKernel {
TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype().with_device(device).verify(q);
TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype().with_device(device).verify(k);
TensorMatcher({D}).with_dtype().with_device(device).verify(q_weight).verify(k_weight);
- TensorMatcher({-1, R}).with_dtype().with_device(device).verify(cos_sin_cache);
+ TensorMatcher({-1, R}).with_dtype().with_device(device).verify(cos_sin_cache);
TensorMatcher({N}).with_dtype(id_type).with_device(device).verify(positions);
const auto num_tokens = static_cast(N.unwrap());
diff --git a/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh b/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh
new file mode 100644
index 000000000..9c7704328
--- /dev/null
+++ b/python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh
@@ -0,0 +1,182 @@
+// CUDA fast path for the Ulysses sequence-parallel output head merge.
+//
+// usp_merge_heads:
+// x [W, S, B, h_local, D] (contiguous, the output all-to-all result)
+// -> out [B, S, W, h_local, D] (contiguous)
+// Replaces `x.permute(2, 1, 0, 3, 4).contiguous()` on the head_dim=2
+// output path of `_usp_output_all_to_all`.
+//
+// A pure copy (no arithmetic), so it is bit-exact with the eager permute by
+// construction. It exists because ATen's generic permute-copy reaches well
+// under half of HBM bandwidth on the packed-DiT shapes, while a single pass
+// with coalesced vectorized stores runs near roofline.
+
+#pragma once
+
+#include // For host dtype helpers and TensorView metadata
+#include // For RuntimeCheck and div_ceil
+
+#include // For CUDA dtype aliases
+#include // For LaunchKernel
+#include // For device::AlignedVector
+
+#include
+
+namespace sglang_usp_relayout {
+
+namespace {
+
+constexpr int kBlockSize = 256;
+constexpr int64_t kMaxGrid = 65535;
+
+inline const char* data_ptr(const tvm::ffi::TensorView& t) {
+ return static_cast(t.data_ptr()) + t.byte_offset();
+}
+
+inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
+ return static_cast(t.data_ptr()) + t.byte_offset();
+}
+
+inline bool aligned16(const void* p) {
+ return (reinterpret_cast(p) & 0xF) == 0;
+}
+
+inline int64_t numel(const tvm::ffi::TensorView& t) {
+ int64_t n = 1;
+ for (int i = 0; i < t.ndim(); ++i) {
+ n *= t.size(i);
+ }
+ return n;
+}
+
+inline int64_t grid_for(int64_t total) {
+ int64_t grid = host::div_ceil(total, static_cast(kBlockSize));
+ if (grid < 1) {
+ grid = 1;
+ }
+ if (grid > kMaxGrid) {
+ grid = kMaxGrid;
+ }
+ return grid;
+}
+
+inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
+ int64_t expected = 1;
+ for (int i = t.ndim() - 1; i >= 0; --i) {
+ if (t.size(i) == 1) {
+ continue;
+ }
+ if (t.stride(i) != expected) {
+ return false;
+ }
+ expected *= t.size(i);
+ }
+ return true;
+}
+
+template
+inline void check_dtype(const tvm::ffi::TensorView& t) {
+ host::RuntimeCheck(host::is_type(t.dtype()), "unexpected dtype for usp_merge_heads tensor");
+}
+
+// out[b, s, w, h, c] = x[w, s, b, h, c]
+template
+__global__ void usp_merge_heads_vec_kernel(
+ T* __restrict__ out,
+ const T* __restrict__ x,
+ int64_t n_vec,
+ int64_t d_vec, // D / kVec
+ int64_t h_local,
+ int64_t batch,
+ int64_t seq,
+ int64_t world) {
+ const int64_t stride = static_cast(gridDim.x) * blockDim.x;
+ for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n_vec; i += stride) {
+ int64_t rest = i;
+ const int64_t c_vec = rest % d_vec;
+ rest /= d_vec;
+ const int64_t h = rest % h_local;
+ rest /= h_local;
+ const int64_t w = rest % world;
+ rest /= world;
+ const int64_t s = rest % seq;
+ const int64_t b = rest / seq;
+
+ const int64_t src_vec = ((((w * seq + s) * batch + b) * h_local) + h) * d_vec + c_vec;
+ device::AlignedVector val;
+ val.load(x, src_vec);
+ val.store(out, i);
+ }
+}
+
+template
+__global__ void usp_merge_heads_scalar_kernel(
+ T* __restrict__ out,
+ const T* __restrict__ x,
+ int64_t total,
+ int64_t head_dim,
+ int64_t h_local,
+ int64_t batch,
+ int64_t seq,
+ int64_t world) {
+ const int64_t stride = static_cast(gridDim.x) * blockDim.x;
+ for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
+ int64_t rest = i;
+ const int64_t c = rest % head_dim;
+ rest /= head_dim;
+ const int64_t h = rest % h_local;
+ rest /= h_local;
+ const int64_t w = rest % world;
+ rest /= world;
+ const int64_t s = rest % seq;
+ const int64_t b = rest / seq;
+
+ out[i] = x[((((w * seq + s) * batch + b) * h_local) + h) * head_dim + c];
+ }
+}
+
+} // namespace
+
+template
+struct UspMergeHeadsKernel {
+ static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) {
+ check_dtype(out);
+ check_dtype(x);
+ host::RuntimeCheck(x.ndim() == 5, "x must be [W, S, B, h_local, D]");
+ host::RuntimeCheck(out.ndim() == 5, "out must be [B, S, W, h_local, D]");
+ for (auto* t : {&x, &out}) {
+ host::RuntimeCheck(t->device().device_type == kDLCUDA, "usp_merge_heads tensors must be CUDA");
+ host::RuntimeCheck(is_dense_contiguous(*t), "usp_merge_heads tensors must be contiguous");
+ }
+ const int64_t world = x.size(0);
+ const int64_t seq = x.size(1);
+ const int64_t batch = x.size(2);
+ const int64_t h_local = x.size(3);
+ const int64_t head_dim = x.size(4);
+ host::RuntimeCheck(
+ out.size(0) == batch && out.size(1) == seq && out.size(2) == world && out.size(3) == h_local &&
+ out.size(4) == head_dim,
+ "out must be the [B, S, W, h_local, D] permutation of x");
+
+ const int64_t total = numel(x);
+ if (total == 0) {
+ return;
+ }
+
+ T* out_ptr = reinterpret_cast(mutable_data_ptr(out));
+ const T* x_ptr = reinterpret_cast(data_ptr(x));
+
+ constexpr int kVec = 16 / sizeof(T);
+ const bool vec_ok = (head_dim % kVec == 0) && aligned16(out_ptr) && aligned16(x_ptr);
+ if (vec_ok) {
+ const int64_t n_vec = total / kVec;
+ host::LaunchKernel(static_cast(grid_for(n_vec)), kBlockSize, out.device())(
+ usp_merge_heads_vec_kernel, out_ptr, x_ptr, n_vec, head_dim / kVec, h_local, batch, seq, world);
+ } else {
+ host::LaunchKernel(static_cast(grid_for(total)), kBlockSize, out.device())(
+ usp_merge_heads_scalar_kernel, out_ptr, x_ptr, total, head_dim, h_local, batch, seq, world);
+ }
+ }
+};
+
+} // namespace sglang_usp_relayout
diff --git a/python/sglang/kernels/jit/csrc/elementwise/activation.cuh b/python/sglang/kernels/jit/csrc/elementwise/activation.cuh
index dc49e58e5..c72ef9f61 100644
--- a/python/sglang/kernels/jit/csrc/elementwise/activation.cuh
+++ b/python/sglang/kernels/jit/csrc/elementwise/activation.cuh
@@ -55,7 +55,13 @@ struct ActivationParams {
uint32_t expert_step;
};
-template
+template <
+ typename T,
+ ActivationKind kAct,
+ bool kUsePDL,
+ bool kFilterExpert,
+ bool kRoundActivation = false,
+ bool kReuseInput = false>
__global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams params) {
using namespace device;
constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
@@ -70,7 +76,7 @@ __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams para
}
const auto offset = tid % num_vecs;
const auto input_offset = token_id * (num_vecs * 2) + offset;
- const auto output_offset = tid;
+ const auto output_offset = kReuseInput ? input_offset : tid;
PDLWaitPrimary();
const auto gate = device::load_as(params.input, input_offset);
const auto up = device::load_as(params.input, input_offset + num_vecs);
@@ -79,9 +85,18 @@ __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams para
for (int i = 0; i < kVecSize; ++i) {
const float gate_f32 = device::cast(gate[i]);
const float up_f32 = device::cast(up[i]);
- out[i] = device::cast(apply_activation_f32(gate_f32) * up_f32);
+ if constexpr (kRoundActivation) {
+ const T activated = device::cast(apply_activation_f32(gate_f32));
+ out[i] = device::cast(device::cast(activated) * up_f32);
+ } else {
+ out[i] = device::cast(apply_activation_f32(gate_f32) * up_f32);
+ }
+ }
+ if constexpr (kReuseInput) {
+ device::store_as(const_cast(params.input), out, output_offset);
+ } else {
+ device::store_as(params.out, out, output_offset);
}
- device::store_as(params.out, out, output_offset);
PDLTriggerSecondary();
}
@@ -117,26 +132,28 @@ struct ActivationKernel {
using kernel_fn_t = decltype(&act_and_mul_kernel);
using unary_kernel_fn_t = decltype(&act_kernel);
- template
- static constexpr kernel_fn_t activation_kernel = act_and_mul_kernel;
+ template
+ static constexpr kernel_fn_t activation_kernel =
+ act_and_mul_kernel;
static_assert(device::kMaxVecBytes % sizeof(T) == 0, "unsupported data type");
- template
+ template
static kernel_fn_t select_kernel(const std::string& type) {
using namespace host;
if (type == "silu") {
- return activation_kernel;
+ return activation_kernel;
} else if (type == "gelu") {
- return activation_kernel;
+ return activation_kernel;
} else if (type == "gelu_tanh") {
- return activation_kernel;
+ return activation_kernel;
} else {
Panic("unsupported activation type: ", type);
}
return nullptr;
}
+ template
static void launch(
const tvm::ffi::TensorView& input,
const tvm::ffi::TensorView& out,
@@ -151,10 +168,11 @@ struct ActivationKernel {
auto device_ = SymbolicDevice{};
device_.set_options();
- TensorMatcher({N, D_out}) //
- .with_dtype()
- .with_device(device_)
- .verify(out);
+ if constexpr (kReuseInput) {
+ TensorMatcher({N, D_out}).with_strides({D_in, 1}).with_dtype().with_device(device_).verify(out);
+ } else {
+ TensorMatcher({N, D_out}).with_dtype().with_device(device_).verify(out);
+ }
TensorMatcher({N, D_in}) //
.with_dtype()
.with_device(device_)
@@ -166,13 +184,16 @@ struct ActivationKernel {
if (num_tokens == 0) return;
RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension");
RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size");
+ if constexpr (kReuseInput) {
+ RuntimeCheck(input.data_ptr() == out.data_ptr(), "in-place activation output must alias input");
+ }
// only get once to avoid overhead
const auto num_total_items = num_tokens * (hidden_size / kVecSize);
RuntimeCheck(num_total_items <= std::numeric_limits::max(), "too many items for 32-bit indexing");
const auto num_blocks = div_ceil(static_cast(num_total_items), kBlockSize);
const auto params = ActivationParams{
.input = input.data_ptr(),
- .out = out.data_ptr(),
+ .out = kReuseInput ? nullptr : out.data_ptr(),
.hidden_dim = hidden_size,
.num_tokens = num_tokens,
.expert_ids = expert_ids,
@@ -180,10 +201,10 @@ struct ActivationKernel {
};
if (expert_ids != nullptr) {
RuntimeCheck(expert_step > 0, "expert_step must be positive");
- const auto kernel = select_kernel(type);
+ const auto kernel = select_kernel(type);
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
} else {
- const auto kernel = select_kernel(type);
+ const auto kernel = select_kernel(type);
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
}
}
@@ -192,6 +213,16 @@ struct ActivationKernel {
launch(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
}
+ static void
+ run_activation_with_rounding(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
+ launch(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
+ }
+
+ static void run_activation_with_rounding_input_inplace(
+ const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
+ launch(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
+ }
+
static void run_activation_filtered(
const tvm::ffi::TensorView input,
const tvm::ffi::TensorView out,
diff --git a/python/sglang/kernels/ops/activation/activation.py b/python/sglang/kernels/ops/activation/activation.py
index 481f74a9c..8fd027cc9 100644
--- a/python/sglang/kernels/ops/activation/activation.py
+++ b/python/sglang/kernels/ops/activation/activation.py
@@ -29,15 +29,26 @@ def _fast_math_flags() -> list[str]:
@cache_once
-def activation_module(dtype: torch.dtype) -> Module:
+def activation_module(dtype: torch.dtype, *, fast_math: bool = True) -> Module:
+ fast_math_flags = _fast_math_flags()
+ if not fast_math and not fast_math_flags:
+ return activation_module(dtype)
args = make_cpp_args(dtype, is_arch_support_pdl())
return load_jit(
- "activation",
+ "activation" if fast_math else "rounded_activation",
*args,
cuda_files=["elementwise/activation.cuh"],
- extra_cuda_cflags=_fast_math_flags(),
+ extra_cuda_cflags=fast_math_flags if fast_math else [],
cuda_wrappers=[
("run_activation", f"ActivationKernel<{args}>::run_activation"),
+ (
+ "run_activation_with_rounding",
+ f"ActivationKernel<{args}>::run_activation_with_rounding",
+ ),
+ (
+ "run_activation_with_rounding_input_inplace",
+ f"ActivationKernel<{args}>::run_activation_with_rounding_input_inplace",
+ ),
(
"run_activation_filtered",
f"ActivationKernel<{args}>::run_activation_filtered",
@@ -65,6 +76,28 @@ def _run_activation_inplace(
module.run_activation(input_2d, out_2d, op_name)
+@register_custom_op(mutates_args=["out"])
+def _run_activation_with_rounding_inplace(
+ op_name: str, input: torch.Tensor, out: torch.Tensor
+) -> None:
+ hidden_size = input.shape[-1] // 2
+ # Fast-math changes FP16 SiLU at eager rounding boundaries on SM90.
+ module = activation_module(input.dtype, fast_math=False)
+ input_2d = input.view(-1, hidden_size * 2)
+ out_2d = out.view(-1, hidden_size)
+ module.run_activation_with_rounding(input_2d, out_2d, op_name)
+
+
+@register_custom_op(mutates_args=["input"])
+def _run_silu_and_mul_with_rounding_inplace(input: torch.Tensor) -> None:
+ hidden_size = input.shape[-1] // 2
+ module = activation_module(input.dtype, fast_math=False)
+ input_2d = input.view(-1, hidden_size * 2)
+ module.run_activation_with_rounding_input_inplace(
+ input_2d, input_2d[:, :hidden_size], "silu"
+ )
+
+
@register_custom_op(mutates_args=["out"])
def _run_activation_filtered_inplace(
op_name: str,
@@ -150,6 +183,23 @@ def silu_and_mul(
return run_activation("silu", input, out, expert_ids, expert_step)
+def silu_and_mul_with_activation_rounding(
+ input: torch.Tensor,
+ out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+ hidden_size = input.shape[-1] // 2
+ if out is None:
+ out = input.new_empty(*input.shape[:-1], hidden_size)
+ _run_activation_with_rounding_inplace("silu", input, out)
+ return out
+
+
+def silu_and_mul_with_activation_rounding_(input: torch.Tensor) -> torch.Tensor:
+ hidden_size = input.shape[-1] // 2
+ _run_silu_and_mul_with_rounding_inplace(input)
+ return input[..., :hidden_size]
+
+
def gelu_and_mul(
input: torch.Tensor,
out: Optional[torch.Tensor] = None,
diff --git a/python/sglang/kernels/ops/diffusion/qknorm_rope.py b/python/sglang/kernels/ops/diffusion/qknorm_rope.py
index 80af7cdf7..80149d5f3 100644
--- a/python/sglang/kernels/ops/diffusion/qknorm_rope.py
+++ b/python/sglang/kernels/ops/diffusion/qknorm_rope.py
@@ -26,8 +26,18 @@ def _jit_qknorm_rope_module(
rope_dim: int,
is_neox: bool,
dtype: torch.dtype,
+ cache_dtype: torch.dtype,
+ round_norm_before_rope: bool,
) -> Module:
- args = make_cpp_args(head_dim, rope_dim, is_neox, is_arch_support_pdl(), dtype)
+ args = make_cpp_args(
+ head_dim,
+ rope_dim,
+ is_neox,
+ is_arch_support_pdl(),
+ dtype,
+ cache_dtype,
+ round_norm_before_rope,
+ )
return load_jit(
"qknorm_rope",
*args,
@@ -43,6 +53,8 @@ def can_use_fused_inplace_qknorm_rope(
rope_dim: int,
is_neox: bool,
dtype: torch.dtype,
+ cache_dtype: torch.dtype = torch.float32,
+ round_norm_before_rope: bool = False,
) -> bool:
if head_dim not in (64, 128, 256):
logger.warning(f"Unsupported head_dim={head_dim} for JIT fused QKNorm+RoPE")
@@ -62,15 +74,29 @@ def can_use_fused_inplace_qknorm_rope(
return False
if is_neox:
rotary_lanes = rope_dim // elems_per_thread
- if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
+ if rotary_lanes < 2 or rotary_lanes % 2:
logger.warning(
- "rope_dim=%s yields invalid rotary_lanes=%s for neox fused QKNorm+RoPE; rotary lane count must be a power of 2",
+ "rope_dim=%s yields invalid rotary_lanes=%s for neox fused QKNorm+RoPE; rotary lane count must be even",
rope_dim,
rotary_lanes,
)
return False
+ if round_norm_before_rope and cache_dtype != dtype:
+ logger.warning(
+ "Exact fused QKNorm+RoPE requires cache dtype %s to match activation dtype %s",
+ cache_dtype,
+ dtype,
+ )
+ return False
try:
- _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, dtype)
+ _jit_qknorm_rope_module(
+ head_dim,
+ rope_dim,
+ is_neox,
+ dtype,
+ cache_dtype,
+ round_norm_before_rope,
+ )
return True
except Exception as e:
logger.warning(f"Failed to load JIT fused QKNorm+RoPE kernel: {e}")
@@ -90,8 +116,16 @@ def fused_inplace_qknorm_rope(
eps: float = 1e-6,
head_dim: int = 0,
rope_dim: int = 0,
+ round_norm_before_rope: bool = False,
) -> None:
head_dim = head_dim or q.size(-1)
rope_dim = rope_dim or cos_sin_cache.size(-1)
- module = _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, q.dtype)
+ module = _jit_qknorm_rope_module(
+ head_dim,
+ rope_dim,
+ is_neox,
+ q.dtype,
+ cos_sin_cache.dtype,
+ round_norm_before_rope,
+ )
module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)
diff --git a/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py b/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py
new file mode 100644
index 000000000..58b23114c
--- /dev/null
+++ b/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py
@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _round_bf16_to_fp32(value):
+ # force the eager BF16 kernel boundary so Triton cannot contract the next add
+ bits = value.to(tl.int32, bitcast=True)
+ rounding_bias = 0x7FFF + ((bits >> 16) & 1)
+ rounded_bits = (bits + rounding_bias) & -65536
+ return rounded_bits.to(tl.float32, bitcast=True)
+
+
+@triton.jit
+def _indexed_scale_shift_bf16_kernel(
+ output_ptr,
+ x_ptr,
+ shift_ptr,
+ scale_ptr,
+ indices_ptr,
+ hidden_size,
+ stride_x_row,
+ stride_shift_row,
+ stride_scale_row,
+ stride_indices,
+ BLOCK_N: tl.constexpr,
+):
+ row = tl.program_id(0)
+ columns = tl.arange(0, BLOCK_N)
+ mask = columns < hidden_size
+ index = tl.load(indices_ptr + row * stride_indices)
+
+ x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(
+ tl.float32
+ )
+ shift = tl.load(
+ shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0
+ ).to(tl.float32)
+ scale = tl.load(
+ scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0
+ ).to(tl.float32)
+
+ one_plus_scale = _round_bf16_to_fp32(1.0 + scale)
+ scaled = _round_bf16_to_fp32(x * one_plus_scale)
+ tl.store(
+ output_ptr + row * stride_x_row + columns,
+ scaled + shift,
+ mask=mask,
+ )
+
+
+@triton.jit
+def _indexed_gate_bf16_kernel(
+ output_ptr,
+ x_ptr,
+ gate_ptr,
+ other_ptr,
+ indices_ptr,
+ hidden_size,
+ stride_x_row,
+ stride_gate_row,
+ stride_other_row,
+ stride_indices,
+ BLOCK_N: tl.constexpr,
+):
+ row = tl.program_id(0)
+ columns = tl.arange(0, BLOCK_N)
+ mask = columns < hidden_size
+ index = tl.load(indices_ptr + row * stride_indices)
+
+ x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(
+ tl.float32
+ )
+ gate = tl.load(
+ gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0
+ ).to(tl.float32)
+ other = tl.load(
+ other_ptr + row * stride_other_row + columns, mask=mask, other=0.0
+ ).to(tl.float32)
+
+ gated = _round_bf16_to_fp32(gate * other)
+ tl.store(
+ output_ptr + row * stride_x_row + columns,
+ x + gated,
+ mask=mask,
+ )
+
+
+def indexed_scale_shift_bf16_(
+ x: torch.Tensor,
+ shift: torch.Tensor,
+ scale: torch.Tensor,
+ indices: torch.Tensor,
+) -> torch.Tensor:
+ rows, hidden_size = x.shape
+ if rows == 0:
+ return x
+ block_n = triton.next_power_of_2(hidden_size)
+ _indexed_scale_shift_bf16_kernel[(rows,)](
+ x,
+ x,
+ shift,
+ scale,
+ indices,
+ hidden_size,
+ x.stride(0),
+ shift.stride(0),
+ scale.stride(0),
+ indices.stride(0),
+ BLOCK_N=block_n,
+ num_warps=8,
+ )
+ return x
+
+
+def indexed_gate_bf16_(
+ x: torch.Tensor,
+ gate: torch.Tensor,
+ other: torch.Tensor,
+ indices: torch.Tensor,
+) -> torch.Tensor:
+ rows, hidden_size = x.shape
+ if rows == 0:
+ return x
+ block_n = triton.next_power_of_2(hidden_size)
+ _indexed_gate_bf16_kernel[(rows,)](
+ x,
+ x,
+ gate,
+ other,
+ indices,
+ hidden_size,
+ x.stride(0),
+ gate.stride(0),
+ other.stride(0),
+ indices.stride(0),
+ BLOCK_N=block_n,
+ num_warps=8,
+ )
+ return x
diff --git a/python/sglang/kernels/ops/diffusion/triton/scale_shift.py b/python/sglang/kernels/ops/diffusion/triton/scale_shift.py
index dff528e87..f5f416c2b 100644
--- a/python/sglang/kernels/ops/diffusion/triton/scale_shift.py
+++ b/python/sglang/kernels/ops/diffusion/triton/scale_shift.py
@@ -5,6 +5,81 @@ import triton.language as tl # type: ignore
from sglang.multimodal_gen.runtime.platforms import current_platform
+@triton.jit
+def _fp32_mul_add_rn(x, scale, residual):
+ """Match separate CUDA FP32 multiply and add rounding (no FMA)."""
+ return tl.inline_asm_elementwise(
+ asm="""{
+ .reg .f32 product;
+ mul.rn.f32 product, $1, $2;
+ add.rn.f32 $0, $3, product;
+ }""",
+ constraints="=f,f,f,f",
+ args=(x, scale, residual),
+ dtype=tl.float32,
+ is_pure=True,
+ pack=1,
+ )
+
+
+@triton.jit
+def _fused_scaled_residual_add_exact_kernel(
+ output_ptr,
+ residual_ptr,
+ x_ptr,
+ scale_ptr,
+ numel: tl.constexpr,
+ width: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+):
+ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ mask = offsets < numel
+ x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32)
+ scale = tl.load(scale_ptr + offsets % width, mask=mask)
+ residual = tl.load(residual_ptr + offsets, mask=mask)
+ output = _fp32_mul_add_rn(x, scale, residual)
+ tl.store(output_ptr + offsets, output, mask=mask)
+
+
+def try_fused_scaled_residual_add_exact(
+ residual: torch.Tensor,
+ x: torch.Tensor,
+ scale: torch.Tensor,
+) -> torch.Tensor | None:
+ """Fuse ``residual + x * scale`` without changing eager FP32 rounding."""
+ if (
+ not current_platform.is_cuda()
+ or torch.is_grad_enabled()
+ or torch.compiler.is_compiling()
+ or residual.dtype != torch.float32
+ or x.dtype not in (torch.float16, torch.bfloat16)
+ or scale.dtype != torch.float32
+ or not residual.is_cuda
+ or residual.device != x.device
+ or residual.device != scale.device
+ or residual.shape != x.shape
+ or scale.shape != (x.shape[-1],)
+ or not residual.is_contiguous()
+ or not x.is_contiguous()
+ or not scale.is_contiguous()
+ or x.numel() == 0
+ ):
+ return None
+
+ output = torch.empty_like(residual)
+ block_size = 1024
+ _fused_scaled_residual_add_exact_kernel[(triton.cdiv(x.numel(), block_size),)](
+ output,
+ residual,
+ x,
+ scale,
+ numel=x.numel(),
+ width=x.shape[-1],
+ BLOCK_SIZE=block_size,
+ )
+ return output
+
+
@triton.jit
def _fused_layernorm_scale_shift_gate_select01_kernel(
output_ptr,
diff --git a/python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py b/python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py
new file mode 100644
index 000000000..a265cf883
--- /dev/null
+++ b/python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _pack_qkv_destination_major_kernel(
+ output_ptr,
+ q_ptr,
+ k_ptr,
+ v_ptr,
+ total_elements,
+ rows,
+ local_heads,
+ head_size,
+ stride_q_row,
+ stride_q_head,
+ stride_k_row,
+ stride_k_head,
+ stride_v_row,
+ stride_v_head,
+ BLOCK_SIZE: tl.constexpr,
+):
+ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ mask = offsets < total_elements
+
+ dim = offsets % head_size
+ head_slot = offsets // head_size
+ local_head = head_slot % local_heads
+ row_slot = head_slot // local_heads
+ row = row_slot % rows
+ destination = row_slot // rows
+ global_head = destination * local_heads + local_head
+
+ q = tl.load(
+ q_ptr + row * stride_q_row + global_head * stride_q_head + dim,
+ mask=mask,
+ )
+ k = tl.load(
+ k_ptr + row * stride_k_row + global_head * stride_k_head + dim,
+ mask=mask,
+ )
+ v = tl.load(
+ v_ptr + row * stride_v_row + global_head * stride_v_head + dim,
+ mask=mask,
+ )
+ output_base = head_slot * (3 * head_size) + dim
+ tl.store(output_ptr + output_base, q, mask=mask)
+ tl.store(output_ptr + output_base + head_size, k, mask=mask)
+ tl.store(output_ptr + output_base + 2 * head_size, v, mask=mask)
+
+
+def pack_qkv_destination_major(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ world_size: int,
+) -> torch.Tensor:
+ rows, global_heads, head_size = q.shape
+ local_heads = global_heads // world_size
+ output = torch.empty(
+ world_size,
+ rows,
+ local_heads,
+ 3 * head_size,
+ dtype=q.dtype,
+ device=q.device,
+ )
+ total_elements = rows * global_heads * head_size
+ if total_elements == 0:
+ return output
+
+ block_size = 1024
+ _pack_qkv_destination_major_kernel[(triton.cdiv(total_elements, block_size),)](
+ output,
+ q,
+ k,
+ v,
+ total_elements,
+ rows,
+ local_heads,
+ head_size,
+ q.stride(0),
+ q.stride(1),
+ k.stride(0),
+ k.stride(1),
+ v.stride(0),
+ v.stride(1),
+ BLOCK_SIZE=block_size,
+ num_warps=8,
+ )
+ return output
diff --git a/python/sglang/kernels/ops/diffusion/usp_relayout.py b/python/sglang/kernels/ops/diffusion/usp_relayout.py
new file mode 100644
index 000000000..edd0cbaac
--- /dev/null
+++ b/python/sglang/kernels/ops/diffusion/usp_relayout.py
@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+ from tvm_ffi.module import Module
+
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+@cache_once
+def _jit_usp_relayout_module(dtype: torch.dtype) -> Module:
+ args = make_cpp_args(dtype)
+ return load_jit(
+ "diffusion_usp_relayout",
+ *args,
+ cuda_files=["diffusion/usp_relayout.cuh"],
+ cuda_wrappers=[
+ (
+ "usp_merge_heads",
+ "sglang_usp_relayout::" f"UspMergeHeadsKernel<{args}>::run",
+ ),
+ ],
+ )
+
+
+def _fake_merge_heads(x: torch.Tensor) -> torch.Tensor:
+ world, seq, batch, h_local, head_dim = x.shape
+ return x.new_empty((batch, seq, world, h_local, head_dim))
+
+
+@register_custom_op(
+ op_name="diffusion_usp_merge_heads",
+ mutates_args=[],
+ fake_impl=_fake_merge_heads,
+)
+def _usp_merge_heads_custom_op(x: torch.Tensor) -> torch.Tensor:
+ world, seq, batch, h_local, head_dim = x.shape
+ out = x.new_empty((batch, seq, world, h_local, head_dim))
+ module = _jit_usp_relayout_module(x.dtype)
+ module.usp_merge_heads(out, x)
+ return out
+
+
+def can_use_usp_merge_heads(x: torch.Tensor) -> bool:
+ return (
+ isinstance(x, torch.Tensor)
+ and torch.version.hip is None
+ and x.is_cuda
+ and x.dtype in _SUPPORTED_DTYPES
+ and x.dim() == 5
+ and x.numel() > 0
+ and x.is_contiguous()
+ )
+
+
+def _usp_merge_heads_cuda(x: torch.Tensor) -> torch.Tensor:
+ """[W, S, B, h_local, D] -> [B, S, W, h_local, D] contiguous.
+
+ Bit-exact single-pass replacement for
+ ``x.permute(2, 1, 0, 3, 4).contiguous()`` on the Ulysses output path.
+ """
+ if not can_use_usp_merge_heads(x):
+ raise RuntimeError("unsupported input for usp_merge_heads CUDA")
+ return _usp_merge_heads_custom_op(x)
+
+
+def usp_merge_heads(x: torch.Tensor) -> torch.Tensor:
+ """Merge Ulysses output heads with an exact eager fallback.
+
+ The backend selection lives here so callers only express the layout
+ transformation. Unsupported devices, layouts, and compiled regions retain
+ the original PyTorch operation.
+ """
+ if not torch.compiler.is_compiling() and can_use_usp_merge_heads(x):
+ return _usp_merge_heads_cuda(x)
+ return x.permute(2, 1, 0, 3, 4).contiguous()
diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md
index ae2418c11..65a8ea68d 100644
--- a/python/sglang/multimodal_gen/README.md
+++ b/python/sglang/multimodal_gen/README.md
@@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus
## Key Features
SGLang Diffusion has the following features:
- - Broad model support: Wan, FastWan, FLUX, Qwen-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3, LingBot World, SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
+ - Broad model support: Wan, FastWan, FLUX, Qwen-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3, MiniMax-H3, LingBot World, SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
- Fast inference speed: empowered by optimized `sgl-kernel` kernels, scheduler/runtime improvements, caching acceleration, and native diffusion hot-path optimizations
- Ease of use: OpenAI-compatible api, CLI, and python sdk support
- Multi-platform support:
diff --git a/python/sglang/multimodal_gen/benchmarks/bench_serving.py b/python/sglang/multimodal_gen/benchmarks/bench_serving.py
index d54a5a3c7..d50f98bdf 100644
--- a/python/sglang/multimodal_gen/benchmarks/bench_serving.py
+++ b/python/sglang/multimodal_gen/benchmarks/bench_serving.py
@@ -111,6 +111,21 @@ def _infer_slo_base_time_ms_from_warmups(
return float(np.median(candidates_ms)) if candidates_ms else None
+def _parse_extra_body(raw: Optional[str]) -> Dict[str, Any]:
+ """Parses --extra-body, which is merged over every generated payload."""
+ if not raw:
+ return {}
+ try:
+ parsed = json.loads(raw)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"--extra-body is not valid JSON: {exc}") from exc
+ if not isinstance(parsed, dict):
+ raise ValueError(
+ f"--extra-body must be a JSON object, got {type(parsed).__name__}."
+ )
+ return parsed
+
+
def _populate_slo_ms_from_warmups(
requests_list: List[RequestFuncInput], warmup_pairs: List[tuple], args
) -> List[RequestFuncInput]:
@@ -496,6 +511,10 @@ async def benchmark(args):
if args.base_url is None:
args.base_url = NetworkAddress(args.host, args.port).to_url()
+ # Parsed before the service wait and the dataset download so a malformed
+ # value fails immediately instead of after minutes of setup.
+ extra_body = _parse_extra_body(args.extra_body)
+
# Wait for service
wait_for_service(args.base_url)
@@ -571,6 +590,13 @@ async def benchmark(args):
requests_list = dataset.get_requests()
logger.info(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
+ if extra_body:
+ logger.info(f"Merging --extra-body into every request: {extra_body}")
+ requests_list = [
+ replace(req, extra_body={**req.extra_body, **extra_body})
+ for req in requests_list
+ ]
+
# Limit concurrency
if args.max_concurrency is not None:
semaphore = asyncio.Semaphore(args.max_concurrency)
@@ -588,10 +614,10 @@ async def benchmark(args):
# Run warmup requests
warmup_pairs: List[tuple] = []
if args.warmup_requests and requests_list:
- # The server always overrides warmup requests to use
- # num_inference_steps=1 (see Req.set_as_warmup), so we match
- # that here to keep the benchmark's SLO estimation consistent.
- warmup_steps = 1
+ # Defaults to 1 to match the server's own boot warmup (see
+ # Req.set_as_warmup) and keep SLO estimation consistent. Raise it
+ # for models that reject a 1-step schedule, such as MiniMax-H3.
+ warmup_steps = args.warmup_inference_steps
logger.info(
f"Running {args.warmup_requests} warmup request(s) with "
f"num_inference_steps={warmup_steps}..."
@@ -828,6 +854,22 @@ if __name__ == "__main__":
default=1,
help="Number of warmup requests to run before measurement.",
)
+ parser.add_argument(
+ "--warmup-inference-steps",
+ type=int,
+ default=1,
+ help="Denoise steps for warmup requests. Raise it for models that "
+ "reject a 1-step schedule.",
+ )
+ parser.add_argument(
+ "--extra-body",
+ type=str,
+ default=None,
+ help="JSON object merged over each JSON request body, for contract "
+ 'fields the generic payload omits (e.g. \'{"task": "t2va"}\' for '
+ "MiniMax-H3). Multipart image requests forward it as an extra_body "
+ "form field instead, which the server may not unpack.",
+ )
parser.add_argument(
"--num-inference-steps",
type=int,
diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py
index 83ef2327c..0e86c6d93 100644
--- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py
+++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py
@@ -12,6 +12,7 @@ from sglang.multimodal_gen.configs.models.dits.lingbot_world import (
LingBotWorldVideoConfig,
)
from sglang.multimodal_gen.configs.models.dits.longlive2 import LongLive2VideoConfig
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import (
@@ -27,6 +28,7 @@ __all__ = [
"Ideogram4DistilledDiTConfig",
"LingBotWorldVideoConfig",
"LongLive2VideoConfig",
+ "MiniMaxH3DiTConfig",
"WanVideoConfig",
"Hunyuan3DDiTConfig",
"MOVAAudioConfig",
diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py
new file mode 100644
index 000000000..d773411e0
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_block
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64
+MINIMAX_H3_ADALN_MODALITY_NUM = 3
+
+
+@dataclass
+class MiniMaxH3DiTArchConfig(DiTArchConfig):
+ _fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
+
+ lora_param_names_mapping: dict = field(default_factory=dict)
+
+ _supported_attention_backends: set[AttentionBackendEnum] = field(
+ default_factory=lambda: {
+ AttentionBackendEnum.FA,
+ AttentionBackendEnum.AITER,
+ AttentionBackendEnum.TORCH_SDPA,
+ }
+ )
+
+ num_layers: int = 50
+ token_refiner_num_layers: int = 2
+ hidden_size: int = 5376
+ num_attention_heads: int = 56
+ attention_head_dim: int = 128
+ ffn_hidden_size: int = 14336
+ latents_dim: int = 24
+ audio_latents_dim: int = 32
+ patch_size: tuple[int, int, int] = (1, 2, 2)
+ text_dim: int = 5120
+ timestep_input_dim: int = 256
+ time_embed_hidden_size: int = 5376
+ time_embed_dim: int = 2688
+ adaln_out_features: int = 18 * 5376
+ final_adaln_out_features: int = 2 * 5376
+ rope_inv_freq_len: int = 16
+ norm_eps: float = 1e-5
+ qk_norm_eps: float = 1e-5
+ final_norm_eps: float = 1e-5
+
+ def __post_init__(self) -> None:
+ super().__post_init__()
+ if isinstance(self.patch_size, list):
+ self.patch_size = tuple(self.patch_size)
+ if len(self.patch_size) != 3:
+ raise ValueError(f"patch_size must have 3 values, got {self.patch_size}.")
+ self.num_channels_latents = self.latents_dim
+
+
+@dataclass
+class MiniMaxH3DiTConfig(DiTConfig):
+ arch_config: MiniMaxH3DiTArchConfig = field(default_factory=MiniMaxH3DiTArchConfig)
+
+
+__all__ = [
+ "MINIMAX_H3_ADALN_MODALITY_NUM",
+ "MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT",
+ "MiniMaxH3DiTArchConfig",
+ "MiniMaxH3DiTConfig",
+]
diff --git a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py
index e97e6d9f7..58c00f45c 100644
--- a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py
+++ b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py
@@ -21,6 +21,10 @@ from sglang.multimodal_gen.configs.models.encoders.ideogram import (
Ideogram4TextEncoderConfig,
)
from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
+from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
+ MiniMaxH3Qwen3VLArchConfig,
+ MiniMaxH3Qwen3VLConfig,
+)
from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig
from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
@@ -36,6 +40,8 @@ __all__ = [
"Flux2MistralTextConfig",
"build_flux2_text_messages",
"LlamaConfig",
+ "MiniMaxH3Qwen3VLArchConfig",
+ "MiniMaxH3Qwen3VLConfig",
"Qwen3TextConfig",
"Qwen3VLConfig",
"T5Config",
diff --git a/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py
new file mode 100644
index 000000000..62d9d3bdc
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native Qwen3-VL encoder configuration for MiniMax H3."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.qwen3vl import (
+ Qwen3VLArchConfig,
+ Qwen3VLConfig,
+)
+
+MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50
+
+
+@dataclass
+class MiniMaxH3Qwen3VLArchConfig(Qwen3VLArchConfig):
+ """The checkpoint is Qwen3-VL-32B, consumed at hidden_states[50]."""
+
+ architectures: list[str] = field(
+ default_factory=lambda: ["MiniMaxH3Qwen3VLEncoder"]
+ )
+ hidden_size: int = 5120
+ intermediate_size: int = 25600
+ num_hidden_layers: int = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+ num_attention_heads: int = 64
+ num_key_value_heads: int = 8
+ head_dim: int = 128
+ text_len: int = 262144
+ hidden_state_skip_layer: int = 0
+
+
+@dataclass
+class MiniMaxH3Qwen3VLConfig(Qwen3VLConfig):
+ arch_config: MiniMaxH3Qwen3VLArchConfig = field(
+ default_factory=MiniMaxH3Qwen3VLArchConfig
+ )
+
+ def post_diffusers_config_update(self) -> None:
+ """Select the in-tree extractor after loading the HF architecture."""
+
+ arch = self.arch_config
+ arch.architectures = ["MiniMaxH3Qwen3VLEncoder"]
+ arch.hidden_size = int(arch.text_config.hidden_size)
+ arch.intermediate_size = int(arch.text_config.intermediate_size)
+ arch.num_attention_heads = int(arch.text_config.num_attention_heads)
+ arch.num_key_value_heads = int(arch.text_config.num_key_value_heads)
+ arch.head_dim = int(arch.text_config.head_dim)
+ arch.num_hidden_layers = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+ arch.text_config.num_hidden_layers = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+ arch.text_config.output_hidden_states = False
+ arch.text_config.use_cache = False
+
+
+__all__ = [
+ "MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER",
+ "MiniMaxH3Qwen3VLArchConfig",
+ "MiniMaxH3Qwen3VLConfig",
+]
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py
index 3438b1b89..f7e1472e4 100644
--- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py
+++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py
@@ -3,6 +3,12 @@
from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig
from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig
from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
+ MiniMaxH3AudioVAEConfig,
+)
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
+ MiniMaxH3VideoVAEConfig,
+)
from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
StableDiffusion3VAEConfig,
)
@@ -11,6 +17,8 @@ from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig
__all__ = [
"DacVAEConfig",
"HunyuanVAEConfig",
+ "MiniMaxH3AudioVAEConfig",
+ "MiniMaxH3VideoVAEConfig",
"StableDiffusion3VAEConfig",
"WanVAEConfig",
"Hunyuan3DVAEConfig",
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py
new file mode 100644
index 000000000..71e3e959a
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_contract import (
+ validate_minimax_h3_vae_latent_stats,
+)
+
+
+@dataclass
+class MiniMaxH3AudioVAEArchConfig(VAEArchConfig):
+ sample_rate: int = 32000
+ latent_channels: int = 32
+ latents_mean: list[float] | None = None
+ latents_std: list[float] | None = None
+ output_channel: int = 2
+
+
+@dataclass
+class MiniMaxH3AudioVAEConfig(VAEConfig):
+ arch_config: MiniMaxH3AudioVAEArchConfig = field(
+ default_factory=MiniMaxH3AudioVAEArchConfig
+ )
+ load_encoder: bool = True
+ load_decoder: bool = True
+
+ def post_init(self) -> None:
+ validate_minimax_h3_vae_latent_stats(
+ self.arch_config,
+ component_name="audio_vae",
+ expected_channels=32,
+ )
+
+
+__all__ = ["MiniMaxH3AudioVAEArchConfig", "MiniMaxH3AudioVAEConfig"]
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_contract.py b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_contract.py
new file mode 100644
index 000000000..417136815
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_contract.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+import math
+from typing import Protocol
+
+
+class MiniMaxH3LatentStatsConfig(Protocol):
+ latent_channels: int
+ latents_mean: list[float] | None
+ latents_std: list[float] | None
+
+
+class MiniMaxH3VAEContractError(ValueError):
+ def __init__(self, component_name: str, detail: str) -> None:
+ super().__init__(f"MiniMax H3 {component_name} {detail}")
+ self.component_name = component_name
+ self.detail = detail
+
+ def __reduce__(self):
+ # BaseException pickles via cls(*args); rebuild from the two ctor args
+ # so the error propagates cleanly across process boundaries.
+ return (type(self), (self.component_name, self.detail))
+
+
+def validate_minimax_h3_vae_latent_stats(
+ arch_config: MiniMaxH3LatentStatsConfig,
+ component_name: str,
+ expected_channels: int,
+) -> None:
+ if arch_config.latent_channels != expected_channels:
+ raise MiniMaxH3VAEContractError(
+ component_name,
+ "latent_channels must be "
+ f"{expected_channels}, got {arch_config.latent_channels!r}",
+ )
+
+ for field_name, values in (
+ ("latents_mean", arch_config.latents_mean),
+ ("latents_std", arch_config.latents_std),
+ ):
+ if values is None:
+ raise MiniMaxH3VAEContractError(
+ component_name,
+ f"config.json missing {field_name}",
+ )
+ if not isinstance(values, list) or not all(
+ isinstance(value, (int, float)) and not isinstance(value, bool)
+ for value in values
+ ):
+ raise MiniMaxH3VAEContractError(
+ component_name,
+ f"config.json {field_name} must be a list of numbers",
+ )
+ if len(values) != expected_channels:
+ raise MiniMaxH3VAEContractError(
+ component_name,
+ f"config.json {field_name} must contain exactly "
+ f"{expected_channels} values, got {len(values)}",
+ )
+ if field_name == "latents_mean" and not all(
+ math.isfinite(value) for value in values
+ ):
+ raise MiniMaxH3VAEContractError(
+ component_name,
+ "config.json latents_mean values must be finite",
+ )
+ if field_name == "latents_std" and not all(
+ math.isfinite(value) and value > 0 for value in values
+ ):
+ raise MiniMaxH3VAEContractError(
+ component_name,
+ "config.json latents_std values must be finite and greater than zero",
+ )
+
+
+__all__ = [
+ "MiniMaxH3VAEContractError",
+ "validate_minimax_h3_vae_latent_stats",
+]
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py
new file mode 100644
index 000000000..d23b6740e
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py
@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_contract import (
+ validate_minimax_h3_vae_latent_stats,
+)
+
+
+@dataclass
+class MiniMaxH3VideoVAEArchConfig(VAEArchConfig):
+ latent_channels: int = 24
+ latents_mean: list[float] | None = None
+ latents_std: list[float] | None = None
+ temporal_compression_ratio: int = 4
+ spatial_compression_ratio: int = 16
+ vae_clip_length: int = 17
+ vae_token_drop: int = 3
+ vae_encoder_tiling: int = 1
+ vae_decoder_tiling: int = 1
+ vae_parallel_tiling: int = 1
+ vae_tile_size: int = 256
+ vae_tile_overlap_min: int = 64
+ vae_chunk_dim: int = -1
+
+
+@dataclass
+class MiniMaxH3VideoVAEConfig(VAEConfig):
+ arch_config: MiniMaxH3VideoVAEArchConfig = field(
+ default_factory=MiniMaxH3VideoVAEArchConfig
+ )
+ load_encoder: bool = True
+ load_decoder: bool = True
+ use_tiling: bool = True
+ use_parallel_tiling: bool = True
+ # The released checkpoint's quality contract uses overlapping latent
+ # tiles. Parallel tiling distributes whole tiles without changing that
+ # recipe. Spatial-shard decode is rejected because validation found output
+ # mismatches on H3.
+ parallel_decode_mode: str = "tiled"
+
+ def resolved_parallel_decode_mode(self) -> str:
+ if self.parallel_decode_mode == "auto":
+ return "tiled"
+ if self.parallel_decode_mode in ("spatial", "spatial_shard"):
+ raise ValueError(
+ "MiniMax H3 rejects spatial-shard VAE decode because it failed "
+ "the released quality contract; use tiled"
+ )
+ if self.parallel_decode_mode == "tiled":
+ return "tiled"
+ if self.parallel_decode_mode == "patch":
+ raise ValueError("MiniMax H3 does not support patch VAE decode; use tiled")
+ raise ValueError(
+ f"unsupported MiniMax H3 VAE parallel decode mode "
+ f"{self.parallel_decode_mode!r}"
+ )
+
+ def post_init(self) -> None:
+ self.resolved_parallel_decode_mode()
+ validate_minimax_h3_vae_latent_stats(
+ self.arch_config,
+ component_name="video_vae",
+ expected_channels=24,
+ )
+
+
+__all__ = ["MiniMaxH3VideoVAEArchConfig", "MiniMaxH3VideoVAEConfig"]
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
index 765dea30a..0867e5ed0 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
@@ -40,6 +40,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
LTX2PipelineConfig,
LTX23PipelineConfig,
)
+from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
+ MiniMaxH3PipelineConfig,
+)
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
@@ -86,4 +89,5 @@ __all__ = [
"LTX23PipelineConfig",
"LingBotWorldCausalDMDConfig",
"LingBotWorldV2CausalDMDConfig",
+ "MiniMaxH3PipelineConfig",
]
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py
index 5edae2a08..616bed20f 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py
@@ -267,6 +267,11 @@ class PipelineConfig:
# return the model-specific config for optimal deployment setting
return ModelDeploymentConfig()
+ def validate_server_args(self, server_args: Any) -> None:
+ """Validate model-owned constraints after server args are normalized."""
+
+ del server_args
+
# Wan2.2 TI2V parameters
boundary_ratio: float | None = None
@@ -394,6 +399,11 @@ class PipelineConfig:
"""
return self.task_type in (ModelTaskType.T2I, ModelTaskType.T2V)
+ def supports_disaggregation(self) -> bool:
+ """Return whether multi-service disaggregated deployment is supported."""
+
+ return True
+
def supports_native_grouped_requests(self):
"""Return whether dynamic batches should run as grouped Req lists."""
return False
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
new file mode 100644
index 000000000..b42b70ea1
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
+from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
+ MiniMaxH3Qwen3VLConfig,
+)
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
+ MiniMaxH3AudioVAEConfig,
+)
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
+ MiniMaxH3VideoVAEConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+ ModelTaskType,
+ PipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+ ModelDeploymentConfig,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+@dataclass
+class MiniMaxH3PipelineConfig(PipelineConfig):
+ """MiniMax H3 native audio-video pipeline configuration."""
+
+ # Canonical H3 materials are prepared by the model-specific stages. The
+ # generic TI2V image resize would both duplicate that work and overwrite
+ # the already-resolved target canvas.
+ skip_input_image_preprocess: bool = True
+ native_only_components = (
+ "text_encoder",
+ "transformer",
+ "video_vae",
+ "audio_vae",
+ )
+ task_type: ModelTaskType = ModelTaskType.TI2V
+ dit_config: MiniMaxH3DiTConfig = field(default_factory=MiniMaxH3DiTConfig)
+ vae_config: MiniMaxH3VideoVAEConfig = field(default_factory=MiniMaxH3VideoVAEConfig)
+ audio_vae_config: MiniMaxH3AudioVAEConfig = field(
+ default_factory=MiniMaxH3AudioVAEConfig
+ )
+ dit_precision: str = "bf16"
+ # The video VAE remains fp32-resident because it also encodes keyframes.
+ # Decode follows the released fp16-autocast recipe unless the user
+ # explicitly disables autocast.
+ vae_precision: str = "fp32"
+ vae_decode_precision: str = "fp16"
+ audio_vae_precision: str = "fp32"
+ text_encoder_configs: tuple[MiniMaxH3Qwen3VLConfig, ...] = field(
+ default_factory=lambda: (MiniMaxH3Qwen3VLConfig(),)
+ )
+ text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
+ text_encoder_extra_args: list[dict] = field(default_factory=lambda: [{}])
+ # The released checkpoint is CFG-distilled and has one positive branch.
+ should_use_guidance: bool = False
+ output_audio_sample_rate: int | None = 32000
+ output_audio_channels: int | None = 2
+ output_av_drift_tolerance_s: float | None = 0.25
+
+ def accepts_audio_input(self) -> bool:
+ return True
+
+ def supports_disaggregation(self) -> bool:
+ return False
+
+ @property
+ def requires_audio_output(self) -> bool:
+ return True
+
+ def get_model_deployment_config(self) -> ModelDeploymentConfig:
+ return ModelDeploymentConfig(
+ speed_mode_enable_torch_compile_by_default=False,
+ keep_resident_min_available_gb=120,
+ keep_resident_components=("dit", "text_encoder", "vae"),
+ auto_enable_cfg_parallel=False,
+ supports_cfg_parallel=False,
+ )
+
+ @staticmethod
+ def _server_arg_value(value):
+ return getattr(value, "value", value)
+
+ def validate_quality_deployment(self, server_args) -> None:
+ """Fail closed unless the resident server matches the measured profile."""
+
+ attention_backend = self._server_arg_value(server_args.attention_backend)
+ attention_backend = (
+ str(attention_backend).strip().lower()
+ if attention_backend is not None
+ else None
+ )
+ capability = current_platform.get_device_capability()
+ capability_int = capability.to_int() if capability is not None else None
+ device_name = (
+ current_platform.get_device_name()
+ if current_platform.is_cuda()
+ else type(current_platform).__name__
+ )
+ model_variant = str(server_args.model_variant or "fl2va").lower()
+ actual = {
+ "attention_backend": attention_backend,
+ "backend": self._server_arg_value(server_args.backend),
+ "component_attention_backends": {},
+ "enable_breakable_cuda_graph": server_args.enable_breakable_cuda_graph,
+ "enable_torch_compile": server_args.enable_torch_compile,
+ "is_dit_layerwise_offload_selected": (
+ server_args.is_dit_layerwise_offload_selected
+ ),
+ "model_variant": model_variant,
+ "num_gpus": server_args.num_gpus,
+ "performance_mode": server_args.performance_mode,
+ "quantization": server_args.quantization,
+ "regional_compile": server_args.regional_compile,
+ "ring_degree": server_args.ring_degree,
+ "sp_degree": server_args.sp_degree,
+ "tp_size": server_args.tp_size,
+ "ulysses_degree": server_args.ulysses_degree,
+ "use_fsdp_inference": server_args.use_fsdp_inference,
+ }
+ actual["component_attention_backends"] = dict(
+ server_args.component_attention_backends or {}
+ )
+ expected = {
+ "attention_backend": {None, "fa"},
+ "backend": {"auto", "sglang"},
+ "component_attention_backends": {},
+ "enable_breakable_cuda_graph": False,
+ "enable_torch_compile": False,
+ "is_dit_layerwise_offload_selected": False,
+ "model_variant": "fl2va",
+ "num_gpus": 4,
+ "performance_mode": "speed",
+ "quantization": None,
+ "regional_compile": False,
+ "ring_degree": 1,
+ "sp_degree": 4,
+ "tp_size": 1,
+ "ulysses_degree": 4,
+ "use_fsdp_inference": False,
+ }
+ mismatches = {
+ name: {"expected": wanted, "actual": actual[name]}
+ for name, wanted in expected.items()
+ if (
+ actual[name] not in wanted
+ if isinstance(wanted, set)
+ else actual[name] != wanted
+ )
+ }
+ if (
+ not current_platform.is_cuda()
+ or "H200" not in device_name.upper()
+ or capability_int != 90
+ ):
+ mismatches["device"] = {
+ "expected": "NVIDIA H200 (compute capability 9.0)",
+ "actual": f"{device_name} (compute capability {capability_int})",
+ }
+ if mismatches:
+ raise ValueError(
+ "MiniMax-H3 approximate quality profiles are validated only for "
+ f"the strict 4xH200 fl2va deployment; mismatches: {mismatches}"
+ )
+
+ def validate_server_args(self, server_args) -> None:
+ # Reject known-inexact VAE modes before any large component download.
+ self.vae_config.resolved_parallel_decode_mode()
+ attention_backend = self._server_arg_value(server_args.attention_backend)
+ if str(attention_backend).strip().lower() == "sage_attn":
+ raise ValueError(
+ "MiniMax-H3 does not support SageAttention: the current packed "
+ "varlen path does not preserve model output"
+ )
+
+ def select_vae_weight_files(
+ self,
+ safetensors_list: list[str],
+ component_model_path: str,
+ component_name: str,
+ vae_precision: str,
+ ) -> list[str]:
+ if component_name == "video_vae":
+ return [os.path.join(component_model_path, "source", "model.safetensors")]
+ return safetensors_list
+
+
+__all__ = ["MiniMaxH3PipelineConfig"]
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py
index bbb4d2767..189f79dc9 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py
@@ -25,6 +25,10 @@ class ModelDeploymentConfig:
auto_enable_cfg_parallel: bool = True
# degree 1 keeps CFG parallel disabled and leaves GPUs available for SP
auto_cfg_parallel_degree_by_num_gpus: tuple[tuple[int, int], ...] = ()
+ # Let performance_mode=speed opt into torch.compile unless the model has
+ # established that the compiled path changes its numerical contract.
+ speed_mode_enable_torch_compile_by_default: bool = True
+ supports_cfg_parallel: bool = True
def get_auto_cfg_parallel_degree(self, num_gpus: int) -> int:
for candidate_num_gpus, cfg_degree in self.auto_cfg_parallel_degree_by_num_gpus:
diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
new file mode 100644
index 000000000..603500183
--- /dev/null
+++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+import math
+import os
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+import msgspec
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+
+_MINIMAX_H3_MAX_SIGNED_SEED = (1 << 63) - 1
+
+
+def _optional_unit_float(value: Any, field_name: str) -> float | None:
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{field_name} must be a number")
+ out = float(value)
+ if out < 0.0 or out > 1.0:
+ raise ValueError(f"{field_name} must be in [0, 1]")
+ return out
+
+
+def _optional_positive_finite_float(value: Any, field_name: str) -> float | None:
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{field_name} must be a number")
+ out = float(value)
+ if not math.isfinite(out) or out <= 0.0:
+ raise ValueError(f"{field_name} must be a positive finite number")
+ return out
+
+
+@dataclass
+class MiniMaxH3SamplingParams(SamplingParams):
+ height: int = 512
+ width: int = 896
+ num_inference_steps: int = 50
+ num_frames: int = field(default=1, init=False)
+ fps: int = field(default=24, init=False)
+ negative_prompt: None = field(default=None, init=False)
+ guidance_scale: float = field(default=1.0, init=False)
+ guidance_scale_2: None = field(default=None, init=False)
+ true_cfg_scale: None = field(default=None, init=False)
+ guidance_rescale: float = field(default=0.0, init=False)
+ cfg_normalization: float = field(default=0.0, init=False)
+ imgvid_cond_noise_aug_for_inference: float | None = None
+ audio_cond_noise_aug_for_inference: float | None = None
+ task: str | None = None
+ conditions: list[dict[str, Any]] | None = None
+ target: dict[str, Any] | None = None
+ audio_flow_shift: float | None = None
+ output_mode: str | None = field(
+ default=None,
+ metadata={"batch_sig_exclude": True},
+ )
+
+ @classmethod
+ def video_request_extra_fields(cls) -> frozenset[str]:
+ return frozenset(
+ {
+ "task",
+ "conditions",
+ "target",
+ "audio_flow_shift",
+ "audio_guidance_scale",
+ "quality",
+ "output_mode",
+ "imgvid_cond_noise_aug_for_inference",
+ "audio_cond_noise_aug_for_inference",
+ }
+ )
+
+ @staticmethod
+ def _video_hooks():
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.video_adapter import (
+ MiniMaxH3VideoModelAdapter,
+ )
+
+ return MiniMaxH3VideoModelAdapter()
+
+ @classmethod
+ def lower_video_request_kwargs(
+ cls,
+ request: Any,
+ kwargs: dict[str, Any],
+ ) -> dict[str, Any]:
+ return cls._video_hooks().lower_video_request_kwargs(request, kwargs)
+
+ def prepare_video_request_for_queue(self, req: Any) -> None:
+ hooks = self._video_hooks()
+ hooks.validate_sampling_params(self)
+ hooks.prepare_for_queue_sync(req)
+
+ def expand_video_request_outputs_for_queue(self, req: Any) -> list[Any]:
+ """Use the same independent-seed grouped path in serve and generate."""
+
+ from sglang.multimodal_gen.runtime.entrypoints.utils import (
+ expand_request_outputs,
+ )
+
+ return expand_request_outputs(req)
+
+ def prepare_synthetic_warmup_request_for_queue(
+ self, req: Any, server_args: Any
+ ) -> None:
+ """Lower generic warmup into one valid native partition request.
+
+ This intentionally calls the existing pre-queue resolver directly
+ instead of the public video admission hook: synthetic warmup disables
+ file delivery, while the public H3 contract correctly requires it.
+ """
+ selected_variant = getattr(server_args, "model_variant", None)
+ if selected_variant is not None:
+ selected_partition = str(selected_variant).strip().lower()
+ else:
+ selected_path = server_args.model_subfolder or server_args.model_path
+ selected_partition = os.path.basename(
+ os.path.normpath(str(selected_path))
+ ).lower()
+
+ if selected_partition == "ref2va":
+ image_path = req.image_path
+ if isinstance(image_path, list):
+ if not image_path:
+ raise ValueError(
+ "MiniMax H3 Ref2VA synthetic warmup requires an image"
+ )
+ image_path = image_path[0]
+ if not isinstance(image_path, str) or not image_path:
+ raise ValueError("MiniMax H3 Ref2VA synthetic warmup requires an image")
+ task = "ref2va"
+ conditions = [
+ {
+ "type": "image",
+ "uri": image_path,
+ "role": "reference",
+ }
+ ]
+ else:
+ task = "t2va"
+ conditions = []
+
+ self.task = task
+ self.conditions = conditions
+ self.target = {
+ "short_edge": 768,
+ "aspect_ratio": "16:9",
+ "duration_seconds": 5.0,
+ }
+ selected_seed = req.seed if isinstance(req.seed, int) else int(req.seed[0])
+ req.extra.update(self.build_request_extra(_seed_override=int(selected_seed)))
+ self._video_hooks().prepare_for_queue_sync(req)
+
+ def project_video_queued_job_fields(self, req: Any) -> dict[str, str]:
+ return self._video_hooks().project_queued_job_fields(req)
+
+ def validate_video_final_outputs(
+ self,
+ output_paths: list[str],
+ req: Any,
+ ) -> dict[str, str]:
+ return self._video_hooks().validate_final_outputs_sync(output_paths, req)
+
+ def cleanup_video_request(self, req: Any) -> None:
+ self._video_hooks().cleanup_request_sync(req)
+
+ def _adjust(self, server_args) -> None:
+ """Apply generic path/output adjustments without deriving time shape.
+
+ The generic helper normally rewrites ``num_frames`` for temporal VAE
+ alignment and GPU sharding. MiniMax H3 resolves that shape from the
+ canonical target during pre-queue admission, so those rewrites must
+ not leak into the offline parameter object. Keep the transport
+ metadata at its internal sentinel values until pre-queue populates the
+ actual batch shape.
+ """
+
+ super()._adjust(server_args)
+ self.fps = 24
+ self.num_frames = 1
+
+ def _validate(self) -> None:
+ self.fps = 24
+ self.num_frames = 1
+ if isinstance(self.target, Mapping):
+ self.target = {
+ field_name: self.target[field_name]
+ for field_name in (
+ "short_edge",
+ "aspect_ratio",
+ "duration_seconds",
+ )
+ if field_name in self.target
+ }
+ super()._validate()
+ _optional_positive_finite_float(self.flow_shift, "flow_shift")
+ _optional_positive_finite_float(self.audio_flow_shift, "audio_flow_shift")
+ if self.enable_frame_interpolation:
+ raise ValueError(
+ "MiniMax H3 does not support enable_frame_interpolation: the "
+ "accepted delivery contract is the canonical 24 fps output"
+ )
+ if self.enable_upscaling:
+ raise ValueError(
+ "MiniMax H3 does not support enable_upscaling: the accepted "
+ "delivery contract is the resolved target canvas"
+ )
+ if self.enable_teacache:
+ raise ValueError(
+ "MiniMax H3 does not support enable_teacache: its packed "
+ "video/audio denoise loop has no lossless TeaCache contract"
+ )
+ if self.rollout:
+ raise ValueError(
+ "MiniMax H3 does not support rollout: its coupled video/audio "
+ "scheduler has no SchedulerRLMixin contract"
+ )
+ if self.return_trajectory_latents or self.return_trajectory_decoded:
+ raise ValueError(
+ "MiniMax H3 does not support trajectory output for its coupled "
+ "video/audio denoise state"
+ )
+ seeds = self.seed if isinstance(self.seed, list) else [self.seed]
+ for seed in seeds:
+ if seed > _MINIMAX_H3_MAX_SIGNED_SEED:
+ raise ValueError(
+ "MiniMax H3 seed must not exceed the signed int64 maximum, "
+ f"got {seed}"
+ )
+ if (
+ isinstance(self.seed, int)
+ and self.seed + self.num_outputs_per_prompt - 1
+ > _MINIMAX_H3_MAX_SIGNED_SEED
+ ):
+ raise ValueError(
+ "MiniMax H3 scalar seed plus output index must fit the "
+ "signed int64 upper bound"
+ )
+
+ def build_request_extra(self, *, _seed_override: int | None = None) -> dict:
+ _optional_unit_float(
+ self.imgvid_cond_noise_aug_for_inference,
+ "imgvid_cond_noise_aug_for_inference",
+ )
+ _optional_unit_float(
+ self.audio_cond_noise_aug_for_inference,
+ "audio_cond_noise_aug_for_inference",
+ )
+ extra = super().build_request_extra()
+ if self.task is not None:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.request_validation import (
+ minimax_h3_validate_canonical_request,
+ )
+
+ extra["minimax_h3_canonical_request"] = (
+ minimax_h3_validate_canonical_request(
+ task=self.task,
+ prompt=self.prompt,
+ conditions=self.conditions,
+ target=self.target,
+ flow_shift=self.flow_shift,
+ audio_flow_shift=self.audio_flow_shift,
+ seed=(
+ _seed_override
+ if _seed_override is not None
+ else self.seed if isinstance(self.seed, int) else None
+ ),
+ )
+ )
+ elif (
+ self.conditions is not None
+ or self.target is not None
+ or self.flow_shift is not None
+ or self.audio_flow_shift is not None
+ ):
+ raise ValueError(
+ "task is required when conditions/target/flow_shift/"
+ "audio_flow_shift are provided"
+ )
+ return extra
+
+ def refresh_request_extra_after_output_expansion(self, req: Any) -> None:
+ """Copy validated canonical identity with the selected scalar Req seed."""
+
+ selected_seed = getattr(req, "seed", None)
+ if isinstance(selected_seed, int):
+ canonical_key = "minimax_h3_canonical_request"
+ canonical = dict(req.extra[canonical_key])
+ canonical["seed"] = selected_seed
+ req.extra[canonical_key] = canonical
+ resolved_plan_key = "minimax_h3_resolved_plan"
+ resolved_plan = req.extra.get(resolved_plan_key)
+ if resolved_plan is not None:
+ req.extra[resolved_plan_key] = msgspec.structs.replace(
+ resolved_plan, seed=selected_seed
+ )
+ else:
+ req.extra.update(self.build_request_extra())
+
+
+__all__ = ["MiniMaxH3SamplingParams"]
diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py
index e280ba2e9..c771c3ce8 100644
--- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py
+++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py
@@ -123,6 +123,10 @@ class SamplingParams:
)
output_quality: str | None = "default"
output_compression: int | None = None
+ # Model-owned, request-scoped approximate acceleration profile. Models
+ # that support it must validate the deployment and workload explicitly.
+ # It intentionally participates in the dynamic-batch signature.
+ quality: str = "lossless"
# Frame interpolation
enable_frame_interpolation: bool = False
@@ -328,6 +332,64 @@ class SamplingParams:
if self.realtime_chunk_size is not None:
req.realtime_chunk_size = self.realtime_chunk_size
+ @classmethod
+ def video_request_extra_fields(cls) -> frozenset[str]:
+ """Declare model-specific multipart video fields accepted by this type."""
+
+ return frozenset()
+
+ @classmethod
+ def lower_video_request_kwargs(
+ cls,
+ request: Any,
+ kwargs: dict[str, Any],
+ ) -> dict[str, Any]:
+ """Adapt generic video-API kwargs before constructing this params type."""
+ del request
+ return kwargs
+
+ def prepare_video_request_for_queue(self, req: Any) -> None:
+ """Resolve model-specific admission facts before a video job is queued."""
+ del req
+
+ def expand_video_request_outputs_for_queue(self, req: Any) -> list[Any] | None:
+ """Return per-output requests when a model owns grouped execution.
+
+ ``None`` preserves the default model-native ``num_outputs`` handling.
+ Models that need the framework's independent-seed request expansion
+ can opt in after their shared pre-queue work has completed.
+ """
+ del req
+ return None
+
+ def prepare_synthetic_warmup_request_for_queue(
+ self, req: Any, server_args: Any
+ ) -> None:
+ """Resolve model-specific facts for one synthetic warmup request."""
+ del req, server_args
+
+ def project_video_queued_job_fields(self, req: Any) -> dict[str, str]:
+ """Return model-resolved fields to publish with the queued video job."""
+ del req
+ return {}
+
+ def validate_video_final_outputs(
+ self,
+ output_paths: list[str],
+ req: Any,
+ ) -> dict[str, str]:
+ """Validate final files and return truthful completion metadata."""
+ del output_paths, req
+ return {}
+
+ def cleanup_video_request(self, req: Any) -> None:
+ """Release request-scoped resources owned by the model integration."""
+ del req
+
+ def refresh_request_extra_after_output_expansion(self, req: Any) -> None:
+ """Refresh request identity after assigning a per-output seed."""
+ del req
+
def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int:
"""Convert output_quality string to compression level."""
if data_type == DataType.ACTION:
@@ -346,6 +408,11 @@ class SamplingParams:
f"prompt_path must be a txt file, got {self.prompt_path!r}"
)
+ if not isinstance(self.quality, str) or not self.quality.strip():
+ raise ValueError(
+ f"quality must be a non-empty string, got {self.quality!r}"
+ )
+
# These are always required to be sane regardless of pipeline.
if (
not isinstance(self.num_outputs_per_prompt, int)
@@ -862,10 +929,20 @@ class SamplingParams:
type=int,
help="Output compression level (0-100, higher means better quality but larger file size)",
)
+ add_argument(
+ "--quality",
+ type=str,
+ help=(
+ "Select a model-owned quality/performance profile. "
+ "Support and validated deployment constraints are model-specific."
+ ),
+ )
add_argument(
"--num-outputs-per-prompt",
+ "--num-outputs",
+ dest="num_outputs_per_prompt",
type=int,
- help="Number of outputs to generate per prompt",
+ help="Number of outputs to generate per prompt (alias: --num-outputs)",
)
add_argument(
"--seed",
diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py
index 93c6a6b54..bb90c805d 100644
--- a/python/sglang/multimodal_gen/registry.py
+++ b/python/sglang/multimodal_gen/registry.py
@@ -37,6 +37,7 @@ from sglang.multimodal_gen.configs.pipeline_configs import (
HunyuanConfig,
LingBotWorldCausalDMDConfig,
LingBotWorldV2CausalDMDConfig,
+ MiniMaxH3PipelineConfig,
WanI2V480PConfig,
WanI2V720PConfig,
WanT2V480PConfig,
@@ -141,6 +142,7 @@ from sglang.multimodal_gen.configs.sample.ltx_2 import (
LTX23HQSamplingParams,
LTX23SamplingParams,
)
+from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
from sglang.multimodal_gen.configs.sample.mova import (
MOVA_360P_SamplingParams,
MOVA_720P_SamplingParams,
@@ -826,6 +828,18 @@ def _register_configs():
lambda hf_id: "mova" in hf_id.lower() and "720p" in hf_id.lower()
],
)
+ register_configs(
+ sampling_param_cls=MiniMaxH3SamplingParams,
+ pipeline_config_cls=MiniMaxH3PipelineConfig,
+ hf_model_paths=[
+ "MiniMaxAI/MiniMax-H3",
+ "MiniMax/MiniMax-H3",
+ ],
+ model_detectors=[
+ lambda model_id: "minimaxh3"
+ in model_id.lower().replace("-", "").replace("_", "")
+ ],
+ )
# FLUX
register_configs(
sampling_param_cls=FluxSamplingParams,
diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/minimax_h3.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/minimax_h3.py
new file mode 100644
index 000000000..240bc763f
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/minimax_h3.py
@@ -0,0 +1,163 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0
+# ==============================================================================
+"""MiniMax H3 breakable CUDA graph packed-prompt padding."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+ MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT,
+)
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
+ prompt_padding as bcg_utils,
+)
+
+
+def is_minimax_h3_transformer(current_model: Any, call_kwargs: dict) -> bool:
+ return (
+ bcg_utils.transformer_class_name_matches(current_model, "minimaxh3")
+ and "prompt_embeds" in call_kwargs
+ and "packed_seq_params" in call_kwargs
+ and "refiner_packed_seq_params" in call_kwargs
+ and "text_pos_info" in call_kwargs
+ )
+
+
+def _position_ids(info: Any) -> torch.Tensor | None:
+ if isinstance(info, dict):
+ ids = info.get("position_ids")
+ else:
+ ids = getattr(info, "position_ids", None)
+ return ids if torch.is_tensor(ids) else None
+
+
+def _replace_position_ids(info: Any, ids: torch.Tensor) -> dict[str, Any]:
+ if isinstance(info, dict):
+ return {**info, "position_ids": ids}
+ # H3 currently passes dictionaries. Avoid mutating an unknown request
+ # object if an alternate frontend supplies one.
+ return {"position_ids": ids}
+
+
+def _replace_psp(
+ psp: Any,
+ *,
+ cu_seqlens_q: torch.Tensor,
+ max_seqlen_q: int,
+) -> dict[str, Any]:
+ if isinstance(psp, dict):
+ return {
+ **psp,
+ "cu_seqlens_q": cu_seqlens_q,
+ "max_seqlen_q": max_seqlen_q,
+ }
+ return {
+ "cu_seqlens_q": cu_seqlens_q,
+ "max_seqlen_q": max_seqlen_q,
+ }
+
+
+def _aligned(value: int) -> int:
+ alignment = MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
+ return (value + alignment - 1) // alignment * alignment
+
+
+def pad_minimax_h3_prompt_kwargs(
+ call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
+) -> dict:
+ prompt = bcg_utils.first_tensor(call_kwargs.get("prompt_embeds"))
+ text_pos = _position_ids(call_kwargs.get("text_pos_info"))
+ img_pos = _position_ids(call_kwargs.get("img_pos_info"))
+ audio_pos = _position_ids(call_kwargs.get("audio_pos_info"))
+ x = call_kwargs.get("x")
+ if (
+ not torch.is_tensor(prompt)
+ or prompt.dim() < 2
+ or not torch.is_tensor(text_pos)
+ or not torch.is_tensor(img_pos)
+ or not torch.is_tensor(audio_pos)
+ or not torch.is_tensor(x)
+ or x.dim() != 3
+ ):
+ return call_kwargs
+
+ text_len = int(prompt.shape[0])
+ bucket = bcg_utils.select_text_bucket(text_len, buckets)
+ if bucket is None:
+ return call_kwargs
+
+ # All used H3 rows are disjoint text, image/video, or audio rows. Derive
+ # the used/media lengths from tensor shapes so padding itself does not
+ # perform a GPU-to-host .item() synchronization on every denoising step.
+ media_rows = int(img_pos.numel()) + int(audio_pos.numel())
+ used = text_len + media_rows
+ source_seq = int(x.shape[1])
+ if source_seq < _aligned(used):
+ return call_kwargs
+
+ out = dict(call_kwargs)
+ # request-local row lists have prompt-dependent shapes, so keep them out
+ # of bucketed BCG signatures and rebuild the rows in the eager break
+ out.pop("local_embedding_layout", None)
+ # Request-static H3 denoising normally carries the live refined-text
+ # length as a host integer to avoid a per-step device scalar read. Host
+ # integers are baked into BCG signatures, however, so different prompt
+ # lengths would miss the same text bucket. Make only the BCG-padded copy a
+ # scalar tensor; the eager embedding break reads its updated replay value.
+ refined_len = out.get("refined_prompt_embeds_length")
+ if refined_len is not None and not torch.is_tensor(refined_len):
+ out["refined_prompt_embeds_length"] = torch.tensor(
+ int(refined_len),
+ dtype=torch.int64,
+ device=prompt.device,
+ )
+ if text_len < bucket:
+ out["prompt_embeds"] = bcg_utils.pad_tensor_dim(prompt, dim=0, target=bucket)
+ # These rows exist only to stabilize the BCG input signature. The
+ # model's eager embedding break trims prompt/text_pos/refiner metadata
+ # back to ``text_len`` before any projection or attention, so their
+ # values never enter the model.
+ dummy_text_pos = torch.arange(
+ used,
+ used + (bucket - text_len),
+ dtype=text_pos.dtype,
+ device=text_pos.device,
+ )
+ out["text_pos_info"] = _replace_position_ids(
+ out["text_pos_info"], torch.cat((text_pos.view(-1), dummy_text_pos))
+ )
+
+ # Do not grow the main packed sequence to media_rows + bucket. Changing
+ # the SP row partition changes GEMM shapes and is measurably non-bitwise
+ # even though dummy rows live in an independent attention segment. A
+ # capture is therefore reusable only inside the request's existing
+ # 64-row packed-sequence alignment group; other groups safely miss the
+ # signature and run eager.
+ packed_cu = torch.tensor([0, used, source_seq], dtype=torch.int32, device=x.device)
+ out["packed_seq_params"] = _replace_psp(
+ out["packed_seq_params"],
+ cu_seqlens_q=packed_cu,
+ # FA accepts an upper bound; keeping this bucket-stable is required
+ # because non-tensor values are baked into the BCG signature.
+ max_seqlen_q=source_seq,
+ )
+ refiner_cu = torch.tensor(
+ [0, text_len, bucket],
+ dtype=torch.int32,
+ device=prompt.device,
+ )
+ out["refiner_packed_seq_params"] = _replace_psp(
+ out["refiner_packed_seq_params"],
+ cu_seqlens_q=refiner_cu,
+ max_seqlen_q=bucket,
+ )
+ return out
+
+
+bcg_utils.register_prompt_padder(
+ is_minimax_h3_transformer, pad_minimax_h3_prompt_kwargs
+)
diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py
index cc9fc0dc9..d0ea0d442 100644
--- a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py
+++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py
@@ -301,6 +301,7 @@ def _ensure_model_padders_registered() -> None:
_model_padders_registered = True
from sglang.multimodal_gen.runtime.breakable_cuda_graph.model_padders import ( # noqa: F401
ideogram,
+ minimax_h3,
qwen_image,
zimage,
)
diff --git a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py
index 1ea75e358..819b8b9da 100644
--- a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py
+++ b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py
@@ -38,6 +38,20 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_gro
_original_similarity = None
+def disable_cache_on_transformer(transformer: torch.nn.Module) -> torch.nn.Module:
+ """Remove Cache-DiT hooks so subsequent requests use the native forward."""
+
+ logger.info("Disabling cache-dit on %s", type(transformer).__name__)
+ target = getattr(transformer, "_sglang_cache_dit_adapter", transformer)
+ cache_dit.disable_cache(target)
+ if target is not transformer:
+ del transformer._sglang_cache_dit_adapter
+ for name in ("_is_parallelized", "_parallelism_config"):
+ if hasattr(transformer, name):
+ delattr(transformer, name)
+ return transformer
+
+
def _patch_cache_dit_similarity():
from cache_dit.caching.cache_contexts import cache_manager
@@ -268,6 +282,7 @@ DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS: dict[str, DualTransformerBlockAdapterSpec]
_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern]] = {
"ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3),
"Krea2Transformer2DModel": ("transformer_blocks", ForwardPattern.Pattern_3),
+ "MiniMaxH3DiTModel": ("blocks", ForwardPattern.Pattern_3),
}
@@ -412,6 +427,8 @@ def enable_cache_on_transformer(
calibrator_config=calibrator_config,
parallelism_config=None,
)
+ if custom_adapter is not None:
+ transformer._sglang_cache_dit_adapter = custom_adapter
if parallelism_config is not None:
context_manager = getattr(transformer, "_context_manager", None)
diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py
index 1c06560e8..70fee44be 100644
--- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py
+++ b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/ipc_a2a.py
@@ -257,8 +257,17 @@ IPC_A2A = IpcA2AState()
def ipc_a2a_ready(group) -> bool:
"""True when the IPC transport is enabled and initialized (initializes
lazily on the first eager call; never inside a graph capture)."""
+ from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+ from sglang.multimodal_gen.runtime.platforms import current_platform
+
if not envs.SGLANG_DIFFUSION_IPC_A2A or IPC_A2A.failed:
return False
+ # TP+Ulysses groups are strided in global-rank order, while this transport
+ # supports the adjacent two-device topology used by TP1+U2. Reject the
+ # transport consistently before lazy initialization so no rank enters IPC
+ # while its peer falls back to NCCL.
+ if not current_platform.is_cuda() or get_tp_world_size() > 1:
+ return False
if IPC_A2A.inited:
return True
if torch.cuda.is_current_stream_capturing():
diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
index cd6e9504c..b296de944 100644
--- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
+++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
@@ -220,6 +220,7 @@ class DiffGenerator:
)
request_groups: list[list[Req]] = []
+ parent_requests: list[tuple[Req, int]] = []
image_paths_per_prompt = self._resolve_image_paths_per_prompt(
prompts, sampling_params_orig.image_path
)
@@ -243,13 +244,32 @@ class DiffGenerator:
sampling_params=sampling_params,
external_trace_header=external_trace_header,
)
- request_groups.append(
- expand_request_outputs(
- req,
- num_prompts=len(prompts),
- prompt_index=i,
+ parent_requests.append((req, i))
+
+ for req, prompt_index in parent_requests:
+ sampling_params = req.sampling_params
+ try:
+ if sampling_params.data_type == DataType.VIDEO:
+ sampling_params.prepare_video_request_for_queue(req)
+ request_groups.append(
+ expand_request_outputs(
+ req,
+ num_prompts=len(prompts),
+ prompt_index=prompt_index,
+ )
)
- )
+ except Exception:
+ if sampling_params.data_type == DataType.VIDEO:
+ sampling_params.cleanup_video_request(req)
+ for prepared_requests in request_groups:
+ if (
+ prepared_requests
+ and prepared_requests[0].data_type == DataType.VIDEO
+ ):
+ prepared_requests[0].sampling_params.cleanup_video_request(
+ prepared_requests[0]
+ )
+ raise
results: list[GenerationResult] = []
total_start_time = time.perf_counter()
@@ -285,6 +305,10 @@ class DiffGenerator:
)
for idx, path in enumerate(output_file_paths):
req = requests[idx]
+ if req.data_type == DataType.VIDEO:
+ req.sampling_params.validate_video_final_outputs(
+ [path], req
+ )
results.append(
GenerationResult(
**self._result_common(
@@ -346,6 +370,11 @@ class DiffGenerator:
for idx in range(len(samples_out)):
req = requests[idx]
+ output_file_path = req.output_file_path(1, 0)
+ if req.data_type == DataType.VIDEO and req.save_output:
+ req.sampling_params.validate_video_final_outputs(
+ [output_file_path], req
+ )
results.append(
GenerationResult(
**self._result_common(
@@ -355,12 +384,23 @@ class DiffGenerator:
frames=frames_out[idx],
audio=audios_out[idx],
prompt_index=global_output_index + idx,
- output_file_path=req.output_file_path(1, 0),
+ output_file_path=output_file_path,
)
)
except Exception as e:
logger.error("Generation failed: %s", e, exc_info=True)
finally:
+ if requests and requests[0].data_type == DataType.VIDEO:
+ try:
+ # Pre-queue resources are shared by the shallow
+ # per-output Req copies, so one idempotent cleanup is
+ # sufficient for the whole parent request.
+ requests[0].sampling_params.cleanup_video_request(requests[0])
+ except Exception:
+ logger.warning(
+ "Failed to clean up model-owned video request resources",
+ exc_info=True,
+ )
global_output_index += len(requests)
total_gen_time = time.perf_counter() - total_start_time
diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py
index c5bbc77da..0517ee733 100644
--- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py
+++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py
@@ -134,6 +134,9 @@ class VideoGenerationsRequest(BaseModel):
diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend
# Performance profiling
perf_dump_path: Optional[str] = None
+ profile: Optional[bool] = False
+ num_profiled_timesteps: Optional[int] = None
+ profile_all_stages: Optional[bool] = False
class VideoListResponse(BaseModel):
diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
index 23ea98e41..880fe0a8d 100644
--- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
+++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
@@ -341,10 +341,14 @@ async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
async def process_generation_batch(
scheduler_client: AsyncSchedulerClient,
batch,
+ *,
+ scheduler_batches=None,
) -> tuple[list[str], OutputBatch]:
total_start_time = time.perf_counter()
with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
- result = await scheduler_client.forward([batch])
+ result = await scheduler_client.forward(
+ scheduler_batches if scheduler_batches is not None else [batch]
+ )
if (
result.output is None
diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
index eb5d1bcd4..639833bcc 100644
--- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
+++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
@@ -83,6 +83,120 @@ def _parse_form_extra_value(value: Any) -> Any:
return value
+_MULTIPART_EXTRA_FORM_FIELDS = (
+ "use_duration_template",
+ "use_resolution_template",
+ "use_system_prompt",
+ "use_guardrails",
+ "guardrails",
+ "video_path",
+ "video_url",
+ "generate_sound",
+ "sound_duration",
+ "condition_frame_indexes",
+ "action_mode",
+ "domain_id",
+ "domain_name",
+ "raw_action_dim",
+ "action_fps",
+ "action",
+ "action_view_point",
+ "action_normalization",
+ "condition_frame_indexes_vision",
+ "condition_video_keep",
+)
+
+
+def _video_sampling_params_cls(server_args) -> type[SamplingParams]:
+ """Resolve the params type selected for the current server."""
+
+ sampling_params_cls = SamplingParams
+ if server_args.pipeline_class_name:
+ from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+ config_classes = get_pipeline_config_classes(server_args.pipeline_class_name)
+ if config_classes is not None:
+ _, sampling_params_cls = config_classes
+ if sampling_params_cls is SamplingParams:
+ from sglang.multimodal_gen.registry import get_model_info
+
+ model_info = get_model_info(
+ server_args.model_path,
+ backend=server_args.backend,
+ model_id=server_args.model_id,
+ )
+ if model_info is not None:
+ sampling_params_cls = model_info.sampling_param_cls
+ return sampling_params_cls
+
+
+def _multipart_extra_form_keys(
+ sampling_params_cls: type[SamplingParams],
+) -> tuple[str, ...]:
+ return tuple(
+ dict.fromkeys(
+ (
+ *VideoGenerationsRequest.model_fields,
+ *_MULTIPART_EXTRA_FORM_FIELDS,
+ *sorted(sampling_params_cls.video_request_extra_fields()),
+ )
+ )
+ )
+
+
+def _filter_multipart_declared_fields(
+ extra_from_form: Dict[str, Any],
+ sampling_params_cls: type[SamplingParams],
+) -> Dict[str, Any]:
+ declared = set(_multipart_extra_form_keys(sampling_params_cls))
+ return {key: value for key, value in extra_from_form.items() if key in declared}
+
+
+def _merge_multipart_extra_form_fields(
+ raw_form: Any,
+ extra_from_form: Dict[str, Any],
+ sampling_params_cls: type[SamplingParams],
+) -> None:
+ for key in _multipart_extra_form_keys(sampling_params_cls):
+ if key in raw_form and key not in extra_from_form:
+ extra_from_form[key] = _parse_form_extra_value(raw_form[key])
+
+
+def _multipart_video_extras(
+ raw_form: Any,
+ *,
+ extra_body: Any,
+ extra_params: Any,
+ sampling_params_cls: type[SamplingParams],
+) -> Dict[str, Any]:
+ """Build and validate multipart extras once for request construction."""
+
+ extra_from_form: Dict[str, Any] = {}
+ if extra_body:
+ try:
+ extra_from_form = flatten_extra_params(json.loads(extra_body))
+ except (json.JSONDecodeError, ValueError, TypeError) as exc:
+ raise HTTPException(
+ status_code=400, detail="extra_body is not valid JSON"
+ ) from exc
+ if extra_params:
+ try:
+ extra_from_form.update(
+ flatten_extra_params({"extra_params": json.loads(extra_params)})
+ )
+ except (json.JSONDecodeError, ValueError, TypeError) as exc:
+ raise HTTPException(
+ status_code=400, detail="extra_params is not valid JSON"
+ ) from exc
+ _merge_multipart_extra_form_fields(
+ raw_form,
+ extra_from_form,
+ sampling_params_cls,
+ )
+ flatten_extra_params(extra_from_form)
+ return _filter_multipart_declared_fields(extra_from_form, sampling_params_cls)
+
+
def _is_probably_video_source(source: Any) -> bool:
content_type = (getattr(source, "content_type", "") or "").lower()
if content_type.startswith("video/"):
@@ -224,45 +338,52 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
server_args.pipeline_config.action_stats_path
)
- return build_sampling_params(
- request_id,
- prompt=request.prompt,
- num_outputs_per_prompt=max(1, min(int(num_outputs), 10)),
- size=request.size,
- width=request.width,
- height=request.height,
- num_frames=num_frames,
- fps=fps,
- image_path=image_path,
- video_path=video_path,
- output_file_name=request_id,
- seed=request.seed,
- generator_device=request.generator_device,
- num_inference_steps=request.num_inference_steps,
- guidance_scale=request.guidance_scale,
- guidance_scale_2=request.guidance_scale_2,
- negative_prompt=request.negative_prompt,
- max_sequence_length=request.max_sequence_length,
- flow_shift=request.flow_shift,
- use_duration_template=_extra_value(request, "use_duration_template"),
- use_resolution_template=_extra_value(request, "use_resolution_template"),
- use_system_prompt=_extra_value(request, "use_system_prompt"),
- use_guardrails=_extra_value(request, "use_guardrails"),
- enable_teacache=request.enable_teacache,
- enable_frame_interpolation=request.enable_frame_interpolation,
- frame_interpolation_exp=request.frame_interpolation_exp,
- frame_interpolation_scale=request.frame_interpolation_scale,
- frame_interpolation_model_path=request.frame_interpolation_model_path,
- enable_upscaling=request.enable_upscaling,
- upscaling_model_path=request.upscaling_model_path,
- upscaling_scale=request.upscaling_scale,
- output_path=request.output_path,
- output_compression=request.output_compression,
- output_quality=request.output_quality,
- perf_dump_path=request.perf_dump_path,
- diffusers_kwargs=request.diffusers_kwargs,
+ kwargs = {
+ "prompt": request.prompt,
+ "num_outputs_per_prompt": max(1, min(int(num_outputs), 10)),
+ "size": request.size,
+ "width": request.width,
+ "height": request.height,
+ "num_frames": num_frames,
+ "fps": fps,
+ "image_path": image_path,
+ "video_path": video_path,
+ "output_file_name": request_id,
+ "seed": request.seed,
+ "generator_device": request.generator_device,
+ "num_inference_steps": request.num_inference_steps,
+ "guidance_scale": request.guidance_scale,
+ "guidance_scale_2": request.guidance_scale_2,
+ "true_cfg_scale": request.true_cfg_scale,
+ "negative_prompt": request.negative_prompt,
+ "max_sequence_length": request.max_sequence_length,
+ "flow_shift": request.flow_shift,
+ "use_duration_template": _extra_value(request, "use_duration_template"),
+ "use_resolution_template": _extra_value(request, "use_resolution_template"),
+ "use_system_prompt": _extra_value(request, "use_system_prompt"),
+ "use_guardrails": _extra_value(request, "use_guardrails"),
+ "enable_teacache": request.enable_teacache,
+ "enable_frame_interpolation": request.enable_frame_interpolation,
+ "frame_interpolation_exp": request.frame_interpolation_exp,
+ "frame_interpolation_scale": request.frame_interpolation_scale,
+ "frame_interpolation_model_path": request.frame_interpolation_model_path,
+ "enable_upscaling": request.enable_upscaling,
+ "upscaling_model_path": request.upscaling_model_path,
+ "upscaling_scale": request.upscaling_scale,
+ "output_path": request.output_path,
+ "output_compression": request.output_compression,
+ "output_quality": request.output_quality,
+ "perf_dump_path": request.perf_dump_path,
+ "profile": request.profile,
+ "num_profiled_timesteps": request.num_profiled_timesteps,
+ "profile_all_stages": request.profile_all_stages,
+ "diffusers_kwargs": request.diffusers_kwargs,
**cosmos3_kwargs,
- )
+ }
+
+ sampling_params_cls = _video_sampling_params_cls(server_args)
+ kwargs = sampling_params_cls.lower_video_request_kwargs(request, kwargs)
+ return build_sampling_params(request_id, **kwargs)
# extract metadata which http_server needs to know
@@ -311,6 +432,7 @@ async def _dispatch_job_async(
job_id: str,
batch: Req,
*,
+ scheduler_batches: list[Req] | None = None,
temp_dirs: list[str] | None = None,
output_persistent: bool = True,
) -> None:
@@ -318,9 +440,30 @@ async def _dispatch_job_async(
try:
save_file_path_list, result = await process_generation_batch(
- async_scheduler_client, batch
+ async_scheduler_client,
+ batch,
+ scheduler_batches=scheduler_batches,
)
save_file_path = save_file_path_list[0]
+ try:
+ final_media_fields = await asyncio.to_thread(
+ batch.sampling_params.validate_video_final_outputs,
+ save_file_path_list,
+ batch,
+ )
+ except Exception:
+ for output_path in save_file_path_list:
+ try:
+ os.remove(output_path)
+ except FileNotFoundError:
+ pass
+ except OSError:
+ logger.warning(
+ "Failed to remove rejected video output %s",
+ output_path,
+ exc_info=True,
+ )
+ raise
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
@@ -343,13 +486,32 @@ async def _dispatch_job_async(
update_fields = add_common_data_to_response(
update_fields, request_id=job_id, result=result
)
+ update_fields.update(final_media_fields)
await VIDEO_STORE.update_fields(job_id, update_fields)
except Exception as e:
logger.error(f"{e}")
await VIDEO_STORE.update_fields(
- job_id, {"status": "failed", "error": {"message": str(e)}}
+ job_id,
+ {
+ "status": "failed",
+ "error": {"message": str(e)},
+ "url": None,
+ "file_path": None,
+ "file_paths": None,
+ "num_outputs": None,
+ },
)
finally:
+ try:
+ await asyncio.to_thread(
+ batch.sampling_params.cleanup_video_request,
+ batch,
+ )
+ except Exception:
+ logger.warning(
+ "Failed to clean up model-owned video request resources",
+ exc_info=True,
+ )
for td in temp_dirs or []:
shutil.rmtree(td, ignore_errors=True)
@@ -376,6 +538,8 @@ async def create_video(
generator_device: Optional[str] = Form("cuda"),
negative_prompt: Optional[str] = Form(None),
guidance_scale: Optional[float] = Form(None),
+ guidance_scale_2: Optional[float] = Form(None),
+ true_cfg_scale: Optional[float] = Form(None),
num_inference_steps: Optional[int] = Form(None),
max_sequence_length: Optional[int] = Form(None),
flow_shift: Optional[float] = Form(None),
@@ -398,6 +562,22 @@ async def create_video(
server_args = get_global_server_args()
task_type = server_args.pipeline_config.task_type
+ is_multipart = "multipart/form-data" in content_type
+ raw_form: Any = None
+ extra_from_form: Dict[str, Any] = {}
+
+ # Parse model-specific multipart metadata before creating request-owned
+ # directories or saving uploads, so malformed JSON leaves no resources.
+ if is_multipart:
+ if not prompt:
+ raise HTTPException(status_code=400, detail="prompt is required")
+ raw_form = await request.form()
+ extra_from_form = _multipart_video_extras(
+ raw_form,
+ extra_body=extra_body,
+ extra_params=extra_params,
+ sampling_params_cls=_video_sampling_params_cls(server_args),
+ )
# Resolve input upload directory (may be a temp dir when saving is disabled)
temp_dirs: list[str] = []
@@ -411,14 +591,11 @@ async def create_video(
# Resolve output directory
effective_output_path = server_args.output_path
output_persistent = True
- if "multipart/form-data" not in content_type:
+ if not is_multipart:
# JSON body may carry a per-request output_path; checked after parsing below
pass
- if "multipart/form-data" in content_type:
- if not prompt:
- raise HTTPException(status_code=400, detail="prompt is required")
-
+ if is_multipart:
video_input_path = None
image_sources = merge_image_input_list(input_reference, reference_url)
if video_reference is not None:
@@ -462,52 +639,10 @@ async def create_video(
status_code=400, detail=f"Failed to process image source: {str(e)}"
)
- # Parse extra_body JSON (if provided in multipart form) to get fps/num_frames overrides
- extra_from_form: Dict[str, Any] = {}
- if extra_body:
- try:
- extra_from_form = flatten_extra_params(json.loads(extra_body))
- except Exception:
- extra_from_form = {}
- if extra_params:
- try:
- extra_from_form.update(
- flatten_extra_params({"extra_params": json.loads(extra_params)})
- )
- except Exception:
- pass
-
def form_value(name: str, value: Any) -> Any:
selected = value if value is not None else extra_from_form.get(name)
return _parse_form_extra_value(selected)
- raw_form = await request.form()
- for key in (
- "use_duration_template",
- "use_resolution_template",
- "use_system_prompt",
- "use_guardrails",
- "guardrails",
- "video_path",
- "video_url",
- "generate_sound",
- "sound_duration",
- "condition_frame_indexes",
- "action_mode",
- "domain_id",
- "domain_name",
- "raw_action_dim",
- "action_fps",
- "action",
- "action_view_point",
- "action_normalization",
- "condition_frame_indexes_vision",
- "condition_video_keep",
- ):
- if key in raw_form and key not in extra_from_form:
- extra_from_form[key] = _parse_form_extra_value(raw_form[key])
- flatten_extra_params(extra_from_form)
-
request_field_names = set(VideoGenerationsRequest.model_fields)
extra_request_fields = {
key: value
@@ -536,6 +671,8 @@ async def create_video(
negative_prompt=form_value("negative_prompt", negative_prompt),
num_inference_steps=form_value("num_inference_steps", num_inference_steps),
guidance_scale=form_value("guidance_scale", guidance_scale),
+ guidance_scale_2=form_value("guidance_scale_2", guidance_scale_2),
+ true_cfg_scale=form_value("true_cfg_scale", true_cfg_scale),
max_sequence_length=form_value("max_sequence_length", max_sequence_length),
flow_shift=form_value("flow_shift", flow_shift),
enable_teacache=form_value("enable_teacache", enable_teacache),
@@ -639,30 +776,59 @@ async def create_video(
try:
sampling_params = _build_video_sampling_params(request_id, req)
except (ValueError, TypeError) as e:
+ for td in temp_dirs:
+ shutil.rmtree(td, ignore_errors=True)
raise HTTPException(status_code=400, detail=str(e))
- job = _video_job_from_sampling(request_id, req, sampling_params)
- await VIDEO_STORE.upsert(request_id, job)
+ batch: Req | None = None
+ scheduler_batches: list[Req] | None = None
+ try:
+ # Build Req for scheduler.
+ trace_headers = extract_trace_headers(request.headers)
+ batch = prepare_request(
+ server_args=server_args,
+ sampling_params=sampling_params,
+ external_trace_header=trace_headers,
+ )
+ # Add diffusers_kwargs if provided.
+ if req.diffusers_kwargs:
+ batch.extra["diffusers_kwargs"] = req.diffusers_kwargs
+ if "max_sequence_length" in req.diffusers_kwargs:
+ batch.max_sequence_length = req.diffusers_kwargs["max_sequence_length"]
+ if "flow_shift" in req.diffusers_kwargs:
+ batch.flow_shift = req.diffusers_kwargs["flow_shift"]
+ await asyncio.to_thread(
+ sampling_params.prepare_video_request_for_queue,
+ batch,
+ )
+ scheduler_batches = sampling_params.expand_video_request_outputs_for_queue(
+ batch
+ )
+ job = _video_job_from_sampling(request_id, req, sampling_params)
+ job.update(sampling_params.project_video_queued_job_fields(batch))
+ await VIDEO_STORE.upsert(request_id, job)
+ except Exception as e:
+ if batch is not None:
+ try:
+ await asyncio.to_thread(sampling_params.cleanup_video_request, batch)
+ except Exception:
+ logger.warning(
+ "Failed to clean up rejected video request resources",
+ exc_info=True,
+ )
+ for td in temp_dirs:
+ shutil.rmtree(td, ignore_errors=True)
+ if isinstance(e, (TypeError, ValueError)):
+ raise HTTPException(status_code=400, detail=str(e)) from e
+ raise
- # Build Req for scheduler
- trace_headers = extract_trace_headers(request.headers)
- batch = prepare_request(
- server_args=server_args,
- sampling_params=sampling_params,
- external_trace_header=trace_headers,
- )
- # Add diffusers_kwargs if provided
- if req.diffusers_kwargs:
- batch.extra["diffusers_kwargs"] = req.diffusers_kwargs
- if "max_sequence_length" in req.diffusers_kwargs:
- batch.max_sequence_length = req.diffusers_kwargs["max_sequence_length"]
- if "flow_shift" in req.diffusers_kwargs:
- batch.flow_shift = req.diffusers_kwargs["flow_shift"]
+ assert batch is not None
# Enqueue the job asynchronously and return immediately
asyncio.create_task(
_dispatch_job_async(
request_id,
batch,
+ scheduler_batches=scheduler_batches,
temp_dirs=temp_dirs or None,
output_persistent=output_persistent,
)
@@ -717,6 +883,21 @@ async def delete_video(video_id: str = Path(...)):
return VideoResponse(**job)
+def _select_video_variant_path(job: dict, variant: str | None) -> str | None:
+ file_paths = job.get("file_paths")
+ if file_paths:
+ try:
+ variant_index = 0 if variant is None else int(variant)
+ except (TypeError, ValueError):
+ return None
+ if 0 <= variant_index < len(file_paths):
+ return file_paths[variant_index]
+ return None
+ if variant not in (None, "0", 0):
+ return None
+ return job.get("file_path")
+
+
@router.get("/{video_id}/content")
async def download_video_content(
video_id: str = Path(...), variant: Optional[str] = Query(None)
@@ -731,9 +912,13 @@ async def download_video_content(
detail=f"Video has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}",
)
- file_path = job.get("file_path")
- if not file_path or not os.path.exists(file_path):
+ file_path = _select_video_variant_path(job, variant)
+ if job.get("status") not in {"completed", "failed"}:
raise HTTPException(status_code=404, detail="Generation is still in-progress")
+ if not file_path or not os.path.exists(file_path):
+ raise HTTPException(
+ status_code=404, detail=f"Video variant {variant} not found"
+ )
media_type = "video/mp4" # default variant
return FileResponse(
diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py
index 2c5b5839c..1494f5847 100644
--- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py
+++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py
@@ -341,6 +341,7 @@ def expand_request_outputs(
req.seed = seeds[0]
req.seeds = None
req.generator = None
+ req.sampling_params.refresh_request_extra_after_output_expansion(req)
return [req]
expanded: list[Req] = []
@@ -365,6 +366,9 @@ def expand_request_outputs(
output_req.output_file_name = _with_output_index_suffix(
req.output_file_name, output_index
)
+ output_req.sampling_params.refresh_request_extra_after_output_expansion(
+ output_req
+ )
output_req.validate()
expanded.append(output_req)
@@ -487,7 +491,7 @@ def _try_save_cuda_video_direct(
if video.shape[0] != 3:
return False
- frames = (video * 255).clamp(0, 255).to(torch.uint8)
+ frames = (video * 255).clamp_(0, 255).to(torch.uint8)
frames = frames.permute(1, 2, 3, 0).contiguous()
num_frames, height, width, _ = frames.shape
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
index 2b7b8eea5..52210c07f 100755
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
@@ -2,6 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
+import importlib
import logging
import os
@@ -15,7 +16,10 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
AttentionMetadataBuilder,
)
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
-from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER_GFX95
+from sglang.multimodal_gen.runtime.platforms.aiter import (
+ USE_AITER_GFX95,
+ USE_AITER_GFX942,
+)
logger = logging.getLogger(__name__)
@@ -205,3 +209,38 @@ class AITerImpl(AttentionImpl):
return_lse=True,
)
return output
+
+ @torch.compiler.disable
+ def forward_varlen(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ ) -> torch.Tensor:
+ del cu_seqlens_host
+ if USE_AITER_GFX942:
+ # The grouped-varlen ASM kernel hangs on H3's ~64K packed
+ # sequences on gfx942; AITER's Triton path handles this shape.
+ attention_func = importlib.import_module(
+ "aiter.ops.triton.attention.mha"
+ ).flash_attn_varlen_func
+ else:
+ attention_func = aiter.flash_attn_varlen_func
+
+ cu_seqlens = cu_seqlens.to(device=query.device, dtype=torch.int32).contiguous()
+ output = attention_func(
+ q=query.contiguous(),
+ k=key.contiguous(),
+ v=value.contiguous(),
+ cu_seqlens_q=cu_seqlens,
+ cu_seqlens_k=cu_seqlens,
+ max_seqlen_q=max_seqlen,
+ max_seqlen_k=max_seqlen,
+ softmax_scale=self.softmax_scale,
+ causal=self.causal,
+ )
+ return output[0] if isinstance(output, tuple) else output
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
index b016f0f81..e0184a0e4 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
@@ -170,6 +170,20 @@ class AttentionImpl(ABC, Generic[T]):
) -> torch.Tensor:
raise NotImplementedError
+ def forward_varlen(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ ) -> torch.Tensor:
+ raise NotImplementedError(
+ f"{type(self).__name__} does not implement packed varlen attention"
+ )
+
def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
return wrap_method_with_debug_kernel_once(
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
index 8f0cfb2f6..21f711852 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
@@ -443,3 +443,28 @@ class FlashAttentionImpl(AttentionImpl):
return out_tensor
raise ValueError(f"flash attention version {fa_ver} is not supported.")
+
+ def forward_varlen(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ ) -> torch.Tensor:
+ del cu_seqlens_host
+ output = flash_attn_varlen_func(
+ query,
+ key,
+ value,
+ cu_seqlens_q=cu_seqlens,
+ cu_seqlens_k=cu_seqlens,
+ max_seqlen_q=max_seqlen,
+ max_seqlen_k=max_seqlen,
+ softmax_scale=self.softmax_scale,
+ causal=self.causal,
+ ver=fa_ver,
+ )
+ return output[0] if isinstance(output, tuple) else output
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
index c117f43f7..901446c0c 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
@@ -94,6 +94,35 @@ class SDPAImpl(AttentionImpl):
output = output.transpose(1, 2)
return output
+ def forward_varlen(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ ) -> torch.Tensor:
+ del max_seqlen
+ bounds = (
+ cu_seqlens_host
+ if cu_seqlens_host is not None
+ else tuple(int(item) for item in cu_seqlens.tolist())
+ )
+ output = torch.empty_like(query)
+ for start, stop in zip(bounds[:-1], bounds[1:]):
+ if start == stop:
+ continue
+ segment = self.forward(
+ query[start:stop].unsqueeze(0),
+ key[start:stop].unsqueeze(0),
+ value[start:stop].unsqueeze(0),
+ None,
+ )
+ output[start:stop].copy_(segment[0])
+ return output
+
class CudnnSDPABackend(SDPABackend):
@staticmethod
@@ -122,7 +151,7 @@ class DynamicCudnnSDPABackend(SDPABackend):
return DynamicCudnnSDPAImpl
-class DynamicCudnnSDPAImpl(AttentionImpl):
+class DynamicCudnnSDPAImpl(SDPAImpl):
def __init__(
self,
num_heads: int,
diff --git a/python/sglang/multimodal_gen/runtime/layers/usp.py b/python/sglang/multimodal_gen/runtime/layers/usp.py
index bd81eb066..f01ba5341 100644
--- a/python/sglang/multimodal_gen/runtime/layers/usp.py
+++ b/python/sglang/multimodal_gen/runtime/layers/usp.py
@@ -8,6 +8,10 @@ import torch.distributed as dist
import torch.distributed._functional_collectives as ft_c
from torch.distributed.tensor.experimental._attention import _cp_options
+from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
+ pack_qkv_destination_major,
+)
+from sglang.kernels.ops.diffusion.usp_relayout import usp_merge_heads
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_sp_group,
get_ulysses_parallel_rank,
@@ -280,6 +284,47 @@ def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
return x
+def _usp_input_all_to_all_packed_qkv(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Exchange 3D Q/K/V with one destination-major Ulysses collective."""
+ world_size = get_ulysses_parallel_world_size()
+ if world_size <= 1:
+ return q, k, v
+
+ assert q.ndim == 3 and q.shape == k.shape == v.shape
+ s_local, h_global, head_size = q.shape
+ assert h_global % world_size == 0
+ h_local = h_global // world_size
+
+ if (
+ q.is_cuda
+ and q.dtype in (torch.float16, torch.bfloat16)
+ and q.dtype == k.dtype == v.dtype
+ and q.stride(-1) == k.stride(-1) == v.stride(-1) == 1
+ and not torch.compiler.is_compiling()
+ ):
+ packed = pack_qkv_destination_major(q, k, v, world_size)
+ else:
+ packed = torch.empty(
+ (world_size, s_local, h_local, 3 * head_size),
+ dtype=q.dtype,
+ device=q.device,
+ )
+ for index, tensor in enumerate((q, k, v)):
+ head_shards = tensor.view(s_local, world_size, h_local, head_size).permute(
+ 1, 0, 2, 3
+ )
+ packed[..., index * head_size : (index + 1) * head_size].copy_(head_shards)
+
+ packed = _usp_all_to_all_single(packed)
+ packed = packed.reshape(s_local * world_size, h_local, 3 * head_size)
+ q, k, v = packed.split(head_size, dim=-1)
+ return q, k, v
+
+
def _usp_input_all_to_all_varlen(
x: torch.Tensor, seq_lens: list[int], head_dim: int = 1
) -> torch.Tensor:
@@ -419,7 +464,7 @@ def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
x = x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, h_global, s_local, d)
else: # head_dim == 2
# Shape transition: [world_size, s_local, b, h_local, d] -> [b, s_local, world_size, h_local, d]
- x = x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, s_local, h_global, d)
+ x = usp_merge_heads(x).reshape(b, s_local, h_global, d)
return x
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
index dbec24882..bb4abc685 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
@@ -113,9 +113,12 @@ class ComponentLoader(ABC):
return {}
def should_raise_customized_load_error(
- self, _server_args: ServerArgs, _component_name: str
+ self, server_args: ServerArgs, component_name: str
) -> bool:
- return False
+ native_only_components = getattr(
+ server_args.pipeline_config, "native_only_components", ()
+ )
+ return component_name in native_only_components
@staticmethod
def _is_component_set_as_layerwise_load(
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
index 870a06744..5818f6b37 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
@@ -54,7 +54,6 @@ class ImageEncoderLoader(TextEncoderLoader):
finalize_encoder_folding(
encoder_config,
server_args.encoder_parallel,
- batched=server_args.batching_max_size > 1,
)
# Always start with local device; load_model will adjust for offload if needed
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
index 66877c392..8c9e85e81 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
@@ -2,7 +2,7 @@ import dataclasses
import glob
import os
import re
-from collections.abc import Generator, Iterable
+from collections.abc import Callable, Generator, Iterable
from contextlib import nullcontext
from typing import cast
@@ -39,6 +39,7 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import (
safetensors_weights_iterator,
)
from sglang.multimodal_gen.runtime.models.encoders.base import (
+ TextEncoder,
finalize_encoder_folding,
get_folding_tp_group,
)
@@ -184,6 +185,7 @@ class TextEncoderLoader(ComponentLoader):
model_name_or_path: str,
fall_back_to_pt: bool,
allow_patterns_overrides: list[str] | None,
+ key_filter: Callable[[str], bool] | None = None,
) -> tuple[str, list[str], bool]:
"""Prepare weights for the model.
@@ -216,7 +218,10 @@ class TextEncoderLoader(ComponentLoader):
if use_safetensors:
hf_weights_files = filter_duplicate_safetensors_files(
- hf_weights_files, hf_folder, index_file
+ hf_weights_files,
+ hf_folder,
+ index_file,
+ key_filter=key_filter,
)
else:
hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
@@ -237,20 +242,39 @@ class TextEncoderLoader(ComponentLoader):
self,
source: "Source",
to_cpu: bool,
+ key_filter: Callable[[str], bool] | None = None,
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""get an iterator for the model weights based on the load format."""
+ source_key_filter: Callable[[str], bool] | None
+ if key_filter is None:
+ source_key_filter = None
+ else:
+
+ def include_source_weight(name: str) -> bool:
+ return key_filter(source.prefix + name)
+
+ source_key_filter = include_source_weight
+
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
source.model_or_path,
source.fall_back_to_pt,
source.allow_patterns_overrides,
+ key_filter=source_key_filter,
)
if use_safetensors:
weights_iterator = safetensors_weights_iterator(
hf_weights_files,
to_cpu=to_cpu,
+ key_filter=source_key_filter,
)
else:
weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
+ if source_key_filter is not None:
+ weights_iterator = (
+ (name, tensor)
+ for name, tensor in weights_iterator
+ if source_key_filter(name)
+ )
# apply the prefix.
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
@@ -261,6 +285,10 @@ class TextEncoderLoader(ComponentLoader):
model_path: str,
to_cpu: bool,
) -> Generator[tuple[str, torch.Tensor], None, None]:
+ key_filter = cast(
+ Callable[[str], bool] | None,
+ getattr(model, "should_materialize_checkpoint_weight", None),
+ )
primary_weights = TextEncoderLoader.Source(
model_path,
prefix="",
@@ -270,6 +298,7 @@ class TextEncoderLoader(ComponentLoader):
yield from self._get_weights_iterator(
primary_weights,
to_cpu,
+ key_filter,
)
secondary_weights = cast(
@@ -280,6 +309,7 @@ class TextEncoderLoader(ComponentLoader):
yield from self._get_weights_iterator(
source,
to_cpu,
+ key_filter,
)
def load_customized(
@@ -314,11 +344,20 @@ class TextEncoderLoader(ComponentLoader):
)
if post_diffusers_config_update is not None:
post_diffusers_config_update()
+ model_cls, _ = ModelRegistry.resolve_model_cls(
+ getattr(encoder_config, "architectures", [])
+ )
# real dims are populated now; resolve fold vs replicate
finalize_encoder_folding(
encoder_config,
server_args.encoder_parallel,
- batched=server_args.batching_max_size > 1,
+ prefer_dp=(
+ server_args.batching_max_size > 1
+ and (server_args.tp_size or 1) == 1
+ and (server_args.dp_size or 1) == 1
+ and issubclass(model_cls, TextEncoder)
+ and model_cls.supports_dp_encode
+ ),
)
encoder_dtype = server_args.pipeline_config.text_encoder_precisions[
encoder_index
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
index c68271c3c..eca50e5f0 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
@@ -28,6 +28,20 @@ _is_npu = is_npu()
logger = init_logger(__name__)
+def _warn_if_expected_param_dtype_missing(
+ model: torch.nn.Module, expected_dtype: torch.dtype | None
+) -> None:
+ if expected_dtype is None:
+ return
+ param_dtypes = {param.dtype for param in model.parameters()}
+ if expected_dtype not in param_dtypes:
+ logger.warning(
+ "Model parameter dtypes do not include expected param dtype, %s vs %s",
+ param_dtypes,
+ expected_dtype,
+ )
+
+
def _server_args_for_transformer_component(
server_args: ServerArgs, component_name: str
) -> ServerArgs:
@@ -89,7 +103,8 @@ class TransformerLoader(ComponentLoader):
# Don't let a quantized load quietly fall back to the unquantized native
# model. That would drop the requested precision and bury the real error.
return (
- component_server_args.transformer_weights_path is not None
+ super().should_raise_customized_load_error(server_args, component_name)
+ or component_server_args.transformer_weights_path is not None
or component_server_args.quantization is not None
)
@@ -185,15 +200,6 @@ class TransformerLoader(ComponentLoader):
for post_load_hook in quant_spec.post_load_hooks:
post_load_hook(model)
- # considering the existent of mixed-precision models (e.g., nunchaku)
- if (
- next(model.parameters()).dtype != quant_spec.param_dtype
- and quant_spec.param_dtype
- ):
- logger.warning(
- "Model dtype does not match expected param dtype, %s vs %s",
- next(model.parameters()).dtype,
- quant_spec.param_dtype,
- )
+ _warn_if_expected_param_dtype_missing(model, quant_spec.param_dtype)
return model
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
index b97ed08f5..e59c87477 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
@@ -142,10 +142,12 @@ class VAELoader(ComponentLoader):
should_offload = self.should_offload(server_args)
target_device = self.target_device(should_offload)
- # Check for auto_map first (custom VAE classes)
+ native_only = component_name in getattr(
+ server_args.pipeline_config, "native_only_components", ()
+ )
auto_map = config.get("auto_map", {})
auto_model_map = auto_map.get("AutoModel")
- if auto_model_map:
+ if auto_model_map and not native_only:
module_path, cls_name = auto_model_map.rsplit(".", 1)
custom_module_file = os.path.join(component_model_path, f"{module_path}.py")
spec = importlib.util.spec_from_file_location("_custom", custom_module_file)
@@ -191,16 +193,18 @@ class VAELoader(ComponentLoader):
for sf_path in safetensors_list:
loaded.update(safetensors_load_file(sf_path))
_backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
- vae.load_state_dict(loaded, strict=False)
+ strict_load = native_only
+ vae.load_state_dict(loaded, strict=strict_load)
- state_keys = set(vae.state_dict().keys())
- loaded_keys = set(loaded.keys())
- missing_keys = sorted(state_keys - loaded_keys)
- unexpected_keys = sorted(loaded_keys - state_keys)
- if missing_keys:
- logger.warning("VAE missing keys: %s", missing_keys)
- if unexpected_keys:
- logger.warning("VAE unexpected keys: %s", unexpected_keys)
+ if not strict_load:
+ state_keys = set(vae.state_dict().keys())
+ loaded_keys = set(loaded.keys())
+ missing_keys = sorted(state_keys - loaded_keys)
+ unexpected_keys = sorted(loaded_keys - state_keys)
+ if missing_keys:
+ logger.warning("VAE missing keys: %s", missing_keys)
+ if unexpected_keys:
+ logger.warning("VAE unexpected keys: %s", unexpected_keys)
if _should_use_channels_last_3d(server_args, component_name):
n = _convert_conv3d_weights_to_channels_last_3d(vae)
diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
index afbfc5437..bd0e175cb 100644
--- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
@@ -20,6 +20,7 @@ from torch.distributed.fsdp import (
FSDPModule,
MixedPrecisionPolicy,
fully_shard,
+ register_fsdp_forward_method,
)
from torch.nn.modules.module import _IncompatibleKeys
@@ -204,7 +205,8 @@ def maybe_load_fsdp_model(
Args:
param_dtype: Data type for model parameters, also used for:
- Model initialization context (set_default_torch_dtype)
- - FSDP mixed precision policy
+ - FSDP mixed precision policy unless the model preserves mixed
+ original parameter dtypes
- Weight loading and casting
reduce_dtype: Data type for gradient reduction in FSDP mixed precision.
strict: If True, enforce strict state dict loading (all keys must match).
@@ -215,8 +217,19 @@ def maybe_load_fsdp_model(
# 1. prepare for loading
default_torch_dtype = param_dtype if param_dtype else torch.bfloat16
+ # Some native models deliberately mix FP32 projections with lower-precision
+ # blocks. FSDP must all-gather those parameters in their original dtypes;
+ # the thread-local compute dtype below remains the requested default.
+ fsdp_param_dtype = (
+ None
+ if fsdp_inference and getattr(model_cls, "_fsdp_mixed_dtype_params", False)
+ else default_torch_dtype
+ )
mp_policy = MixedPrecisionPolicy(
- default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
+ param_dtype=fsdp_param_dtype,
+ reduce_dtype=reduce_dtype,
+ output_dtype=output_dtype,
+ cast_forward_inputs=False,
)
set_mixed_precision_policy(
@@ -279,6 +292,8 @@ def maybe_load_fsdp_model(
fsdp_shard_conditions=getattr(model, "_fsdp_shard_conditions", None),
pin_cpu_memory=pin_cpu_memory,
)
+ if callable(getattr(model, "refine_prompt_embeds", None)):
+ register_fsdp_forward_method(model, "refine_prompt_embeds")
param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
index d501cd5ec..8128f7cd0 100644
--- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
@@ -601,7 +601,12 @@ def _resolve_quant_config(
# in source dtype and are quantized in
# process_weights_after_loading.
quant_cls = get_quantization_config(server_args.quantization)
- return quant_cls()
+ quant_kwargs = {}
+ if server_args.quantization in {"fp8", "mxfp4"}:
+ quant_kwargs["ignored_layers"] = getattr(
+ server_args, "quantization_ignored_layers", None
+ )
+ return quant_cls(**quant_kwargs)
quant_config = get_quant_config(hf_config, component_model_path)
if quant_config is None and server_args.transformer_weights_path:
diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py
index 2d172af6d..4d48eb1d4 100644
--- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py
+++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py
@@ -65,7 +65,10 @@ def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
# So, we use the index_file to
# look up which safetensors files should be used.
def filter_duplicate_safetensors_files(
- hf_weights_files: list[str], hf_folder: str, index_file: str
+ hf_weights_files: list[str],
+ hf_folder: str,
+ index_file: str,
+ key_filter: Callable[[str], bool] | None = None,
) -> list[str]:
# model.safetensors.index.json is a mapping from keys in the
# torch state_dict to safetensors file holding that weight.
@@ -79,6 +82,9 @@ def filter_duplicate_safetensors_files(
weight_map = json.load(f)["weight_map"]
weight_files_in_index = set()
for weight_name in weight_map:
+ # remove only shards whose indexed tensors are all filtered
+ if key_filter is not None and not key_filter(weight_name):
+ continue
weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
# Filter out any fields that are not found in the index file.
hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
index db2a0a761..aaf14edd3 100644
--- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
@@ -107,6 +107,21 @@ class _ExpandedOutputParts:
trajectory_decoded_parts: list[list[torch.Tensor]] | None = None
+def _worker_cpu_intra_op_threads(num_gpus: int) -> int | None:
+ """CPU intra-op thread budget for one of `num_gpus` co-located workers.
+
+ torch defaults the intra-op pool to every host core in every worker, so
+ co-located workers oversubscribe the host num_gpus-fold and any CPU op
+ past the ~32k-element parallel grain pays pool wakeup contention instead
+ of microseconds (measured 500x on request-static packed layouts). An
+ explicit OMP_NUM_THREADS keeps deployer intent (returns None).
+ """
+ if "OMP_NUM_THREADS" in os.environ:
+ return None
+ cpu_count = os.cpu_count() or 1
+ return max(1, min(16, cpu_count // max(1, num_gpus)))
+
+
class GPUWorker(GPUWorkerPostTrainingMixin):
"""
A worker that executes the model on a single GPU.
@@ -198,6 +213,9 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
def init_device_and_model(self) -> None:
"""Initialize the device and load the model."""
torch.get_device_module().set_device(self.local_rank)
+ intra_op_threads = _worker_cpu_intra_op_threads(self.server_args.num_gpus)
+ if intra_op_threads is not None:
+ torch.set_num_threads(intra_op_threads)
# Set environment variables for distributed initialization
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = str(self.master_port)
diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
index 8b7734068..494ab4947 100644
--- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
+++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
@@ -67,8 +67,12 @@ class LayerwiseOffloadManager:
)
self.copy_stream = torch.get_device_module().Stream()
+ # ``named_parameters()`` is relative to ``model``, just like the path in
+ # ``layers_attr_str``. Anchor the match so a manager for top-level
+ # ``blocks`` cannot also capture an unrelated nested list such as
+ # ``token_refiner.blocks`` whose forward hooks run at a different time.
self._layer_name_re = re.compile(
- rf"(^|\.){re.escape(layers_attr_str)}\.(\d+)(\.|$)"
+ rf"^{re.escape(layers_attr_str)}\.(?P\d+)(\.|$)"
)
# layer_idx -> {dtype: consolidated_pinned_cpu_tensor}
@@ -99,7 +103,7 @@ class LayerwiseOffloadManager:
if not m:
return None
try:
- return int(m.group(2))
+ return int(m.group("layer_idx"))
except Exception:
return None
@@ -612,6 +616,10 @@ class LayerwiseOffloadableModuleMixin:
self.layerwise_offload_managers = []
named_modules = dict(self.named_modules())
configured_layer_names = []
+ # These legacy tuning knobs are explicitly DiT-scoped. Auxiliary
+ # components still support layerwise streaming, but their layers run
+ # once per component use and get no reuse benefit from DiT residency.
+ dit_tuning_enabled = self.layerwise_offload_dit_group_enabled
for layer_name in self.layer_names:
module_list = named_modules.get(layer_name)
if not isinstance(module_list, (torch.nn.ModuleList, torch.nn.Sequential)):
@@ -620,14 +628,17 @@ class LayerwiseOffloadableModuleMixin:
continue
num_layers = len(module_list)
- if server_args.dit_offload_prefetch_size < 1.0:
- prefetch_size = 1 + int(
- round(server_args.dit_offload_prefetch_size * (num_layers - 1))
- )
+ prefetch_value = (
+ server_args.dit_offload_prefetch_size if dit_tuning_enabled else 0.0
+ )
+ if prefetch_value < 1.0:
+ prefetch_size = 1 + int(round(prefetch_value * (num_layers - 1)))
else:
- prefetch_size = int(server_args.dit_offload_prefetch_size)
+ prefetch_size = int(prefetch_value)
- resident_value = server_args.dit_layerwise_resident_layers
+ resident_value = (
+ server_args.dit_layerwise_resident_layers if dit_tuning_enabled else 0.0
+ )
if resident_value <= 0:
resident_layers = 0
elif resident_value < 1.0:
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
new file mode 100644
index 000000000..0abc01624
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -0,0 +1,1692 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 packed-token DiT.
+
+Native SGLang implementation of the MiniMax H3 audio-video DiT. The forward
+contract accepts packed inference keyword arguments and returns packed logits.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from sglang.kernels.ops.activation.activation import (
+ silu_and_mul_with_activation_rounding_,
+)
+from sglang.kernels.ops.diffusion.qknorm_rope import (
+ can_use_fused_inplace_qknorm_rope,
+ fused_inplace_qknorm_rope,
+)
+from sglang.kernels.ops.diffusion.triton.indexed_modulation import (
+ indexed_gate_bf16_,
+ indexed_scale_shift_bf16_,
+)
+from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+ MINIMAX_H3_ADALN_MODALITY_NUM,
+ MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT,
+ MiniMaxH3DiTArchConfig,
+ MiniMaxH3DiTConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+ get_tp_world_size,
+ tensor_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.layers.linear import (
+ ColumnParallelLinear,
+ MergedColumnParallelLinear,
+ RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+ QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+ LayerwiseOffloadableModuleMixin,
+ is_layerwise_offloaded_module,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
+from sglang.multimodal_gen.runtime.platforms import (
+ AttentionBackendEnum,
+ current_platform,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+ eager_on_graph,
+)
+
+_ARCH_DEFAULTS = MiniMaxH3DiTArchConfig()
+_BF16_DTYPE = torch.bfloat16
+_FP32_DTYPE = torch.float32
+
+_MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER = (
+ "video_patch_proj.weight",
+ "video_patch_proj.bias",
+ "audio_patch_proj.weight",
+ "audio_patch_proj.bias",
+ "time_embedder.proj_in.weight",
+ "time_embedder.proj_in.bias",
+ "time_embedder.proj_out.weight",
+ "time_embedder.proj_out.bias",
+ "final_layer.video_out.weight",
+ "final_layer.video_out.bias",
+ "final_layer.audio_out.weight",
+ "final_layer.audio_out.bias",
+)
+MINIMAX_H3_FP32_PARAM_NAMES = frozenset(_MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER)
+MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
+
+
+def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
+ if key not in kwargs or kwargs[key] is None:
+ raise ValueError(f"MiniMaxH3DiTModel.forward requires kwarg {key!r}")
+ return kwargs[key]
+
+
+# The exhaustive keyword contract of MiniMaxH3DiTModel.forward. Anything not
+# listed here is rejected with a TypeError before any tensor work starts.
+_FORWARD_SUPPORTED_KWARGS = frozenset(
+ {
+ "x",
+ "audio_x",
+ "img_position_ids",
+ "rope_cache",
+ "unique_timesteps",
+ "inverse_indices",
+ "update_mask",
+ "update_audio_mask",
+ "token_tags",
+ "block_token_tags",
+ "block_combined_indices",
+ "skip_mask_out_condition",
+ "prompt_embeds",
+ "refined_prompt_embeds_length",
+ "img_pos_info",
+ "audio_pos_info",
+ "text_pos_info",
+ "img_pos_for_infer_output_info",
+ "local_embedding_layout",
+ "packed_seq_params",
+ "refiner_packed_seq_params",
+ }
+)
+
+
+def _ulysses_ctx() -> tuple[int, int]:
+ """(world_size, rank) of the Ulysses sequence-parallel group.
+
+ Returns (1, 0) when model parallelism is not initialized (unit tests /
+ single-process debug paths init tp=1 sp=1 which also yields ws=1).
+ """
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_ulysses_parallel_rank,
+ get_ulysses_parallel_world_size,
+ model_parallel_is_initialized,
+ )
+
+ if not model_parallel_is_initialized():
+ return 1, 0
+ return get_ulysses_parallel_world_size(), get_ulysses_parallel_rank()
+
+
+def _ring_world_size() -> int:
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_ring_parallel_world_size,
+ model_parallel_is_initialized,
+ )
+
+ if not model_parallel_is_initialized():
+ return 1
+ return get_ring_parallel_world_size()
+
+
+def _reorder_grouped_qkv_to_qkv(
+ weight: torch.Tensor,
+ *,
+ num_query_groups: int,
+ heads_per_group: int,
+ head_dim: int,
+) -> torch.Tensor:
+ per_group = (heads_per_group + 2) * head_dim
+ expected_out = num_query_groups * per_group
+ if weight.shape[0] != expected_out:
+ raise ValueError(
+ "qkv weight has incompatible output dim for grouped checkpoint layout: "
+ f"got {tuple(weight.shape)}, expected first dim {expected_out}."
+ )
+
+ rest_shape = weight.shape[1:]
+ grouped = weight.reshape(num_query_groups, per_group, *rest_shape)
+ q, k, v = torch.split(
+ grouped,
+ [heads_per_group * head_dim, head_dim, head_dim],
+ dim=1,
+ )
+ return torch.cat(
+ [
+ q.reshape(num_query_groups * heads_per_group * head_dim, *rest_shape),
+ k.reshape(num_query_groups * head_dim, *rest_shape),
+ v.reshape(num_query_groups * head_dim, *rest_shape),
+ ],
+ dim=0,
+ )
+
+
+def _copy_grouped_qkv_tp_shard(
+ param: torch.Tensor,
+ loaded_weight: torch.Tensor,
+ *,
+ num_query_groups: int,
+ head_dim: int,
+ tp_rank: int,
+ tp_size: int,
+) -> bool:
+ """Copy a dense MHA checkpoint directly into its TP-local Q/K/V rows."""
+ if (
+ tp_size <= 0
+ or not 0 <= tp_rank < tp_size
+ or num_query_groups % tp_size
+ or getattr(param, "output_dim", None) != 0
+ or getattr(param, "is_sharded_weight", False)
+ or getattr(param, "packed_dim", None) is not None
+ or param.dtype != _BF16_DTYPE
+ or loaded_weight.dtype != _BF16_DTYPE
+ or not param.is_contiguous()
+ or not loaded_weight.is_contiguous()
+ ):
+ return False
+
+ expected_rows = num_query_groups * 3 * head_dim
+ local_groups = num_query_groups // tp_size
+ rest_shape = loaded_weight.shape[1:]
+ if loaded_weight.shape[0] != expected_rows or tuple(param.shape) != (
+ 3 * local_groups * head_dim,
+ *rest_shape,
+ ):
+ return False
+
+ grouped = loaded_weight.view(num_query_groups, 3, head_dim, *rest_shape)
+ grouped = grouped.narrow(0, tp_rank * local_groups, local_groups)
+ target = param.data.view(3, local_groups, head_dim, *rest_shape)
+ for index in range(3):
+ target[index].copy_(grouped[:, index])
+ return True
+
+
+def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> nn.RMSNorm:
+ # RMSNorm uses fp32 accumulation with bf16 inputs and outputs.
+ # torch.nn.RMSNorm upcasts reduced-precision inputs for the variance
+ # reduction, matching that accumulation semantic.
+ return nn.RMSNorm(size, eps=eps, dtype=dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = torch.chunk(x, 2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def _modulate_scale_shift(
+ x: torch.Tensor,
+ shift: torch.Tensor,
+ scale: torch.Tensor,
+ indices: torch.Tensor,
+ *,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """Apply indexed affine modulation, reusing disposable CUDA BF16 input."""
+ # Apply per-index affine modulation: x * (1 + scale[idx]) + shift[idx].
+ if (
+ x.is_cuda
+ and x.dtype == _BF16_DTYPE
+ and dtype == _BF16_DTYPE
+ and shift.dtype == _BF16_DTYPE
+ and scale.dtype == _BF16_DTYPE
+ and x.is_contiguous()
+ ):
+ return indexed_scale_shift_bf16_(x, shift, scale, indices)
+ return (
+ x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)
+ ).to(dtype)
+
+
+def _modulate_gate(
+ x: torch.Tensor,
+ gate: torch.Tensor,
+ other: torch.Tensor,
+ indices: torch.Tensor,
+ *,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """Apply indexed gated residual, reusing disposable CUDA BF16 input."""
+ # Apply the per-index gated residual: x + gate[idx] * other.
+ if (
+ x.is_cuda
+ and x.dtype == _BF16_DTYPE
+ and dtype == _BF16_DTYPE
+ and gate.dtype == _BF16_DTYPE
+ and other.dtype == _BF16_DTYPE
+ and x.is_contiguous()
+ and other.is_contiguous()
+ ):
+ return indexed_gate_bf16_(x, gate, other, indices)
+ return (x + gate.index_select(0, indices) * other).to(dtype)
+
+
+def _silu_mul(hidden: torch.Tensor, *, reuse_input: bool) -> torch.Tensor:
+ if (
+ reuse_input
+ and hidden.is_cuda
+ and hidden.dtype == _BF16_DTYPE
+ and hidden.is_contiguous()
+ and hidden.shape[-1] % 16 == 0
+ ):
+ return silu_and_mul_with_activation_rounding_(hidden)
+ gate, up = hidden.chunk(2, dim=-1)
+ return nn.functional.silu(gate) * up
+
+
+def _apply_qk_norm(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ q_norm: nn.RMSNorm,
+ k_norm: nn.RMSNorm,
+ head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if (
+ q.is_cuda
+ and q.dtype == _BF16_DTYPE
+ and q.dtype == k.dtype == q_norm.weight.dtype == k_norm.weight.dtype
+ and q.stride(-1) == k.stride(-1) == 1
+ and q.stride(-2) == k.stride(-2) == head_dim
+ and q_norm.eps == k_norm.eps
+ and not torch.compiler.is_compiling()
+ ):
+ fused_inplace_qknorm(
+ q,
+ k,
+ q_norm.weight,
+ k_norm.weight,
+ eps=q_norm.eps,
+ head_dim=head_dim,
+ )
+ return q, k
+ return q_norm(q), k_norm(k)
+
+
+class MiniMaxH3Rope(nn.Module):
+ """3D rope over (t, h, w); rotates 96 of 128 head dims (rotary_percent 0.75).
+
+ Frequency layout concatenates temporal, height, and width embeddings twice,
+ with 16 frequencies per axis (inv_freq = base^-(arange(0,32,2)/32)).
+ """
+
+ def __init__(self, inv_freq_len: int) -> None:
+ super().__init__()
+ self.register_buffer(
+ "inv_freq",
+ torch.empty(inv_freq_len, dtype=_FP32_DTYPE),
+ persistent=True,
+ )
+
+ def forward(self, img_position_ids: torch.Tensor) -> torch.Tensor:
+ """img_position_ids: [1, S, 3] (t, h, w) -> freqs [S, rot_dim=96]."""
+ if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1:
+ raise ValueError(
+ "img_position_ids must be [1, S, 3], got "
+ f"{list(img_position_ids.shape)}"
+ )
+ pos = img_position_ids[0].to(_FP32_DTYPE) # [S, 3]
+ per_axis = pos.unsqueeze(-1) * self.inv_freq.view(1, 1, -1) # [S, 3, 16]
+ t_f, h_f, w_f = per_axis.unbind(dim=1) # each [S, 16]
+ half = torch.cat((t_f, h_f, w_f), dim=-1) # [S, 48]
+ return torch.cat((half, half), dim=-1) # [S, 96]
+
+
+def _rope_cos_sin_cache(freqs: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor:
+ """Build the activation-dtype cos|sin cache for fused Q/K RoPE."""
+ half = freqs.shape[-1] // 2
+ return (
+ torch.cat(
+ (torch.cos(freqs[:, :half]), torch.sin(freqs[:, :half])),
+ dim=-1,
+ )
+ .to(dtype=dtype, copy=False)
+ .contiguous()
+ )
+
+
+def _apply_rope_cos_sin(
+ x: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> torch.Tensor:
+ """Rotate the first cached RoPE dims; pass the remaining head dims through."""
+ rot_dim = cos.shape[-1]
+ x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
+ x_rot = (x_rot * cos) + (_rotate_half(x_rot) * sin)
+ return torch.cat((x_rot, x_pass), dim=-1)
+
+
+def _apply_rope_qk(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos_sin_cache: torch.Tensor,
+ positions: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if not q.is_cuda:
+ half = cos_sin_cache.shape[-1] // 2
+ cos_half, sin_half = cos_sin_cache.split(half, dim=-1)
+ cos = torch.cat((cos_half, cos_half), dim=-1).unsqueeze(1)
+ sin = torch.cat((sin_half, sin_half), dim=-1).unsqueeze(1)
+ return (
+ _apply_rope_cos_sin(q, cos, sin),
+ _apply_rope_cos_sin(k, cos, sin),
+ )
+
+ from sgl_kernel import rotary_embedding as apply_sgl_kernel_rotary_embedding
+
+ apply_sgl_kernel_rotary_embedding(
+ positions,
+ q.view(q.shape[0], -1),
+ k.view(k.shape[0], -1),
+ q.shape[-1],
+ cos_sin_cache,
+ True,
+ )
+ return q, k
+
+
+class MiniMaxH3TimeEmbedder(nn.Module):
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ *,
+ prefix: str,
+ ) -> None:
+ super().__init__()
+ self.frequency_embedding_size = arch.timestep_input_dim
+ self.proj_in = ColumnParallelLinear(
+ arch.timestep_input_dim,
+ arch.time_embed_hidden_size,
+ bias=True,
+ gather_output=False,
+ params_dtype=_FP32_DTYPE,
+ quant_config=None,
+ prefix=f"{prefix}.proj_in",
+ )
+ self.proj_out = RowParallelLinear(
+ arch.time_embed_hidden_size,
+ arch.time_embed_dim,
+ bias=True,
+ input_is_parallel=True,
+ params_dtype=_FP32_DTYPE,
+ quant_config=None,
+ prefix=f"{prefix}.proj_out",
+ )
+ self.register_buffer("_frequency_cache", None, persistent=False)
+
+ def forward(self, t: torch.Tensor) -> torch.Tensor:
+ """t: [M] -> [M, time_embed_dim] fp32.
+
+ The sinusoidal embedding stays fp32 throughout and concatenates cosine
+ values before sine values.
+ """
+ half = self.frequency_embedding_size // 2
+ freqs = self._frequency_cache
+ if freqs is None or freqs.device != t.device:
+ # Construct this on the execution device once so the values keep
+ # the established CUDA numerics without repeating arange/exp on
+ # every denoise step.
+ freqs = torch.exp(
+ -math.log(10000.0)
+ * torch.arange(half, dtype=_FP32_DTYPE, device=t.device)
+ / half
+ )
+ self._frequency_cache = freqs
+ args = t.to(_FP32_DTYPE)[:, None] * freqs[None]
+ t_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ hidden, _ = self.proj_in(t_freq)
+ hidden = nn.functional.silu(hidden)
+ out, _ = self.proj_out(hidden)
+ return out
+
+
+def _minimax_h3_attention_core_impl(
+ attention: MiniMaxH3Attention,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ cu_seqlens_host: tuple[int, ...] | None,
+ max_seqlen: int,
+ ulysses_active: bool,
+) -> torch.Tensor:
+ """Dynamic varlen attention and Ulysses collectives.
+
+ This is the narrow BCG break point: projections, normalization, RoPE,
+ residuals, and MLPs remain captured while the dynamic packed attention
+ kernel and sequence-parallel collectives execute eagerly.
+ """
+
+ if ulysses_active:
+ from sglang.multimodal_gen.runtime.layers.usp import (
+ _usp_input_all_to_all_packed_qkv,
+ _usp_output_all_to_all,
+ )
+
+ q, k, v = _usp_input_all_to_all_packed_qkv(q, k, v)
+
+ if attention._attention_impl is None:
+ attention._set_attention_backend(
+ get_attn_backend(
+ attention.head_dim,
+ q.dtype,
+ supported_attention_backends=attention._supported_attention_backends,
+ )
+ )
+ out = attention._attention_impl.forward_varlen(
+ q,
+ k,
+ v,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ cu_seqlens_host=cu_seqlens_host,
+ )
+ if ulysses_active:
+ out = _usp_output_all_to_all(out[None], head_dim=2)[0]
+ return out
+
+
+_minimax_h3_attention_core_bcg = eager_on_graph(True)(_minimax_h3_attention_core_impl)
+
+
+class MiniMaxH3Attention(nn.Module):
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ bcg_breakpoint: bool = True,
+ ) -> None:
+ super().__init__()
+ self.bcg_breakpoint = bcg_breakpoint
+ self.tp_size = get_tp_world_size()
+ if arch.num_attention_heads % self.tp_size:
+ raise ValueError(
+ "MiniMax H3 attention heads must be divisible by TP size: "
+ f"{arch.num_attention_heads} % {self.tp_size} != 0"
+ )
+ self.total_num_heads = arch.num_attention_heads
+ self.num_heads = self.total_num_heads // self.tp_size
+ self.head_dim = arch.attention_head_dim
+ self.inner_dim = self.total_num_heads * self.head_dim
+ self.local_inner_dim = self.num_heads * self.head_dim
+ self.softmax_scale = self.head_dim**-0.5
+ self._supported_attention_backends = arch._supported_attention_backends
+ self._attention_impl = None
+ # The checkpoint stores one fused qkv tensor. Each logical Q/K/V
+ # matrix must be sharded independently; a plain ColumnParallelLinear
+ # would instead slice across the concatenated tensor and is incorrect
+ # for TP > 1.
+ self.qkv_proj = MergedColumnParallelLinear(
+ arch.hidden_size,
+ [self.inner_dim] * 3,
+ bias=False,
+ gather_output=False,
+ params_dtype=_BF16_DTYPE,
+ quant_config=quant_config,
+ prefix=f"{prefix}.qkv_proj",
+ )
+ self._install_qkv_weight_loader(arch)
+ self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
+ self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
+ # cache width covers cos/sin for temporal, height, and width frequencies
+ rope_dim = 6 * arch.rope_inv_freq_len
+ self._use_fused_qknorm_rope = (
+ current_platform.is_cuda()
+ and can_use_fused_inplace_qknorm_rope(
+ arch.attention_head_dim,
+ rope_dim,
+ True,
+ _BF16_DTYPE,
+ cache_dtype=_BF16_DTYPE,
+ round_norm_before_rope=True,
+ )
+ )
+ self.out_proj = RowParallelLinear(
+ self.inner_dim,
+ arch.hidden_size,
+ bias=False,
+ input_is_parallel=True,
+ params_dtype=_BF16_DTYPE,
+ quant_config=quant_config,
+ prefix=f"{prefix}.out_proj",
+ )
+
+ def _set_attention_backend(self, backend) -> None:
+ impl_cls = backend.get_impl_cls()
+ self._attention_impl = impl_cls(
+ num_heads=self.num_heads,
+ head_size=self.head_dim,
+ causal=False,
+ softmax_scale=self.softmax_scale,
+ num_kv_heads=self.num_heads,
+ )
+
+ def _install_qkv_weight_loader(self, arch: MiniMaxH3DiTArchConfig) -> None:
+ weight = self.qkv_proj.weight
+ base_loader = weight.weight_loader
+
+ def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+ # The grouped checkpoint layout is
+ # [num_query_groups, q_per_group + k + v] before splitting.
+ # MiniMax H3 uses MHA, so checkpoint rows are per-head [q, k, v],
+ # while SGLang stores [q_all, k_all, v_all].
+ if _copy_grouped_qkv_tp_shard(
+ param,
+ loaded_weight,
+ num_query_groups=arch.num_attention_heads,
+ head_dim=arch.attention_head_dim,
+ tp_rank=self.qkv_proj.tp_rank,
+ tp_size=self.tp_size,
+ ):
+ return
+ reordered = _reorder_grouped_qkv_to_qkv(
+ loaded_weight,
+ num_query_groups=arch.num_attention_heads,
+ heads_per_group=1,
+ head_dim=arch.attention_head_dim,
+ )
+ base_loader(param, reordered)
+
+ if hasattr(weight, "_weight_loader"):
+ weight._weight_loader = _weight_loader
+ else:
+ weight.weight_loader = _weight_loader
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ *,
+ rope_cache: tuple[torch.Tensor, torch.Tensor] | None,
+ cu_seqlens: torch.Tensor,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ max_seqlen: int,
+ ulysses_active: bool = False,
+ ) -> torch.Tensor:
+ """x: [T, hidden] packed thd rows -> [T, hidden].
+
+ Operation order: fused qkv projection -> per-head q/k RMSNorm -> RoPE
+ on q/k -> variable-length non-causal flash attention -> output projection.
+
+ With Ulysses sequence parallelism, x holds this rank's row shard;
+ qkv/norm/RoPE run locally, an all-to-all trades sequence for heads.
+ Each rank attends the full sequence with heads/world_size local heads,
+ so cu_seqlens retains global packed-document semantics. The inverse
+ all-to-all restores the row shard before the output projection.
+ """
+ total = x.shape[0]
+ qkv, _ = self.qkv_proj(x)
+ q, k, v = qkv.split(self.local_inner_dim, dim=-1)
+ q = q.view(total, self.num_heads, self.head_dim)
+ k = k.view(total, self.num_heads, self.head_dim)
+ v = v.view(total, self.num_heads, self.head_dim)
+ if rope_cache is None:
+ q, k = _apply_qk_norm(
+ q,
+ k,
+ self.q_norm,
+ self.k_norm,
+ self.head_dim,
+ )
+ else:
+ cos_sin_cache, positions = rope_cache
+ if self._use_fused_qknorm_rope and not torch.compiler.is_compiling():
+ fused_inplace_qknorm_rope(
+ q,
+ k,
+ self.q_norm.weight,
+ self.k_norm.weight,
+ cos_sin_cache,
+ positions,
+ is_neox=True,
+ eps=self.q_norm.eps,
+ head_dim=self.head_dim,
+ rope_dim=cos_sin_cache.shape[-1],
+ round_norm_before_rope=True,
+ )
+ else:
+ q, k = _apply_qk_norm(
+ q,
+ k,
+ self.q_norm,
+ self.k_norm,
+ self.head_dim,
+ )
+ q, k = _apply_rope_qk(q, k, cos_sin_cache, positions)
+
+ attention_core = (
+ _minimax_h3_attention_core_bcg
+ if self.bcg_breakpoint
+ else _minimax_h3_attention_core_impl
+ )
+ out = attention_core(
+ self,
+ q,
+ k,
+ v,
+ cu_seqlens=cu_seqlens,
+ cu_seqlens_host=cu_seqlens_host,
+ max_seqlen=max_seqlen,
+ ulysses_active=ulysses_active,
+ )
+ out = out.reshape(total, self.num_heads * self.head_dim)
+ out, _ = self.out_proj(out)
+ return out
+
+
+class MiniMaxH3MLP(nn.Module):
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ ) -> None:
+ super().__init__()
+ # As with qkv, gate and up are two independently sharded logical
+ # matrices even though the checkpoint stores them fused.
+ self.fc1 = MergedColumnParallelLinear(
+ arch.hidden_size,
+ [arch.ffn_hidden_size] * 2,
+ bias=False,
+ gather_output=False,
+ params_dtype=_BF16_DTYPE,
+ quant_config=quant_config,
+ prefix=f"{prefix}.fc1",
+ )
+ # Chunk the fused fc1 output as [gate, up], then compute
+ # silu(gate) * up.
+ self.fc2 = RowParallelLinear(
+ arch.ffn_hidden_size,
+ arch.hidden_size,
+ bias=False,
+ input_is_parallel=True,
+ params_dtype=_BF16_DTYPE,
+ quant_config=quant_config,
+ prefix=f"{prefix}.fc2",
+ )
+ self.reuse_fc1_activation = quant_config is None
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ hidden, _ = self.fc1(x)
+ hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation)
+ out, _ = self.fc2(hidden)
+ return out
+
+
+class MiniMaxH3AdalnProj(nn.Module):
+ """SiLU + zero-init linear over unique condition embeddings.
+
+ Per block, three modalities each produce six H-wide vectors:
+ [M, t_dim] -> [M, 3*6H] -> view(M*3, 6H) -> chunk(6).
+ The final layer uses one modality and produces two H-wide vectors:
+ [M, t_dim] -> [M, 2H] -> chunk(2).
+ """
+
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ out_features: int,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ expand_ratio: int,
+ modality_num: int,
+ ) -> None:
+ super().__init__()
+ if out_features != expand_ratio * arch.hidden_size * modality_num:
+ raise ValueError(
+ "adaln out_features mismatch: "
+ f"{out_features} != {expand_ratio}*{arch.hidden_size}*{modality_num}"
+ )
+ self.expand_ratio = expand_ratio
+ self.modality_num = modality_num
+ self.hidden_size = arch.hidden_size
+ self.linear = ColumnParallelLinear(
+ arch.time_embed_dim,
+ out_features,
+ bias=True,
+ gather_output=False,
+ params_dtype=_BF16_DTYPE,
+ quant_config=quant_config,
+ prefix=f"{prefix}.linear",
+ )
+
+ def project_local(self, adaln_input: torch.Tensor) -> torch.Tensor:
+ x, _ = self.linear(adaln_input)
+ return x
+
+ def split_output(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
+ m = x.shape[0]
+ x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
+ return tuple(x.chunk(self.expand_ratio, dim=-1))
+
+ def forward(self, adaln_input: torch.Tensor) -> tuple[torch.Tensor, ...]:
+ """adaln_input: SiLU(t_emb) BF16 -> expand_ratio tensors of [M*modality_num, H]."""
+ x = self.project_local(adaln_input)
+ if get_tp_world_size() > 1:
+ x = tensor_model_parallel_all_gather(x)
+ return self.split_output(x)
+
+
+class MiniMaxH3TokenRefinerBlock(nn.Module):
+ """Standard pre-norm transformer block without AdaLN or RoPE."""
+
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ ) -> None:
+ super().__init__()
+ self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
+ self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps)
+ # The whole dynamic text-refiner/scatter phase is one BCG break point;
+ # do not nest per-attention graph breaks inside it.
+ self.attn = MiniMaxH3Attention(
+ arch,
+ quant_config,
+ prefix=f"{prefix}.attn",
+ bcg_breakpoint=False,
+ )
+ self.mlp = MiniMaxH3MLP(arch, quant_config, prefix=f"{prefix}.mlp")
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ max_seqlen: int,
+ ) -> torch.Tensor:
+ x = x + self.attn(
+ self.norm1(x),
+ rope_cache=None,
+ cu_seqlens=cu_seqlens,
+ cu_seqlens_host=cu_seqlens_host,
+ max_seqlen=max_seqlen,
+ )
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class MiniMaxH3TokenRefiner(nn.Module):
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ ) -> None:
+ super().__init__()
+ self.blocks = nn.ModuleList(
+ [
+ MiniMaxH3TokenRefinerBlock(
+ arch,
+ quant_config,
+ prefix=f"{prefix}.blocks.{index}",
+ )
+ for index in range(arch.token_refiner_num_layers)
+ ]
+ )
+ self.final_norm = _norm(arch.hidden_size, eps=arch.final_norm_eps)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ *,
+ cu_seqlens: torch.Tensor,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ max_seqlen: int,
+ ) -> torch.Tensor:
+ for block in self.blocks:
+ x = block(
+ x,
+ cu_seqlens=cu_seqlens,
+ cu_seqlens_host=cu_seqlens_host,
+ max_seqlen=max_seqlen,
+ )
+ return self.final_norm(x)
+
+
+class MiniMaxH3DiTBlock(nn.Module):
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ ) -> None:
+ super().__init__()
+ self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
+ self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps)
+ self.attn = MiniMaxH3Attention(
+ arch,
+ quant_config,
+ prefix=f"{prefix}.attn",
+ )
+ self.mlp = MiniMaxH3MLP(arch, quant_config, prefix=f"{prefix}.mlp")
+ self.adaln_proj = MiniMaxH3AdalnProj(
+ arch,
+ arch.adaln_out_features,
+ quant_config,
+ prefix=f"{prefix}.adaln_proj",
+ expand_ratio=6,
+ modality_num=MINIMAX_H3_ADALN_MODALITY_NUM,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ *,
+ adaln_input: torch.Tensor,
+ combined_indices: torch.Tensor,
+ rope_cache: tuple[torch.Tensor, torch.Tensor],
+ cu_seqlens: torch.Tensor,
+ cu_seqlens_host: tuple[int, ...] | None = None,
+ max_seqlen: int,
+ ulysses_active: bool = False,
+ adaln_params: tuple[torch.Tensor, ...] | None = None,
+ ) -> torch.Tensor:
+ """x: [T, H]; adaln_input: [M, t_dim]; combined_indices: [T]
+ (= inverse_indices * modality_num + token_tags.clamp(min=0)).
+
+ Each block computes AdaLN parameters once, then applies
+ norm1 -> scale/shift -> attention -> gated residual, followed by
+ norm2 -> scale/shift -> MLP -> gated residual.
+ """
+ if adaln_params is None:
+ adaln_params = self.adaln_proj(adaln_input)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln_params
+
+ residual = x
+ h = self.norm1(x)
+ h = _modulate_scale_shift(
+ h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE
+ )
+ h = self.attn(
+ h,
+ rope_cache=rope_cache,
+ cu_seqlens=cu_seqlens,
+ cu_seqlens_host=cu_seqlens_host,
+ max_seqlen=max_seqlen,
+ ulysses_active=ulysses_active,
+ )
+ x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE)
+
+ residual = x
+ h = self.norm2(x)
+ h = _modulate_scale_shift(
+ h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE
+ )
+ h = self.mlp(h)
+ return _modulate_gate(
+ residual, gate_mlp, h, combined_indices, dtype=_BF16_DTYPE
+ )
+
+
+class MiniMaxH3FinalLayer(nn.Module):
+ def __init__(
+ self,
+ arch: MiniMaxH3DiTArchConfig,
+ quant_config: QuantizationConfig | None,
+ *,
+ prefix: str,
+ ) -> None:
+ super().__init__()
+ video_patch_dim = (
+ arch.latents_dim
+ * arch.patch_size[0]
+ * arch.patch_size[1]
+ * arch.patch_size[2]
+ )
+ self.norm = _norm(arch.hidden_size, eps=arch.final_norm_eps)
+ self.adaln_proj = MiniMaxH3AdalnProj(
+ arch,
+ arch.final_adaln_out_features,
+ quant_config,
+ prefix=f"{prefix}.adaln_proj",
+ expand_ratio=2,
+ modality_num=1,
+ )
+ self.video_out = ColumnParallelLinear(
+ arch.hidden_size,
+ video_patch_dim,
+ bias=True,
+ gather_output=False,
+ params_dtype=_FP32_DTYPE,
+ quant_config=None,
+ prefix=f"{prefix}.video_out",
+ )
+ self.audio_out = ColumnParallelLinear(
+ arch.hidden_size,
+ arch.audio_latents_dim,
+ bias=True,
+ gather_output=False,
+ params_dtype=_FP32_DTYPE,
+ quant_config=None,
+ prefix=f"{prefix}.audio_out",
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ *,
+ adaln_input: torch.Tensor,
+ inverse_indices: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Project all rows into TP-local video/audio output shards.
+
+ Apply single-modality shift/scale AdaLN to the final normalized
+ activations, cast to fp32, then apply both output heads to all rows.
+ The model gathers output columns only after selecting live media rows,
+ preserving the GEMM shape while reducing collective payload.
+ """
+ shift, scale = self.adaln_proj(adaln_input)
+ h = self.norm(x)
+ h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE)
+ # Preserve full precision through both final output projections.
+ h = h.to(_FP32_DTYPE)
+ video, _ = self.video_out(h)
+ audio, _ = self.audio_out(h)
+ return video, audio
+
+
+class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
+ _fsdp_shard_conditions = _ARCH_DEFAULTS._fsdp_shard_conditions
+ # parameters mix fp32 (patch projections, timestep embedder, and output
+ # heads) with bf16 blocks; FSDP must gather in each parameter's own dtype
+ _fsdp_mixed_dtype_params = True
+ _compile_conditions = _ARCH_DEFAULTS._compile_conditions
+ _supported_attention_backends = _ARCH_DEFAULTS._supported_attention_backends
+ param_names_mapping = _ARCH_DEFAULTS.param_names_mapping
+ reverse_param_names_mapping = _ARCH_DEFAULTS.reverse_param_names_mapping
+ lora_param_names_mapping = _ARCH_DEFAULTS.lora_param_names_mapping
+
+ def _can_batch_block_adaln(self) -> bool:
+ return (
+ get_tp_world_size() > 1
+ and not torch.compiler.is_compiling()
+ and not envs.SGLANG_CACHE_DIT_ENABLED
+ and not hasattr(self, "_sglang_cache_dit_adapter")
+ and not is_layerwise_offloaded_module(self)
+ and all(type(block) is MiniMaxH3DiTBlock for block in self.blocks)
+ )
+
+ def _validate_tp_config(
+ self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int
+ ) -> None:
+ if tp_size <= 0:
+ raise ValueError("TP size must be positive.")
+ if arch.num_attention_heads <= 0:
+ raise ValueError("num_attention_heads must be positive.")
+ if arch.hidden_size <= 0:
+ raise ValueError("hidden_size must be positive.")
+ if arch.attention_head_dim <= 0:
+ raise ValueError("attention_head_dim must be positive.")
+ if arch.ffn_hidden_size <= 0:
+ raise ValueError("ffn_hidden_size must be positive.")
+ for name, value in (
+ ("num_attention_heads", arch.num_attention_heads),
+ ("hidden_size", arch.hidden_size),
+ ("ffn_hidden_size", arch.ffn_hidden_size),
+ ("time_embed_hidden_size", arch.time_embed_hidden_size),
+ ("adaln_out_features", arch.adaln_out_features),
+ ("final_adaln_out_features", arch.final_adaln_out_features),
+ ("video_patch_output_dim", arch.latents_dim * math.prod(arch.patch_size)),
+ ("audio_patch_output_dim", arch.audio_latents_dim),
+ ):
+ if value % tp_size:
+ raise ValueError(
+ f"MiniMax H3 {name}={value} must be divisible by "
+ f"TP size {tp_size}."
+ )
+
+ @staticmethod
+ def _validate_sequence_parallel_config(
+ *,
+ arch: MiniMaxH3DiTArchConfig,
+ tp_size: int,
+ ulysses_size: int,
+ ring_size: int,
+ ) -> None:
+ if ulysses_size <= 0:
+ raise ValueError("MiniMax H3 Ulysses size must be positive.")
+ if ring_size != 1:
+ raise NotImplementedError(
+ "MiniMax H3 packed multi-segment attention does not support "
+ "Ring or mixed USP. Set --ring-degree 1 and use Ulysses "
+ "sequence parallelism."
+ )
+ local_heads = arch.num_attention_heads // tp_size
+ if local_heads % ulysses_size:
+ raise ValueError(
+ f"MiniMax H3 TP-local heads {local_heads} must be divisible by "
+ f"Ulysses size {ulysses_size} (total heads="
+ f"{arch.num_attention_heads}, TP={tp_size})."
+ )
+ if MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT % ulysses_size:
+ raise ValueError(
+ "MiniMax H3 packed sequence alignment "
+ f"{MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT} must be divisible by "
+ f"Ulysses size {ulysses_size}. Choose a Ulysses size that "
+ "divides both the TP-local attention heads and the packed "
+ "sequence alignment."
+ )
+
+ def __init__(
+ self,
+ config: MiniMaxH3DiTConfig,
+ hf_config: dict[str, Any],
+ quant_config: QuantizationConfig | None = None,
+ ) -> None:
+ super().__init__(config=config, hf_config=hf_config)
+ arch = config.arch_config
+ self.arch = arch
+ self.hidden_size = arch.hidden_size
+ self.num_attention_heads = arch.num_attention_heads
+ self.num_channels_latents = arch.latents_dim
+ tp_size = get_tp_world_size()
+ ulysses_size, _ = _ulysses_ctx()
+ self._validate_tp_config(arch=arch, tp_size=tp_size)
+ self._validate_sequence_parallel_config(
+ arch=arch,
+ tp_size=tp_size,
+ ulysses_size=ulysses_size,
+ ring_size=_ring_world_size(),
+ )
+
+ self.video_patch_proj = ColumnParallelLinear(
+ arch.latents_dim
+ * arch.patch_size[0]
+ * arch.patch_size[1]
+ * arch.patch_size[2],
+ arch.hidden_size,
+ bias=True,
+ gather_output=True,
+ params_dtype=_FP32_DTYPE,
+ quant_config=None,
+ prefix="video_patch_proj",
+ )
+ self.audio_patch_proj = ColumnParallelLinear(
+ arch.audio_latents_dim,
+ arch.hidden_size,
+ bias=True,
+ gather_output=True,
+ params_dtype=_FP32_DTYPE,
+ quant_config=None,
+ prefix="audio_patch_proj",
+ )
+ self.condition_proj = ColumnParallelLinear(
+ arch.text_dim,
+ arch.hidden_size,
+ bias=True,
+ gather_output=True,
+ params_dtype=_BF16_DTYPE,
+ quant_config=quant_config,
+ prefix="condition_proj",
+ )
+ self.time_embedder = MiniMaxH3TimeEmbedder(
+ arch,
+ prefix="time_embedder",
+ )
+ self.rope = MiniMaxH3Rope(arch.rope_inv_freq_len)
+ self.token_refiner = MiniMaxH3TokenRefiner(
+ arch,
+ quant_config,
+ prefix="token_refiner",
+ )
+ self.blocks = nn.ModuleList(
+ [
+ MiniMaxH3DiTBlock(
+ arch,
+ quant_config,
+ prefix=f"blocks.{index}",
+ )
+ for index in range(arch.num_layers)
+ ]
+ )
+ self.layer_names = ["blocks"]
+ self.final_layer = MiniMaxH3FinalLayer(
+ arch,
+ quant_config,
+ prefix="final_layer",
+ )
+ self._resolved_attention_backend: AttentionBackendEnum | None = None
+ self._mark_missing_params_required()
+
+ def _resolve_attention_backend_once(self) -> None:
+ if self._resolved_attention_backend is not None:
+ return
+ backend = get_attn_backend(
+ self.arch.attention_head_dim,
+ _BF16_DTYPE,
+ supported_attention_backends=self._supported_attention_backends,
+ )
+ for module in self.modules():
+ if isinstance(module, MiniMaxH3Attention):
+ module._set_attention_backend(backend)
+ self._resolved_attention_backend = backend.get_enum()
+
+ def _mark_missing_params_required(self) -> None:
+ for _, param in self.named_parameters():
+ param.missing_param_init = "error"
+
+ def post_load_weights(self) -> None:
+ for name in _MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER:
+ param = self.get_parameter(name)
+ if param.dtype != _FP32_DTYPE:
+ raise ValueError(
+ f"{name} must stay fp32 after load, got {param.dtype}."
+ )
+ # assign=True loading may re-register this persistent buffer as a parameter
+ rope_inv_freq = self.rope.inv_freq
+ if rope_inv_freq.dtype != _FP32_DTYPE:
+ raise ValueError(
+ f"rope.inv_freq must stay fp32 after load, got {rope_inv_freq.dtype}."
+ )
+
+ @staticmethod
+ def _pos_ids(pos_info: Any, key: str) -> torch.Tensor:
+ if isinstance(pos_info, dict):
+ ids = pos_info.get("position_ids")
+ else:
+ ids = getattr(pos_info, "position_ids", None)
+ if ids is None:
+ raise ValueError(f"{key}.position_ids is required")
+ return ids.view(-1).to(torch.long)
+
+ @staticmethod
+ def _psp_field(psp: Any, key: str, field: str) -> Any:
+ if isinstance(psp, dict):
+ value = psp.get(field)
+ else:
+ value = getattr(psp, field, None)
+ if value is None:
+ raise ValueError(f"{key}.{field} is required")
+ return value
+
+ @staticmethod
+ def _psp_optional_field(psp: Any, field: str) -> Any:
+ if isinstance(psp, dict):
+ return psp.get(field)
+ return getattr(psp, field, None)
+
+ def refine_prompt_embeds(
+ self,
+ prompt_embeds: torch.Tensor,
+ refiner_cu_seqlens: torch.Tensor,
+ *,
+ device: torch.device,
+ ) -> torch.Tensor:
+ """Project and refine request-static text conditioning once."""
+ text_len = int(refiner_cu_seqlens[1].item())
+ if text_len <= 0 or text_len > int(prompt_embeds.shape[0]):
+ raise ValueError(
+ "refiner cu_seqlens live text length must be in "
+ f"[1, {int(prompt_embeds.shape[0])}], got {text_len}"
+ )
+ text_rows = prompt_embeds[:text_len].to(device=device, dtype=_BF16_DTYPE)
+ true_refiner_cu = torch.stack(
+ (
+ refiner_cu_seqlens[0],
+ refiner_cu_seqlens[1],
+ refiner_cu_seqlens[1],
+ )
+ )
+ text_embed, _ = self.condition_proj(text_rows)
+ return self.token_refiner(
+ text_embed,
+ cu_seqlens=true_refiner_cu,
+ cu_seqlens_host=(0, text_len, text_len),
+ max_seqlen=text_len,
+ )
+
+ def build_rope_cache(
+ self,
+ img_position_ids: torch.Tensor,
+ *,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Build request-static RoPE inputs for this Ulysses rank."""
+ if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1:
+ raise ValueError(
+ "img_position_ids must be [1, S, 3], got "
+ f"{list(img_position_ids.shape)}"
+ )
+ seq_len = int(img_position_ids.shape[1])
+ sp_ws, sp_rank = _ulysses_ctx()
+ if seq_len % sp_ws:
+ raise ValueError(
+ f"packed seq_len {seq_len} not divisible by ulysses world size {sp_ws}"
+ )
+ local_seq_len = seq_len // sp_ws
+ row_start = sp_rank * local_seq_len
+ rope_freqs = self.rope(
+ img_position_ids[:, row_start : row_start + local_seq_len]
+ ).to(device)
+ return (
+ _rope_cos_sin_cache(rope_freqs, dtype=_BF16_DTYPE),
+ torch.arange(
+ local_seq_len,
+ device=device,
+ dtype=torch.long,
+ ),
+ )
+
+ @eager_on_graph(True)
+ def _embed(
+ self,
+ *,
+ x: torch.Tensor,
+ audio_x: torch.Tensor,
+ text_embeddings_selected: torch.Tensor,
+ unique_timesteps: torch.Tensor,
+ img_pos: torch.Tensor,
+ audio_pos: torch.Tensor,
+ text_pos: torch.Tensor,
+ refiner_cu_seqlens: torch.Tensor,
+ refiner_max_seqlen: int,
+ row_start: int,
+ row_stop: int,
+ device: torch.device,
+ refined_prompt_embeds_length: int | torch.Tensor | None = None,
+ local_embedding_layout: dict[str, torch.Tensor | int] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Build embeddings for one contiguous block-stack row shard.
+
+ Returns (decoder_input [S_local, H] bf16, t_emb [M, t_dim] fp32).
+ """
+ # BCG pads the prompt tensor only to stabilize its input signature.
+ # Raw-input callers recover the live length from refiner metadata;
+ # request-static refined inputs carry it as a host integer and avoid a
+ # per-step device scalar read. Running the refiner at the bucketed M
+ # dimension changes GEMM selection and is not bitwise equivalent.
+ if refined_prompt_embeds_length is None:
+ text_len = int(refiner_cu_seqlens[1].item())
+ elif torch.is_tensor(refined_prompt_embeds_length):
+ # BCG turns this request-varying host constant into a scalar input
+ # so different live lengths can replay one padded-text signature.
+ # _embed is an eager graph break, so this value is read outside
+ # captured CUDA graphs.
+ text_len = int(refined_prompt_embeds_length.item())
+ else:
+ text_len = int(refined_prompt_embeds_length)
+ if text_len <= 0 or text_len > int(text_embeddings_selected.shape[0]):
+ raise ValueError(
+ "refiner cu_seqlens live text length must be in "
+ f"[1, {int(text_embeddings_selected.shape[0])}], got {text_len}"
+ )
+ text_pos = text_pos[:text_len]
+ if refined_prompt_embeds_length is not None:
+ text_embed = text_embeddings_selected[:text_len].to(
+ device=device, dtype=_BF16_DTYPE
+ )
+ if int(text_embed.shape[-1]) != self.hidden_size:
+ raise ValueError(
+ "refined prompt embeddings must have hidden width "
+ f"{self.hidden_size}, got {int(text_embed.shape[-1])}"
+ )
+ else:
+ text_embed = self.refine_prompt_embeds(
+ text_embeddings_selected,
+ refiner_cu_seqlens,
+ device=device,
+ )
+
+ local_seq_len = row_stop - row_start
+ trusted_layout = local_embedding_layout is not None
+ if trusted_layout:
+ used_len = text_len + int(img_pos.numel()) + int(audio_pos.numel())
+ local_live_rows = min(max(used_len - row_start, 0), local_seq_len)
+ embeddings = torch.empty(
+ (local_seq_len, self.hidden_size), device=device, dtype=_BF16_DTYPE
+ )
+ if local_live_rows < local_seq_len:
+ embeddings[local_live_rows:].zero_()
+ else:
+ # Direct model callers do not provide the serving-time partition
+ # contract. Preserve their historical zero-fill/add semantics for
+ # sparse or overlapping row maps.
+ embeddings = torch.zeros(
+ (local_seq_len, self.hidden_size), device=device, dtype=_BF16_DTYPE
+ )
+
+ if local_embedding_layout is None:
+ text_source_ids = torch.nonzero(
+ (text_pos >= row_start) & (text_pos < row_stop),
+ as_tuple=False,
+ ).view(-1)
+ text_row_ids = text_pos.index_select(0, text_source_ids) - row_start
+ img_global_ids = img_pos.index_select(
+ 0,
+ torch.nonzero(
+ (img_pos >= row_start) & (img_pos < row_stop),
+ as_tuple=False,
+ ).view(-1),
+ )
+ img_row_ids = img_global_ids - row_start
+ audio_global_ids = audio_pos.index_select(
+ 0,
+ torch.nonzero(
+ (audio_pos >= row_start) & (audio_pos < row_stop),
+ as_tuple=False,
+ ).view(-1),
+ )
+ audio_row_ids = audio_global_ids - row_start
+ else:
+ text_source_start = int(local_embedding_layout["text_source_start"])
+ text_source_stop = int(local_embedding_layout["text_source_stop"])
+ img_global_ids = local_embedding_layout["img_global_ids"]
+ img_row_ids = local_embedding_layout["img_row_ids"]
+ audio_global_ids = local_embedding_layout["audio_global_ids"]
+ audio_row_ids = local_embedding_layout["audio_row_ids"]
+
+ write_rows = embeddings.index_copy_ if trusted_layout else embeddings.index_add_
+ if trusted_layout:
+ text_rows = text_source_stop - text_source_start
+ if text_rows:
+ embeddings[:text_rows].copy_(
+ text_embed[text_source_start:text_source_stop]
+ )
+ elif text_row_ids.numel():
+ write_rows(
+ 0,
+ text_row_ids,
+ text_embed.index_select(0, text_source_ids).to(_BF16_DTYPE),
+ )
+
+ # latent embedders stay fp32; only rows owned by this SP rank are
+ # projected, then cast during scattering into the bf16 sequence
+ if img_row_ids.numel():
+ x_rows = (
+ x.view(-1, x.shape[-1]).index_select(0, img_global_ids).to(_FP32_DTYPE)
+ )
+ video_embed, _ = self.video_patch_proj(x_rows)
+ write_rows(
+ 0,
+ img_row_ids,
+ video_embed.to(_BF16_DTYPE),
+ )
+
+ if audio_row_ids.numel():
+ audio_rows = (
+ audio_x.view(-1, audio_x.shape[-1])
+ .index_select(0, audio_global_ids)
+ .to(_FP32_DTYPE)
+ )
+ audio_embed, _ = self.audio_patch_proj(audio_rows)
+ write_rows(
+ 0,
+ audio_row_ids,
+ audio_embed.to(_BF16_DTYPE),
+ )
+
+ t_emb = self.time_embedder(unique_timesteps)
+ return embeddings, t_emb
+
+ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
+ """Packed inference forward.
+
+ Keyword names follow the checkpoint's serving contract.
+ Returns `(video_logits, audio_logits)` from rows selected by
+ `img_pos_for_infer_output_info` and `audio_pos_info`, with condition
+ rows zeroed by update masks.
+ """
+ # Strict keyword contract: refuse any kwarg forward does not consume.
+ unexpected = sorted(set(kwargs) - _FORWARD_SUPPORTED_KWARGS)
+ if unexpected:
+ raise TypeError(
+ "MiniMaxH3DiTModel.forward received unexpected kwargs: "
+ f"{unexpected}; supported kwargs: "
+ f"{sorted(_FORWARD_SUPPORTED_KWARGS)}"
+ )
+
+ x = _required_kwarg(kwargs, "x")
+ audio_x = _required_kwarg(kwargs, "audio_x")
+ img_position_ids = _required_kwarg(kwargs, "img_position_ids")
+ unique_timesteps = _required_kwarg(kwargs, "unique_timesteps")
+ inverse_indices = (
+ _required_kwarg(kwargs, "inverse_indices").view(-1).to(torch.long)
+ )
+ update_mask = _required_kwarg(kwargs, "update_mask")
+ block_token_tags = kwargs.get("block_token_tags")
+ token_tags = kwargs.get("token_tags")
+ if block_token_tags is None:
+ token_tags = _required_kwarg(kwargs, "token_tags").view(-1).to(torch.long)
+ else:
+ block_token_tags = block_token_tags.view(-1).to(torch.long)
+ token_tags = None
+ skip_mask_out_condition = bool(kwargs.get("skip_mask_out_condition", False))
+
+ text_selected = _required_kwarg(kwargs, "prompt_embeds")
+
+ img_pos = self._pos_ids(_required_kwarg(kwargs, "img_pos_info"), "img_pos_info")
+ audio_pos = self._pos_ids(
+ _required_kwarg(kwargs, "audio_pos_info"), "audio_pos_info"
+ )
+ text_pos = self._pos_ids(
+ _required_kwarg(kwargs, "text_pos_info"),
+ "text_pos_info",
+ )
+ infer_out_pos = self._pos_ids(
+ _required_kwarg(kwargs, "img_pos_for_infer_output_info"),
+ "img_pos_for_infer_output_info",
+ )
+
+ psp = _required_kwarg(kwargs, "packed_seq_params")
+ cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(
+ torch.int32
+ )
+ raw_cu_seqlens_host = self._psp_optional_field(psp, "cu_seqlens_q_host")
+ cu_seqlens_host = tuple(
+ int(value)
+ for value in (
+ cu_seqlens.tolist()
+ if raw_cu_seqlens_host is None
+ else raw_cu_seqlens_host
+ )
+ )
+ max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
+ refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
+ refiner_cu = self._psp_field(
+ refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q"
+ ).to(torch.int32)
+ refiner_max = int(
+ self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")
+ )
+
+ if x.dim() != 3 or x.shape[0] != 1:
+ raise ValueError(f"x must be [1, S, C], got {list(x.shape)}")
+ seq_len = int(x.shape[1])
+ if token_tags is not None and token_tags.shape[0] != seq_len:
+ raise ValueError(
+ "token_tags must cover the full packed sequence "
+ f"({seq_len}), got {token_tags.shape[0]}."
+ )
+ if inverse_indices.shape[0] != seq_len:
+ raise ValueError(
+ f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}"
+ )
+ device = x.device
+ self._resolve_attention_backend_once()
+ if _ring_world_size() != 1:
+ raise NotImplementedError(
+ "MiniMax H3 packed multi-segment attention requires "
+ "--ring-degree 1; Ring and mixed USP are unsupported."
+ )
+
+ sp_ws, sp_rank = _ulysses_ctx()
+ local_seq_len = seq_len
+ if sp_ws > 1:
+ if seq_len % sp_ws:
+ raise ValueError(
+ f"packed seq_len {seq_len} not divisible by ulysses "
+ f"world size {sp_ws}"
+ )
+ local_heads = self.num_attention_heads // get_tp_world_size()
+ if local_heads % sp_ws:
+ raise ValueError(
+ f"TP-local heads {local_heads} not divisible by Ulysses "
+ f"world size {sp_ws} (total heads={self.num_attention_heads}, "
+ f"TP={get_tp_world_size()})"
+ )
+ local_seq_len = seq_len // sp_ws
+ row_start = sp_rank * local_seq_len
+ row_stop = row_start + local_seq_len
+
+ # RoPE and latent projections are row-local before Ulysses exchanges
+ # sequence for heads inside attention. Serving normally prepares the
+ # request-static cache once; direct model callers use this fallback.
+ rope_cache = kwargs.get("rope_cache")
+ if rope_cache is None:
+ rope_freqs = self.rope(img_position_ids[:, row_start:row_stop]).to(device)
+ rope_cache = (
+ _rope_cos_sin_cache(rope_freqs, dtype=_BF16_DTYPE),
+ torch.arange(
+ local_seq_len,
+ device=device,
+ dtype=torch.long,
+ ),
+ )
+ img_pos = img_pos.to(device)
+ audio_pos = audio_pos.to(device)
+ text_pos = text_pos.to(device)
+
+ decoder_input, t_emb = self._embed(
+ x=x,
+ audio_x=audio_x,
+ text_embeddings_selected=text_selected,
+ unique_timesteps=unique_timesteps.view(-1).to(device),
+ img_pos=img_pos,
+ audio_pos=audio_pos,
+ text_pos=text_pos,
+ refiner_cu_seqlens=refiner_cu.to(device),
+ refiner_max_seqlen=refiner_max,
+ row_start=row_start,
+ row_stop=row_stop,
+ device=device,
+ refined_prompt_embeds_length=kwargs.get("refined_prompt_embeds_length"),
+ local_embedding_layout=kwargs.get("local_embedding_layout"),
+ )
+ # request-step AdaLN input shared by all blocks
+ adaln_input = nn.functional.silu(t_emb).to(_BF16_DTYPE)
+ inverse_indices = inverse_indices.to(device)
+ block_inverse = inverse_indices[row_start:row_stop]
+ if block_token_tags is None:
+ assert token_tags is not None
+ token_tags = token_tags.to(device)
+ block_token_tags = token_tags[row_start:row_stop].clamp(min=0)
+ else:
+ block_token_tags = block_token_tags.to(device)
+ if block_token_tags.shape[0] != local_seq_len:
+ raise ValueError(
+ "block_token_tags must cover the rank-local packed sequence "
+ f"({local_seq_len}), got {block_token_tags.shape[0]}."
+ )
+ block_combined = kwargs.get("block_combined_indices")
+ if block_combined is None:
+ block_combined = torch.add(
+ block_token_tags,
+ block_inverse,
+ alpha=MINIMAX_H3_ADALN_MODALITY_NUM,
+ )
+
+ hidden = decoder_input
+ cu_seqlens = cu_seqlens.to(device)
+ block_adaln_params = None
+ if self._can_batch_block_adaln():
+ local_adaln = torch.stack(
+ [block.adaln_proj.project_local(adaln_input) for block in self.blocks]
+ )
+ gathered_adaln = tensor_model_parallel_all_gather(local_adaln)
+ block_adaln_params = tuple(
+ block.adaln_proj.split_output(output)
+ for block, output in zip(self.blocks, gathered_adaln)
+ )
+ # With Ulysses sequence parallelism, shard rows across the group for
+ # the block stack. Attention trades sequence for heads internally;
+ # everything else, including the final layer, is row-local.
+ for index, block in enumerate(self.blocks):
+ hidden = block(
+ hidden,
+ adaln_input=adaln_input,
+ combined_indices=block_combined,
+ rope_cache=rope_cache,
+ cu_seqlens=cu_seqlens,
+ cu_seqlens_host=cu_seqlens_host,
+ max_seqlen=max_seqlen,
+ ulysses_active=sp_ws > 1,
+ adaln_params=(
+ None if block_adaln_params is None else block_adaln_params[index]
+ ),
+ )
+ video_logits, audio_logits = self.final_layer(
+ hidden,
+ adaln_input=adaln_input,
+ inverse_indices=block_inverse,
+ )
+ if sp_ws > 1:
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_sp_group,
+ )
+
+ video_width = video_logits.shape[-1]
+ logits = get_sp_group().all_gather(
+ torch.cat((video_logits, audio_logits), dim=-1), dim=0
+ )
+ video_logits, audio_logits = logits.split(
+ (video_width, logits.shape[-1] - video_width), dim=-1
+ )
+
+ # Preserve the full-row output GEMM (and therefore its numerical
+ # contract), but defer TP column gathers until after dead text/padding
+ # rows have been removed. For hybrid TP+Ulysses, the preceding SP row
+ # gather also carries only the TP-local output width.
+ video_logits = video_logits.index_select(0, infer_out_pos.to(device))
+ audio_logits = audio_logits.index_select(0, audio_pos.to(device))
+ if get_tp_world_size() > 1:
+ video_logits = tensor_model_parallel_all_gather(video_logits)
+ audio_logits = tensor_model_parallel_all_gather(audio_logits)
+ if not skip_mask_out_condition:
+ update_mask = update_mask.view(-1).to(device)
+ if update_mask.shape[0] != video_logits.shape[0]:
+ raise ValueError(
+ "update_mask length mismatch: "
+ f"{update_mask.shape[0]} != {video_logits.shape[0]}"
+ )
+ video_logits = video_logits * update_mask.unsqueeze(-1)
+ # Audio has no condition rows in the supported tasks, so its
+ # derived update mask is all ones. Honor an explicit mask when
+ # provided.
+ update_audio_mask = kwargs.get("update_audio_mask")
+ if update_audio_mask is not None:
+ audio_logits = audio_logits * update_audio_mask.view(-1).unsqueeze(-1)
+ return video_logits, audio_logits
+
+
+EntryClass = MiniMaxH3DiTModel
+
+__all__ = [
+ "MINIMAX_H3_FP32_BUFFER_NAMES",
+ "MINIMAX_H3_FP32_PARAM_NAMES",
+ "MiniMaxH3DiTModel",
+ "_reorder_grouped_qkv_to_qkv",
+]
diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py
index 08fc03174..33727bd95 100644
--- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py
+++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py
@@ -123,11 +123,11 @@ def encoder_dp_worthwhile(
def finalize_encoder_folding(
- config: EncoderConfig, policy: str = "auto", batched: bool = False
+ config: EncoderConfig, policy: str = "auto", prefer_dp: bool = False
) -> None:
"""resolve fold-vs-replicate once real dims are known (post update_model_arch,
pre construction); folding shards the weights, so it rules out dp for the
- lifetime of the loaded model. `batched` is the batching ceiling being > 1."""
+ lifetime of the loaded model. `prefer_dp` means the runtime can engage dp."""
if config.parallel_folding_mode is None:
return
group = get_folding_tp_group(config)
@@ -138,7 +138,7 @@ def finalize_encoder_folding(
# a batched encode prefers dp (one all_gather) over folding (an
# all_reduce per layer), so leave a dp-capable encoder unsharded
keep = (
- not (batched and encoder_dp_capable(config))
+ not (prefer_dp and encoder_dp_capable(config))
and encoder_folding_worthwhile(config, group.world_size)
and group_has_measured_topology(group)
)
diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py
new file mode 100644
index 000000000..b7ff05638
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py
@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native, TP-foldable Qwen3-VL layer-50 encoder for MiniMax H3."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
+ MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
+ MiniMaxH3Qwen3VLConfig,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
+from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel
+
+MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
+_LAYER_WEIGHT_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.")
+
+
+def _is_unconsumed_checkpoint_weight(name: str) -> bool:
+ """Weights intentionally absent from the layer-50 feature extractor."""
+
+ if name == "lm_head.weight" or name.startswith("model.language_model.norm."):
+ return True
+ match = _LAYER_WEIGHT_RE.match(name)
+ return bool(match and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER)
+
+
+class MiniMaxH3Qwen3VLEncoder(TextEncoder):
+ """Qwen3-VL-32B multimodal backbone ending at hidden_states[50].
+
+ The component loader builds and loads this module under the encoder-folding
+ TP group. A TP=1/SP=8 DiT deployment therefore shards the encoder over all
+ eight otherwise-idle ranks during encoding.
+ """
+
+ supports_dp_encode = True
+
+ @staticmethod
+ def should_materialize_checkpoint_weight(name: str) -> bool:
+ return (
+ "rotary_emb.inv_freq" not in name
+ and not _is_unconsumed_checkpoint_weight(name)
+ )
+
+ def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None:
+ super().__init__(config)
+ arch = config.arch_config
+ selected_layer = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
+ if int(arch.text_config.num_hidden_layers) != selected_layer:
+ raise ValueError(
+ "MiniMax H3 Qwen3-VL config must be trimmed to "
+ f"{selected_layer} language layers before construction"
+ )
+ self.model = Qwen3VLModel(arch, use_tensor_parallel=True)
+ # H3 consumes the unnormalized output immediately after layer 49.
+ self.model.language_model.norm = nn.Identity()
+ self.image_token_id = int(arch.image_token_id)
+ self.video_token_id = int(arch.video_token_id)
+ self.selected_lm_layer = selected_layer
+ self.hidden_dim = MINIMAX_H3_QWEN3VL_HIDDEN_DIM
+
+ @property
+ def device(self) -> torch.device:
+ return next(self.parameters()).device
+
+ @torch.no_grad()
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ position_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ **kwargs: Any,
+ ) -> BaseEncoderOutput:
+ outputs = self.model(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ use_cache=False,
+ **kwargs,
+ )
+ return BaseEncoderOutput(last_hidden_state=outputs.last_hidden_state)
+
+ @torch.no_grad()
+ def encode_ids(
+ self,
+ input_ids: torch.Tensor,
+ *,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.Tensor | None = None,
+ pixel_values_videos: torch.Tensor | None = None,
+ video_grid_thw: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ if input_ids.dim() != 1:
+ raise ValueError(f"input_ids must be 1-D, got {list(input_ids.shape)}")
+ if (pixel_values is None) != (image_grid_thw is None):
+ raise ValueError("pixel_values and image_grid_thw must be given together")
+ if (pixel_values_videos is None) != (video_grid_thw is None):
+ raise ValueError(
+ "pixel_values_videos and video_grid_thw must be given together"
+ )
+
+ host_ids = input_ids.to(device="cpu", dtype=torch.long)[None]
+ host_image_grid_thw = (
+ image_grid_thw.to(device="cpu", dtype=torch.long)
+ if image_grid_thw is not None
+ else None
+ )
+ host_video_grid_thw = (
+ video_grid_thw.to(device="cpu", dtype=torch.long)
+ if video_grid_thw is not None
+ else None
+ )
+ position_ids = None
+ if host_image_grid_thw is not None or host_video_grid_thw is not None:
+ position_ids, _ = self.model.get_rope_index(
+ host_ids,
+ host_image_grid_thw,
+ host_video_grid_thw,
+ attention_mask=torch.ones_like(host_ids),
+ )
+ ids = host_ids.to(self.device)
+ call_kwargs: dict[str, Any] = {
+ "input_ids": ids,
+ "attention_mask": torch.ones_like(ids),
+ "output_attentions": False,
+ "output_hidden_states": False,
+ "return_dict": True,
+ "use_cache": False,
+ }
+ if position_ids is not None:
+ call_kwargs["position_ids"] = position_ids.to(self.device)
+ if pixel_values is not None:
+ call_kwargs["pixel_values"] = pixel_values.to(self.device, torch.bfloat16)
+ call_kwargs["image_grid_thw"] = host_image_grid_thw
+ if pixel_values_videos is not None:
+ call_kwargs["pixel_values_videos"] = pixel_values_videos.to(
+ self.device, torch.bfloat16
+ )
+ call_kwargs["video_grid_thw"] = host_video_grid_thw
+
+ hidden = self.model(**call_kwargs).last_hidden_state[0].to(torch.bfloat16)
+ expected_shape = [int(ids.shape[1]), self.hidden_dim]
+ if list(hidden.shape) != expected_shape:
+ raise ValueError(
+ f"unexpected hidden shape {list(hidden.shape)}, "
+ f"expected {expected_shape}"
+ )
+ return hidden
+
+ def load_weights(
+ self,
+ weights: Iterable[tuple[str, torch.Tensor]],
+ ) -> set[str]:
+ params = dict(self.named_parameters(remove_duplicate=False))
+ loaded: set[str] = set()
+ for name, loaded_weight in weights:
+ if not self.should_materialize_checkpoint_weight(name):
+ continue
+ param = params.get(name)
+ if param is None:
+ raise KeyError(
+ f"Unexpected MiniMax H3 Qwen3-VL checkpoint weight: {name}"
+ )
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
+ try:
+ weight_loader(param, loaded_weight.to(param.dtype))
+ except Exception as exc:
+ raise RuntimeError(
+ "Failed to load MiniMax H3 Qwen3-VL weight "
+ f"{name!r}: checkpoint={tuple(loaded_weight.shape)}, "
+ f"parameter={tuple(param.shape)}"
+ ) from exc
+ loaded.add(name)
+ return loaded
+
+
+EntryClass = MiniMaxH3Qwen3VLEncoder
+
+__all__ = ["MiniMaxH3Qwen3VLEncoder"]
diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py
index a91fdd8dd..a74a03f38 100644
--- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py
+++ b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py
@@ -4,9 +4,10 @@ from transformers import (
Cache,
DynamicCache,
)
-from transformers.masking_utils import create_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.utils import TransformersKwargs, is_torchdynamo_compiling
+from transformers.utils.generic import is_flash_attention_requested
+from transformers.vision_utils import get_vision_cu_seqlens, get_vision_position_ids
from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig
from sglang.multimodal_gen.runtime.distributed import (
@@ -208,7 +209,11 @@ class Qwen3VLTextAttention(nn.Module):
super().__init__()
self.config = config
self.layer_idx = layer_idx
- self.head_dim = config.hidden_size // config.num_attention_heads
+ self.head_dim = (
+ int(config.head_dim)
+ if getattr(config, "head_dim", None) is not None
+ else config.hidden_size // config.num_attention_heads
+ )
self.total_num_heads = config.num_attention_heads
self.total_num_key_value_heads = config.num_key_value_heads
tp_size = _tp_world_size() if use_tensor_parallel else 1
@@ -582,14 +587,6 @@ class Qwen3VLTextModel(nn.Module):
else:
text_position_ids = position_ids[0]
- attention_mask = create_causal_mask(
- config=self.config,
- inputs_embeds=inputs_embeds,
- attention_mask=attention_mask,
- past_key_values=past_key_values,
- position_ids=text_position_ids,
- )
-
hidden_states = inputs_embeds
# create position embeddings to be shared across the decoder layers
@@ -651,7 +648,8 @@ class Qwen3VLTextModel(nn.Module):
):
visual_pos_masks = visual_pos_masks.to(hidden_states.device)
visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
- local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
+ local_this = hidden_states[visual_pos_masks, :]
+ local_this.add_(visual_embeds)
hidden_states[visual_pos_masks, :] = local_this
return hidden_states
@@ -664,10 +662,13 @@ class Qwen3VLModel(nn.Module):
config: Qwen3VLConfig
_no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
- def __init__(self, config):
+ def __init__(self, config, *, use_tensor_parallel: bool = False):
super().__init__()
self.visual = Qwen3VLVisionModel._from_config(config.vision_config)
- self.language_model = Qwen3VLTextModel(config.text_config)
+ self.language_model = Qwen3VLTextModel(
+ config.text_config,
+ use_tensor_parallel=use_tensor_parallel,
+ )
self.rope_deltas = None # cache rope_deltas here
self.config = config
@@ -868,6 +869,25 @@ class Qwen3VLModel(nn.Module):
# Same implementation as for images
return self.get_image_features(pixel_values_videos, video_grid_thw)
+ def _get_flat_visual_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ grid_thw: Optional[torch.LongTensor],
+ ):
+ pixel_values = pixel_values.type(self.visual.dtype)
+ vision_kwargs = {}
+ if grid_thw is not None and grid_thw.device.type == "cpu":
+ if not is_flash_attention_requested(self.visual.config):
+ vision_kwargs = {
+ "position_ids": get_vision_position_ids(
+ grid_thw, self.visual.spatial_merge_size
+ ).to(pixel_values.device),
+ "cu_seqlens": get_vision_cu_seqlens(grid_thw),
+ }
+ grid_thw = grid_thw.to(pixel_values.device)
+ visual_out = self.visual(pixel_values, grid_thw=grid_thw, **vision_kwargs)
+ return visual_out.pooler_output, visual_out.deepstack_features
+
def get_image_features(
self,
pixel_values: torch.FloatTensor,
@@ -882,10 +902,9 @@ class Qwen3VLModel(nn.Module):
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
The temporal, height and width of feature shape of each image in LLM.
"""
- pixel_values = pixel_values.type(self.visual.dtype)
- visual_out = self.visual(pixel_values, grid_thw=image_grid_thw)
- image_embeds = visual_out.pooler_output
- deepstack_image_embeds = visual_out.deepstack_features
+ image_embeds, deepstack_image_embeds = self._get_flat_visual_features(
+ pixel_values, image_grid_thw
+ )
split_sizes = (
image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2
).tolist()
@@ -996,35 +1015,40 @@ class Qwen3VLModel(nn.Module):
return_dict if return_dict is not None else self.config.use_return_dict
)
- if inputs_embeds is None:
+ inputs_embeds_owned = inputs_embeds is None
+ if inputs_embeds_owned:
inputs_embeds = self.get_input_embeddings()(input_ids)
image_mask = None
video_mask = None
if pixel_values is not None:
- image_embeds, deepstack_image_embeds = self.get_image_features( # long
+ image_embeds, deepstack_image_embeds = self._get_flat_visual_features(
pixel_values, image_grid_thw
)
- image_embeds = torch.cat(image_embeds, dim=0).to(
- inputs_embeds.device, inputs_embeds.dtype
- )
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
image_mask, _ = self.get_placeholder_mask(
input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
)
- inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+ if inputs_embeds_owned:
+ inputs_embeds.masked_scatter_(image_mask, image_embeds)
+ else:
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+ inputs_embeds_owned = True
if pixel_values_videos is not None:
- video_embeds, deepstack_video_embeds = self.get_video_features(
+ video_embeds, deepstack_video_embeds = self._get_flat_visual_features(
pixel_values_videos, video_grid_thw
)
- video_embeds = torch.cat(video_embeds, dim=0).to(
- inputs_embeds.device, inputs_embeds.dtype
- )
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
_, video_mask = self.get_placeholder_mask(
input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
)
- inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+ if inputs_embeds_owned:
+ inputs_embeds.masked_scatter_(video_mask, video_embeds)
+ else:
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+ inputs_embeds_owned = True
visual_pos_masks = None
deepstack_visual_embeds = None
@@ -1040,8 +1064,8 @@ class Qwen3VLModel(nn.Module):
deepstack_image_embeds, deepstack_video_embeds
):
embed_joint = img_embed.new_zeros(
- visual_pos_masks.sum(), img_embed.shape[-1]
- ).to(img_embed.device)
+ img_embed.shape[0] + vid_embed.shape[0], img_embed.shape[-1]
+ )
embed_joint[image_mask_joint, :] = img_embed
embed_joint[video_mask_joint, :] = vid_embed
deepstack_visual_embeds.append(embed_joint)
diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py
new file mode 100644
index 000000000..6c067b0a7
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_minimax_h3_euler_ancestral.py
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+
+
+def _require_finite_tensor(tensor: torch.Tensor, name: str) -> None:
+ if not bool(torch.isfinite(tensor).all().item()):
+ raise ValueError(f"{name} must be finite")
+
+
+def _validate_unit_timestep(timestep: torch.Tensor, name: str) -> None:
+ if not isinstance(timestep, torch.Tensor):
+ raise ValueError(f"{name} must be a torch.Tensor")
+ if not torch.is_floating_point(timestep):
+ raise ValueError(f"{name} must be a floating point tensor")
+ _require_finite_tensor(timestep, name)
+ out_of_range = (timestep < 0) | (timestep > 1)
+ if bool(out_of_range.any().item()):
+ raise ValueError(f"{name} must be in [0, 1]")
+
+
+def _validate_sigma(value: float, name: str) -> float:
+ sigma = float(value)
+ if not math.isfinite(sigma):
+ raise ValueError(f"{name} must be finite")
+ if sigma < 0.0:
+ raise ValueError(f"{name} must be non-negative")
+ return sigma
+
+
+def _validate_timestep_sigma_pair(
+ timestep: torch.Tensor,
+ sigma_curr: float,
+ name: str,
+) -> float:
+ _validate_unit_timestep(timestep, f"{name}_timestep")
+ sigma = _validate_sigma(sigma_curr, f"{name}_sigma_curr")
+ expected = 1.0 - timestep.detach().to(dtype=torch.float32)
+ actual = torch.full_like(expected, sigma)
+ if not torch.allclose(actual, expected, rtol=1e-5, atol=1e-5):
+ raise ValueError(f"{name}_sigma_curr must equal 1 - {name}_timestep")
+ return sigma
+
+
+def minimax_h3_rf_v_to_x0(
+ xt: torch.Tensor,
+ v: torch.Tensor,
+ timestep: torch.Tensor,
+) -> torch.Tensor:
+ if xt.shape != v.shape:
+ raise ValueError(f"xt and v shapes must match, got {xt.shape} vs {v.shape}")
+ if not torch.is_floating_point(xt):
+ raise ValueError("xt must be a floating point tensor")
+ if not torch.is_floating_point(v):
+ raise ValueError("v must be a floating point tensor")
+ _require_finite_tensor(xt, "xt")
+ _require_finite_tensor(v, "v")
+ _validate_unit_timestep(timestep, "timestep")
+ x0 = _minimax_h3_rf_v_to_x0(xt, v, timestep)
+ _require_finite_tensor(x0, "x0")
+ return x0
+
+
+def _minimax_h3_rf_v_to_x0(
+ xt: torch.Tensor,
+ v: torch.Tensor,
+ timestep: torch.Tensor,
+) -> torch.Tensor:
+ cond_t = timestep.to(device=xt.device, dtype=xt.dtype)
+ while cond_t.ndim < xt.ndim:
+ cond_t = cond_t.unsqueeze(-1)
+ sigma_t = 1 - cond_t
+ return xt + sigma_t * v
+
+
+def minimax_h3_euler_eta0_step(
+ state: torch.Tensor,
+ denoised: torch.Tensor,
+ *,
+ sigma_curr: float,
+ sigma_next: float,
+) -> torch.Tensor:
+ if state.shape != denoised.shape:
+ raise ValueError(
+ f"state and denoised shapes must match, got {state.shape} vs "
+ f"{denoised.shape}"
+ )
+ if not torch.is_floating_point(state):
+ raise ValueError("state must be a floating point tensor")
+ if not torch.is_floating_point(denoised):
+ raise ValueError("denoised must be a floating point tensor")
+ _require_finite_tensor(state, "state")
+ _require_finite_tensor(denoised, "denoised")
+ sigma_curr = _validate_sigma(sigma_curr, "sigma_curr")
+ sigma_next = _validate_sigma(sigma_next, "sigma_next")
+ if sigma_curr == 0.0 and sigma_next != 0.0:
+ raise ValueError("sigma_next must be 0 when sigma_curr is 0")
+ out = _minimax_h3_euler_eta0_step(
+ state,
+ denoised,
+ sigma_curr=sigma_curr,
+ sigma_next=sigma_next,
+ )
+ _require_finite_tensor(out, "euler_eta0_step output")
+ return out
+
+
+def _minimax_h3_euler_eta0_step(
+ state: torch.Tensor,
+ denoised: torch.Tensor,
+ *,
+ sigma_curr: float,
+ sigma_next: float,
+ sigma_ratio: torch.Tensor | None = None,
+) -> torch.Tensor:
+ if sigma_curr == 0.0:
+ return state
+ compute_dtype = torch.float32
+ if state.dtype not in (torch.float16, torch.bfloat16):
+ compute_dtype = state.dtype
+ if sigma_ratio is None:
+ sigma_curr_t = state.new_tensor(sigma_curr, dtype=compute_dtype)
+ sigma_next_t = state.new_tensor(sigma_next, dtype=compute_dtype)
+ ratio = sigma_next_t / sigma_curr_t
+ else:
+ ratio = sigma_ratio.to(device=state.device, dtype=compute_dtype)
+ out = ratio * state.to(dtype=compute_dtype) + (1.0 - ratio) * denoised.to(
+ dtype=compute_dtype
+ )
+ return out.to(dtype=state.dtype)
+
+
+class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
+ def __init__(self, **config: Any) -> None:
+ if config:
+ raise ValueError(
+ f"{type(self).__name__} does not accept config fields: "
+ f"{sorted(config)}"
+ )
+
+ def set_shift(self, _flow_shift: float) -> None:
+ """Ignore flow shift, matching the previous loader-specific path."""
+
+ def step_denoising(
+ self,
+ *,
+ input_visual_latent: torch.Tensor,
+ input_audio_latent: torch.Tensor,
+ timestep: torch.Tensor,
+ noise_pred_visual: torch.Tensor,
+ noise_pred_audio: torch.Tensor,
+ sigma_curr: float,
+ sigma_next: float,
+ video_timestep: torch.Tensor | None = None,
+ audio_timestep: torch.Tensor | None = None,
+ video_sigma_curr: float | None = None,
+ video_sigma_next: float | None = None,
+ audio_sigma_curr: float | None = None,
+ audio_sigma_next: float | None = None,
+ ) -> dict[str, torch.Tensor]:
+ visual_timestep = timestep if video_timestep is None else video_timestep
+ audio_timestep = timestep if audio_timestep is None else audio_timestep
+ visual_sigma_curr = sigma_curr if video_sigma_curr is None else video_sigma_curr
+ visual_sigma_next = sigma_next if video_sigma_next is None else video_sigma_next
+ audio_sigma_curr = sigma_curr if audio_sigma_curr is None else audio_sigma_curr
+ audio_sigma_next = sigma_next if audio_sigma_next is None else audio_sigma_next
+ visual_sigma_curr = _validate_timestep_sigma_pair(
+ visual_timestep,
+ visual_sigma_curr,
+ "video",
+ )
+ audio_sigma_curr = _validate_timestep_sigma_pair(
+ audio_timestep,
+ audio_sigma_curr,
+ "audio",
+ )
+
+ denoised_visual = minimax_h3_rf_v_to_x0(
+ input_visual_latent,
+ noise_pred_visual,
+ visual_timestep,
+ )
+ denoised_audio = minimax_h3_rf_v_to_x0(
+ input_audio_latent,
+ noise_pred_audio,
+ audio_timestep,
+ )
+ return {
+ "output_visual_latent": minimax_h3_euler_eta0_step(
+ input_visual_latent,
+ denoised_visual,
+ sigma_curr=visual_sigma_curr,
+ sigma_next=visual_sigma_next,
+ ),
+ "output_audio_latent": minimax_h3_euler_eta0_step(
+ input_audio_latent,
+ denoised_audio,
+ sigma_curr=audio_sigma_curr,
+ sigma_next=audio_sigma_next,
+ ),
+ }
+
+
+EntryClass = MiniMaxH3EulerAncestralEta0SchedulerAdapter
+
+__all__ = [
+ "MiniMaxH3EulerAncestralEta0SchedulerAdapter",
+ "minimax_h3_euler_eta0_step",
+ "minimax_h3_rf_v_to_x0",
+]
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3.py
new file mode 100644
index 000000000..26f256b47
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3.py
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
+ MiniMaxH3AudioVAEConfig,
+)
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
+ MiniMaxH3VideoVAEConfig,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+ LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.minimax_h3_audio_vae import (
+ DacAudioVAE,
+)
+from sglang.multimodal_gen.runtime.models.vaes.minimax_h3_video_vae import (
+ AutoencoderKLLegacy,
+)
+
+
+class MiniMaxH3VideoVAE(AutoencoderKLLegacy, LayerwiseOffloadableModuleMixin):
+ layerwise_offload_dit_group_enabled = False
+ # EncoderFCN3D indexes its down containers instead of calling them, so they
+ # cannot host layerwise hooks. Keep the small encoder resident.
+ layer_names = ["decoder.transformer_blocks"]
+
+ def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None:
+ arch = config.arch_config
+ parallel_decode_mode = config.resolved_parallel_decode_mode()
+ use_tiled_decode = config.use_tiling and parallel_decode_mode == "tiled"
+ super().__init__(
+ in_channels=3,
+ out_ch=3,
+ ch=128,
+ embed_dim=24,
+ z_channels=24,
+ use_3d_conv=True,
+ zq_ch_encoder=None,
+ zq_ch_decoder=None,
+ num_res_blocks=2,
+ num_res_blocks_decoder=None,
+ ch_mult=[1, 2, 2, 4, 4, 8],
+ space_down=[2, 2, 2, 2, 1, 1],
+ space_up=[1, 2, 2, 2, 2, 1],
+ time_down=[1, 2, 2, 1, 1, 1],
+ time_up=None,
+ padding_mode="reflect",
+ padding_mode_t=None,
+ use_t_isolated_gn=True,
+ causal_encoder=True,
+ causal_decoder=False,
+ use_vit_decoder=True,
+ vit_decoder_kwargs={
+ "dim_head": 64,
+ "ffn_activation_fn": "silu",
+ "ffn_use_gated": True,
+ "heads": 32,
+ "norm_affine": True,
+ "norm_type": "rms_norm",
+ "num_layers": 36,
+ "qk_norm_affine": False,
+ "qk_norm_type": "rms_norm",
+ "rope_dim_ratio": 0.75,
+ "rope_theta": 100.0,
+ },
+ shift_factor=0.0,
+ scaling_factor=1.0,
+ pixel_norm_type="imagenet",
+ clip_length=arch.vae_clip_length,
+ token_drop=arch.vae_token_drop,
+ encoder_tiling=bool(arch.vae_encoder_tiling),
+ decoder_tiling=use_tiled_decode,
+ parallel_tiling=use_tiled_decode
+ and config.use_parallel_decode
+ and config.use_parallel_tiling
+ and bool(arch.vae_parallel_tiling),
+ tile_size=int(arch.vae_tile_size),
+ tile_overlap_min=int(arch.vae_tile_overlap_min),
+ encoder_parallel=False,
+ decoder_parallel=False,
+ chunk_dim=int(arch.vae_chunk_dim),
+ )
+ self.sglang_config = config
+ self.use_parallel_decode = config.use_parallel_decode
+ self.parallel_decode_mode = parallel_decode_mode
+
+ def prepare_decoder_autocast_weights(self, dtype) -> int:
+ return self.decoder.prepare_autocast_linear_weights(dtype)
+
+
+class MiniMaxH3AudioVAE(DacAudioVAE, LayerwiseOffloadableModuleMixin):
+ layerwise_offload_dit_group_enabled = False
+ # BigVGAN stores each executable upsampler inside a one-element ModuleList.
+ # The outer ``decoder.ups`` containers are indexed but never called, so hooks
+ # must target the inner lists whose ConvTranspose1d modules run forward.
+ layer_names = [
+ "encoder.block",
+ *(f"decoder.ups.{index}" for index in range(7)),
+ "decoder.resblocks",
+ ]
+
+ def __init__(self, config: MiniMaxH3AudioVAEConfig) -> None:
+ super().__init__(
+ encoder_dim=64,
+ encoder_rates=[2, 4, 4, 5, 5],
+ latent_dim=2048,
+ decoder_dim=1024,
+ decoder_rates=[5, 5, 2, 2, 2, 2, 2],
+ sample_rate=32000,
+ vae_latent_channels=32,
+ attn_proj=True,
+ decoder_type="bigvgan",
+ )
+ self.config = config
+
+
+EntryClass = [MiniMaxH3VideoVAE, MiniMaxH3AudioVAE]
+
+__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"]
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/__init__.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/__init__.py
new file mode 100644
index 000000000..94070ad45
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/__init__.py
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .audio_vae import DacAudioVAE
+
+__all__ = ["DacAudioVAE"]
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/alias_free.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/alias_free.py
new file mode 100644
index 000000000..de8c266f1
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/alias_free.py
@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+if "sinc" in dir(torch):
+ sinc = torch.sinc
+else:
+ # This code is adopted from adefossez's julius.core.sinc under the MIT License
+ # https://adefossez.github.io/julius/julius/core.html
+ def sinc(x: torch.Tensor):
+ """
+ Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+ __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+ """
+ return torch.where(
+ x == 0,
+ torch.tensor(1.0, device=x.device, dtype=x.dtype),
+ torch.sin(math.pi * x) / math.pi / x,
+ )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+def kaiser_sinc_filter1d(
+ cutoff, half_width, kernel_size
+): # return filter [1,1,kernel_size]
+ even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+
+ # For kaiser window
+ delta_f = 4 * half_width
+ A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+ if A > 50.0:
+ beta = 0.1102 * (A - 8.7)
+ elif A >= 21.0:
+ beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+ else:
+ beta = 0.0
+ window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+ # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+ if even:
+ time = torch.arange(-half_size, half_size) + 0.5
+ else:
+ time = torch.arange(kernel_size) - half_size
+ if cutoff == 0:
+ filter_ = torch.zeros_like(time)
+ else:
+ filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+ """
+ Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal.
+ """
+ filter_ /= filter_.sum()
+ filter = filter_.view(1, 1, kernel_size)
+
+ return filter
+
+
+class LowPassFilter1d(nn.Module):
+ def __init__(
+ self,
+ cutoff=0.5,
+ half_width=0.6,
+ stride: int = 1,
+ padding: bool = True,
+ padding_mode: str = "replicate",
+ kernel_size: int = 12,
+ ):
+ """
+ kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible.
+ """
+ super().__init__()
+ if cutoff < -0.0:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+ self.kernel_size = kernel_size
+ self.even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = stride
+ self.padding = padding
+ self.padding_mode = padding_mode
+ filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+ self.register_buffer("filter", filter)
+
+ # Input [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ if self.padding:
+ x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+ out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+ return out
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = (
+ int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ )
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = (
+ self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+ )
+ filter = kaiser_sinc_filter1d(
+ cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ kernel_size=self.kernel_size,
+ )
+ self.register_buffer("filter", filter)
+
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ x = F.pad(x, (self.pad, self.pad), mode="replicate")
+ x = F.conv_transpose1d(
+ x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+ )
+ x.mul_(self.ratio)
+ x = x[..., self.pad_left : -self.pad_right]
+
+ return x
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = (
+ int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ )
+ self.lowpass = LowPassFilter1d(
+ cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ stride=ratio,
+ kernel_size=self.kernel_size,
+ )
+
+ def forward(self, x):
+ xx = self.lowpass(x)
+
+ return xx
+
+
+class Activation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ):
+ super().__init__()
+ self.up_ratio = up_ratio
+ self.down_ratio = down_ratio
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ def forward(self, x):
+ x = self.upsample(x)
+ x = self.act(x)
+ x = self.downsample(x)
+
+ return x
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py
new file mode 100644
index 000000000..161c56269
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/audio_vae.py
@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: Apache-2.0
+# DAC-lineage audio VAE: waveform encoder + BigVGAN decoder (inference-only bundle).
+import math
+from typing import List
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.functional import scaled_dot_product_attention
+from torch.nn.utils.parametrizations import weight_norm
+
+from .bigvgan import AttrDict, BigVGAN
+
+
+class GeGluMlp(nn.Module):
+ def __init__(self, in_features, hidden_features):
+ super().__init__()
+ self.norm = nn.LayerNorm(in_features)
+ self.act = nn.GELU(approximate="tanh")
+ self.w0 = nn.Linear(in_features, hidden_features)
+ self.w1 = nn.Linear(in_features, hidden_features)
+ self.w2 = nn.Linear(hidden_features, in_features)
+
+ def forward(self, x):
+ x = self.norm(x)
+ x = self.act(self.w0(x)).mul_(self.w1(x))
+ x = self.w2(x)
+ return x
+
+
+class CausalAttention(nn.Module):
+ def __init__(self, in_dim, out_dim, num_heads):
+ super().__init__()
+ if in_dim > out_dim:
+ # assert in_dim // num_heads == out_dim
+ self.head_dim = in_dim // num_heads
+ self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False)
+ self.q_bias = nn.Parameter(torch.zeros(in_dim))
+ self.v_bias = nn.Parameter(torch.zeros(in_dim))
+ self.register_buffer("zero_k_bias", torch.zeros(in_dim))
+ else:
+ # assert out_dim // num_heads == in_dim
+ self.head_dim = out_dim // num_heads
+ self.qkv = nn.Linear(in_dim, out_dim * 3, bias=False)
+ self.q_bias = nn.Parameter(torch.zeros(out_dim))
+ self.v_bias = nn.Parameter(torch.zeros(out_dim))
+ self.register_buffer("zero_k_bias", torch.zeros(out_dim))
+
+ self.in_dim = in_dim
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.scale = self.head_dim**-0.5
+ self.proj = nn.Linear(out_dim, out_dim)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ B, N, C = x.shape
+ qkv = F.linear(
+ input=x,
+ weight=self.qkv.weight,
+ bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)),
+ )
+ q, k, v = (
+ qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
+ .permute(2, 0, 3, 1, 4)
+ .unbind(0)
+ )
+
+ x = scaled_dot_product_attention(
+ q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True
+ )
+
+ if self.in_dim > self.out_dim:
+ x = torch.mean(x, dim=1)
+ if self.in_dim // self.num_heads != self.out_dim:
+ x = nn.functional.adaptive_avg_pool1d(x, self.out_dim)
+ else:
+ x = x.transpose(1, 2).reshape(B, N, -1)
+ x = self.proj(x)
+ return x
+
+
+class AttnProjection(nn.Module):
+ def __init__(
+ self, in_dim, out_dim, num_heads, norm_layer=nn.LayerNorm, mlp_ratio=2
+ ):
+ super().__init__()
+ assert out_dim % in_dim == 0 or in_dim % out_dim == 0
+ self.in_dim = in_dim
+ self.out_dim = out_dim
+ self.norm1 = norm_layer(in_dim)
+ self.attn = CausalAttention(in_dim, out_dim, num_heads)
+ self.proj = nn.Linear(in_dim, out_dim)
+ self.norm3 = norm_layer(in_dim)
+
+ self.norm2 = norm_layer(out_dim)
+ hidden_dim = int(out_dim * mlp_ratio)
+ self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim)
+ # self.mlp = FeedForward(out_dim, out_dim)
+
+ def forward(self, x):
+ x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x)))
+ return self.mlp(self.norm2(x)).add_(x)
+
+
+def WNConv1d(*args, **kwargs):
+ return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+@torch.jit.script
+def snake(x, alpha):
+ shape = x.shape
+ x = x.reshape(shape[0], shape[1], -1)
+ x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+ x = x.reshape(shape)
+ return x
+
+
+class Snake1d(nn.Module):
+ def __init__(self, channels):
+ super().__init__()
+ self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+ def forward(self, x):
+ return snake(x, self.alpha)
+
+
+class ResidualUnit(nn.Module):
+ def __init__(self, dim: int = 16, dilation: int = 1):
+ super().__init__()
+ pad = ((7 - 1) * dilation) // 2
+ self.block = nn.Sequential(
+ Snake1d(dim),
+ WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+ Snake1d(dim),
+ WNConv1d(dim, dim, kernel_size=1),
+ )
+
+ def forward(self, x):
+ y = self.block(x)
+ pad = (x.shape[-1] - y.shape[-1]) // 2
+ if pad > 0:
+ x = x[..., pad:-pad]
+ return x + y
+
+
+class EncoderBlock(nn.Module):
+ def __init__(self, dim: int = 16, stride: int = 1):
+ super().__init__()
+ self.block = nn.Sequential(
+ ResidualUnit(dim // 2, dilation=1),
+ ResidualUnit(dim // 2, dilation=3),
+ ResidualUnit(dim // 2, dilation=9),
+ Snake1d(dim // 2),
+ WNConv1d(
+ dim // 2,
+ dim,
+ kernel_size=2 * stride,
+ stride=stride,
+ padding=math.ceil(stride / 2),
+ ),
+ )
+
+ def forward(self, x):
+ return self.block(x)
+
+
+class Encoder(nn.Module):
+ def __init__(
+ self,
+ d_model: int = 64,
+ strides: list = [2, 4, 8, 8],
+ d_latent: int = 64,
+ ):
+ super().__init__()
+ # Create first convolution
+ self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+ # Create EncoderBlocks that double channels as they downsample by `stride`
+ for stride in strides:
+ d_model *= 2
+ self.block += [EncoderBlock(d_model, stride=stride)]
+
+ # Create last convolution
+ self.block += [
+ Snake1d(d_model),
+ WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
+ ]
+
+ # Wrap black into nn.Sequential
+ self.block = nn.Sequential(*self.block)
+ self.enc_dim = d_model
+
+ def forward(self, x):
+ return self.block(x)
+
+
+class DacAudioVAE(nn.Module):
+ def __init__(
+ self,
+ encoder_dim: int = 64,
+ encoder_rates: List[int] = [2, 4, 8, 8],
+ latent_dim: int = None,
+ decoder_dim: int = 1536,
+ decoder_rates: List[int] = [8, 8, 4, 2],
+ sample_rate: int = 44100,
+ vae_latent_channels: int = 64,
+ attn_proj: bool = False,
+ decoder_type: str = "bigvgan",
+ ):
+ super().__init__()
+
+ self.encoder_dim = encoder_dim
+ self.encoder_rates = encoder_rates
+ self.decoder_dim = decoder_dim
+ self.decoder_rates = decoder_rates
+ self.sample_rate = sample_rate
+ self.attn_proj = attn_proj
+ self.decoder_type = decoder_type
+
+ if latent_dim is None:
+ latent_dim = encoder_dim * (2 ** len(encoder_rates))
+
+ self.latent_dim = latent_dim
+
+ self.hop_length = np.prod(encoder_rates)
+ self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
+
+ if latent_dim % vae_latent_channels == 0:
+ self.attn_proj_dim = vae_latent_channels
+ else:
+ # smallest power of two >= vae_latent_channels
+ self.attn_proj_dim = 2 ** int(np.ceil(np.log2(vae_latent_channels)))
+
+ self.mean_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1)
+ self.logs_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1)
+
+ self.dec_in_proj = nn.Conv1d(vae_latent_channels, latent_dim, 1)
+
+ if self.decoder_type == "bigvgan":
+ if sample_rate == 16000:
+ bigvgan_conf = {
+ "resblock": "1",
+ "num_mels": latent_dim,
+ "upsample_rates": [5, 5, 2, 2, 2, 2],
+ "upsample_kernel_sizes": [9, 9, 4, 4, 4, 4],
+ "upsample_initial_channel": decoder_dim,
+ "resblock_kernel_sizes": [3, 7, 11],
+ "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ "use_tanh_at_final": False,
+ "use_bias_at_final": False,
+ "activation": "snakebeta",
+ "snake_logscale": True,
+ }
+ elif sample_rate == 32000:
+ bigvgan_conf = {
+ "resblock": "1",
+ "num_mels": latent_dim,
+ "upsample_rates": [5, 5, 2, 2, 2, 2, 2],
+ "upsample_kernel_sizes": [9, 9, 4, 4, 4, 4, 4],
+ "upsample_initial_channel": decoder_dim,
+ "resblock_kernel_sizes": [3, 7, 11],
+ "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+ "use_tanh_at_final": False,
+ "use_bias_at_final": False,
+ "activation": "snakebeta",
+ "snake_logscale": True,
+ }
+ else:
+ raise ValueError(f"Invalid sample_rate: {sample_rate}")
+
+ h = AttrDict(**bigvgan_conf)
+ self.decoder = BigVGAN(h)
+ else:
+ raise ValueError(f"Invalid decoder type: {self.decoder_type}")
+
+ if self.attn_proj:
+ self.pre_block = AttnProjection(latent_dim, self.attn_proj_dim, num_heads=8)
+
+ self.sample_rate = sample_rate
+
+ def preprocess(self, audio_data, sample_rate):
+ if sample_rate is None:
+ sample_rate = self.sample_rate
+
+ length = audio_data.shape[-1]
+ right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+ if right_pad:
+ audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+ return audio_data
+
+ def decode(self, z: torch.Tensor):
+ """Decode given latent codes and return audio data
+
+ Parameters
+ ----------
+ z : Tensor[B x D x T]
+ Continuous latent representation
+
+ Returns
+ -------
+ Tensor[B x 1 x length]
+ Decoded audio data.
+ """
+ z = self.dec_in_proj(z)
+ return self.decoder(z)
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/bigvgan.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/bigvgan.py
new file mode 100644
index 000000000..77ccb3870
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_audio_vae/bigvgan.py
@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2024 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+
+import torch
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d, Parameter
+from torch.nn.utils.parametrizations import weight_norm
+
+from .alias_free import Activation1d
+
+
+def get_padding(kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+
+# Adapted from https://github.com/EdwardDixon/snake under the MIT license.
+@torch.jit.script
+def snakebeta(x, alpha, beta):
+ shape = x.shape
+ x = x.reshape(shape[0], shape[1], -1)
+ x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+ x = x.reshape(shape)
+ return x
+
+
+class SnakeBeta(nn.Module):
+ def __init__(
+ self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+ ):
+ super(SnakeBeta, self).__init__()
+ self.in_features = in_features
+
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale:
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ self.beta = Parameter(torch.zeros(in_features) * alpha)
+ else:
+ self.alpha = Parameter(torch.ones(in_features) * alpha)
+ self.beta = Parameter(torch.ones(in_features) * alpha)
+
+ self.alpha.requires_grad = alpha_trainable
+ self.beta.requires_grad = alpha_trainable
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ x = snakebeta(x, alpha, beta)
+ return x
+
+
+class AttrDict(dict):
+ def __init__(self, *args, **kwargs):
+ super(AttrDict, self).__init__(*args, **kwargs)
+ self.__dict__ = self
+
+
+class AMPBlock1(torch.nn.Module):
+ """
+ AMPBlock applies SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+ AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
+
+ Args:
+ h (AttrDict): Hyperparameters.
+ channels (int): Number of convolution channels.
+ kernel_size (int): Size of the convolution kernel. Default is 3.
+ dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+ activation (str): Activation function type. Must be 'snakebeta'.
+ """
+
+ def __init__(
+ self,
+ h: AttrDict,
+ channels: int,
+ kernel_size: int = 3,
+ dilation: tuple = (1, 3, 5),
+ activation: str = None,
+ ):
+ super().__init__()
+
+ self.h = h
+
+ self.convs1 = nn.ModuleList(
+ [
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ dilation=d,
+ padding=get_padding(kernel_size, d),
+ )
+ )
+ for d in dilation
+ ]
+ )
+
+ self.convs2 = nn.ModuleList(
+ [
+ weight_norm(
+ Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ stride=1,
+ dilation=1,
+ padding=get_padding(kernel_size, 1),
+ )
+ )
+ for _ in range(len(dilation))
+ ]
+ )
+
+ self.num_layers = len(self.convs1) + len(
+ self.convs2
+ ) # Total number of conv layers
+
+ if activation == "snakebeta":
+ self.activations = nn.ModuleList(
+ [
+ Activation1d(
+ activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)
+ )
+ for _ in range(self.num_layers)
+ ]
+ )
+ else:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ def forward(self, x):
+ activation_iter = iter(self.activations)
+ for c1, c2 in zip(self.convs1, self.convs2):
+ a1 = next(activation_iter)
+ a2 = next(activation_iter)
+ xt = a1(x)
+ xt = c1(xt)
+ xt = a2(xt)
+ xt = c2(xt)
+ x = xt.add_(x)
+
+ return x
+
+
+class BigVGAN(torch.nn.Module):
+ """
+ BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
+
+ Args:
+ h (AttrDict): Hyperparameters.
+ """
+
+ def __init__(self, h: AttrDict):
+ super().__init__()
+ self.h = h
+
+ self.num_kernels = len(h.resblock_kernel_sizes)
+ self.num_upsamples = len(h.upsample_rates)
+
+ # Pre-conv
+ self.conv_pre = weight_norm(
+ Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
+ )
+
+ # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+ if h.resblock == "1":
+ resblock_class = AMPBlock1
+ else:
+ raise ValueError(
+ f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}"
+ )
+
+ # Transposed conv-based upsamplers. does not apply anti-aliasing
+ self.ups = nn.ModuleList()
+ for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+ self.ups.append(
+ nn.ModuleList(
+ [
+ weight_norm(
+ ConvTranspose1d(
+ h.upsample_initial_channel // (2**i),
+ h.upsample_initial_channel // (2 ** (i + 1)),
+ k,
+ u,
+ padding=(k - u) // 2,
+ )
+ )
+ ]
+ )
+ )
+
+ # Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+ self.resblocks = nn.ModuleList()
+ for i in range(len(self.ups)):
+ ch = h.upsample_initial_channel // (2 ** (i + 1))
+ for j, (k, d) in enumerate(
+ zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
+ ):
+ self.resblocks.append(
+ resblock_class(h, ch, k, d, activation=h.activation)
+ )
+
+ # Post-conv
+ if h.activation != "snakebeta":
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+ activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+
+ self.activation_post = Activation1d(activation=activation_post)
+
+ # Whether to use bias for the final conv_post. Default to True for backward compatibility
+ self.use_bias_at_final = h.get("use_bias_at_final", True)
+ self.conv_post = weight_norm(
+ Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+ )
+
+ # Final tanh activation. Defaults to True for backward compatibility
+ self.use_tanh_at_final = h.get("use_tanh_at_final", True)
+
+ def forward(self, x):
+ # Pre-conv
+ x = self.conv_pre(x)
+
+ for i in range(self.num_upsamples):
+ # Upsampling
+ for i_up in range(len(self.ups[i])):
+ x = self.ups[i][i_up](x)
+ # AMP blocks
+ xs = None
+ for j in range(self.num_kernels):
+ if xs is None:
+ xs = self.resblocks[i * self.num_kernels + j](x)
+ else:
+ xs += self.resblocks[i * self.num_kernels + j](x)
+ x = xs.div_(self.num_kernels)
+
+ # Post-conv
+ x = self.activation_post(x)
+ x = self.conv_post(x)
+ # Final tanh activation
+ if self.use_tanh_at_final:
+ x.tanh_()
+ else:
+ x.clamp_(min=-1.0, max=1.0) # Bound the output to [-1, 1]
+
+ return x
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/__init__.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/__init__.py
new file mode 100644
index 000000000..28f7bde2d
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/__init__.py
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .klvae import AutoencoderKLLegacy
+
+__all__ = ["AutoencoderKLLegacy"]
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/attention.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/attention.py
new file mode 100644
index 000000000..22e880d5d
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/attention.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# Attention module for the MiniMax H3 visual VAE (inference-only bundle).
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from diffusers.utils import logging
+
+from .flash import flash_attn
+from .vit_utils import _env_flag, apply_rotary_pos_emb_qk
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _vit_norm_input(module, hidden_states):
+ if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1"):
+ return hidden_states.float()
+ weight = getattr(module, "weight", None)
+ return hidden_states.to(getattr(weight, "dtype", hidden_states.dtype))
+
+
+def _apply_qk_norm(module, hidden_states):
+ if (
+ _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1")
+ and isinstance(module, (nn.LayerNorm, nn.RMSNorm))
+ and getattr(module, "weight", None) is None
+ and getattr(module, "bias", None) is None
+ and hidden_states.is_cuda
+ and hidden_states.dtype in (torch.float16, torch.bfloat16)
+ and not torch.is_grad_enabled()
+ and not torch.compiler.is_compiling()
+ ):
+ # CUDA LayerNorm/RMSNorm accumulates half/bfloat16 inputs in FP32.
+ # With no affine parameters its half output is bit-identical to the
+ # released FP32-norm-then-cast recipe, without two full-tensor casts.
+ with torch.autocast("cuda", enabled=False):
+ return module(hidden_states)
+ return module(_vit_norm_input(module, hidden_states)).to(hidden_states.dtype)
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ heads,
+ dim_head,
+ embed_dim: Optional[int] = None,
+ qk_norm_type: Optional[str] = None,
+ qk_norm_affine: bool = False,
+ bias: bool = True,
+ out_bias: Optional[bool] = None,
+ eps: float = 1e-5,
+ **kwargs,
+ ):
+ super().__init__()
+ self.dim_head = dim_head
+ self.heads = heads
+ self.attn_inner_dim = dim_head * heads
+ self.embed_dim = embed_dim if embed_dim is not None else self.attn_inner_dim
+
+ out_bias = out_bias if out_bias is not None else bias
+
+ if qk_norm_type is None:
+ self.norm_q = None
+ self.norm_k = None
+ elif qk_norm_type == "layer_norm":
+ self.norm_q = nn.LayerNorm(
+ dim_head, eps=eps, elementwise_affine=qk_norm_affine
+ )
+ self.norm_k = nn.LayerNorm(
+ dim_head, eps=eps, elementwise_affine=qk_norm_affine
+ )
+ elif qk_norm_type == "rms_norm":
+ self.norm_q = nn.RMSNorm(
+ dim_head, eps=eps, elementwise_affine=qk_norm_affine
+ )
+ self.norm_k = nn.RMSNorm(
+ dim_head, eps=eps, elementwise_affine=qk_norm_affine
+ )
+ else:
+ raise ValueError(
+ f"unknown qk_norm_type: {qk_norm_type}. Should be None,'layer_norm','rms_norm'"
+ )
+
+ self.to_qkv = nn.Linear(self.embed_dim, self.attn_inner_dim * 3, bias=bias)
+
+ self.to_out = nn.Linear(self.attn_inner_dim, self.embed_dim, bias=out_bias)
+
+ if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0):
+ logger.warning(f"Unused kwargs: {kwargs}")
+
+ def _perform_attention(self, query, key, value, pack_info):
+ cu_seqlens = pack_info.get("cu_seqlens", None)
+ mask_mod = pack_info.get("mask_mod", None)
+ block_sparse = pack_info.get("block_sparse", None)
+ valid_seq_len = pack_info.get("valid_seq_len", None)
+
+ if cu_seqlens is not None:
+ raise NotImplementedError(
+ "varlen attention is not supported in this inference-only bundle"
+ )
+
+ padded_seq_len = query.shape[1]
+ if valid_seq_len is not None:
+ valid_seq_len = int(valid_seq_len)
+ if not 0 < valid_seq_len <= padded_seq_len:
+ raise ValueError(
+ "valid_seq_len must be in (0, padded_seq_len], got "
+ f"{valid_seq_len} for padded_seq_len={padded_seq_len}"
+ )
+ query = query[:, :valid_seq_len]
+ key = key[:, :valid_seq_len]
+ value = value[:, :valid_seq_len]
+
+ if mask_mod is not None:
+ hidden_states = flash_attn(
+ query,
+ key,
+ value,
+ mask_mod=mask_mod,
+ block_sparse=block_sparse,
+ )
+ else:
+ hidden_states = flash_attn(
+ query,
+ key,
+ value,
+ )
+
+ if valid_seq_len is not None and valid_seq_len < padded_seq_len:
+ hidden_states = torch.cat(
+ [
+ hidden_states,
+ hidden_states.new_zeros(
+ hidden_states.shape[0],
+ padded_seq_len - valid_seq_len,
+ hidden_states.shape[2],
+ hidden_states.shape[3],
+ ),
+ ],
+ dim=1,
+ )
+
+ return hidden_states
+
+ def perform_attention(self, query, key, value, pack_info={}):
+ return self._perform_attention(query, key, value, pack_info)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ rotary_pos_emb: Optional[torch.Tensor] = None,
+ pack_info: dict = {},
+ ) -> torch.Tensor:
+ batch_size, seq_len, _ = hidden_states.shape
+
+ qkv = self.to_qkv(hidden_states)
+ qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head)
+ query, key, value = torch.chunk(qkv, 3, dim=-1)
+
+ if self.norm_q is not None:
+ query = _apply_qk_norm(self.norm_q, query)
+ if self.norm_k is not None:
+ key = _apply_qk_norm(self.norm_k, key)
+
+ if rotary_pos_emb is not None:
+ query, key = apply_rotary_pos_emb_qk(query, key, rotary_pos_emb)
+
+ hidden_states = self.perform_attention(query, key, value, pack_info)
+
+ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
+ hidden_states = self.to_out(hidden_states)
+
+ return hidden_states
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py
new file mode 100644
index 000000000..b1f92236e
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py
@@ -0,0 +1,281 @@
+# SPDX-License-Identifier: Apache-2.0
+# Transformer building blocks for the MiniMax H3 visual VAE ViT decoder.
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+
+from sglang.kernels.ops.activation.activation import (
+ silu_and_mul_with_activation_rounding,
+)
+from sglang.kernels.ops.diffusion.triton.scale_shift import (
+ try_fused_scaled_residual_add_exact,
+)
+
+from .attention import Attention
+from .vit_utils import _env_flag, _vit_torch_compile_kwargs
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _vit_norm_input(module, hidden_states):
+ if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1"):
+ return hidden_states.float()
+ return hidden_states.to(getattr(module.weight, "dtype", hidden_states.dtype))
+
+
+def _scaled_residual_add(residual, x, scale):
+ fused = try_fused_scaled_residual_add_exact(residual, x, scale)
+ return residual + x * scale if fused is None else fused
+
+
+class FeedForward(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ dim_out: Optional[int] = None,
+ mult: int = 4,
+ activation_fn: str = "silu",
+ bias: bool = True,
+ use_gated: bool = True,
+ glu_balanced: bool = False,
+ ):
+ super().__init__()
+ ratio = 2 / 3 if (use_gated and glu_balanced) else 1
+ inner_dim = round(dim * mult * ratio)
+ dim_out = dim_out if dim_out is not None else dim
+ self.use_gated = use_gated
+
+ if use_gated:
+ self.w1 = nn.Linear(dim, inner_dim * 2, bias=bias)
+ else:
+ self.w1 = nn.Linear(dim, inner_dim, bias=bias)
+
+ if activation_fn == "silu":
+ self.act_fn = nn.SiLU()
+ elif activation_fn == "gelu":
+ self.act_fn = nn.GELU()
+ elif activation_fn == "gelu-approximate":
+ self.act_fn = nn.GELU(approximate="tanh")
+ else:
+ raise ValueError(f"Unsupported activation function: {activation_fn}")
+
+ self.w2 = nn.Linear(inner_dim, dim_out, bias=bias)
+ self._compile_forward_enabled = _env_flag(
+ "MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE", "0"
+ )
+ self._compile_forward_fatal = _env_flag(
+ "MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE_FATAL", "0"
+ )
+ self._compiled_forward = None
+
+ def _forward_impl(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.w1(hidden_states)
+
+ if self.use_gated:
+ if (
+ isinstance(self.act_fn, nn.SiLU)
+ and hidden_states.is_cuda
+ and hidden_states.dtype in (torch.float16, torch.bfloat16)
+ and hidden_states.is_contiguous()
+ and hidden_states.shape[-1] % 32 == 0
+ ):
+ hidden_states = silu_and_mul_with_activation_rounding(hidden_states)
+ else:
+ gate, hidden_states = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.act_fn(gate).mul_(hidden_states)
+ else:
+ hidden_states = self.act_fn(hidden_states)
+
+ hidden_states = self.w2(hidden_states)
+ return hidden_states
+
+ def _get_forward_impl(self):
+ if not self._compile_forward_enabled:
+ return self._forward_impl
+ if self._compiled_forward is not None:
+ return self._compiled_forward
+ if not hasattr(torch, "compile"):
+ message = (
+ "torch.compile is unavailable; falling back to eager ViT FeedForward"
+ )
+ if self._compile_forward_fatal:
+ raise RuntimeError(message)
+ logger.warning(f"[ViTFeedForward] {message}")
+ self._compile_forward_enabled = False
+ return self._forward_impl
+
+ kwargs = _vit_torch_compile_kwargs(
+ "MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE"
+ )
+ try:
+ self._compiled_forward = torch.compile(self._forward_impl, **kwargs)
+ logger.info(f"[ViTFeedForward] torch.compile enabled kwargs={kwargs}")
+ except Exception as exc:
+ if self._compile_forward_fatal:
+ raise
+ logger.warning(
+ f"[ViTFeedForward] torch.compile setup failed: {type(exc).__name__}: {exc}; "
+ "falling back to eager"
+ )
+ self._compile_forward_enabled = False
+ self._compiled_forward = None
+ return self._forward_impl
+ return self._compiled_forward
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ forward_impl = self._get_forward_impl()
+ try:
+ return forward_impl(hidden_states)
+ except Exception as exc:
+ if (
+ self._compile_forward_enabled
+ and self._compiled_forward is not None
+ and forward_impl is self._compiled_forward
+ and not self._compile_forward_fatal
+ ):
+ logger.warning(
+ f"[ViTFeedForward] compiled forward failed: {type(exc).__name__}: {exc}; "
+ "disabling compile and retrying eager"
+ )
+ self._compile_forward_enabled = False
+ self._compiled_forward = None
+ return self._forward_impl(hidden_states)
+ raise
+
+
+class RotaryEmbeddingND(nn.Module):
+ def __init__(self, dim, rotary_base=10000, n_dim=3, use_angle=False):
+ super().__init__()
+ self.dim = dim
+ self.n_dim = n_dim
+
+ if dim % (2 * n_dim) != 0:
+ raise ValueError(
+ f"head_dim {dim} must be divisible by 2 * n_dim {2 * n_dim}"
+ )
+
+ if use_angle:
+ self.angle_scale = 2.0 * math.pi
+ else:
+ self.angle_scale = 1.0
+
+ inv_freq = 1 / rotary_base ** torch.arange(
+ 0, 1, 2 * n_dim / dim, dtype=torch.float32
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, img_ids):
+ B, N, D = img_ids.shape
+ if D != self.n_dim:
+ raise ValueError(f"Expected {self.n_dim} dimensions, got {D}")
+
+ with torch.autocast("cuda", enabled=False):
+ angles = (
+ self.angle_scale
+ * img_ids[:, :, :, None]
+ * self.inv_freq.to(img_ids.device)[None, None, None, :]
+ )
+ angles = angles.flatten(2, 3)
+ angles = angles.tile(2)
+ angles = angles.unsqueeze(2)
+
+ cos = torch.cos(angles)
+ sin = torch.sin(angles)
+
+ return cos.to(dtype=img_ids.dtype), sin.to(dtype=img_ids.dtype)
+
+
+@maybe_allow_in_graph
+class TransformerBlock(nn.Module):
+ def __init__(
+ self,
+ heads: int,
+ dim_head: int,
+ embed_dim: Optional[int] = None,
+ ffn_glu_balanced: bool = False,
+ norm_type: str = "layer_norm",
+ norm_affine: bool = True,
+ qk_norm_type: str = "rms_norm",
+ qk_norm_affine: bool = False,
+ ffn_activation_fn: str = "silu",
+ ffn_use_gated: bool = True,
+ use_scale: bool = True,
+ bias: bool = True,
+ eps: float = 1e-5,
+ **kwargs,
+ ):
+ super().__init__()
+ dim = embed_dim if embed_dim is not None else dim_head * heads
+ self.use_scale = use_scale
+
+ if norm_type == "layer_norm":
+ norm_class = nn.LayerNorm
+ elif norm_type == "rms_norm":
+ norm_class = nn.RMSNorm
+ else:
+ raise ValueError(f"unknown norm_type {norm_type}")
+
+ self.norm1 = norm_class(
+ dim,
+ elementwise_affine=norm_affine,
+ eps=eps,
+ )
+ self.attn = Attention(
+ heads=heads,
+ dim_head=dim_head,
+ embed_dim=dim,
+ qk_norm_type=qk_norm_type,
+ qk_norm_affine=qk_norm_affine,
+ bias=bias,
+ eps=eps,
+ **kwargs,
+ )
+ if use_scale:
+ self.scale1 = nn.Parameter(torch.zeros(dim))
+
+ self.norm2 = norm_class(
+ dim,
+ elementwise_affine=norm_affine,
+ eps=eps,
+ )
+ self.ff = FeedForward(
+ dim=dim,
+ activation_fn=ffn_activation_fn,
+ bias=bias,
+ use_gated=ffn_use_gated,
+ glu_balanced=ffn_glu_balanced,
+ )
+ if use_scale:
+ self.scale2 = nn.Parameter(torch.zeros(dim))
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ rotary_pos_emb: Optional[torch.FloatTensor] = None,
+ pack_info: dict = {},
+ ):
+ norm_hidden_states = self.norm1(_vit_norm_input(self.norm1, hidden_states)).to(
+ hidden_states.dtype
+ )
+ attn_output = self.attn(norm_hidden_states, rotary_pos_emb, pack_info)
+ if self.use_scale:
+ hidden_states = _scaled_residual_add(
+ hidden_states, attn_output, self.scale1
+ )
+ else:
+ hidden_states = hidden_states + attn_output
+
+ norm_hidden_states = self.norm2(_vit_norm_input(self.norm2, hidden_states)).to(
+ hidden_states.dtype
+ )
+ ff_output = self.ff(norm_hidden_states)
+ if self.use_scale:
+ hidden_states = _scaled_residual_add(hidden_states, ff_output, self.scale2)
+ else:
+ hidden_states = hidden_states + ff_output
+
+ return hidden_states
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/conv.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/conv.py
new file mode 100644
index 000000000..98a2e483d
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/conv.py
@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# 3D convolution for the MiniMax H3 visual VAE.
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BaseConv3d(nn.Conv3d):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=1,
+ padding=0,
+ bias=True,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ ):
+ super().__init__(
+ in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ bias=bias,
+ padding_mode=padding_mode,
+ )
+ padding_mode = "constant" if padding_mode == "zeros" else padding_mode
+ padding_mode_t = "constant" if padding_mode_t == "zeros" else padding_mode_t
+ self.pad_mode = padding_mode
+ self.pad_mode_t = padding_mode_t or ("constant" if causal else "replicate")
+ self.causal = causal
+
+ def _apply_temporal_padding(self, x):
+ B, C, D, H, W = x.shape
+ if D > 1:
+ pad_size = (
+ 0,
+ 0,
+ 0,
+ 0,
+ self.padding[0] * 2 if self.causal else self.padding[0],
+ 0 if self.causal else self.padding[0],
+ )
+ return F.pad(x, pad_size, mode=self.pad_mode_t)
+ else:
+ if self.pad_mode_t == "constant":
+ assert self.causal, "Zeros padding is only supported for causal mode"
+ return F.pad(
+ x,
+ (0, 0, 0, 0, self.kernel_size[0] - 1, 0),
+ mode="constant",
+ )
+ else:
+ return x.expand(-1, -1, self.kernel_size[0], -1, -1)
+
+ def _apply_padding(self, x):
+ if sum(self.padding) == 0:
+ return x
+
+ x = F.pad(
+ x,
+ (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 0, 0),
+ mode=self.pad_mode,
+ )
+
+ x = self._apply_temporal_padding(x)
+ return x
+
+ def forward(self, x):
+ if sum(self.padding) == 0:
+ return super().forward(x)
+
+ x = self._apply_padding(x)
+ return F.conv3d(
+ x,
+ self.weight,
+ self.bias,
+ stride=self.stride,
+ padding=0,
+ dilation=self.dilation,
+ )
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/flash.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/flash.py
new file mode 100644
index 000000000..3be87d12c
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/flash.py
@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# Torch-native attention implemented with PyTorch SDPA instead of FA4/CUTLASS.
+import os
+from contextlib import nullcontext
+
+import torch
+import torch.nn.functional as F
+
+_BLOCK_CAUSAL_MASK_MOD_CACHE = {}
+
+
+def _auto_sdpa_backend_name() -> str | None:
+ """Return the ROCm-only correctness fallback for H3 video-VAE SDPA."""
+ if torch.version.hip is None:
+ return None
+
+ from sglang.srt.utils import is_gfx95_supported
+
+ # Fused ROCm SDPA corrupts the dense ViT decode on gfx950. Keep every
+ # non-gfx950 platform, including CUDA, on PyTorch's unchanged auto path.
+ return "math" if is_gfx95_supported() else None
+
+
+_AUTO_SDPA_BACKEND = _auto_sdpa_backend_name()
+
+
+def _as_bool_mask(mask, *, device):
+ if not isinstance(mask, torch.Tensor):
+ mask = torch.as_tensor(mask, device=device)
+ return mask.to(device=device, dtype=torch.bool)
+
+
+def _ensure_nonempty_rows(mask):
+ if mask.numel() == 0 or mask.shape[-1] == 0:
+ return mask
+ empty = ~mask.any(dim=-1)
+ mask[..., 0] |= empty
+ return mask
+
+
+def _sdpa_kernel_context():
+ backend_name = os.environ.get("MINIMAX_H3_TORCH_SDPA_BACKEND", "auto").lower()
+ if backend_name in {"", "auto", "default"}:
+ backend_name = _AUTO_SDPA_BACKEND
+ if backend_name is None:
+ return nullcontext()
+
+ from torch.nn.attention import SDPBackend, sdpa_kernel
+
+ backends = {
+ "math": SDPBackend.MATH,
+ "flash": SDPBackend.FLASH_ATTENTION,
+ "flash_attention": SDPBackend.FLASH_ATTENTION,
+ "efficient": SDPBackend.EFFICIENT_ATTENTION,
+ "mem_efficient": SDPBackend.EFFICIENT_ATTENTION,
+ "cudnn": SDPBackend.CUDNN_ATTENTION,
+ "cudnn_attention": SDPBackend.CUDNN_ATTENTION,
+ }
+ if backend_name not in backends:
+ raise ValueError(
+ "MINIMAX_H3_TORCH_SDPA_BACKEND must be one of "
+ f"{sorted([*backends, 'auto', 'default'])}, got {backend_name!r}"
+ )
+ return sdpa_kernel(backends=[backends[backend_name]])
+
+
+def _sdpa_attention(query, key, value, causal=False, attn_mask=None):
+ # query/key/value arrive as [B, S, H, D]; PyTorch SDPA expects
+ # [B, H, S, D].
+ q = query.transpose(1, 2)
+ k = key.transpose(1, 2)
+ v = value.transpose(1, 2)
+ if attn_mask is not None and attn_mask.dim() == 3:
+ attn_mask = attn_mask.unsqueeze(0)
+ with _sdpa_kernel_context():
+ out = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ attn_mask=attn_mask,
+ dropout_p=0.0,
+ is_causal=causal,
+ )
+ return out.transpose(1, 2).nan_to_num(0.0)
+
+
+def _mask_mod_to_dense(mask_mod, batch, heads, q_len, kv_len, device, aux_tensors=None):
+ q_idx = torch.arange(q_len, device=device).view(q_len, 1)
+ kv_idx = torch.arange(kv_len, device=device).view(1, kv_len)
+ dense = torch.empty((batch, heads, q_len, kv_len), dtype=torch.bool, device=device)
+ for b in range(batch):
+ b_idx = torch.tensor(b, device=device)
+ for h in range(heads):
+ h_idx = torch.tensor(h, device=device)
+ mask = mask_mod(b_idx, h_idx, q_idx, kv_idx, None, aux_tensors)
+ dense[b, h] = _as_bool_mask(mask, device=device)
+ return _ensure_nonempty_rows(dense)
+
+
+#########################################################
+# Block causal attention
+#########################################################
+
+
+def make_block_causal_mask_mod(num_tokens, block_size, num_special=0, suffix=False):
+ if num_tokens < 0:
+ raise ValueError(f"num_tokens must be non-negative, got {num_tokens}")
+ if block_size <= 0:
+ raise ValueError(f"block_size must be positive, got {block_size}")
+ if num_special < 0:
+ raise ValueError(f"num_special must be non-negative, got {num_special}")
+
+ cache_key = (num_tokens, block_size, num_special, suffix)
+ if cache_key in _BLOCK_CAUSAL_MASK_MOD_CACHE:
+ return _BLOCK_CAUSAL_MASK_MOD_CACHE[cache_key]
+
+ if suffix:
+
+ def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors):
+ del b, h, seqlen_info, aux_tensors
+ q_is_special = q_idx >= num_tokens
+ kv_is_special = kv_idx >= num_tokens
+ return (
+ q_is_special
+ | kv_is_special
+ | (q_idx // block_size >= kv_idx // block_size)
+ )
+
+ else:
+
+ def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors):
+ del b, h, seqlen_info, aux_tensors
+ q_is_special = q_idx < num_special
+ kv_is_special = kv_idx < num_special
+ q_block_idx = (q_idx - num_special) // block_size
+ kv_block_idx = (kv_idx - num_special) // block_size
+ return q_is_special | kv_is_special | (q_block_idx >= kv_block_idx)
+
+ mask_mod.block_sparse_cache_key = (
+ "block_causal",
+ num_tokens,
+ block_size,
+ num_special,
+ suffix,
+ )
+ _BLOCK_CAUSAL_MASK_MOD_CACHE[cache_key] = mask_mod
+ return mask_mod
+
+
+#########################################################
+# Public entry point
+#########################################################
+
+
+@torch.compiler.disable
+def flash_attn(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ causal: bool = False,
+ mask_mod=None,
+ block_sparse=None,
+ aux_tensors=None,
+) -> torch.Tensor:
+ use_masked = mask_mod is not None or block_sparse is not None
+
+ if block_sparse is not None and mask_mod is None:
+ raise ValueError("block_sparse requires mask_mod")
+ if causal and mask_mod is not None:
+ raise ValueError(
+ "causal must be encoded in mask_mod when using masked attention"
+ )
+ if aux_tensors is not None and not use_masked:
+ raise ValueError("aux_tensors is only supported with masked attention")
+
+ if use_masked:
+ batch, q_len, heads, _ = query.shape
+ kv_len = key.shape[1]
+ dense_mask = _mask_mod_to_dense(
+ mask_mod,
+ batch,
+ heads,
+ q_len,
+ kv_len,
+ query.device,
+ aux_tensors=aux_tensors,
+ )
+ return _sdpa_attention(query, key, value, attn_mask=dense_mask)
+
+ return _sdpa_attention(query, key, value, causal=causal)
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/klvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/klvae.py
new file mode 100644
index 000000000..9aafb6b3a
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/klvae.py
@@ -0,0 +1,1297 @@
+# SPDX-License-Identifier: Apache-2.0
+# MiniMax H3 visual VAE: 3D causal CNN encoder + ViT3D decoder (inference-only bundle).
+import math
+import os
+from typing import List, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models import ModelMixin
+from diffusers.utils import logging
+from PIL import Image
+
+from sglang.multimodal_gen.runtime.distributed import (
+ get_decode_parallel_group_coordinator,
+ get_decode_parallel_rank,
+ get_decode_parallel_world_size,
+ model_parallel_is_initialized,
+)
+
+from .processor import (
+ VAEProcessor,
+ get_denormalize_transform,
+ get_normalize_transform,
+)
+from .vae_cnn import EncoderFCN3D
+from .vae_vit import ViT3DDecoder
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _resolve_temporal_cat_dtype():
+ raw = (
+ os.environ.get("MINIMAX_H3_VAE_DECODER_TEMPORAL_CAT_DTYPE", "").strip().lower()
+ )
+ if raw in ("", "0", "false", "no", "off", "none", "keep", "default"):
+ return None
+ mapping = {
+ "fp16": torch.float16,
+ "float16": torch.float16,
+ "half": torch.float16,
+ "bf16": torch.bfloat16,
+ "bfloat16": torch.bfloat16,
+ "fp32": torch.float32,
+ "float32": torch.float32,
+ }
+ if raw not in mapping:
+ raise ValueError(
+ "MINIMAX_H3_VAE_DECODER_TEMPORAL_CAT_DTYPE must be one of "
+ "fp16|bf16|fp32|keep, got %r" % raw
+ )
+ return mapping[raw]
+
+
+def _resolve_temporal_stream_cat():
+ raw = (
+ os.environ.get("MINIMAX_H3_VAE_DECODER_STREAM_TEMPORAL_CAT", "1")
+ .strip()
+ .lower()
+ )
+ return raw not in ("0", "false", "no", "off", "disable", "disabled")
+
+
+def get_tile_parallel_state():
+ if not dist.is_initialized() or not model_parallel_is_initialized():
+ return 0, 1
+ return get_decode_parallel_rank(), get_decode_parallel_world_size()
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters, upcast_fp32=True):
+ if upcast_fp32:
+ parameters = parameters.to(dtype=torch.float32)
+
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.std = self.logvar.mul(0.5).exp_()
+
+ @torch.compiler.disable
+ def sample(self, generator=None):
+ noise = torch.randn(self.mean.shape, generator=generator)
+ noise = noise.to(device=self.parameters.device)
+ return noise.mul_(self.std).add_(self.mean)
+
+
+class ClsTokenAggregator:
+ def __init__(self, vae_model):
+ self.vae = vae_model
+ self.cls_tokens = []
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ if self.cls_tokens and hasattr(self.vae.encoder, "loss_info"):
+ self.vae.encoder.loss_info["cls_token"] = torch.stack(
+ self.cls_tokens, dim=0
+ ).mean(dim=0)
+ return False
+
+ def collect(self):
+ if (
+ hasattr(self.vae.encoder, "loss_info")
+ and "cls_token" in self.vae.encoder.loss_info
+ ):
+ self.cls_tokens.append(self.vae.encoder.loss_info["cls_token"].clone())
+
+ def collect_stacked(self, num_tiles, sample_batch_size):
+ if (
+ hasattr(self.vae.encoder, "loss_info")
+ and "cls_token" in self.vae.encoder.loss_info
+ ):
+ cls_token = self.vae.encoder.loss_info["cls_token"]
+ cls_token = cls_token.unflatten(0, (num_tiles, sample_batch_size))
+ self.cls_tokens.extend(token.clone() for token in cls_token)
+
+
+class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ Abstract shared base for the MiniMax H3 visual VAE.
+
+ This class only carries the shared inference machinery (temporal
+ chunking, tiling, encode/decode entry points). Instantiate the concrete
+ subclass ``AutoencoderKLLegacy`` via ``from_pretrained`` instead.
+ """
+
+ _compilable_modules = ["encoder", "decoder"]
+ _deprecated_kwargs = [
+ "clip_length",
+ "token_drop",
+ "isolated_first_frame",
+ "isolated_last_frame",
+ "isolated_key_frame",
+ "encoder_tiling",
+ "decoder_tiling",
+ "parallel_tiling",
+ "stack_tiling",
+ "tile_size",
+ "tile_overlap_min",
+ "decoder_tile_size",
+ "decoder_tile_overlap_min",
+ "latent_patch_size",
+ "crop_mode",
+ "encoder_parallel",
+ "decoder_parallel",
+ "chunk_dim",
+ ] # legacy config keys accepted by from_pretrained for checkpoint compatibility
+
+ def setup_forward(self, **kwargs):
+ self.clip_length = kwargs.get("clip_length", 17)
+ self.token_drop = kwargs.get("token_drop", 0)
+ self.frame_drop = self.token_drop * self.vae_ratio_t
+ self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t
+ self.tokens_chunk_size = math.ceil(self.clip_length / self.vae_ratio_t)
+ self.token_overlap = (-self.token_drop) % self.tokens_chunk_size
+ self.frame_overlap = max(
+ self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0
+ )
+ self.isolated_first_frame = kwargs.get("isolated_first_frame", False)
+ self.isolated_last_frame = kwargs.get("isolated_last_frame", False)
+ self.isolated_key_frame = kwargs.get("isolated_key_frame", False)
+
+ self.encoder_tiling = kwargs.get("encoder_tiling", False)
+ self.decoder_tiling = kwargs.get("decoder_tiling", False)
+ self.stack_tiling = kwargs.get("stack_tiling", False)
+ self.tile_size = kwargs.get("tile_size", 256)
+ self.tile_overlap_min = kwargs.get("tile_overlap_min", 64)
+ self.decoder_tile_size = kwargs.get("decoder_tile_size", self.tile_size)
+ self.decoder_tile_overlap_min = kwargs.get(
+ "decoder_tile_overlap_min", self.tile_overlap_min
+ )
+ self._blend_weight_cache = {}
+ self.latent_patch_size = kwargs.get("latent_patch_size", 1)
+ self.crop_mode = kwargs.get("crop_mode", "top_left")
+ self.pixel_norm_type = kwargs.get("pixel_norm_type", "imagenet")
+
+ if kwargs.get("encoder_parallel", False) or kwargs.get(
+ "decoder_parallel", False
+ ):
+ raise ValueError(
+ "MiniMax H3 VAE spatial sharding is unsupported; use complete-tile "
+ "parallel decode instead"
+ )
+ parallel_tiling = kwargs.get("parallel_tiling", False)
+ if hasattr(self, "parallel_tiling") and parallel_tiling != self.parallel_tiling:
+ logger.warning(
+ "Do not support changing parallel tiling after initialization"
+ )
+ else:
+ self.parallel_tiling = parallel_tiling
+
+ processor_kwargs = {
+ "vae_ratio": self.vae_ratio,
+ "vae_ratio_t": self.vae_ratio_t,
+ "clip_length": self.clip_length,
+ "frame_overlap": self.frame_overlap,
+ "token_overlap": self.token_overlap,
+ "tokens_chunk_size": self.tokens_chunk_size,
+ "isolated_last_frame": self.isolated_last_frame,
+ "latent_patch_size": self.latent_patch_size,
+ "crop_mode": self.crop_mode,
+ "pixel_norm_type": self.pixel_norm_type,
+ "transform": self.transform,
+ "transform_rev": self.transform_rev,
+ "use_3d_conv": self.use_3d_conv,
+ }
+ if hasattr(self, "processor"):
+ for key, value in processor_kwargs.items():
+ setattr(self.processor, key, value)
+ else:
+ self.processor = VAEProcessor(**processor_kwargs)
+
+ def split_tiles(self, input_len, is_decoder=False):
+ tile_size = self.decoder_tile_size if is_decoder else self.tile_size
+ tile_overlap_min = (
+ self.decoder_tile_overlap_min if is_decoder else self.tile_overlap_min
+ )
+
+ if tile_size >= input_len:
+ return [0], [input_len], []
+
+ N = math.ceil(input_len / tile_size)
+ while True:
+ overlaps = [tile_overlap_min] * (N - 1)
+ remaining = tile_size * N - sum(overlaps) - input_len
+
+ if remaining < 0:
+ N += 1
+ else:
+ break
+
+ remaining_units = remaining // self.vae_ratio
+ for i in range(remaining_units):
+ overlaps[i % (N - 1)] += self.vae_ratio
+
+ tile_start_idx = [0]
+ for i in range(N - 1):
+ tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i])
+
+ tile_len = [tile_size] * N
+ return tile_start_idx, tile_len, overlaps
+
+ def blend(
+ self, a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int
+ ) -> torch.Tensor:
+ blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)
+
+ cache_key = (blend_extent, b.device, b.dtype)
+ weights = self._blend_weight_cache.get(cache_key)
+ if weights is None:
+ positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype)
+ weights = (1 - positions / blend_extent, positions / blend_extent)
+ self._blend_weight_cache[cache_key] = weights
+ weight_a, weight_b = weights
+
+ shape = [1] * a.ndim
+ shape[dim] = blend_extent
+ weight_a = weight_a.view(shape)
+ weight_b = weight_b.view(shape)
+
+ slice_a = [slice(None)] * a.ndim
+ slice_a[dim] = slice(-blend_extent, None)
+ a_overlap = a[tuple(slice_a)]
+
+ slice_b = [slice(None)] * b.ndim
+ slice_b[dim] = slice(0, blend_extent)
+ b_overlap = b[tuple(slice_b)]
+
+ blended = a_overlap * weight_a
+ blended.add_(b_overlap * weight_b)
+
+ if blend_extent < b.shape[dim]:
+ slice_b_rest = [slice(None)] * b.ndim
+ slice_b_rest[dim] = slice(blend_extent, None)
+ b_rest = b[tuple(slice_b_rest)]
+ return torch.cat([blended, b_rest], dim=dim)
+ else:
+ return blended
+
+ def _assemble_tiles(self, rows, y_overlap, x_overlap):
+ output_height = sum(
+ row[0].shape[-2] - (y_overlap[i] if i < len(rows) - 1 else 0)
+ for i, row in enumerate(rows)
+ )
+ output_width = sum(
+ tile.shape[-1] - (x_overlap[j] if j < len(rows[0]) - 1 else 0)
+ for j, tile in enumerate(rows[0])
+ )
+ output = rows[0][0].new_empty(
+ (*rows[0][0].shape[:-2], output_height, output_width)
+ )
+
+ y_offset = 0
+ # released VAE blends vertically before horizontally; preserve that
+ # order while writing each cropped tile into the final tensor
+ for i, row in enumerate(rows):
+ row_height = row[0].shape[-2] - (y_overlap[i] if i < len(rows) - 1 else 0)
+ x_offset = 0
+ for j, tile in enumerate(row):
+ if i > 0:
+ tile = self.blend(rows[i - 1][j], tile, y_overlap[i - 1], dim=-2)
+ if j > 0:
+ tile = self.blend(row[j - 1], tile, x_overlap[j - 1], dim=-1)
+ if i < len(rows) - 1:
+ tile = tile[..., : -y_overlap[i], :]
+ if j < len(row) - 1:
+ tile = tile[..., :, : -x_overlap[j]]
+
+ tile_height, tile_width = tile.shape[-2:]
+ output[
+ ...,
+ y_offset : y_offset + tile_height,
+ x_offset : x_offset + tile_width,
+ ].copy_(tile)
+ x_offset += tile_width
+ y_offset += row_height
+
+ return output
+
+ def _all_gather_tiled_results(self, tasks, num_tiles):
+ tile_rank, tile_world_size = get_tile_parallel_state()
+
+ if not tasks:
+ raise ValueError(f"Found empty tasks on tile rank {tile_rank}")
+
+ expected_tasks = (
+ num_tiles - tile_rank + tile_world_size - 1
+ ) // tile_world_size
+ if len(tasks) != expected_tasks:
+ raise ValueError(
+ f"Expected {expected_tasks} tiled tasks on rank {tile_rank}, "
+ f"got {len(tasks)} for num_tiles={num_tiles}, "
+ f"world_size={tile_world_size}"
+ )
+
+ max_tasks = (num_tiles + tile_world_size - 1) // tile_world_size
+ if len(tasks) == max_tasks:
+ stacked = torch.stack(tasks, dim=0)
+ else:
+ stacked = tasks[0].new_empty((max_tasks, *tasks[0].shape))
+ torch.stack(tasks, dim=0, out=stacked[: len(tasks)])
+ stacked[len(tasks) :].zero_()
+
+ # Round-robin tile ownership makes every rank's task count known from
+ # num_tiles and world size. Pad only the leading task dimension and use a
+ # single equal-shape all-gather; the previous path paid for a barrier,
+ # a shape all-gather, and then a padded data all-gather per temporal
+ # clip.
+ gathered = get_decode_parallel_group_coordinator().all_gather(
+ stacked, separate_tensors=True
+ )
+
+ results = [None] * num_tiles
+ for rank, rank_tensors in enumerate(gathered):
+ num_rank_tasks = (num_tiles - rank + tile_world_size - 1) // tile_world_size
+ for k in range(num_rank_tasks):
+ global_idx = k * tile_world_size + rank
+ results[global_idx] = rank_tensors[k]
+
+ return results
+
+ def _local_tile_indices(self, num_tiles, tile_rank, tile_world_size):
+ return list(range(tile_rank, num_tiles, tile_world_size))
+
+ def _run_tile_tasks(
+ self, tiles, tile_indices, forward_fn, stack_tiling, cls_agg=None
+ ):
+ if stack_tiling and tile_indices:
+ sample_batch_size = tiles[0].shape[0]
+ tile_batch = torch.cat([tiles[idx] for idx in tile_indices], dim=0)
+ output_batch = forward_fn(tile_batch)
+ output_tiles = output_batch.unflatten(
+ 0, (len(tile_indices), sample_batch_size)
+ ).unbind(dim=0)
+ if cls_agg is not None:
+ cls_agg.collect_stacked(len(tile_indices), sample_batch_size)
+ return list(output_tiles)
+
+ tasks = []
+ for idx in tile_indices:
+ tasks.append(forward_fn(tiles[idx]))
+ if cls_agg is not None:
+ cls_agg.collect()
+ return tasks
+
+ def tiled_encode(self, x):
+ if self.parallel_tiling: # Fast online encoding for large videos
+ tile_rank, tile_world_size = get_tile_parallel_state()
+ else:
+ tile_rank, tile_world_size = 0, 1
+
+ height, width = x.shape[-2], x.shape[-1]
+ y_idx, y_len, y_overlap = self.split_tiles(height, False)
+ x_idx, x_len, x_overlap = self.split_tiles(width, False)
+
+ i_max, j_max = len(y_idx), len(x_idx)
+ num_tiles = i_max * j_max
+ if tile_world_size > num_tiles:
+ # Every rank executes this replicated path. When the canvas has
+ # fewer tiles than ranks, run locally instead of assigning
+ # empty task lists that cannot participate in the tensor gather.
+ tile_rank, tile_world_size = 0, 1
+
+ x_tiles = []
+ for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
+ for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
+ tile = x[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
+ x_tiles.append(tile)
+
+ with ClsTokenAggregator(self) as agg:
+ local_tile_indices = self._local_tile_indices(
+ num_tiles, tile_rank, tile_world_size
+ )
+ stack_tiling = self.stack_tiling and not (
+ self.training and getattr(self.encoder, "mask_enabled", False)
+ )
+ encoded_tasks = self._run_tile_tasks(
+ x_tiles, local_tile_indices, self.encode, stack_tiling, agg
+ )
+
+ if tile_world_size > 1:
+ all_encoded = self._all_gather_tiled_results(encoded_tasks, num_tiles)
+ if agg.cls_tokens:
+ agg.cls_tokens = self._all_gather_tiled_results(
+ agg.cls_tokens, num_tiles
+ )
+ else:
+ all_encoded = encoded_tasks
+
+ rows = [[None for _ in range(j_max)] for _ in range(i_max)]
+ for idx, encoded in enumerate(all_encoded):
+ i, j = idx // j_max, idx % j_max
+ rows[i][j] = encoded.to(x.device)
+
+ latent_y_overlap = [
+ tile_overlap // self.vae_ratio for tile_overlap in y_overlap
+ ]
+ latent_x_overlap = [
+ tile_overlap // self.vae_ratio for tile_overlap in x_overlap
+ ]
+
+ z = self._assemble_tiles(rows, latent_y_overlap, latent_x_overlap)
+
+ return z
+
+ def tiled_decode(self, z):
+ if self.parallel_tiling: # Fast online decoding for large videos
+ tile_rank, tile_world_size = get_tile_parallel_state()
+ else:
+ tile_rank, tile_world_size = 0, 1
+
+ height, width = (
+ z.shape[-2] * self.vae_ratio,
+ z.shape[-1] * self.vae_ratio,
+ )
+ y_idx, y_len, y_overlap = self.split_tiles(height, True)
+ x_idx, x_len, x_overlap = self.split_tiles(width, True)
+
+ i_max, j_max = len(y_idx), len(x_idx)
+ num_tiles = i_max * j_max
+ if tile_world_size > num_tiles:
+ # See tiled_encode: small canvases fall back to replicated local
+ # tiling so no decode rank is assigned an empty task list.
+ tile_rank, tile_world_size = 0, 1
+
+ z_tiles = []
+ for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
+ i_pos, i_len = (
+ i_pos // self.vae_ratio,
+ i_len // self.vae_ratio,
+ )
+ for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
+ j_pos, j_len = (j_pos // self.vae_ratio, j_len // self.vae_ratio)
+ tile = z[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
+ z_tiles.append(tile)
+
+ local_tile_indices = self._local_tile_indices(
+ num_tiles, tile_rank, tile_world_size
+ )
+ stack_tiling = self.stack_tiling and not (
+ self.training and getattr(self.decoder, "mask_enabled", False)
+ )
+ decoded_tasks = self._run_tile_tasks(
+ z_tiles, local_tile_indices, self.decode, stack_tiling
+ )
+
+ if tile_world_size > 1:
+ all_decoded = self._all_gather_tiled_results(decoded_tasks, num_tiles)
+ else:
+ all_decoded = decoded_tasks
+
+ rows = [[None for _ in range(j_max)] for _ in range(i_max)]
+ for idx, decoded in enumerate(all_decoded):
+ i, j = idx // j_max, idx % j_max
+ rows[i][j] = decoded.to(z.device)
+
+ dec = self._assemble_tiles(rows, y_overlap, x_overlap)
+ return dec
+
+ def _adaptive_encode(self, x):
+ if self.encoder_tiling:
+ return self.tiled_encode(x)
+ else:
+ return self.encode(x)
+
+ def _adaptive_decode(self, z):
+ if self.decoder_tiling:
+ return self.tiled_decode(z)
+ else:
+ return self.decode(z)
+
+ def trim_code(self, z, target_codes):
+ if target_codes < z.shape[2]:
+ if self.causal_encoder:
+ z = z[:, :, -target_codes:, :, :]
+ else:
+ start_frame = (z.shape[2] - target_codes) // 2
+ z = z[:, :, start_frame : start_frame + target_codes, :, :]
+ return z
+
+ def trim_output(self, dec, target_frames):
+ if target_frames < dec.shape[2]:
+ if self.causal_encoder: # This is defined by encoder, not decoder
+ dec = dec[:, :, -target_frames:, :, :]
+ else:
+ start_frame = (dec.shape[2] - target_frames) // 2
+ dec = dec[:, :, start_frame : start_frame + target_frames, :, :]
+ return dec
+
+ def encode_temporal(self, x):
+ offset_frame = (
+ 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0
+ )
+
+ frame_num = x.shape[2]
+ pad_size = (offset_frame - frame_num) % self.clip_length
+ padded_frame_num = frame_num + pad_size
+ num_chunks = (padded_frame_num - offset_frame) // self.clip_length
+
+ z_list = []
+ for i in range(num_chunks):
+ start_idx = i * self.clip_length + offset_frame
+ end_idx = (i + 1) * self.clip_length + offset_frame
+ clip_x = x[:, :, start_idx : min(end_idx, frame_num), :, :]
+ if end_idx > frame_num:
+ pad_frames = x[:, :, -1:].expand(-1, -1, end_idx - frame_num, -1, -1)
+ clip_x = torch.cat([clip_x, pad_frames], dim=2)
+
+ if self.isolated_key_frame:
+ key_frame = clip_x[:, :, :1, :, :]
+ z_key = self._adaptive_encode(key_frame)
+
+ if clip_x.shape[2] > 1:
+ video_frames = clip_x[:, :, 1:, :, :]
+ z_video = self._adaptive_encode(video_frames)
+ z = torch.cat([z_key, z_video], dim=2)
+ else:
+ z = z_key
+ else:
+ z = self._adaptive_encode(clip_x)
+
+ z_list.append(z)
+
+ z = torch.cat(z_list, dim=2)
+ if self.token_drop > 0:
+ z = z[:, :, : -self.token_drop]
+
+ if self.isolated_first_frame:
+ input_first_frame = x[:, :, :1, :, :]
+ z_first_frame = self._adaptive_encode(input_first_frame)
+
+ if self.frame_pre_padding == 0:
+ z = torch.cat([z_first_frame, z], dim=2)
+ else:
+ z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2)
+
+ if self.isolated_last_frame:
+ last_frame_idx = padded_frame_num - self.frame_drop + offset_frame
+ if last_frame_idx >= frame_num:
+ input_last_frame = x[:, :, -1:, :, :]
+ else:
+ input_last_frame = x[:, :, last_frame_idx : last_frame_idx + 1, :, :]
+ z_last_frame = self._adaptive_encode(input_last_frame)
+ z = torch.cat([z, z_last_frame], dim=2)
+
+ return z
+
+ def _decode_temporal_pad_frames(self, z, pad_tokens):
+ if pad_tokens <= 0:
+ return 0
+ intra_tail = self.clip_length % self.vae_ratio_t
+ if intra_tail == 0:
+ return int(pad_tokens) * int(self.vae_ratio_t)
+
+ z_len_before_pad = z.shape[2] - pad_tokens
+ return sum(
+ (
+ intra_tail
+ if (z_len_before_pad + k) % self.tokens_chunk_size == 0
+ else self.vae_ratio_t
+ )
+ for k in range(pad_tokens)
+ )
+
+ def _decode_temporal_output_frame_plan(
+ self, z, z_head, z_tail, num_chunks, pad_tokens
+ ):
+ chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+ split_count = int(self.token_drop > 0) + 1
+ total_frames = 0
+ final_overlap_frames = 0
+
+ if z_head is not None:
+ total_frames += 1
+
+ for i in range(num_chunks):
+ t_start_idx = i * self.tokens_chunk_size
+ t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
+ clip_token_len = max(
+ 0, min(t_end_idx, z.shape[2]) - min(t_start_idx, z.shape[2])
+ )
+ if i == 0 and z_head is not None:
+ clip_token_len += z_head.shape[2]
+ if i == num_chunks - 1 and z_tail is not None:
+ clip_token_len += z_tail.shape[2]
+
+ clip_frame_len = clip_token_len * self.vae_ratio_t
+ if i == 0 and z_head is not None:
+ clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t)
+ if i == num_chunks - 1 and z_tail is not None:
+ clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t)
+
+ for j in range(split_count):
+ f_start_idx = j * chunk_dec
+ f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len)
+ chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding)
+ if j == 0:
+ total_frames += chunk_frames
+ else:
+ final_overlap_frames = chunk_frames
+
+ total_frames += final_overlap_frames
+ if z_tail is not None:
+ total_frames += 1
+
+ pad_frames = self._decode_temporal_pad_frames(z, pad_tokens)
+ return int(total_frames), int(pad_frames), int(total_frames - pad_frames)
+
+ def _decode_temporal_streaming(
+ self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype
+ ):
+ total_frames, pad_frames, output_frames = (
+ self._decode_temporal_output_frame_plan(
+ z, z_head, z_tail, num_chunks, pad_tokens
+ )
+ )
+ if output_frames <= 0:
+ raise ValueError(
+ f"decode_temporal streaming planned non-positive output_frames={output_frames} "
+ f"total_frames={total_frames} pad_frames={pad_frames}"
+ )
+
+ chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+ split_count = int(self.token_drop > 0) + 1
+ dec = None
+ dec_overlap = None
+ write_pos = 0
+ logical_frames = 0
+ dropped_frames = 0
+ decoded_count = 0
+
+ def write_part(part):
+ nonlocal dec, write_pos, logical_frames, dropped_frames
+ part_frames = int(part.shape[2])
+ if part_frames <= 0:
+ return
+ logical_frames += part_frames
+ if dec is None:
+ out_shape = list(part.shape)
+ out_shape[2] = output_frames
+ dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
+
+ remaining = int(dec.shape[2]) - write_pos
+ copy_frames = min(part_frames, max(0, remaining))
+ if copy_frames > 0:
+ dec[:, :, write_pos : write_pos + copy_frames, :, :].copy_(
+ part[:, :, :copy_frames, :, :]
+ )
+ write_pos += copy_frames
+ dropped_frames += part_frames - copy_frames
+
+ for i in range(num_chunks):
+ t_start_idx = i * self.tokens_chunk_size
+ t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
+ clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
+
+ if i == 0 and z_head is not None:
+ clip_z = torch.cat([z_head, clip_z], dim=2)
+
+ if i == num_chunks - 1 and z_tail is not None:
+ clip_z = torch.cat([clip_z, z_tail], dim=2)
+
+ clip_dec = self._adaptive_decode(clip_z)
+ decoded_count += 1
+ if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
+ clip_dec = clip_dec.to(temporal_cat_dtype)
+ if clip_dec.device != z.device:
+ clip_dec = clip_dec.to(z.device)
+
+ dec_tail = None
+ if i == 0 and z_head is not None:
+ write_part(
+ clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :]
+ )
+ clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :]
+
+ if i == num_chunks - 1 and z_tail is not None:
+ dec_tail = clip_dec[:, :, -1:, :, :]
+ clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :]
+
+ for j in range(split_count):
+ f_start_idx = j * chunk_dec
+ f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
+ clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
+ clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]
+
+ if j == 0:
+ if dec_overlap is not None:
+ clip_dec_chunk = self.blend(
+ dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
+ )
+ dec_overlap = None
+ write_part(clip_dec_chunk)
+ else:
+ # Break the view's reference to the full decoded clip so earlier
+ # temporal chunks can be released before the final output exists.
+ dec_overlap = clip_dec_chunk.contiguous()
+
+ if i == num_chunks - 1:
+ if dec_overlap is not None:
+ write_part(dec_overlap)
+ dec_overlap = None
+ if dec_tail is not None:
+ write_part(dec_tail)
+
+ del clip_dec, clip_z
+
+ if dec is None:
+ raise RuntimeError("decode_temporal streaming produced no output tensor")
+ if (
+ logical_frames != total_frames
+ or dropped_frames != pad_frames
+ or write_pos != output_frames
+ ):
+ raise RuntimeError(
+ "decode_temporal streaming frame plan mismatch: "
+ f"logical_frames={logical_frames} total_frames={total_frames} "
+ f"dropped_frames={dropped_frames} pad_frames={pad_frames} "
+ f"write_pos={write_pos} output_frames={output_frames}"
+ )
+
+ return dec
+
+ def decode_temporal(self, z):
+ chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+
+ isolated_token_num = 0
+ if self.isolated_first_frame and self.frame_pre_padding == 0:
+ isolated_token_num = isolated_token_num + 1
+ if self.isolated_last_frame:
+ isolated_token_num = isolated_token_num + 1
+
+ pseudo_total_tokens = z.shape[2] - isolated_token_num + self.token_drop
+
+ pad_tokens = 0
+ remainder = pseudo_total_tokens % self.tokens_chunk_size
+ if remainder != 0:
+ if self.training:
+ raise ValueError(f"Temporal token length {z.shape[2]} is wrong!")
+ else:
+ pad_tokens = self.tokens_chunk_size - remainder
+ pseudo_total_tokens = pseudo_total_tokens + pad_tokens
+
+ pseudo_num_chunks = pseudo_total_tokens // self.tokens_chunk_size
+ num_chunks = pseudo_num_chunks - int(self.token_drop > 0)
+
+ z_head = None
+ if self.isolated_first_frame and self.frame_pre_padding == 0:
+ z_head = z[:, :, :1, :, :]
+ z = z[:, :, 1:, :, :]
+
+ z_tail = None
+ if self.isolated_last_frame:
+ z_tail = z[:, :, -1:, :, :]
+ z = z[:, :, :-1, :, :]
+
+ if pad_tokens > 0:
+ pad_z = z[:, :, -1:, :, :].expand(-1, -1, pad_tokens, -1, -1)
+ z = torch.cat([z, pad_z], dim=2)
+
+ temporal_cat_dtype = _resolve_temporal_cat_dtype()
+ if not self.training and _resolve_temporal_stream_cat():
+ return self._decode_temporal_streaming(
+ z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype
+ )
+
+ decoded_tasks = []
+ for i in range(num_chunks):
+ t_start_idx = i * self.tokens_chunk_size
+ t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
+ clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
+
+ if i == 0 and z_head is not None:
+ clip_z = torch.cat([z_head, clip_z], dim=2)
+
+ if i == num_chunks - 1 and z_tail is not None:
+ clip_z = torch.cat([clip_z, z_tail], dim=2)
+
+ clip_dec = self._adaptive_decode(clip_z)
+ if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
+ clip_dec = clip_dec.to(temporal_cat_dtype)
+
+ decoded_tasks.append((i, clip_dec))
+
+ clip_dec_list = [clip_dec.to(z.device) for _, clip_dec in decoded_tasks]
+
+ dec_list = []
+ dec_overlap = None
+
+ dec_head = None
+ if z_head is not None:
+ dec_head = clip_dec_list[0][
+ :, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :
+ ]
+ clip_dec_list[0] = clip_dec_list[0][:, :, self.vae_ratio_t :, :, :]
+
+ dec_tail = None
+ if z_tail is not None:
+ dec_tail = clip_dec_list[-1][:, :, -1:, :, :]
+ clip_dec_list[-1] = clip_dec_list[-1][:, :, : -self.vae_ratio_t, :, :]
+
+ if dec_head is not None:
+ dec_list.append(dec_head)
+
+ for i in range(num_chunks):
+ for j in range(int(self.token_drop > 0) + 1):
+ clip_dec = clip_dec_list[i]
+
+ f_start_idx = j * chunk_dec
+ f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
+ clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
+ clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]
+
+ if j == 0:
+ if dec_overlap is not None:
+ clip_dec_chunk = self.blend(
+ dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
+ )
+ dec_list.append(clip_dec_chunk)
+ else:
+ dec_overlap = clip_dec_chunk
+
+ if dec_overlap is not None:
+ dec_list.append(dec_overlap)
+
+ if dec_tail is not None:
+ dec_list.append(dec_tail)
+
+ dec = torch.cat(dec_list, dim=2)
+
+ pad_frames = self._decode_temporal_pad_frames(z, pad_tokens)
+ if pad_frames > 0:
+ dec = dec[:, :, :-pad_frames, :, :]
+
+ return dec
+
+ def decode_base(self, z, frame_num=None, process_image=False):
+ if process_image or not self.use_3d_conv:
+ if not self.use_3d_conv and z.ndim == 5:
+ z = z.squeeze(2)
+
+ recon = self._adaptive_decode(z)
+ else:
+ recon = self.decode_temporal(z)
+
+ if self.use_3d_conv:
+ if frame_num is not None:
+ target_frames = frame_num
+ else:
+ target_frames = recon.shape[2]
+
+ recon = self.trim_output(recon, target_frames)
+ if process_image:
+ recon = recon.squeeze(2)
+
+ return recon
+
+ #########################################################
+ # following methods are for inference
+ #########################################################
+
+ @torch.no_grad()
+ def encode_images(
+ self,
+ images: Union[List[np.ndarray], List[torch.Tensor]],
+ transform_input: bool = False,
+ use_fp16_latent: bool = False,
+ verbose: bool = False,
+ ) -> List[torch.Tensor]:
+ """encode images into latents
+
+ Args:
+ images (Union[List[np.ndarray], List[torch.Tensor]]):
+ List of images, single input will be wrapped in a list.
+ If input is a list of np.ndarray, it should be in shape B * (H, W, 3), dtype uint8.
+ If input is a list of torch.Tensor, it should be in shape B * (3, H, W), dtype float32.
+ transform_input (bool, optional):
+ Whether to transform input using ImageNet std/mean. Defaults to False.
+ If input is a list of np.ndarray, it will always be set to True.
+ use_fp16_latent (bool, optional):
+ Whether to use fp16 latent. Defaults to False.
+ verbose (bool, optional):
+ Whether to print debug information. Defaults to False.
+
+ Returns:
+ List[torch.Tensor]:
+ List of image latents.
+ If self.use_3d_conv is True, it should be in shape B * (D, 1, H', W').
+ Otherwise, it should be in shape B * (D, H', W').
+ """
+
+ images = self.processor._ensure_list(images)
+ runtime_owned = False
+
+ if isinstance(images[0], Image.Image):
+ images = [np.array(image) for image in images]
+
+ if isinstance(images[0], np.ndarray):
+ device = next(self.parameters()).device
+ images = self.processor.convert_numpy_to_tensor(images, device)
+ images = torch.split(images, 1, dim=0)
+ transform_input = True
+ runtime_owned = True
+
+ if transform_input:
+ images = [
+ image.unsqueeze(0) if image.ndim == 3 else image for image in images
+ ]
+ images = [
+ self.processor.transform_tensor(image, runtime_owned=runtime_owned)
+ for image in images
+ ]
+
+ prepared = []
+ for image_tensor in images:
+ if image_tensor.ndim == 3:
+ image_tensor = image_tensor.unsqueeze(0)
+ _, _, h, w = image_tensor.shape
+ new_h, new_w = self.processor._align_to_total_patch_size(h, w)
+ image_tensor = self.processor._crop_to_align(image_tensor, new_h, new_w)
+ prepared.append(image_tensor)
+
+ if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1:
+ stacked = torch.cat(prepared, dim=0)
+ if verbose:
+ logger.info(f"batch encode input shape {tuple(stacked.shape)}")
+ all_latents = self.encode_base(stacked, True)
+ image_latents = [
+ all_latents[i].contiguous() for i in range(all_latents.shape[0])
+ ]
+ else:
+ image_latents = []
+ for image_tensor in prepared:
+ if verbose:
+ logger.info(f"input shape {tuple(image_tensor.shape)}")
+ image_latent = self.encode_base(image_tensor, True)
+ image_latents.append(image_latent.squeeze(0).contiguous())
+
+ if use_fp16_latent:
+ image_latents = [lat.to(torch.float16) for lat in image_latents]
+
+ if verbose:
+ for lat in image_latents:
+ logger.info(f"image latent shape {tuple(lat.shape)}")
+
+ return image_latents
+
+ @torch.no_grad()
+ def encode_videos(
+ self,
+ videos: Union[List[np.ndarray], List[torch.Tensor]],
+ transform_input: bool = False,
+ use_fp16_latent: bool = False,
+ verbose: bool = False,
+ encode_prefix: bool = False,
+ ) -> List[torch.Tensor]:
+ """encode videos into latents
+
+ Args:
+ videos (Union[List[np.ndarray], List[torch.Tensor]]):
+ List of videos, single input will be wrapped in a list.
+ If input is a list of np.ndarray, it should be in shape B * (T, H, W, 3), dtype uint8.
+ If input is a list of torch.Tensor, it should be in shape B * (3, T, H, W), dtype float32.
+ transform_input (bool, optional):
+ Whether to transform input using ImageNet std/mean. Defaults to False.
+ If input is a list of np.ndarray, it will always be set to True.
+ use_fp16_latent (bool, optional):
+ Whether to use fp16 latent. Defaults to False.
+ verbose (bool, optional):
+ Whether to print debug information. Defaults to False.
+ encode_prefix (bool, optional):
+ Continuation (prefix) mode: prepend normalized
+ black frames to token alignment, append black frames to chunk
+ alignment, encode with token_drop disabled, then discard only
+ the trailing padding tokens. Returns both latents and leading
+ pad-frame counts. Defaults to False.
+
+ Returns:
+ List[torch.Tensor]:
+ List of video latents, shape B * (D, T', H', W').
+ With encode_prefix=True, returns
+ (List[torch.Tensor], List[int]).
+ """
+
+ videos = self.processor._ensure_list(videos)
+ runtime_owned = False
+
+ if isinstance(videos[0], np.ndarray):
+ device = next(self.parameters()).device
+ videos = [
+ self.processor.convert_numpy_to_tensor(video, device)
+ for video in videos
+ ]
+ transform_input = True
+ runtime_owned = True
+
+ if transform_input:
+ videos = [
+ self.processor.transform_tensor(video, runtime_owned=runtime_owned)
+ for video in videos
+ ]
+ videos = [video.transpose(0, 1) for video in videos]
+
+ if encode_prefix:
+ if self.isolated_last_frame:
+ raise ValueError("encode_prefix does not support isolated_last_frame")
+
+ video_latents = []
+ prefix_pad_frames = []
+ for video in videos:
+ if video.ndim == 4:
+ video = video.unsqueeze(0)
+ _, _, _, h, w = video.shape
+ new_h, new_w = self.processor._align_to_total_patch_size(h, w)
+ video = self.processor._crop_to_align(
+ video, new_h, new_w, is_video=True
+ )
+
+ model_alignment = (
+ self.token_drop,
+ self.frame_drop,
+ self.token_overlap,
+ self.frame_overlap,
+ )
+ processor_alignment = (
+ self.processor.token_overlap,
+ self.processor.frame_overlap,
+ )
+ self.token_drop = 0
+ self.frame_drop = 0
+ self.token_overlap = 0
+ self.frame_overlap = 0
+ self.processor.token_overlap = 0
+ self.processor.frame_overlap = 0
+ try:
+ orig_frames = video.shape[2]
+ leading, trailing, drop_tokens = (
+ self.processor.align_video_length_2pass(orig_frames)
+ )
+ _, _, _, cropped_h, cropped_w = video.shape
+ if leading > 0:
+ black = self.processor.transform(
+ video.new_zeros(leading, 3, cropped_h, cropped_w)
+ )
+ black = black.unsqueeze(0).permute(0, 2, 1, 3, 4)
+ video = torch.cat([black, video], dim=2)
+ if trailing > 0:
+ black = self.processor.transform(
+ video.new_zeros(trailing, 3, cropped_h, cropped_w)
+ )
+ black = black.unsqueeze(0).permute(0, 2, 1, 3, 4)
+ video = torch.cat([video, black], dim=2)
+
+ if verbose:
+ logger.info(
+ f"[encode_prefix] {orig_frames} frames -> "
+ f"pad leading={leading}, trailing={trailing} -> "
+ f"{video.shape[2]} frames"
+ )
+
+ video_latent = self.encode_base(video, False)
+ if drop_tokens > 0:
+ video_latent = video_latent[:, :, :-drop_tokens, :, :]
+ prefix_pad_frames.append(leading)
+ finally:
+ (
+ self.token_drop,
+ self.frame_drop,
+ self.token_overlap,
+ self.frame_overlap,
+ ) = model_alignment
+ (
+ self.processor.token_overlap,
+ self.processor.frame_overlap,
+ ) = processor_alignment
+
+ video_latents.append(video_latent.squeeze(0).contiguous())
+
+ if use_fp16_latent:
+ video_latents = [lat.to(torch.float16) for lat in video_latents]
+ if verbose:
+ for latent in video_latents:
+ logger.info(f"video latent shape {tuple(latent.shape)}")
+ return video_latents, prefix_pad_frames
+
+ prepared = []
+ for video in videos:
+ if video.ndim == 4:
+ video = video.unsqueeze(0)
+ used_frame_length = self.processor.get_suitable_video_length(
+ video.shape[2], verbose
+ )
+ _, _, _, h, w = video.shape
+ new_h, new_w = self.processor._align_to_total_patch_size(h, w)
+ video = video[:, :, :used_frame_length, :, :]
+ video = self.processor._crop_to_align(video, new_h, new_w, is_video=True)
+ prepared.append(video)
+
+ if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1:
+ stacked = torch.cat(prepared, dim=0)
+ if verbose:
+ logger.info(f"batch encode input shape {tuple(stacked.shape)}")
+ all_latents = self.encode_base(stacked, False)
+ video_latents = [
+ all_latents[i].contiguous() for i in range(all_latents.shape[0])
+ ]
+ else:
+ video_latents = []
+ for video in prepared:
+ if verbose:
+ logger.info(f"input shape {tuple(video.shape)}")
+ video_latent = self.encode_base(video, False)
+ video_latents.append(video_latent.squeeze(0).contiguous())
+
+ if use_fp16_latent:
+ video_latents = [lat.to(torch.float16) for lat in video_latents]
+
+ if verbose:
+ for lat in video_latents:
+ logger.info(f"video latent shape {tuple(lat.shape)}")
+
+ return video_latents
+
+
+# ============================================================================
+# Legacy CNN VAE
+# ============================================================================
+
+
+class AutoencoderKLLegacy(AutoencoderKL):
+ r"""
+ A VAE model (legacy CNN-based) for encoding pixels into latents and decoding latent representations into pixels.
+ """
+
+ @register_to_config
+ def __init__(
+ self,
+ in_channels=3,
+ out_ch=3,
+ ch=128,
+ embed_dim=16,
+ z_channels=16,
+ use_3d_conv=False,
+ # cnn vae
+ zq_ch_encoder=None,
+ zq_ch_decoder=None,
+ num_res_blocks=2,
+ num_res_blocks_decoder=None,
+ ch_mult=[1, 2, 2, 4, 4, 8],
+ space_down=[2, 2, 2, 2, 1, 1],
+ space_up=[1, 2, 2, 2, 2, 1],
+ time_down=None,
+ time_up=None,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ use_t_isolated_gn=False,
+ causal_encoder=True,
+ causal_decoder=True,
+ use_vit_decoder=False,
+ vit_decoder_kwargs=None,
+ # stats
+ shift_factor=0.0,
+ scaling_factor=1.0,
+ # pixel normalization
+ pixel_norm_type="imagenet",
+ # others
+ **kwargs,
+ ):
+ ModelMixin.__init__(self) # NOTE: avoid wrong @register_to_config
+
+ if not use_3d_conv or not use_vit_decoder:
+ raise NotImplementedError(
+ "this release only supports use_3d_conv=True with use_vit_decoder=True"
+ )
+
+ self.transform = get_normalize_transform(pixel_norm_type)
+ self.transform_rev = get_denormalize_transform(pixel_norm_type)
+
+ self.use_3d_conv = use_3d_conv
+ self.causal_encoder = causal_encoder
+ self.causal_decoder = causal_decoder
+ self.slidedec = self.causal_encoder and not self.causal_decoder
+
+ # some registered parameters for simplicity
+ self.vae_ratio = int(np.cumprod(space_down)[-1])
+ self.vae_ratio_t = int(np.cumprod(time_down)[-1]) if time_down else 1
+ self.config["vae_ratio"] = self.vae_ratio
+ self.config["vae_ratio_t"] = self.vae_ratio_t
+
+ # Configure inference-time chunking and tiling.
+ self.setup_forward(**kwargs)
+
+ # init encoder
+ encoder_config = {
+ "double_z": True,
+ "z_channels": z_channels,
+ "zq_ch": zq_ch_encoder,
+ "in_channels": in_channels,
+ "ch": ch,
+ "num_res_blocks": num_res_blocks,
+ "ch_mult": ch_mult,
+ "space_down": space_down,
+ "time_down": time_down,
+ "padding_mode": padding_mode,
+ "padding_mode_t": padding_mode_t,
+ "causal": causal_encoder,
+ "use_t_isolated_gn": use_t_isolated_gn,
+ }
+ self.encoder = EncoderFCN3D(**encoder_config)
+
+ # init pointwise quant/post_quant conv
+ self.quant_conv = nn.Conv3d(z_channels * 2, 2 * embed_dim, 1)
+ self.post_quant_conv = nn.Conv3d(embed_dim, z_channels, 1)
+
+ self.use_vit_decoder = use_vit_decoder
+
+ # init decoder
+ vit_kwargs = {
+ "patch_size": self.vae_ratio,
+ "in_channels": z_channels,
+ "out_channels": out_ch,
+ **(vit_decoder_kwargs or {}),
+ }
+ vit_kwargs.setdefault("patch_size_t", self.vae_ratio_t)
+ vit_kwargs.setdefault("t_causal", causal_decoder)
+ self.decoder = ViT3DDecoder(**vit_kwargs)
+
+ @torch.no_grad()
+ def encode(self, x):
+ return self.quant_conv(self.encoder(x))
+
+ @torch.no_grad()
+ def decode(self, z):
+ z2 = self.post_quant_conv(z)
+ if self.use_vit_decoder:
+ return self.decoder(z2)
+ return self.decoder(z2, z)
+
+ def encode_base(self, input, process_image=False):
+ if self.use_3d_conv and input.ndim == 4:
+ input = input.unsqueeze(2)
+
+ if process_image or not self.use_3d_conv:
+ moments = self._adaptive_encode(input)
+ else:
+ moments = self.encode_temporal(input)
+
+ z = DiagonalGaussianDistribution(moments).sample()
+
+ if process_image and self.use_3d_conv:
+ z = self.trim_code(z, 1)
+
+ return z
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/norm.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/norm.py
new file mode 100644
index 000000000..c1ad34d89
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/norm.py
@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+# Torch-native normalization for the MiniMax H3 visual VAE.
+import math
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .conv import BaseConv3d
+
+
+def _validate_activation(activation):
+ valid_activations = {"identity", "silu", "relu"}
+ if activation not in valid_activations:
+ raise ValueError(
+ f"Unsupported activation: {activation}. Supported: {valid_activations}"
+ )
+
+
+def _apply_activation(x, activation):
+ _validate_activation(activation)
+ if activation == "identity":
+ return x
+ if activation == "silu":
+ return F.silu(x)
+ return F.relu(x)
+
+
+def _merge_time_to_batch(x):
+ batch, channels, depth, height, width = x.shape
+ return (
+ x.permute(0, 2, 1, 3, 4)
+ .contiguous()
+ .view(batch * depth, channels, 1, height, width)
+ )
+
+
+def _split_time_from_batch(x, batch):
+ batch_depth, channels, _, height, width = x.shape
+ depth = batch_depth // batch
+ return (
+ x.view(batch, depth, channels, height, width)
+ .permute(0, 2, 1, 3, 4)
+ .contiguous()
+ )
+
+
+def fused_group_norm(x, num_groups, weight, bias, eps=1e-5, activation="silu"):
+ out = F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps)
+ return _apply_activation(out, activation)
+
+
+def fused_spatial_norm(
+ f,
+ num_groups,
+ norm_weight,
+ norm_bias,
+ dynamic_scale,
+ dynamic_bias,
+ eps=1e-5,
+ activation="silu",
+):
+ norm_f = F.group_norm(
+ f,
+ num_groups,
+ weight=norm_weight,
+ bias=norm_bias,
+ eps=eps,
+ )
+ out = norm_f * dynamic_scale + dynamic_bias
+ return _apply_activation(out, activation)
+
+
+class DummyAffine(torch.nn.Module):
+ def __init__(self, num_channels, affine=True):
+ super().__init__()
+ if affine:
+ self.weight = torch.nn.Parameter(torch.ones(num_channels))
+ self.bias = torch.nn.Parameter(torch.zeros(num_channels))
+ else:
+ self.register_parameter("weight", None)
+ self.register_parameter("bias", None)
+
+ def forward(self, input):
+ if self.weight is None:
+ return input
+ shape = [1, -1] + [1] * (input.dim() - 2)
+ return input * self.weight.view(*shape) + self.bias.view(*shape)
+
+
+class FusedGroupNorm3D(torch.nn.Module):
+ """Compatibility wrapper implemented with native PyTorch ops."""
+
+ def __init__(
+ self,
+ num_groups,
+ num_channels,
+ eps=1e-5,
+ affine=True,
+ activation="silu",
+ cond_channels=None,
+ use_t_isolated_gn=False,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ ):
+ super().__init__()
+ _validate_activation(activation)
+ self.num_groups = num_groups
+ self.num_channels = num_channels
+ self.eps = eps
+ self.affine = affine
+ self.activation = activation
+ self.use_t_isolated_gn = use_t_isolated_gn
+
+ if cond_channels is not None:
+ self.use_spatial_affine = True
+ self.norm_layer = DummyAffine(num_channels, affine=affine)
+ self.conv_y = BaseConv3d(
+ cond_channels,
+ num_channels,
+ kernel_size=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+ self.conv_b = BaseConv3d(
+ cond_channels,
+ num_channels,
+ kernel_size=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+ else:
+ self.use_spatial_affine = False
+ if self.affine:
+ self.weight = torch.nn.Parameter(torch.ones(num_channels))
+ self.bias = torch.nn.Parameter(torch.zeros(num_channels))
+ else:
+ self.register_parameter("weight", None)
+ self.register_parameter("bias", None)
+
+ def forward(self, f, cond=None):
+ need_reshape = self.use_t_isolated_gn and f.dim() == 5
+ batch = f.shape[0] if need_reshape else None
+ f_size = f.shape[-3:]
+ if need_reshape:
+ f = _merge_time_to_batch(f)
+
+ if self.use_spatial_affine:
+ scale = self.conv_y(cond)
+ bias = self.conv_b(cond)
+ if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:
+ scale = F.interpolate(scale, size=f_size, mode="nearest")
+ bias = F.interpolate(bias, size=f_size, mode="nearest")
+ if need_reshape:
+ scale = _merge_time_to_batch(scale)
+ bias = _merge_time_to_batch(bias)
+ out = fused_spatial_norm(
+ f,
+ self.num_groups,
+ self.norm_layer.weight,
+ self.norm_layer.bias,
+ scale,
+ bias,
+ self.eps,
+ self.activation,
+ )
+ else:
+ if cond is not None:
+ raise NotImplementedError("Dynamic affine is not defined")
+ weight = self.weight if self.affine else None
+ bias = self.bias if self.affine else None
+ out = fused_group_norm(
+ f, self.num_groups, weight, bias, self.eps, self.activation
+ )
+
+ if need_reshape:
+ out = _split_time_from_batch(out, batch)
+ return out
+
+
+class TemporalIsolatedGroupNorm(nn.GroupNorm):
+ def forward(self, input):
+ if input.dim() == 5:
+ batch = input.shape[0]
+ input = _merge_time_to_batch(input)
+ output = super().forward(input)
+ return _split_time_from_batch(output, batch)
+ return super().forward(input)
+
+
+class SpatialNorm3D(nn.Module):
+ def __init__(
+ self,
+ f_channels,
+ zq_channels,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ use_t_isolated_gn=False,
+ ):
+ super().__init__()
+ norm_cls = TemporalIsolatedGroupNorm if use_t_isolated_gn else nn.GroupNorm
+ self.norm_layer = norm_cls(
+ num_groups=32, num_channels=f_channels, eps=1e-6, affine=True
+ )
+
+ self.conv_y = BaseConv3d(
+ zq_channels,
+ f_channels,
+ kernel_size=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+ self.conv_b = BaseConv3d(
+ zq_channels,
+ f_channels,
+ kernel_size=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+
+ def forward(self, f, zq):
+ f_size = f.shape[-3:]
+ norm_f = self.norm_layer(f)
+ scale = self.conv_y(zq)
+ bias = self.conv_b(zq)
+
+ if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:
+ scale = F.interpolate(scale, size=f_size, mode="nearest")
+ bias = F.interpolate(bias, size=f_size, mode="nearest")
+
+ return norm_f * scale + bias
+
+
+def get_spatial_norm_3d(
+ num_channels,
+ cond_channels,
+ *,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ use_t_isolated_gn=False,
+):
+ if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":
+ return FusedGroupNorm3D(
+ num_groups=32,
+ num_channels=num_channels,
+ eps=1e-6,
+ affine=True,
+ cond_channels=cond_channels,
+ use_t_isolated_gn=use_t_isolated_gn,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+ return SpatialNorm3D(
+ num_channels,
+ cond_channels,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ use_t_isolated_gn=use_t_isolated_gn,
+ )
+
+
+def get_group_norm_3d(num_channels, use_t_isolated_gn=False):
+ if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":
+ return FusedGroupNorm3D(
+ num_groups=32,
+ num_channels=num_channels,
+ eps=1e-6,
+ affine=True,
+ use_t_isolated_gn=use_t_isolated_gn,
+ )
+
+ norm_cls = TemporalIsolatedGroupNorm if use_t_isolated_gn else nn.GroupNorm
+ return norm_cls(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py
new file mode 100644
index 000000000..3fa00d9a7
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/processor.py
@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# Tensor pre/post-processing for the MiniMax H3 visual VAE.
+import math
+from typing import Tuple
+
+import numpy as np
+import torch
+from diffusers.utils import logging
+from einops import rearrange
+from torchvision.transforms import Normalize
+
+NORM_CONFIGS = {
+ "imagenet": {
+ "mean": (0.485, 0.456, 0.406),
+ "std": (0.229, 0.224, 0.225),
+ },
+ "simple": {
+ "mean": (0.5, 0.5, 0.5),
+ "std": (0.5, 0.5, 0.5),
+ },
+ "raw": {
+ "mean": (0.0, 0.0, 0.0),
+ "std": (1.0, 1.0, 1.0),
+ },
+}
+
+
+def get_norm_constants(
+ norm_type: str = "imagenet",
+) -> Tuple[Tuple[float, ...], Tuple[float, ...]]:
+ if norm_type not in NORM_CONFIGS:
+ raise ValueError(
+ f"Unknown norm_type: {norm_type}. Must be one of {list(NORM_CONFIGS.keys())}"
+ )
+ config = NORM_CONFIGS[norm_type]
+ return config["mean"], config["std"]
+
+
+def get_normalize_transform(
+ norm_type: str = "imagenet", *, inplace: bool = False
+) -> Normalize:
+ mean, std = get_norm_constants(norm_type)
+ return Normalize(mean, std, inplace=inplace)
+
+
+def get_denormalize_transform(norm_type: str = "imagenet") -> Normalize:
+ mean, std = get_norm_constants(norm_type)
+ inv_mean = tuple(-m / s for m, s in zip(mean, std))
+ inv_std = tuple(1.0 / s for s in std)
+ return Normalize(inv_mean, inv_std)
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class VAEProcessor:
+
+ def __init__(
+ self,
+ *,
+ vae_ratio,
+ vae_ratio_t,
+ clip_length,
+ frame_overlap,
+ token_overlap,
+ tokens_chunk_size,
+ isolated_last_frame,
+ latent_patch_size,
+ crop_mode,
+ pixel_norm_type="imagenet",
+ transform=None,
+ transform_rev=None,
+ use_3d_conv=False,
+ ):
+ self.vae_ratio = vae_ratio
+ self.vae_ratio_t = vae_ratio_t
+ self.clip_length = clip_length
+ self.frame_overlap = frame_overlap
+ self.token_overlap = token_overlap
+ self.tokens_chunk_size = tokens_chunk_size
+ self.isolated_last_frame = isolated_last_frame
+ self.latent_patch_size = latent_patch_size
+ self.crop_mode = crop_mode
+ self.transform = transform or get_normalize_transform(pixel_norm_type)
+ self._runtime_owned_transform = (
+ get_normalize_transform(pixel_norm_type, inplace=True)
+ if transform is None
+ else None
+ )
+ self.transform_rev = transform_rev or get_denormalize_transform(pixel_norm_type)
+ self.use_3d_conv = use_3d_conv
+
+ def _ensure_list(self, data):
+ return data if isinstance(data, list) else [data]
+
+ def _align_to_total_patch_size(self, h, w):
+ total_patch_size = self.latent_patch_size * self.vae_ratio
+ new_h = (h // total_patch_size) * total_patch_size
+ new_w = (w // total_patch_size) * total_patch_size
+ return new_h, new_w
+
+ def _crop_to_align(self, tensor, new_h, new_w, is_video=False):
+ if is_video:
+ _, _, _, h, w = tensor.shape
+ else:
+ _, _, h, w = tensor.shape
+
+ if self.crop_mode == "center":
+ top = (h - new_h) // 2
+ left = (w - new_w) // 2
+ else:
+ top = 0
+ left = 0
+
+ if is_video:
+ return tensor[:, :, :, top : top + new_h, left : left + new_w]
+ else:
+ return tensor[:, :, top : top + new_h, left : left + new_w]
+
+ def _align_target_token(self, T, mode):
+ intra_tail = self.clip_length % self.vae_ratio_t
+ min_frames = intra_tail or self.vae_ratio_t
+ full_chunks = T // self.clip_length
+ remainder = T % self.clip_length
+
+ if remainder == 0:
+ return max(T, min_frames)
+
+ if mode == "pad":
+ aligned_r = (
+ math.ceil((remainder - intra_tail) / self.vae_ratio_t)
+ * self.vae_ratio_t
+ + intra_tail
+ )
+ if aligned_r > self.clip_length:
+ return (full_chunks + 1) * self.clip_length + intra_tail
+ return full_chunks * self.clip_length + aligned_r
+ else: # trim
+ k = (remainder - intra_tail) // self.vae_ratio_t
+ if k >= 0:
+ target = (
+ full_chunks * self.clip_length + k * self.vae_ratio_t + intra_tail
+ )
+ return max(target, min_frames)
+ elif full_chunks > 0:
+ return full_chunks * self.clip_length
+ else:
+ return min_frames
+
+ def _align_target(self, T, mode, granularity):
+ if granularity == "chunk":
+ step = self.clip_length
+ tail = self.frame_overlap
+ if self.isolated_last_frame:
+ tail += 1
+
+ k = math.ceil((T - tail) / step) if mode == "pad" else (T - tail) // step
+ return max(k, 1) * step + tail
+
+ isolated_extra = 1 if self.isolated_last_frame else 0
+ return self._align_target_token(T - isolated_extra, mode) + isolated_extra
+
+ def align_video_length(self, video_length, mode="pad", granularity="chunk"):
+ target = self._align_target(video_length, mode, granularity)
+ delta = target - video_length
+ if delta > 0 and mode == "trim":
+ raise ValueError(
+ f"Cannot trim {video_length} frames to valid length {target}: "
+ f"not enough frames (granularity={granularity})"
+ )
+ return delta
+
+ def align_video_length_2pass(self, video_length):
+ """Return the leading/trailing frame pads and trailing latent drop.
+
+ This is the continuation-prefix (2-pass) alignment. The caller temporarily disables the model's normal token
+ drop and keeps these mirrored processor fields at zero.
+ """
+ if self.isolated_last_frame:
+ raise ValueError(
+ "align_video_length_2pass does not support isolated_last_frame"
+ )
+ if self.token_overlap != 0 or self.frame_overlap != 0:
+ raise ValueError("align_video_length_2pass requires token_drop=0 alignment")
+
+ leading = self.align_video_length(video_length, mode="pad", granularity="token")
+ token_aligned = video_length + leading
+ trailing = self.align_video_length(
+ token_aligned, mode="pad", granularity="chunk"
+ )
+
+ if trailing > 0:
+ intra_tail = self.clip_length % self.vae_ratio_t
+ full_chunks = token_aligned // self.clip_length
+ remainder = token_aligned % self.clip_length
+ real_tokens = full_chunks * self.tokens_chunk_size
+ if remainder > 0:
+ real_tokens += (remainder - intra_tail) // self.vae_ratio_t + 1
+ drop_tokens = self.get_latent_length(token_aligned + trailing) - real_tokens
+ else:
+ drop_tokens = 0
+
+ return leading, trailing, drop_tokens
+
+ def get_suitable_video_length(self, video_length, verbose=False):
+ used_frame_length = video_length + self.align_video_length(
+ video_length, mode="trim", granularity="chunk"
+ )
+ if verbose:
+ logger.info(
+ f"Pick first {used_frame_length} frames from {video_length}-frame video"
+ )
+ return used_frame_length
+
+ def get_latent_length(self, video_length):
+ tail_frame = self.frame_overlap
+ tail_token = self.token_overlap
+ if self.isolated_last_frame:
+ tail_frame += 1
+ tail_token += 1
+
+ video_length = self.get_suitable_video_length(video_length)
+ latent_length = (
+ int((video_length - tail_frame) // self.clip_length)
+ * self.tokens_chunk_size
+ + tail_token
+ )
+ return latent_length
+
+ def transform_tensor(self, tensor, *, runtime_owned=False):
+ B, T = None, None
+ if tensor.ndim == 5:
+ if tensor.shape[2] == 3:
+ tensor = tensor.transpose(1, 2)
+ B, _, T, _, _ = tensor.shape
+ tensor = rearrange(tensor, "b c t h w -> (b t) c h w")
+ elif tensor.ndim == 4:
+ if tensor.shape[0] == 3:
+ tensor = tensor.transpose(0, 1)
+ elif tensor.ndim == 3:
+ tensor = tensor.unsqueeze(0)
+ else:
+ raise ValueError(f"Unsupported tensor shape: {tensor.shape}")
+
+ transform = (
+ self._runtime_owned_transform
+ if runtime_owned and self._runtime_owned_transform is not None
+ else self.transform
+ )
+ tensor = transform(tensor)
+
+ if B is not None and T is not None:
+ tensor = rearrange(tensor, "(b t) c h w -> b c t h w", b=B, t=T)
+
+ return tensor.contiguous()
+
+ def revert_tensor(self, tensor):
+ B, T = None, None
+ if self.use_3d_conv:
+ tensor = tensor.unsqueeze(2) if tensor.ndim == 4 else tensor
+ B, _, T, _, _ = tensor.shape
+ tensor = rearrange(tensor, "b c t h w -> (b t) c h w")
+ tensor_rev = self.transform_rev(tensor).clamp_(0, 1)
+ if B is not None:
+ tensor_rev = rearrange(tensor_rev, "(b t) c h w -> b c t h w", b=B, t=T)
+ return tensor_rev.contiguous()
+
+ @staticmethod
+ def convert_numpy_to_tensor(numpy_array, device=None):
+ if isinstance(numpy_array, list):
+ numpy_array = np.stack(numpy_array, axis=0)
+ tensor = torch.from_numpy(numpy_array)
+ # Keep decoded uint8 pixels compact across the host-to-device copy.
+ # Casting the full video on CPU quadruples both the temporary host
+ # allocation and transfer volume for no loss of information.
+ if device is not None:
+ tensor = tensor.to(device)
+ tensor = tensor.permute(0, 3, 1, 2)
+ return tensor.to(torch.float32).div_(255.0)
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_cnn.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_cnn.py
new file mode 100644
index 000000000..1f72c67f7
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_cnn.py
@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# 3D causal CNN encoder for the MiniMax H3 visual VAE (inference-only bundle).
+import os
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .conv import BaseConv3d
+from .norm import get_group_norm_3d, get_spatial_norm_3d
+
+# ============================================================================
+# 3D CNN Components
+# ============================================================================
+
+
+def norm_silu(x, norm, cond=None):
+ if cond is None:
+ return F.silu(norm(x), inplace=True)
+ else:
+ return F.silu(norm(x, cond), inplace=True)
+
+
+class Downsample3D(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ time_stride=1,
+ space_stride=2,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ ):
+ super().__init__()
+ self.time_stride = time_stride
+ self.space_stride = space_stride
+
+ assert time_stride in [1, 2]
+ assert space_stride in [1, 2, 3]
+
+ self.conv = BaseConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ padding=(1, 0, 0),
+ stride=(time_stride, space_stride, space_stride),
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+ self.causal = self.conv.causal
+ self.pad_mode = self.conv.pad_mode
+
+ def forward(self, x):
+ if self.space_stride == 2:
+ pad = (0, 1, 0, 1, 0, 0)
+ x = F.pad(x, pad, mode=self.pad_mode)
+ return self.conv(x)
+
+
+class ResnetBlock3D(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels=None,
+ zq_ch=None,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ use_t_isolated_gn=False,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+
+ self.use_fused_norm = (
+ os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true"
+ )
+
+ if zq_ch is None:
+ self.norm1 = get_group_norm_3d(
+ in_channels, use_t_isolated_gn=use_t_isolated_gn
+ )
+ self.norm2 = get_group_norm_3d(
+ out_channels, use_t_isolated_gn=use_t_isolated_gn
+ )
+ else:
+ self.norm1 = get_spatial_norm_3d(
+ in_channels,
+ zq_ch,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ use_t_isolated_gn=use_t_isolated_gn,
+ )
+ self.norm2 = get_spatial_norm_3d(
+ out_channels,
+ zq_ch,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ use_t_isolated_gn=use_t_isolated_gn,
+ )
+
+ self.conv1 = BaseConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+
+ self.conv2 = BaseConv3d(
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+
+ if self.in_channels != self.out_channels:
+ self.nin_shortcut = BaseConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+
+ def forward(self, x, zq=None):
+ h = x
+
+ if self.use_fused_norm:
+ h = self.norm1(h, zq)
+ else:
+ h = norm_silu(h, self.norm1, zq)
+
+ h = self.conv1(h)
+
+ if self.use_fused_norm:
+ h = self.norm2(h, zq)
+ else:
+ h = norm_silu(h, self.norm2, zq)
+
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ x = self.nin_shortcut(x)
+
+ return h.add_(x)
+
+
+class EncoderFCN3D(nn.Module):
+ def __init__(
+ self,
+ ch,
+ ch_mult,
+ space_down,
+ time_down,
+ num_res_blocks,
+ in_channels,
+ z_channels,
+ double_z=False,
+ zq_ch=None,
+ padding_mode="zeros",
+ padding_mode_t=None,
+ causal=True,
+ use_t_isolated_gn=False,
+ ):
+ super().__init__()
+ self.ch = ch
+ self.num_levels = len(ch_mult)
+
+ if isinstance(num_res_blocks, int):
+ self.num_res_blocks = [num_res_blocks] * self.num_levels
+ else:
+ self.num_res_blocks = num_res_blocks
+
+ self.space_down_factors = space_down
+ self.time_down_factors = time_down
+ self.in_channels = in_channels
+
+ self.use_fused_norm = (
+ os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true"
+ )
+
+ block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
+ block_in = [block_mid[0]] + block_mid[:-1]
+ block_out = block_mid
+
+ conv_kwargs = dict(
+ padding_mode=padding_mode,
+ padding_mode_t=padding_mode_t,
+ causal=causal,
+ )
+
+ self.conv_in = BaseConv3d(
+ in_channels, block_in[0], kernel_size=3, padding=1, **conv_kwargs
+ )
+
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_levels):
+ down = nn.Module()
+
+ down.block = nn.ModuleList()
+ for i in range(self.num_res_blocks[i_level]):
+ down.block.append(
+ ResnetBlock3D(
+ in_channels=block_in[i_level] if i == 0 else block_mid[i_level],
+ out_channels=block_mid[i_level],
+ zq_ch=zq_ch,
+ use_t_isolated_gn=use_t_isolated_gn,
+ **conv_kwargs,
+ )
+ )
+
+ if space_down[i_level] * time_down[i_level] > 1:
+ down.downsample = Downsample3D(
+ block_mid[i_level],
+ block_out[i_level],
+ time_stride=time_down[i_level],
+ space_stride=space_down[i_level],
+ **conv_kwargs,
+ )
+ else:
+ if block_out[i_level] != block_mid[i_level]:
+ down.downsample = BaseConv3d(
+ block_mid[i_level],
+ block_out[i_level],
+ kernel_size=1,
+ **conv_kwargs,
+ )
+
+ self.down.append(down)
+
+ if zq_ch is None:
+ self.norm_out = get_group_norm_3d(
+ block_out[-1], use_t_isolated_gn=use_t_isolated_gn
+ )
+ else:
+ self.norm_out = get_spatial_norm_3d(
+ block_out[-1],
+ zq_ch,
+ use_t_isolated_gn=use_t_isolated_gn,
+ **conv_kwargs,
+ )
+
+ self.conv_out = BaseConv3d(
+ block_out[-1],
+ 2 * z_channels if double_z else z_channels,
+ kernel_size=3,
+ padding=1,
+ **conv_kwargs,
+ )
+
+ def forward(self, x, zq=None):
+ h = self.conv_in(x)
+ for i_level in range(self.num_levels):
+ for i_block in range(self.num_res_blocks[i_level]):
+ h = self.down[i_level].block[i_block](h, zq)
+ if hasattr(self.down[i_level], "downsample"):
+ h = self.down[i_level].downsample(h)
+
+ if self.use_fused_norm:
+ h = self.norm_out(h, zq)
+ else:
+ h = norm_silu(h, self.norm_out, zq)
+
+ h = self.conv_out(h)
+ return h
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py
new file mode 100644
index 000000000..802f4ec6a
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py
@@ -0,0 +1,374 @@
+# SPDX-License-Identifier: Apache-2.0
+# ViT3D decoder for the MiniMax H3 visual VAE (inference-only bundle).
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import logging
+
+from .base_module import RotaryEmbeddingND, TransformerBlock
+from .flash import make_block_causal_mask_mod
+from .vit_utils import create_token_ids, prepare_rotary_pos_emb
+
+logger = logging.get_logger(__name__)
+
+
+def _linear_with_module_dtype(linear, tensor, out_dtype=None):
+ weight = getattr(linear, "weight", None)
+ target_dtype = getattr(weight, "dtype", tensor.dtype)
+ output = linear(tensor.to(target_dtype))
+ if out_dtype is not None and output.dtype != out_dtype:
+ output = output.to(out_dtype)
+ return output
+
+
+def _pack_tensors_3d(tensors, patch_size, patch_size_t):
+ batch_size, num_channels_tensors, temporal, height, width = tensors.shape
+
+ tensors = tensors.view(
+ batch_size,
+ num_channels_tensors,
+ temporal // patch_size_t,
+ patch_size_t,
+ height // patch_size,
+ patch_size,
+ width // patch_size,
+ patch_size,
+ )
+ tensors = tensors.permute(0, 2, 4, 6, 1, 3, 5, 7)
+ tensors = tensors.reshape(
+ batch_size,
+ (temporal // patch_size_t) * (height // patch_size) * (width // patch_size),
+ num_channels_tensors * patch_size_t * patch_size * patch_size,
+ )
+ return tensors
+
+
+def _unpack_tensors_3d(tensors, patch_size, patch_size_t, temporal, height, width):
+ batch_size, num_patches, channels = tensors.shape
+ num_channels_tensors = channels // (patch_size_t * patch_size * patch_size)
+
+ tensors = tensors.view(
+ batch_size,
+ temporal // patch_size_t,
+ height // patch_size,
+ width // patch_size,
+ num_channels_tensors,
+ patch_size_t,
+ patch_size,
+ patch_size,
+ )
+ tensors = tensors.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
+ tensors = tensors.reshape(batch_size, num_channels_tensors, temporal, height, width)
+ return tensors
+
+
+class ViTBase(ModelMixin, ConfigMixin):
+ """Base class for ViT Encoder and Decoder with common functionality."""
+
+ _no_split_modules = ["TransformerBlock"]
+
+ def _init_weights(self):
+ def basic_init(m):
+ if isinstance(m, nn.Linear):
+ nn.init.xavier_uniform_(m.weight)
+ if m.bias is not None:
+ nn.init.constant_(m.bias, 0)
+
+ self.apply(basic_init)
+
+ def init_mask_config(self, dim, is_3d=False):
+ self._mask_dim = dim
+ self._mask_is_3d = is_3d
+ self.register_buffer("mask_token", torch.zeros(1, 1, dim))
+
+ def set_mask_config(self, mask_config):
+ self.mask_prob = mask_config.get("mask_prob", 0.0)
+ self.mask_enabled = self.mask_prob > 0
+ self.mask_style = mask_config.get("mask_style", "replace")
+ if self.mask_enabled and self.mask_style == "drop" and self.mask_prob < 1.0:
+ logger.warning("mask_style='drop' with mask_prob < 1.0")
+ if self._mask_is_3d:
+ self.temporal_scale_range = mask_config.get(
+ "temporal_scale_range", (0.3, 0.5)
+ )
+ self.spatial_scale_range = mask_config.get(
+ "spatial_scale_range", (0.1, 0.25)
+ )
+ self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.75)
+ self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.95)
+ else:
+ self.spatial_scale_range = mask_config.get(
+ "spatial_scale_range", (0.15, 0.15)
+ )
+ self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.5)
+ self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.75)
+ self.aspect_ratio_range = mask_config.get("aspect_ratio_range", (0.75, 1.5))
+ self.max_retries = mask_config.get("max_retries", 100)
+ if (
+ self.mask_enabled
+ and self.mask_style == "drop"
+ and getattr(self, "t_causal", False)
+ ):
+ logger.warning("mask_style='drop' with t_causal may cause issues")
+ if self.mask_enabled and "mask_token" in self._buffers:
+ del self._buffers["mask_token"]
+ self.mask_token = nn.Parameter(torch.randn(1, 1, self._mask_dim) * 0.02)
+
+ def init_suffix_tokens(self, dim, num_register_tokens, has_cls_token=True):
+ self.num_register_tokens = num_register_tokens
+ if num_register_tokens > 0:
+ self.register_tokens = nn.Parameter(
+ torch.randn(1, num_register_tokens, dim) * 0.02
+ )
+ else:
+ self.register_tokens = None
+ if has_cls_token:
+ self.cls_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02)
+
+ def apply_mask_preprocess(self, hidden_states, img_ids, patch_dims, num_suffix):
+ if self.training and self.mask_enabled:
+ raise NotImplementedError(
+ "mask modeling is not supported in this inference-only bundle"
+ )
+ return hidden_states, img_ids
+
+ def forward_transformer_blocks(self, hidden_states, rotary_pos_emb, pack_info=None):
+ if pack_info is None:
+ pack_info = {}
+ for block in self.transformer_blocks:
+ hidden_states = block(hidden_states, rotary_pos_emb, pack_info)
+ return hidden_states
+
+ def apply_mask_postprocess(self, hidden_states, num_patches):
+ if self.training and self.mask_enabled and self.mask_style == "drop":
+ raise NotImplementedError(
+ "mask modeling is not supported in this inference-only bundle"
+ )
+ return hidden_states
+
+
+class ViT3DDecoder(ViTBase):
+ """Vision Transformer Video Decoder using TransformerBlock."""
+
+ @register_to_config
+ def __init__(
+ self,
+ patch_size: int = 16,
+ patch_size_t: int = 4,
+ t_causal: bool = False,
+ in_channels: int = 16,
+ out_channels: int = 3,
+ num_layers: int = 24,
+ heads: int = 16,
+ dim_head: int = 64,
+ norm_type: str = "layer_norm",
+ norm_affine: bool = True,
+ qk_norm_type: str = None,
+ qk_norm_affine: bool = False,
+ ffn_activation_fn: str = "gelu",
+ ffn_use_gated: bool = False,
+ rope_theta: float = 100.0,
+ rope_dim_ratio: float = 1.0,
+ bias: bool = True,
+ eps: float = 1e-5,
+ num_register_tokens: int = 4,
+ mask_config: dict = {},
+ **kwargs,
+ ):
+ super().__init__()
+
+ dim = heads * dim_head
+ rope_apply_dim = int(dim_head * rope_dim_ratio)
+
+ self.pos_embed = RotaryEmbeddingND(
+ rope_apply_dim, rope_theta, n_dim=3, use_angle=True
+ )
+
+ self.x_embedder = nn.Linear(in_channels, dim)
+
+ self.init_suffix_tokens(dim, num_register_tokens, has_cls_token=False)
+
+ self.t_causal = t_causal
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ TransformerBlock(
+ heads=heads,
+ dim_head=dim_head,
+ norm_type=norm_type,
+ norm_affine=norm_affine,
+ qk_norm_type=qk_norm_type,
+ qk_norm_affine=qk_norm_affine,
+ ffn_activation_fn=ffn_activation_fn,
+ ffn_use_gated=ffn_use_gated,
+ bias=bias,
+ eps=eps,
+ **kwargs,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.norm_out = nn.LayerNorm(dim, elementwise_affine=norm_affine, eps=eps)
+ patch_dim = out_channels * patch_size_t * patch_size * patch_size
+ self.proj_out = nn.Linear(dim, patch_dim)
+
+ self.init_mask_config(dim, is_3d=True)
+ self.set_mask_config(mask_config)
+
+ self._rotary_pos_emb_cache = None
+ self._autocast_linear_dtype = None
+
+ if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0):
+ logger.warning(f"Unused kwargs: {kwargs}")
+
+ def _apply(self, fn, recurse=True):
+ result = super()._apply(fn, recurse=recurse)
+ self._rotary_pos_emb_cache = None
+ self._autocast_linear_dtype = None
+ return result
+
+ def prepare_autocast_linear_weights(self, dtype: torch.dtype) -> int:
+ """Keep decoder-block linear weights in their autocast compute dtype.
+
+ PyTorch autocast does not cache casts for these frozen parameters, so
+ tiled decode otherwise converts every FP32 weight and bias once per
+ block invocation. Persisting the rounded values is numerically
+ equivalent to the per-call autocast conversion. The embedding and
+ output projections stay FP32 because their calls explicitly disable
+ autocast.
+ """
+
+ if dtype not in (torch.float16, torch.bfloat16):
+ raise ValueError(
+ "MiniMax H3 decoder autocast weights require fp16 or bf16, "
+ f"got {dtype}"
+ )
+ if self._autocast_linear_dtype == dtype:
+ return 0
+
+ converted = 0
+ for block in self.transformer_blocks:
+ for linear in (
+ block.attn.to_qkv,
+ block.attn.to_out,
+ block.ff.w1,
+ block.ff.w2,
+ ):
+ if linear.weight.dtype != dtype:
+ linear.to(dtype=dtype)
+ converted += 1
+ self._autocast_linear_dtype = dtype
+ return converted
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ B, C, latent_T, latent_H, latent_W = x.shape
+ patch_size = self.config.patch_size
+ patch_size_t = self.config.patch_size_t
+ num_suffix = 1 + self.num_register_tokens
+
+ hidden_states = _pack_tensors_3d(x, 1, 1)
+ latent_size = (latent_T, latent_H, latent_W)
+
+ with torch.autocast("cuda", enabled=False):
+ hidden_states = _linear_with_module_dtype(
+ self.x_embedder, hidden_states, hidden_states.dtype
+ )
+
+ num_patches = hidden_states.shape[1]
+
+ tokens = [hidden_states]
+
+ if self.register_tokens is not None:
+ register_tokens = self.register_tokens.expand(B, -1, -1)
+ tokens.append(register_tokens)
+
+ cls_token = torch.zeros_like(hidden_states[:, 0:1, :])
+ tokens.append(cls_token)
+ hidden_states = torch.cat(tokens, dim=1)
+
+ patch_dims = [latent_T, latent_H, latent_W]
+ rotary_dtype = (
+ torch.get_autocast_dtype("cuda")
+ if x.is_cuda and torch.is_autocast_enabled("cuda")
+ else hidden_states.dtype
+ )
+ cache_enabled = (
+ not self.training
+ and not self.mask_enabled
+ and not torch.compiler.is_compiling()
+ )
+ cache_key = (
+ B,
+ latent_T,
+ latent_H,
+ latent_W,
+ num_suffix,
+ x.device,
+ x.dtype,
+ rotary_dtype,
+ )
+ cache_record = self._rotary_pos_emb_cache if cache_enabled else None
+ cache_hit = cache_record is not None and cache_record[0] == cache_key
+ if cache_hit:
+ img_ids = cache_record[1]
+ else:
+ img_ids = create_token_ids(latent_size, x.device, x.dtype).expand(B, -1, -1)
+ suffix_ids = torch.zeros(
+ (B, num_suffix, 3), device=x.device, dtype=img_ids.dtype
+ )
+ img_ids = torch.cat([img_ids, suffix_ids], dim=1)
+
+ hidden_states, img_ids = self.apply_mask_preprocess(
+ hidden_states, img_ids, patch_dims, num_suffix
+ )
+ cache_img_ids = img_ids
+
+ pack_info = {}
+ if self.t_causal:
+ spatial_size = latent_H * latent_W
+ mask_mod = make_block_causal_mask_mod(
+ num_tokens=num_patches,
+ block_size=spatial_size,
+ suffix=True,
+ )
+ pack_info["mask_mod"] = mask_mod
+
+ if cache_hit:
+ rotary_pos_emb = cache_record[2]
+ else:
+ rotary_pos_emb = prepare_rotary_pos_emb(
+ self.pos_embed(img_ids),
+ dtype=rotary_dtype,
+ )
+ if cache_enabled:
+ self._rotary_pos_emb_cache = (
+ cache_key,
+ cache_img_ids,
+ rotary_pos_emb,
+ )
+
+ for block in self.transformer_blocks:
+ hidden_states = block(hidden_states, rotary_pos_emb, pack_info)
+
+ hidden_states = self.norm_out(hidden_states)
+
+ hidden_states = self.apply_mask_postprocess(hidden_states, num_patches)
+
+ with torch.autocast("cuda", enabled=False):
+ output = _linear_with_module_dtype(
+ self.proj_out, hidden_states, hidden_states.dtype
+ )
+
+ output = output[:, :num_patches, :]
+
+ video_t = latent_size[0] * patch_size_t
+ video_h = latent_size[1] * patch_size
+ video_w = latent_size[2] * patch_size
+ output = _unpack_tensors_3d(
+ output, patch_size, patch_size_t, video_t, video_h, video_w
+ )
+
+ return output
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vit_utils.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vit_utils.py
new file mode 100644
index 000000000..000c9b6c8
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vit_utils.py
@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+# ViT runtime helpers for the MiniMax H3 visual VAE.
+import os
+from collections.abc import Sequence
+from typing import Tuple
+
+import torch
+from diffusers.utils import logging
+
+
+def _env_flag(name, default="0"):
+ value = os.environ.get(name, default)
+ return str(value).strip().lower() in ("1", "true", "yes", "on")
+
+
+def _env_optional_bool(name, default=""):
+ value = str(os.environ.get(name, default)).strip().lower()
+ if value in ("", "default", "auto", "none", "unset"):
+ return None
+ return value not in ("0", "false", "no", "off", "disabled")
+
+
+def _vit_torch_compile_kwargs(prefix):
+ kwargs = {}
+ backend = os.environ.get(f"{prefix}_BACKEND", "inductor").strip()
+ mode = os.environ.get(f"{prefix}_MODE", "reduce-overhead").strip()
+ if backend and backend.lower() not in ("default", "none"):
+ kwargs["backend"] = backend
+ if mode and mode.lower() not in ("default", "none"):
+ kwargs["mode"] = mode
+ kwargs["fullgraph"] = _env_flag(f"{prefix}_FULLGRAPH", "0")
+ dynamic = _env_optional_bool(f"{prefix}_DYNAMIC")
+ if dynamic is not None:
+ kwargs["dynamic"] = dynamic
+ return kwargs
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def create_token_ids(
+ patch_dims, device, dtype, id_type="length_normalized", flatten=True
+):
+ coords_list = []
+
+ if isinstance(id_type, str):
+ id_type_list = [id_type] * len(patch_dims)
+ elif isinstance(id_type, list):
+ id_type_list = id_type
+ if len(id_type_list) != len(patch_dims):
+ raise ValueError("id_type list must match patch_dims")
+ else:
+ raise ValueError("id_type must be a string or a list")
+
+ if "area_normalized" in id_type_list or id_type == "area_normalized":
+ raise NotImplementedError(
+ "area_normalized id_type is not supported in this inference-only bundle"
+ )
+
+ for _dim_size, _id_type in zip(patch_dims, id_type_list):
+ if isinstance(_dim_size, torch.Tensor):
+ coords_list.append(_dim_size.to(device=device, dtype=dtype))
+ continue
+
+ if _id_type == "length_normalized":
+ coords = torch.arange(0.5, _dim_size, dtype=dtype, device=device)
+ coords = coords / _dim_size
+ coords = 2.0 * coords - 1.0
+ else:
+ coords = torch.arange(_dim_size, dtype=dtype, device=device)
+
+ coords_list.append(coords)
+
+ coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1)
+ if flatten:
+ coords = coords.flatten(0, len(patch_dims) - 1)
+
+ return coords.unsqueeze(0)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = torch.chunk(x, 2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb_impl(
+ t: torch.Tensor, rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor]
+) -> torch.Tensor:
+ cos, sin = rotary_pos_emb[:2]
+
+ if cos.dim() != 4:
+ raise ValueError(f"cos must be [B, N, 1, D], got {cos.shape}")
+
+ cos = cos.to(t.dtype)
+ sin = sin.to(t.dtype)
+
+ rot_dim = cos.shape[-1]
+ t_dim = t.shape[-1]
+
+ if rot_dim < t_dim:
+ t_rot, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+ scaled = t_rot * cos
+ scaled.add_(_rotate_half(t_rot) * sin)
+ t_rot = scaled
+ t = torch.cat((t_rot, t_pass), dim=-1)
+ else:
+ scaled = t * cos
+ scaled.add_(_rotate_half(t) * sin)
+ t = scaled
+
+ return t
+
+
+def prepare_rotary_pos_emb(
+ rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor],
+ *,
+ dtype: torch.dtype,
+) -> tuple[torch.Tensor, ...]:
+ """Prebuild the native Q/K rotary cache once per ViT decoder forward."""
+ cos, sin = rotary_pos_emb
+ if (
+ not cos.is_cuda
+ or dtype not in (torch.float16, torch.bfloat16)
+ or cos.shape != sin.shape
+ or cos.dim() != 4
+ or cos.shape[0] != 1
+ or cos.shape[2] != 1
+ or cos.shape[-1] % 2
+ or _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0")
+ ):
+ return cos, sin
+
+ cos = cos.to(dtype=dtype)
+ sin = sin.to(dtype=dtype)
+ half = cos.shape[-1] // 2
+ # RotaryEmbeddingND repeats each half. The native kernel consumes the
+ # compact NeoX cache [cos_half | sin_half].
+ cache = torch.cat(
+ (cos[0, :, 0, :half], sin[0, :, 0, :half]),
+ dim=-1,
+ ).contiguous()
+ positions = torch.arange(
+ cos.shape[1],
+ dtype=torch.long,
+ device=cos.device,
+ )
+ return cos, sin, cache, positions
+
+
+_COMPILED_APPLY_ROTARY_POS_EMB = None
+_APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = False
+
+
+def _get_apply_rotary_pos_emb_impl():
+ global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED
+ if _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED or not _env_flag(
+ "MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0"
+ ):
+ return _apply_rotary_pos_emb_impl
+ if _COMPILED_APPLY_ROTARY_POS_EMB is not None:
+ return _COMPILED_APPLY_ROTARY_POS_EMB
+ if not hasattr(torch, "compile"):
+ message = (
+ "torch.compile is unavailable; falling back to eager ViT rotary embedding"
+ )
+ if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"):
+ raise RuntimeError(message)
+ logger.warning(f"[ViTRope] {message}")
+ _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True
+ return _apply_rotary_pos_emb_impl
+
+ kwargs = _vit_torch_compile_kwargs("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE")
+ try:
+ _COMPILED_APPLY_ROTARY_POS_EMB = torch.compile(
+ _apply_rotary_pos_emb_impl, **kwargs
+ )
+ logger.info(f"[ViTRope] torch.compile enabled kwargs={kwargs}")
+ except Exception as exc:
+ if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"):
+ raise
+ logger.warning(
+ f"[ViTRope] torch.compile setup failed: {type(exc).__name__}: {exc}; "
+ "falling back to eager"
+ )
+ _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True
+ _COMPILED_APPLY_ROTARY_POS_EMB = None
+ return _apply_rotary_pos_emb_impl
+ return _COMPILED_APPLY_ROTARY_POS_EMB
+
+
+def apply_rotary_pos_emb(
+ t: torch.Tensor, rotary_pos_emb: Sequence[torch.Tensor]
+) -> torch.Tensor:
+ global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED
+ fn = _get_apply_rotary_pos_emb_impl()
+ try:
+ return fn(t, rotary_pos_emb)
+ except Exception as exc:
+ if fn is _COMPILED_APPLY_ROTARY_POS_EMB and not _env_flag(
+ "MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"
+ ):
+ logger.warning(
+ f"[ViTRope] compiled call failed: {type(exc).__name__}: {exc}; "
+ "disabling compile and retrying eager"
+ )
+ _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True
+ _COMPILED_APPLY_ROTARY_POS_EMB = None
+ return _apply_rotary_pos_emb_impl(t, rotary_pos_emb)
+ raise
+
+
+def apply_rotary_pos_emb_qk(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ rotary_pos_emb: Sequence[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Apply the exact native NeoX rotary kernel to Q/K together when possible."""
+ if (
+ len(rotary_pos_emb) == 4
+ and query.is_cuda
+ and query.shape == key.shape
+ and query.dtype == key.dtype
+ and query.dtype in (torch.float16, torch.bfloat16)
+ and query.dim() == 4
+ and query.shape[0] == 1
+ and not torch.compiler.is_compiling()
+ ):
+ _, _, cache, positions = rotary_pos_emb
+ if (
+ cache.is_cuda
+ and cache.dtype == query.dtype
+ and cache.dim() == 2
+ and cache.shape[0] == query.shape[1]
+ and cache.shape[1] <= query.shape[-1]
+ and positions.is_cuda
+ and positions.shape == (query.shape[1],)
+ ):
+ from sgl_kernel import rotary_embedding
+
+ query = query.contiguous()
+ key = key.contiguous()
+ rotary_embedding(
+ positions,
+ query.view(query.shape[1], -1),
+ key.view(key.shape[1], -1),
+ query.shape[-1],
+ cache,
+ True,
+ )
+ return query, key
+
+ return (
+ apply_rotary_pos_emb(query, rotary_pos_emb),
+ apply_rotary_pos_emb(key, rotary_pos_emb),
+ )
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py
new file mode 100644
index 000000000..28cc4b9a2
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py
@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
+ MiniMaxH3PipelineConfig,
+)
+from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+ ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.stages import InputValidationStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3 import (
+ MiniMaxH3AudioEncodingStage,
+ MiniMaxH3DecodingStage,
+ MiniMaxH3DenoisingStage,
+ MiniMaxH3LatentPreparationStage,
+ MiniMaxH3TextEncodingStage,
+ MiniMaxH3TimestepPreparationStage,
+ MiniMaxH3VisualEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
+ MiniMaxH3PartitionAdmissionStage,
+ MiniMaxH3ReleaseMetadata,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class MiniMaxH3Pipeline(LoRAPipeline, ComposedPipelineBase):
+ pipeline_name = "MiniMaxH3Pipeline"
+ default_model_subfolder = "FL2VA"
+ is_video_pipeline = True
+ pipeline_config_cls = MiniMaxH3PipelineConfig
+ sampling_params_cls = MiniMaxH3SamplingParams
+ _required_config_modules = [
+ "processor",
+ "text_encoder",
+ "tokenizer",
+ "video_vae",
+ "audio_vae",
+ # scheduler intentionally absent: model_index carries scheduler=null;
+ # per-modality sigma schedules are generated in TimestepPreparation
+ # from the task profile, and the loop scheduler math lives in
+ # scheduling_minimax_h3_euler_ancestral (stages accept scheduler=None).
+ "transformer",
+ ]
+
+ @staticmethod
+ def model_subfolder_for_variant(variant: str) -> str:
+ if not isinstance(variant, str) or not variant.strip():
+ raise ValueError("MiniMax H3 model variant must be a non-empty string")
+ normalized = variant.strip().lower()
+ subfolders = {
+ "fl2va": "FL2VA",
+ "ref2va": "Ref2VA",
+ }
+ try:
+ return subfolders[normalized]
+ except KeyError as exc:
+ raise ValueError(
+ f"unsupported MiniMax H3 model variant {variant!r}; "
+ f"supported: {sorted(subfolders)!r}"
+ ) from exc
+
+ def _load_config(self):
+ model_variant = self.server_args.model_variant
+ if model_variant is not None:
+ semantic_subfolder = self.model_subfolder_for_variant(model_variant)
+ explicit_subfolder = self.server_args.model_subfolder
+ if (
+ explicit_subfolder is not None
+ and explicit_subfolder.strip().lower() != semantic_subfolder.lower()
+ ):
+ raise ValueError(
+ "MiniMax H3 --model-variant and --model-subfolder select "
+ f"different weight partitions: variant={model_variant!r} maps to "
+ f"{semantic_subfolder!r}, model_subfolder="
+ f"{explicit_subfolder!r}"
+ )
+ self.server_args.model_subfolder = semantic_subfolder
+ model_index = super()._load_config()
+ self.release_metadata = MiniMaxH3ReleaseMetadata.from_model_index(model_index)
+ if (
+ model_variant is not None
+ and self.release_metadata.partition != model_variant.strip().lower()
+ ):
+ raise ValueError(
+ "MiniMax H3 loaded checkpoint partition does not match "
+ f"--model-variant {model_variant!r}"
+ )
+ return model_index
+
+ def validate_disagg_role(self, role: RoleType) -> None:
+ if role != RoleType.MONOLITHIC:
+ raise ValueError(
+ "MiniMaxH3Pipeline only supports monolithic deployment; "
+ f"disaggregation role {role.value!r} is not supported"
+ )
+
+ def create_pipeline_stages(self, server_args: ServerArgs) -> None:
+ # Per-model sigma override from model_index.json; contract tests
+ # construct the pipeline without model_path, hence the guard.
+ release_metadata = getattr(self, "release_metadata", None)
+ sigma_shift_scales = (
+ release_metadata.sigma_shift_scales
+ if release_metadata is not None
+ else None
+ )
+ self.add_stage(InputValidationStage())
+ if release_metadata is not None:
+ self.add_stage(MiniMaxH3PartitionAdmissionStage(release_metadata))
+ self.add_stage(
+ MiniMaxH3TextEncodingStage(
+ text_encoder=self.get_module("text_encoder"),
+ tokenizer=self.get_module("tokenizer"),
+ processor=self.get_module("processor"),
+ )
+ )
+ self.add_stage(
+ MiniMaxH3VisualEncodingStage(
+ video_vae=self.get_module("video_vae"),
+ vae_arch_config=server_args.pipeline_config.vae_config.arch_config,
+ )
+ )
+ self.add_stage(
+ MiniMaxH3AudioEncodingStage(
+ audio_vae=self.get_module("audio_vae"),
+ vae_arch_config=server_args.pipeline_config.audio_vae_config.arch_config,
+ )
+ )
+ self.add_stage(MiniMaxH3LatentPreparationStage())
+ self.add_stage(
+ MiniMaxH3TimestepPreparationStage(
+ sigma_shift_scales=sigma_shift_scales,
+ )
+ )
+ self.add_stage(
+ MiniMaxH3DenoisingStage(
+ transformer=self.get_module("transformer"),
+ pipeline=self,
+ )
+ )
+ self.add_stage(
+ MiniMaxH3DecodingStage(
+ video_vae=self.get_module("video_vae"),
+ audio_vae=self.get_module("audio_vae"),
+ )
+ )
+
+
+EntryClass = MiniMaxH3Pipeline
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
index d25fa3b28..095454229 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
@@ -82,6 +82,7 @@ class ComposedPipelineBase(ABC):
# the name of the pipeline it associated with, in diffusers
pipeline_name: str
+ default_model_subfolder: str | None = None
def is_lora_effective(self):
return False
@@ -172,7 +173,32 @@ class ComposedPipelineBase(ABC):
self.modules[module_name] = module
def _load_config(self) -> dict[str, Any]:
- model_path = maybe_download_model(self.model_path, force_diffusers_model=True)
+ model_subfolder = self.server_args.model_subfolder
+ if model_subfolder is None and not os.path.isfile(
+ os.path.join(self.model_path, "model_index.json")
+ ):
+ model_subfolder = self.default_model_subfolder
+
+ if model_subfolder is None:
+ model_path = maybe_download_model(
+ self.model_path, force_diffusers_model=True
+ )
+ else:
+ model_subfolder = os.path.normpath(model_subfolder)
+ if (
+ os.path.isabs(model_subfolder)
+ or model_subfolder == ".."
+ or model_subfolder.startswith(f"..{os.sep}")
+ ):
+ raise ValueError(
+ f"model_subfolder must stay inside the model repository: {model_subfolder!r}"
+ )
+ model_root = maybe_download_model(
+ self.model_path,
+ allow_patterns=[f"{model_subfolder}/**"],
+ )
+ model_path = os.path.join(model_root, model_subfolder)
+
self.model_path = model_path
logger.info("Model path: %s", model_path)
config = verify_model_config_and_directory(model_path)
@@ -444,13 +470,26 @@ class ComposedPipelineBase(ABC):
component_load_specs: list[ComponentLoadSpec] = []
# enqueue only real weight loads (e.g., scheduler, tokenizer is excluded); skipped/provided modules keep old handling
- for index, (
- module_name,
- (
- transformers_or_diffusers,
- architecture,
- ),
- ) in enumerate(model_index.items()):
+ for index, (module_name, component_spec) in enumerate(model_index.items()):
+ # Diffusers uses JSON null for unavailable optional components.
+ # Check before unpacking the normal [library, architecture] pair.
+ if component_spec is None:
+ logger.warning(
+ "Module %s in model_index.json has null value, removing from required_config_modules",
+ module_name,
+ )
+ if module_name in self.required_config_modules:
+ self.required_config_modules.remove(module_name)
+ continue
+ if (
+ not isinstance(component_spec, (list, tuple))
+ or len(component_spec) != 2
+ ):
+ raise ValueError(
+ f"Module {module_name!r} in model_index.json must be null or "
+ f"a [library, architecture] pair, got {component_spec!r}"
+ )
+ transformers_or_diffusers, architecture = component_spec
if transformers_or_diffusers is None:
logger.warning(
"Module %s in model_index.json has null value, removing from required_config_modules",
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
index d739ebe0f..793381272 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
@@ -161,13 +161,23 @@ class DecodingStage(PipelineStage):
def scale_and_shift(self, latents: torch.Tensor, server_args):
return scale_and_shift_latents(latents, server_args, self.vae)
- def _get_vae_decode_fn(self, vae, server_args: ServerArgs):
+ def _get_vae_decode_fn(
+ self,
+ vae,
+ server_args: ServerArgs,
+ *,
+ decode_fn=None,
+ compiled_callable: ActiveTargetCompiledCallable | None = None,
+ ):
+ decode_fn = decode_fn or vae.decode
if not server_args.enable_torch_compile or not isinstance(vae, nn.Module):
- return vae.decode
+ return decode_fn
+
+ compiled_callable = compiled_callable or self._compiled_vae_decode
will_compile = (
- self._compiled_vae_decode.target_id != id(vae)
- or self._compiled_vae_decode.compiled_module is None
+ compiled_callable.target_id != id(vae)
+ or compiled_callable.compiled_module is None
)
if current_platform.is_npu():
compile_kwargs = build_torch_compile_kwargs(mode=None)
@@ -183,8 +193,8 @@ class DecodingStage(PipelineStage):
if will_compile:
logger.info("Compiling VAE decode with mode: %s", mode)
- return self._compiled_vae_decode.get_or_compile(
- vae, vae.decode, compile_kwargs=compile_kwargs
+ return compiled_callable.get_or_compile(
+ vae, decode_fn, compile_kwargs=compile_kwargs
)
@torch.no_grad()
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py
index b3acb2319..47dbd89d5 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py
@@ -30,6 +30,7 @@ class StageDedupMixin:
deduplicated_output_fields: ClassVar[tuple[str, ...]] = ()
deduplicated_tensor_tree_output_fields: ClassVar[tuple[str, ...]] = ()
deduplicated_deepcopy_output_fields: ClassVar[tuple[str, ...]] = ()
+ deduplicated_extra_output_keys: ClassVar[tuple[str, ...]] = ()
deduplicated_extra_tensor_tree_output_keys: ClassVar[tuple[str, ...]] = ()
def run_grouped_requests(
@@ -60,6 +61,7 @@ class StageDedupMixin:
cls.deduplicated_output_fields
or cls.deduplicated_tensor_tree_output_fields
or cls.deduplicated_deepcopy_output_fields
+ or cls.deduplicated_extra_output_keys
or cls.deduplicated_extra_tensor_tree_output_keys
)
@@ -109,8 +111,8 @@ class StageDedupMixin:
tensor references, which is the low-overhead path for read-only outputs
such as embeddings. Tensor-tree fields recursively clone tensors.
Deepcopy fields are for mutable request-local runtime objects, such as
- scheduler instances. Extra keys clone selected ``Req.extra`` entries
- without replacing the destination extra dict.
+ scheduler instances. Extra output keys follow the same shallow-copy
+ contract, while extra tensor-tree keys recursively clone tensors.
"""
for field in self.deduplicated_output_fields:
setattr(dst, field, self.copy_stage_output(getattr(src, field)))
@@ -118,6 +120,9 @@ class StageDedupMixin:
setattr(dst, field, self.clone_tensor_tree(getattr(src, field)))
for field in self.deduplicated_deepcopy_output_fields:
setattr(dst, field, deepcopy(getattr(src, field)))
+ for key in self.deduplicated_extra_output_keys:
+ if key in src.extra:
+ dst.extra[key] = self.copy_stage_output(src.extra[key])
for key in self.deduplicated_extra_tensor_tree_output_keys:
if key in src.extra:
dst.extra[key] = self.clone_tensor_tree(src.extra[key])
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
index 63b33bda4..fbc8a722e 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
@@ -76,6 +76,9 @@ from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_c
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
ComponentUse,
)
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
+ is_fsdp_managed_module,
+)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
LayerwiseOffloadableModuleMixin,
is_layerwise_offloaded_module,
@@ -227,7 +230,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
num_attention_heads = (
self.server_args.pipeline_config.dit_config.num_attention_heads
)
- attn_head_size = hidden_size // num_attention_heads
+ attn_head_size = getattr(
+ self.server_args.pipeline_config.dit_config,
+ "attention_head_dim",
+ hidden_size // num_attention_heads,
+ )
# torch compile
# list of offloaded dit modules if torch compile is enabled. cleared after compile and warmup
@@ -341,10 +348,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
not args.enable_torch_compile
or not args.offload_during_compile
or not args.warmup
- # a subclass with its own forward would never run the restore
- or type(self).forward is not DenoisingStage.forward
+ or not self._owns_compile_warmup_lifecycle()
or args.use_fsdp_inference
- or envs.SGLANG_CACHE_DIT_ENABLED
+ or self._cache_dit_requested()
or not isinstance(module, LayerwiseOffloadableModuleMixin)
or is_layerwise_offloaded_module(module)
):
@@ -353,6 +359,15 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
if is_layerwise_offloaded_module(module):
self._offloaded_dit_modules_for_compile.append(module)
+ def _owns_compile_warmup_lifecycle(self) -> bool:
+ """Whether ``forward`` enters ``_offload_for_torch_compile_warmup``.
+
+ Custom denoising loops opt in explicitly after wiring the same restore
+ lifecycle. This keeps the safety guard without silently disabling the
+ optimization solely because a model overrides ``forward``.
+ """
+ return type(self).forward is DenoisingStage.forward
+
def _move_resident_components_for_warmup(self) -> list[torch.nn.Module]:
"""Move resident non-DiT components off-device while the warmup
denoising (the compile/autotune peak) runs; forward() moves them back."""
@@ -365,6 +380,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
if (
isinstance(module, torch.nn.Module)
and id(module) not in dit_ids
+ and not is_fsdp_managed_module(module)
and not is_layerwise_offloaded_module(module)
):
param = next(module.parameters(), None)
@@ -386,7 +402,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
self.server_args, "enable_torch_compile", False
) or not isinstance(module, nn.Module):
return
- if envs.SGLANG_CACHE_DIT_ENABLED and not self._cache_dit_enabled:
+ if self._cache_dit_requested() and not self._cache_dit_enabled:
logger.debug("Deferring torch.compile until cache-dit is enabled")
return
if self._torch_compile_registry.is_compiled(module):
@@ -434,6 +450,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
def _cache_dit_dual_model_name(self) -> str:
return "wan2.2"
+ def _cache_dit_requested(self) -> bool:
+ return envs.SGLANG_CACHE_DIT_ENABLED
+
def _cache_dit_secondary_uses_primary_config(self) -> bool:
return False
@@ -604,7 +623,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
# Keep cache-dit disabled for ordinary warmup, but allow torch.compile
# warmup to mount cache-dit before Dynamo traces the transformer.
- if not envs.SGLANG_CACHE_DIT_ENABLED:
+ if not self._cache_dit_requested():
return
if batch.is_warmup and not getattr(
self.server_args, "enable_torch_compile", False
@@ -692,9 +711,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
logger.info(
"cache-dit enabled on transformer (steps=%d, Fn=%d, Bn=%d, rdt=%.3f)",
primary_num_steps,
- envs.SGLANG_CACHE_DIT_FN,
- envs.SGLANG_CACHE_DIT_BN,
- envs.SGLANG_CACHE_DIT_RDT,
+ primary_config.Fn_compute_blocks,
+ primary_config.Bn_compute_blocks,
+ primary_config.residual_diff_threshold,
)
self._cache_dit_enabled = True
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/__init__.py
new file mode 100644
index 000000000..a85586468
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/__init__.py
@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""MiniMax H3-specific pipeline stages."""
+
+from .stages.audio_encoding import MiniMaxH3AudioEncodingStage
+from .stages.decoding import MiniMaxH3DecodingStage
+from .stages.denoising import MiniMaxH3DenoisingStage
+from .stages.latent_preparation import MiniMaxH3LatentPreparationStage
+from .stages.text_encoding import MiniMaxH3TextEncodingStage
+from .stages.timestep_preparation import MiniMaxH3TimestepPreparationStage
+from .stages.visual_encoding import MiniMaxH3VisualEncodingStage
+
+__all__ = [
+ "MiniMaxH3AudioEncodingStage",
+ "MiniMaxH3DecodingStage",
+ "MiniMaxH3DenoisingStage",
+ "MiniMaxH3LatentPreparationStage",
+ "MiniMaxH3TextEncodingStage",
+ "MiniMaxH3TimestepPreparationStage",
+ "MiniMaxH3VisualEncodingStage",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/canvas.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/canvas.py
new file mode 100644
index 000000000..df4453d14
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/canvas.py
@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 keyframe target-canvas preparation.
+
+Geometry behavior:
+- auto-aspect canvases delegate to the shared adaptive v2 shape resolver;
+- cover-crop: aspect-preserving max-scale LANCZOS resize + center crop,
+ upscaling refused unless explicitly allowed.
+
+Both the Qwen presentation (pixel_values) and the visual-condition tokenizer consume
+the SAME prepared canvas image, so preparation
+is cached per request in batch.extra.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+)
+
+MINIMAX_H3_CANVAS_MULTIPLE = 32
+MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY = "minimax_h3_prepared_keyframes"
+
+
+def minimax_h3_cover_crop_plan(
+ *,
+ source_width: int,
+ source_height: int,
+ target_width: int,
+ target_height: int,
+ allow_upscale: bool,
+) -> dict[str, Any]:
+ """Deterministic aspect-preserving cover-crop transform."""
+ if source_width <= 0 or source_height <= 0:
+ raise ValueError("cover_crop requires positive source width/height")
+ scale = max(
+ target_width / float(source_width), target_height / float(source_height)
+ )
+ if scale > 1.0 and not allow_upscale:
+ raise ValueError(
+ "target_canvas cover_crop would upscale the source; set "
+ f"allow_upscale=true (source={source_width}x{source_height}, "
+ f"target={target_width}x{target_height})"
+ )
+ resized_width = max(target_width, int(round(source_width * scale)))
+ resized_height = max(target_height, int(round(source_height * scale)))
+ left = max(0, (resized_width - target_width) // 2)
+ top = max(0, (resized_height - target_height) // 2)
+ return {
+ "scale": scale,
+ "resized_size": (resized_width, resized_height),
+ "crop_box": (left, top, left + target_width, top + target_height),
+ }
+
+
+def minimax_h3_prepare_keyframe_canvas(
+ image: Any,
+ *,
+ target_width: int,
+ target_height: int,
+ allow_upscale: bool = False,
+) -> Any:
+ """Prepare a PIL image onto the target canvas.
+
+ Identity (no resample) when the image already IS the canvas.
+ """
+ from PIL import Image
+
+ image = image.convert("RGB")
+ if image.size == (target_width, target_height):
+ return image
+ plan = minimax_h3_cover_crop_plan(
+ source_width=image.size[0],
+ source_height=image.size[1],
+ target_width=target_width,
+ target_height=target_height,
+ allow_upscale=allow_upscale,
+ )
+ resized = image.resize(plan["resized_size"], Image.Resampling.LANCZOS)
+ return resized.crop(plan["crop_box"])
+
+
+def minimax_h3_stretch_keyframe_canvas(
+ image: Any,
+ *,
+ target_width: int,
+ target_height: int,
+) -> Any:
+ """Stretch the FL first frame directly onto the resolved target canvas."""
+
+ from PIL import Image
+
+ image = image.convert("RGB")
+ if image.size == (target_width, target_height):
+ return image
+ return image.resize((target_width, target_height), Image.Resampling.LANCZOS)
+
+
+def _keyframe_materials(plan: Any) -> list[Any]:
+ return [m for m in plan.materials if m.material_chain == "image.target_canvas"]
+
+
+def _keyframe_canvas_size(shape: Any) -> tuple[int, int]:
+ geometry = str(shape["geometry"])
+ if geometry != "resolved_v2":
+ raise ValueError(
+ "fl2va keyframe preparation requires pre-queue resolved_v2 "
+ f"geometry, got {geometry!r}"
+ )
+ return int(shape["width"]), int(shape["height"])
+
+
+def _validate_keyframe_materials(plan: Any, keyframes: list[Any]) -> tuple[int, ...]:
+ if str(plan.task) != "fl2va":
+ raise ValueError("keyframe target-canvas materials require plan.task='fl2va'")
+ semantic_indices = tuple(material.frame_index for material in keyframes)
+ if semantic_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+ raise ValueError(
+ "fl2va keyframes must use one of the ordered frame_index signatures "
+ f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {semantic_indices!r}"
+ )
+ frame_count = plan.shape.get("frame_count")
+ if isinstance(frame_count, bool) or not isinstance(frame_count, int):
+ raise ValueError("fl2va keyframe preparation requires an integer frame_count")
+ if frame_count <= 1:
+ raise ValueError("fl2va keyframe preparation requires frame_count > 1")
+ expected_pixels = tuple(
+ frame_count - 1 if index == -1 else index for index in semantic_indices
+ )
+ resolved_pixels = tuple(material.resolved_frame_index for material in keyframes)
+ if resolved_pixels != expected_pixels:
+ raise ValueError(
+ "fl2va keyframe resolved_frame_index values disagree with semantic "
+ f"anchors: expected {expected_pixels!r}, got {resolved_pixels!r}"
+ )
+ return semantic_indices
+
+
+def minimax_h3_prepared_keyframes(batch: Any, plan: Any) -> dict[str, Any]:
+ """Resolve + prepare one or two fl2va keyframes once per request.
+
+ The target canvas is shared across keyframes and must already be frozen by
+ the pre-queue probe/resolve hook.
+ Top-level ``image`` / ``canvas_width`` / ``canvas_height`` keys mirror the
+ first-keyframe payload for compatibility; per-keyframe entries live under
+ ``images``.
+ """
+ keyframes = _keyframe_materials(plan)
+ semantic_indices = _validate_keyframe_materials(plan, keyframes)
+ cached = batch.extra.get(MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY)
+ if cached is not None:
+ cached_indices = tuple(cached.get("semantic_frame_indices") or ())
+ cached_images = cached.get("images") or ()
+ if cached_indices != semantic_indices or len(cached_images) != len(keyframes):
+ raise ValueError(
+ "cached fl2va keyframe preparation disagrees with the resolved plan"
+ )
+ return cached
+
+ canvas_w, canvas_h = _keyframe_canvas_size(plan.shape)
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.prequeue import (
+ MINIMAX_H3_PROBE_FACTS_EXTRA_KEY,
+ MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY,
+ )
+
+ probe_facts = batch.extra.get(MINIMAX_H3_PROBE_FACTS_EXTRA_KEY)
+ material_shapes = batch.extra.get(MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY)
+ for material in keyframes:
+ condition_index = int(material.condition_index)
+ facts = (
+ probe_facts.get(condition_index) if isinstance(probe_facts, dict) else None
+ )
+ material_shape = (
+ material_shapes.get(condition_index)
+ if isinstance(material_shapes, dict)
+ else None
+ )
+ if not isinstance(facts, dict) or not isinstance(material_shape, dict):
+ raise ValueError(
+ "fl2va keyframe preparation requires cached pre-queue probe and "
+ f"shape facts for conditions[{condition_index}]"
+ )
+ if (
+ int(material_shape.get("width") or 0),
+ int(material_shape.get("height") or 0),
+ ) != (canvas_w, canvas_h):
+ raise ValueError(
+ "fl2va keyframe material shape disagrees with the resolved target: "
+ f"condition={condition_index}, material="
+ f"{material_shape.get('width')}x{material_shape.get('height')}, "
+ f"target={canvas_w}x{canvas_h}"
+ )
+
+ from PIL import Image, ImageOps
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_localize_material_uri,
+ )
+
+ entries: list[dict[str, Any]] = []
+ for keyframe_index, material in enumerate(keyframes):
+ image_path = minimax_h3_localize_material_uri(
+ batch,
+ material.uri,
+ condition_type=material.condition_type,
+ condition_index=int(material.condition_index),
+ )
+ with Image.open(image_path) as source_image:
+ image = ImageOps.exif_transpose(source_image)
+ # A request's first semantic keyframe is its geometry anchor, including
+ # the single-image last-frame-only signature [-1]. Only the second image in the
+ # two-keyframe FL extension is a follower and receives cover-crop.
+ prepared_image = (
+ minimax_h3_stretch_keyframe_canvas(
+ image,
+ target_width=canvas_w,
+ target_height=canvas_h,
+ )
+ if keyframe_index == 0
+ else minimax_h3_prepare_keyframe_canvas(
+ image,
+ target_width=canvas_w,
+ target_height=canvas_h,
+ allow_upscale=True,
+ )
+ )
+ entries.append(
+ {
+ "image": prepared_image,
+ "canvas_width": canvas_w,
+ "canvas_height": canvas_h,
+ "condition_index": int(material.condition_index),
+ "frame_index": (
+ None if material.frame_index is None else int(material.frame_index)
+ ),
+ "resolved_frame_index": (
+ None
+ if material.resolved_frame_index is None
+ else int(material.resolved_frame_index)
+ ),
+ }
+ )
+
+ payload = {
+ "image": entries[0]["image"],
+ "canvas_width": canvas_w,
+ "canvas_height": canvas_h,
+ "images": entries,
+ "semantic_frame_indices": [
+ int(item["frame_index"])
+ for item in entries
+ if item.get("frame_index") is not None
+ ],
+ "pixel_frame_indices": [
+ int(item["resolved_frame_index"])
+ for item in entries
+ if item.get("resolved_frame_index") is not None
+ ],
+ "frame_count": (
+ int(plan.shape["frame_count"])
+ if plan.shape.get("frame_count") is not None
+ else None
+ ),
+ }
+ batch.extra[MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY] = payload
+ return payload
+
+
+__all__ = [
+ "MINIMAX_H3_CANVAS_MULTIPLE",
+ "MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY",
+ "minimax_h3_cover_crop_plan",
+ "minimax_h3_prepare_keyframe_canvas",
+ "minimax_h3_prepared_keyframes",
+ "minimax_h3_stretch_keyframe_canvas",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py
new file mode 100644
index 000000000..3a5e038a1
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/condition_noise.py
@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 visual/audio condition-noise augmentation.
+
+The request's condition timestep is applied to both the tensor value and the
+DiT timestep. Tokenizer artifacts remain clean
+and reusable; this module materializes the fixed noised anchors immediately
+before the denoise loop.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
+ minimax_h3_patchify_video_latent,
+)
+
+# Channel-major packed audio rows always carry a stereo layout.
+MINIMAX_H3_AUDIO_COND_CHANNELS = 2
+
+
+def minimax_h3_imgvid_cond_noise_aug_rows(
+ clean_rows: torch.Tensor,
+ *,
+ condition_shapes: Sequence[Sequence[int]],
+ target_latent_t: int,
+ imgvid_cond_num_frames: int,
+ seed: int,
+ noise_aug: float,
+) -> torch.Tensor:
+ """Apply the imgvid-condition RF noise recipe to packed clean rows.
+
+ ``condition_shapes`` contains ``(latent_t, latent_h, latent_w)`` in packed
+ visual-condition order. A new CPU generator with the same row seed is
+ created for every condition. Under the dependent-noise policy, each draw
+ uses the target temporal length plus the template's imgvid-condition frame
+ count, then slices the prefix matching the current condition.
+ """
+
+ noise_aug = float(noise_aug)
+ if not 0.0 <= noise_aug <= 1.0:
+ raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}")
+ if noise_aug == 1.0:
+ return clean_rows
+ if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 96:
+ raise ValueError(
+ "clean imgvid condition rows must have shape [n, 96], got "
+ f"{list(clean_rows.shape)}"
+ )
+
+ target_latent_t = int(target_latent_t)
+ imgvid_cond_num_frames = int(imgvid_cond_num_frames)
+ if target_latent_t <= 0:
+ raise ValueError(f"target_latent_t must be positive, got {target_latent_t}")
+ if imgvid_cond_num_frames <= 0:
+ raise ValueError(
+ "imgvid_cond_num_frames must be positive when condition rows exist, "
+ f"got {imgvid_cond_num_frames}"
+ )
+
+ parsed_shapes: list[tuple[int, int, int]] = []
+ expected_rows = 0
+ for raw_shape in condition_shapes:
+ if len(raw_shape) != 3:
+ raise ValueError(
+ "each imgvid condition shape must be (latent_t, latent_h, latent_w), "
+ f"got {list(raw_shape)}"
+ )
+ latent_t, latent_h, latent_w = (int(value) for value in raw_shape)
+ if latent_t <= 0 or latent_h <= 0 or latent_w <= 0:
+ raise ValueError(
+ f"imgvid condition shape must be positive, got {list(raw_shape)}"
+ )
+ if latent_h % 2 or latent_w % 2:
+ raise ValueError(
+ "imgvid condition spatial dimensions must be divisible by 2, "
+ f"got {(latent_t, latent_h, latent_w)}"
+ )
+ parsed_shapes.append((latent_t, latent_h, latent_w))
+ expected_rows += latent_t * (latent_h // 2) * (latent_w // 2)
+ if not parsed_shapes:
+ raise ValueError("condition_shapes must not be empty")
+ if int(clean_rows.shape[0]) != expected_rows:
+ raise ValueError(
+ f"clean imgvid condition rows {int(clean_rows.shape[0])} != "
+ f"shape-derived rows {expected_rows}"
+ )
+
+ out: list[torch.Tensor] = []
+ row_offset = 0
+ timestep = torch.tensor(noise_aug, dtype=torch.float32, device=clean_rows.device)
+ for latent_t, latent_h, latent_w in parsed_shapes:
+ full_t = target_latent_t + imgvid_cond_num_frames
+ if full_t < latent_t:
+ raise ValueError(
+ f"condition latent_t {latent_t} exceeds the noise draw "
+ f"length {full_t}"
+ )
+ generator = torch.Generator(device="cpu").manual_seed(int(seed))
+ noise = torch.randn(
+ 1,
+ 24,
+ full_t,
+ latent_h,
+ latent_w,
+ generator=generator,
+ dtype=torch.float32,
+ device="cpu",
+ )[:, :, :latent_t]
+ noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]).to(
+ device=clean_rows.device, dtype=torch.float32
+ )
+ row_count = int(noise_rows.shape[0])
+ clean_part = clean_rows[row_offset : row_offset + row_count].to(torch.float32)
+ out.append(timestep * clean_part + (1.0 - timestep) * noise_rows)
+ row_offset += row_count
+ return (out[0] if len(out) == 1 else torch.cat(out, dim=0)).contiguous()
+
+
+def minimax_h3_audio_cond_noise_aug_rows(
+ clean_rows: torch.Tensor,
+ *,
+ condition_audio_t: Sequence[int],
+ seed: int,
+ noise_aug: float,
+) -> torch.Tensor:
+ """Apply the audio-condition RF noise recipe to packed clean rows.
+
+ ``condition_audio_t`` contains the latent T of each audio-bearing
+ condition in canonical request order. Noise is drawn per condition
+ element, with a fresh CPU generator seeded with ``seed + 1`` for every
+ element. Consequently each condition restarts the
+ same RNG stream; concatenating the rows and drawing once would be
+ numerically different for ordered multi-reference requests.
+
+ The mix is intentionally evaluated on CPU in fp32 before the packed rows
+ are transferred to the DiT device.
+ """
+
+ noise_aug = float(noise_aug)
+ if not 0.0 <= noise_aug <= 1.0:
+ raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}")
+ if noise_aug == 1.0:
+ return clean_rows
+ if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 32:
+ raise ValueError(
+ "clean audio condition rows must have shape [n, 32], got "
+ f"{list(clean_rows.shape)}"
+ )
+
+ audio_channels = MINIMAX_H3_AUDIO_COND_CHANNELS
+ parsed_audio_t = [int(value) for value in condition_audio_t]
+ if not parsed_audio_t:
+ raise ValueError("condition_audio_t must not be empty")
+ if any(value <= 0 for value in parsed_audio_t):
+ raise ValueError(
+ f"condition audio latent lengths must be positive, got {parsed_audio_t}"
+ )
+ expected_rows = audio_channels * sum(parsed_audio_t)
+ if int(clean_rows.shape[0]) != expected_rows:
+ raise ValueError(
+ f"clean audio condition rows {int(clean_rows.shape[0])} != "
+ f"shape-derived rows {expected_rows}"
+ )
+
+ out: list[torch.Tensor] = []
+ row_offset = 0
+ timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu")
+ for audio_t in parsed_audio_t:
+ row_count = audio_channels * audio_t
+ clean_part = (
+ clean_rows[row_offset : row_offset + row_count]
+ .detach()
+ .to(device="cpu", dtype=torch.float32)
+ )
+ generator = torch.Generator(device="cpu").manual_seed(int(seed) + 1)
+ noise = torch.randn(
+ clean_part.shape,
+ generator=generator,
+ dtype=torch.float32,
+ device="cpu",
+ )
+ out.append(timestep * clean_part + (1.0 - timestep) * noise)
+ row_offset += row_count
+ rows = out[0] if len(out) == 1 else torch.cat(out, dim=0)
+ return rows.to(device=clean_rows.device, dtype=torch.float32).contiguous()
+
+
+__all__ = [
+ "minimax_h3_audio_cond_noise_aug_rows",
+ "minimax_h3_imgvid_cond_noise_aug_rows",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/constants.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/constants.py
new file mode 100644
index 000000000..1f59adf35
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/constants.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+# Direct-encode text embeddings: {"positive":
+# {"hidden_states": Tensor[text_len, 5120] bf16 cpu, "text_len": int}}
+MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY = "minimax_h3_text_embeddings"
+# Direct keyframe encode: {"rows": Tensor[n_rows, 96] fp32 cpu,
+# "latent_h": int, "latent_w": int, "canvas_height": int,
+# "canvas_width": int, "keyframes": [...],
+# "semantic_frame_indices": [...], "pixel_frame_indices": [...]}
+MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY = "minimax_h3_keyframe_cond_rows"
+# Direct sigma schedules: {"video": [float], "audio": [float]}
+MINIMAX_H3_SIGMAS_EXTRA_KEY = "minimax_h3_sigmas"
+# Direct denoise state: {"initial_video_rows", "initial_audio_rows",
+# "latent_t", "latent_h", "latent_w", "audio_t"}
+MINIMAX_H3_DENOISE_STATE_EXTRA_KEY = "minimax_h3_denoise_state"
+# ref2va direct reference encodes.
+MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY = "minimax_h3_reference_image_rows"
+MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY = "minimax_h3_reference_audio_rows"
+MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY = "minimax_h3_reference_video_rows"
+MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY = "minimax_h3_prepared_reference_video"
+
+MINIMAX_H3_SUPPORTED_FPS = 24
+MINIMAX_H3_MIN_DURATION_SECONDS = 4.0
+MINIMAX_H3_MAX_DURATION_SECONDS = 15.0
+
+# The distilled checkpoint has exactly one positive denoise branch.
+MINIMAX_H3_DEFAULT_BRANCHES: tuple = ({"name": "cond_1"},)
+
+# Audited 4xH200 T2VA profiles. The tuple is
+# (warmup steps, residual-difference threshold, max consecutive cached steps).
+MINIMAX_H3_QUALITY_PROFILES: dict[str, tuple[int, float, int] | None] = {
+ "lossless": None,
+ "high": (4, 0.04, 1),
+ "medium": (4, 0.12, 3),
+ "low": (4, 0.24, 3),
+}
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py
new file mode 100644
index 000000000..54715b2b8
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py
@@ -0,0 +1,519 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 cfg-distilled full denoise loop.
+
+Per step, the positive presentation is forwarded exactly once. Video and audio
+target rows chain through the Euler-eta0 update while visual and audio condition
+rows stay pinned to their noised step-0 anchors.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, nullcontext
+from typing import Any, Callable
+
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+ MINIMAX_H3_ADALN_MODALITY_NUM,
+)
+
+MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
+# ref2va audio reference anchor timestep
+MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0
+# Packed row widths: video rows are [1,2,2]-patchified 24-channel latents
+# (24 * 1 * 2 * 2 = 96); audio rows carry the 32-dim audio latent.
+MINIMAX_H3_VIDEO_ROW_WIDTH = 96
+MINIMAX_H3_AUDIO_ROW_WIDTH = 32
+
+
+@torch.inference_mode()
+def _minimax_h3_update_target_rows_(
+ state: torch.Tensor,
+ velocity: torch.Tensor,
+ *,
+ sigma_t: torch.Tensor,
+ sigma_curr: float,
+ sigma_ratio: torch.Tensor,
+ one_minus_sigma_ratio: torch.Tensor,
+ denoised_scratch: torch.Tensor,
+) -> None:
+ torch.mul(sigma_t, velocity, out=denoised_scratch)
+ torch.add(state, denoised_scratch, out=denoised_scratch)
+ if sigma_curr == 0.0:
+ return
+ torch.mul(one_minus_sigma_ratio, denoised_scratch, out=velocity)
+ torch.mul(sigma_ratio, state, out=state)
+ torch.add(state, velocity, out=state)
+
+
+def _ulysses_ctx() -> tuple[int, int]:
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_ulysses_parallel_rank,
+ get_ulysses_parallel_world_size,
+ model_parallel_is_initialized,
+ )
+
+ if not model_parallel_is_initialized():
+ return 1, 0
+ return get_ulysses_parallel_world_size(), get_ulysses_parallel_rank()
+
+
+def _build_local_embedding_layout(
+ *,
+ seq_len: int,
+ text_pos: torch.Tensor,
+ img_pos: torch.Tensor,
+ audio_pos: torch.Tensor,
+ world_size: int,
+ rank: int,
+ device: torch.device,
+) -> dict[str, torch.Tensor | int]:
+ if seq_len % world_size:
+ raise ValueError(
+ f"packed seq_len {seq_len} not divisible by Ulysses world size "
+ f"{world_size}"
+ )
+ local_seq_len = seq_len // world_size
+ row_start = rank * local_seq_len
+ row_stop = row_start + local_seq_len
+
+ def local_ids(pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ source_ids = torch.nonzero(
+ (pos >= row_start) & (pos < row_stop),
+ as_tuple=False,
+ ).view(-1)
+ return source_ids.to(device), pos.index_select(0, source_ids).to(device)
+
+ text_source_start = min(row_start, int(text_pos.shape[0]))
+ text_source_stop = min(row_stop, int(text_pos.shape[0]))
+ _, img_global_ids = local_ids(img_pos)
+ _, audio_global_ids = local_ids(audio_pos)
+ return {
+ "text_source_start": text_source_start,
+ "text_source_stop": text_source_stop,
+ "img_global_ids": img_global_ids,
+ "img_row_ids": img_global_ids - row_start,
+ "audio_global_ids": audio_global_ids,
+ "audio_row_ids": audio_global_ids - row_start,
+ }
+
+
+class MiniMaxH3DenoiseBranch:
+ """Static per-branch state: packed layout + fixed forward kwargs.
+
+ `packed` is a minimax_h3_packed_sequence(...) result (or equivalent layout
+ dict); `text_embeddings` is the branch's [text_len, 5120] hidden states;
+ `token_tags` must already carry any fl2va vision-span overrides.
+ """
+
+ def __init__(
+ self,
+ *,
+ packed: dict[str, torch.Tensor],
+ text_embeddings: torch.Tensor,
+ token_tags: torch.Tensor,
+ device: torch.device,
+ ) -> None:
+ seq_len = int(packed["seq_len"])
+ self.seq_len = seq_len
+ self.img_pos = packed["img_pos"].view(-1).to(torch.long)
+ self.audio_pos = packed["audio_pos"].view(-1).to(torch.long)
+ self.update_mask = packed["update_mask"].view(-1).to(torch.bool)
+ # ref2va: audio_pos may include reference-audio anchor rows
+ # (audio_update_mask False); absent means all rows are targets.
+ if "audio_update_mask" in packed:
+ self.audio_update_mask = packed["audio_update_mask"].view(-1).to(torch.bool)
+ else:
+ self.audio_update_mask = torch.ones(
+ self.audio_pos.shape[0], dtype=torch.bool
+ )
+ # Packed H3 layouts place reference rows before the generated suffix
+ # for both modalities. Keep the suffix boundary once so the denoise
+ # hot path can update contiguous views instead of gathering and
+ # scattering the same static index tensors on every step.
+ self.video_target_start = int((~self.update_mask).sum())
+ self.audio_target_start = int((~self.audio_update_mask).sum())
+ self.video_target_slice = slice(self.video_target_start, None)
+ self.audio_target_slice = slice(self.audio_target_start, None)
+ text_pos = packed["text_pos"].view(-1).to(torch.long)
+ text_len = int(text_pos.shape[0])
+ if list(text_embeddings.shape)[0] != text_len:
+ raise ValueError(
+ f"text_embeddings rows {list(text_embeddings.shape)} != "
+ f"packed text_len {text_len}"
+ )
+ if int(token_tags.view(-1).shape[0]) != seq_len:
+ raise ValueError(
+ f"token_tags length {int(token_tags.view(-1).shape[0])} != "
+ f"seq_len {seq_len}"
+ )
+ cu = packed["cu_seqlens"].to(torch.int32)
+ self.img_pos_dev = self.img_pos.to(device)
+ self.audio_pos_dev = self.audio_pos.to(device)
+ self.update_mask_dev = self.update_mask.to(device)
+ self.audio_update_mask_dev = self.audio_update_mask.to(device)
+ # Resolve the remaining step-static packed-sequence and anchor row
+ # sets once, keeping nonzero-driven work out of the hot loop.
+ self.img_cond_seq_idx = self.img_pos_dev[~self.update_mask_dev]
+ self.img_target_seq_idx = self.img_pos_dev[self.update_mask_dev]
+ self.audio_target_seq_idx = self.audio_pos_dev[self.audio_update_mask_dev]
+ self.audio_ref_seq_idx = self.audio_pos_dev[~self.audio_update_mask_dev]
+ self.cond_row_idx = torch.nonzero(~self.update_mask_dev).view(-1)
+ self.audio_ref_row_idx = torch.nonzero(~self.audio_update_mask_dev).view(-1)
+ # rows that keep the video timestep each step: text, padding, and
+ # video target rows — everything the three overwrite sets do not cover
+ self.n_video_timestep_rows = (
+ seq_len
+ - int(self.img_cond_seq_idx.numel())
+ - int(self.audio_target_seq_idx.numel())
+ - int(self.audio_ref_seq_idx.numel())
+ )
+ # persistent packed-row buffers; every img/audio position is fully
+ # rewritten by index_copy_ on the first forward_kwargs() call, then
+ # only the target-row subset each step after (condition/reference
+ # rows never change post-priming -- see forward_kwargs).
+ self._x_buffer_primed = False
+ self.x_buffer = torch.zeros(
+ 1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device
+ )
+ self.audio_x_buffer = torch.zeros(
+ 1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device
+ )
+ text_pos_dev = text_pos.to(device)
+ ulysses_world_size, ulysses_rank = _ulysses_ctx()
+ token_tags_host = token_tags.view(-1).to(dtype=torch.long)
+ local_seq_len = seq_len // ulysses_world_size
+ local_row_start = ulysses_rank * local_seq_len
+ local_row_stop = local_row_start + local_seq_len
+ self.local_row_slice = slice(local_row_start, local_row_stop)
+ self.block_token_tags = (
+ token_tags_host[local_row_start:local_row_stop].clamp(min=0).to(device)
+ )
+ self.static_kwargs: dict[str, Any] = {
+ # Cast the fp64 position grid on the host. MiniMaxH3Rope casts it to
+ # fp32 as its first op anyway, so the values are identical; doing the
+ # cast on CPU also avoids platforms that cannot execute fp64 on
+ # device (e.g. Iluvatar CoreX returns zeros for device fp64).
+ "img_position_ids": packed["img_position_ids"][None]
+ .to(torch.float32)
+ .to(device),
+ "update_mask": self.update_mask_dev,
+ "block_token_tags": self.block_token_tags,
+ "skip_mask_out_condition": True,
+ "prompt_embeds": text_embeddings.to(device),
+ "img_pos_info": {"position_ids": self.img_pos_dev},
+ "audio_pos_info": {"position_ids": self.audio_pos_dev},
+ "text_pos_info": {"position_ids": text_pos_dev},
+ "img_pos_for_infer_output_info": {"position_ids": self.img_target_seq_idx},
+ "local_embedding_layout": _build_local_embedding_layout(
+ seq_len=seq_len,
+ text_pos=text_pos,
+ img_pos=self.img_pos,
+ audio_pos=self.audio_pos,
+ world_size=ulysses_world_size,
+ rank=ulysses_rank,
+ device=device,
+ ),
+ "packed_seq_params": {
+ "cu_seqlens_q": cu.to(device),
+ "cu_seqlens_q_host": tuple(int(value) for value in cu.tolist()),
+ "max_seqlen_q": int(cu[1]),
+ },
+ "refiner_packed_seq_params": {
+ "cu_seqlens_q": torch.tensor(
+ [0, text_len, text_len], dtype=torch.int32, device=device
+ ),
+ "cu_seqlens_q_host": (0, text_len, text_len),
+ "max_seqlen_q": text_len,
+ },
+ }
+
+ def forward_kwargs(
+ self,
+ *,
+ video_rows: torch.Tensor,
+ audio_rows: torch.Tensor,
+ step_timesteps: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+ ) -> dict[str, Any]:
+ x = self.x_buffer
+ audio_x = self.audio_x_buffer
+ if not self._x_buffer_primed:
+ # First step: condition/reference rows have just been pinned into
+ # video_rows/audio_rows (see minimax_h3_denoise_loop) and never
+ # change again, so this is the only step that needs the full
+ # img/audio extent written into the persistent buffers.
+ x[0].index_copy_(0, self.img_pos_dev, video_rows)
+ audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows)
+ self._x_buffer_primed = True
+ else:
+ # Later steps: only the target-row subset changed since the
+ # buffers were primed; rewriting condition/reference rows again
+ # would just copy the same bytes already sitting there.
+ x[0].index_copy_(
+ 0, self.img_target_seq_idx, video_rows[self.video_target_slice]
+ )
+ audio_x[0].index_copy_(
+ 0, self.audio_target_seq_idx, audio_rows[self.audio_target_slice]
+ )
+ unique_timesteps, inverse_indices, block_combined_indices = step_timesteps
+ return {
+ **self.static_kwargs,
+ "x": x,
+ "audio_x": audio_x,
+ "unique_timesteps": unique_timesteps,
+ "inverse_indices": inverse_indices,
+ "block_combined_indices": block_combined_indices,
+ }
+
+ def _expand_step_timesteps(
+ self,
+ *,
+ t_video: float,
+ t_audio: float,
+ imgvid_cond_timestep: float,
+ audio_ref_cond_timestep: float,
+ inverse_indices_by_pattern: dict[tuple[int, ...], torch.Tensor],
+ block_combined_by_pattern: dict[tuple[int, ...], torch.Tensor],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Build step-local timestep and AdaLN index tensors.
+
+ Packed-sequence timestep semantics: non-media rows (text and padding)
+ inherit the current video timestep, condition rows pin their noise-aug
+ timesteps. The row->timestep layout is step-static, so instead of
+ materializing the full timestep tensor and paying a device-syncing
+ torch.unique per step, the at-most-four candidate values are deduped
+ in fp32 on the host — torch.unique on the candidate tensor keeps exact
+ fp32 collision semantics — and inverse indices are index_fill'ed from
+ the static position sets.
+ """
+ candidates: list[float] = []
+ fill_groups: list[tuple[torch.Tensor, int]] = []
+ base_slot = -1
+ if self.n_video_timestep_rows > 0:
+ base_slot = len(candidates)
+ candidates.append(float(t_video))
+ for seq_idx, value in (
+ (self.img_cond_seq_idx, imgvid_cond_timestep),
+ (self.audio_target_seq_idx, t_audio),
+ (self.audio_ref_seq_idx, audio_ref_cond_timestep),
+ ):
+ if seq_idx.numel() > 0:
+ fill_groups.append((seq_idx, len(candidates)))
+ candidates.append(float(value))
+ unique_cpu, slot_to_unique = torch.unique(
+ torch.tensor(candidates, dtype=torch.float32),
+ sorted=True,
+ return_inverse=True,
+ )
+ device = self.img_pos_dev.device
+ base_index = int(slot_to_unique[base_slot]) if base_slot >= 0 else 0
+ pattern = tuple(slot_to_unique.tolist())
+ inverse_indices = inverse_indices_by_pattern.get(pattern)
+ if inverse_indices is None:
+ inverse_indices = torch.full(
+ (self.seq_len,), base_index, dtype=torch.long, device=device
+ )
+ for seq_idx, slot in fill_groups:
+ inverse_indices.index_fill_(0, seq_idx, int(slot_to_unique[slot]))
+ inverse_indices_by_pattern[pattern] = inverse_indices
+ block_combined = block_combined_by_pattern.get(pattern)
+ if block_combined is None:
+ block_combined = torch.add(
+ self.block_token_tags,
+ inverse_indices[self.local_row_slice],
+ alpha=MINIMAX_H3_ADALN_MODALITY_NUM,
+ )
+ block_combined_by_pattern[pattern] = block_combined
+ return unique_cpu.to(device), inverse_indices, block_combined
+
+ def prepare_timestep_plan(
+ self,
+ *,
+ video_timesteps: list[float],
+ audio_timesteps: list[float],
+ imgvid_cond_noise_aug: float,
+ audio_ref_cond_noise_aug: float,
+ ) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+ """Stage every step's packed timestep state before denoising."""
+ if len(video_timesteps) != len(audio_timesteps):
+ raise ValueError("video/audio timestep plans must have equal length")
+ inverse_indices_by_pattern: dict[tuple[int, ...], torch.Tensor] = {}
+ block_combined_by_pattern: dict[tuple[int, ...], torch.Tensor] = {}
+ return [
+ self._expand_step_timesteps(
+ t_video=t_video,
+ t_audio=t_audio,
+ imgvid_cond_timestep=max(t_video, imgvid_cond_noise_aug),
+ audio_ref_cond_timestep=max(t_audio, audio_ref_cond_noise_aug),
+ inverse_indices_by_pattern=inverse_indices_by_pattern,
+ block_combined_by_pattern=block_combined_by_pattern,
+ )
+ for t_video, t_audio in zip(video_timesteps, audio_timesteps)
+ ]
+
+
+def minimax_h3_denoise_loop(
+ *,
+ model: Any,
+ model_forward: (
+ Callable[[Any, dict[str, Any], int], tuple[torch.Tensor, torch.Tensor]] | None
+ ) = None,
+ positive: MiniMaxH3DenoiseBranch,
+ initial_video_rows: torch.Tensor,
+ initial_audio_rows: torch.Tensor,
+ keyframe_cond_rows: torch.Tensor | None,
+ audio_ref_rows: torch.Tensor | None = None,
+ sigmas_video: list[float],
+ sigmas_audio: list[float],
+ device: torch.device,
+ imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP,
+ audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
+ on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
+ step_profiler: Callable[[int], AbstractContextManager] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Run the full denoise loop; returns final (video_rows, audio_rows).
+
+ ``initial_video_rows`` covers all image rows of the positive layout. For a
+ conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to
+ pin those rows across every step. The model's raw positive velocity is the
+ update signal; MiniMax H3 only supports cfg-distilled checkpoints.
+ ``model_forward`` is the native-stage hook for residency/BCG runners and
+ receives the zero-based loop step; the default keeps this helper
+ independently testable with a plain callable.
+ """
+ if len(sigmas_video) != len(sigmas_audio):
+ raise ValueError("video/audio sigma schedules must have equal length")
+ if len(sigmas_video) < 2:
+ raise ValueError("sigma schedules need at least 2 entries")
+ n_cond = positive.video_target_start
+ if keyframe_cond_rows is None:
+ if n_cond != 0:
+ raise ValueError(
+ f"layout has {n_cond} cond rows but keyframe_cond_rows is None"
+ )
+ else:
+ if int(keyframe_cond_rows.shape[0]) != n_cond:
+ raise ValueError(
+ f"keyframe_cond_rows {int(keyframe_cond_rows.shape[0])} != "
+ f"layout cond rows {n_cond}"
+ )
+ video_rows = initial_video_rows.to(device=device, dtype=torch.float32, copy=True)
+ audio_rows = initial_audio_rows.to(device=device, dtype=torch.float32, copy=True)
+ if int(video_rows.shape[0]) != int(positive.img_pos.shape[0]):
+ raise ValueError(
+ f"initial video rows {int(video_rows.shape[0])} != positive layout "
+ f"rows {int(positive.img_pos.shape[0])}"
+ )
+ if int(audio_rows.shape[0]) != int(positive.audio_pos.shape[0]):
+ raise ValueError(
+ f"initial audio rows {int(audio_rows.shape[0])} != positive layout "
+ f"rows {int(positive.audio_pos.shape[0])}"
+ )
+ n_audio_ref = positive.audio_target_start
+ if audio_ref_rows is None:
+ if n_audio_ref != 0:
+ raise ValueError(
+ f"layout has {n_audio_ref} audio ref rows but audio_ref_rows is None"
+ )
+ audio_anchor = None
+ else:
+ if int(audio_ref_rows.shape[0]) != n_audio_ref:
+ raise ValueError(
+ f"audio_ref_rows {int(audio_ref_rows.shape[0])} != layout "
+ f"audio ref rows {n_audio_ref}"
+ )
+ audio_anchor = audio_ref_rows.to(device=device, dtype=torch.float32)
+ cond_anchor = (
+ keyframe_cond_rows.to(device=device, dtype=torch.float32)
+ if keyframe_cond_rows is not None
+ else None
+ )
+ if cond_anchor is not None:
+ video_rows.index_copy_(0, positive.cond_row_idx, cond_anchor)
+ if audio_anchor is not None:
+ audio_rows.index_copy_(0, positive.audio_ref_row_idx, audio_anchor)
+
+ num_steps = len(sigmas_video) - 1
+ video_target_slice = positive.video_target_slice
+ audio_target_slice = positive.audio_target_slice
+ video_timesteps = [1.0 - sigma for sigma in sigmas_video[:-1]]
+ audio_timesteps = [1.0 - sigma for sigma in sigmas_audio[:-1]]
+ # One H2D copy per schedule, preserving the previous Python-float
+ # subtraction followed by fp32 conversion.
+ video_step_t = torch.tensor(video_timesteps, dtype=torch.float32, device=device)
+ audio_step_t = torch.tensor(audio_timesteps, dtype=torch.float32, device=device)
+ timestep_plan = positive.prepare_timestep_plan(
+ video_timesteps=video_timesteps,
+ audio_timesteps=audio_timesteps,
+ imgvid_cond_noise_aug=float(imgvid_cond_noise_aug_for_inference),
+ audio_ref_cond_noise_aug=float(audio_cond_noise_aug_for_inference),
+ )
+ # match the scheduler's device-fp32 math once, then reuse one denoised
+ # scratch per modality instead of allocating intermediates every step
+ video_sigmas = torch.tensor(sigmas_video, dtype=torch.float32, device=device)
+ audio_sigmas = torch.tensor(sigmas_audio, dtype=torch.float32, device=device)
+ video_sigma_ratios = video_sigmas[1:] / video_sigmas[:-1]
+ audio_sigma_ratios = audio_sigmas[1:] / audio_sigmas[:-1]
+ video_sigma_t = 1.0 - video_step_t
+ audio_sigma_t = 1.0 - audio_step_t
+ video_one_minus_sigma_ratios = 1.0 - video_sigma_ratios
+ audio_one_minus_sigma_ratios = 1.0 - audio_sigma_ratios
+ video_denoised_scratch = torch.empty_like(video_rows[video_target_slice])
+ audio_denoised_scratch = torch.empty_like(audio_rows[audio_target_slice])
+ for step in range(num_steps):
+ step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
+ with step_cm:
+ s_v = sigmas_video[step]
+ s_a = sigmas_audio[step]
+
+ fk = positive.forward_kwargs(
+ video_rows=video_rows,
+ audio_rows=audio_rows,
+ step_timesteps=timestep_plan[step],
+ )
+ with torch.inference_mode():
+ if model_forward is None:
+ v_video, v_audio = model(**fk)
+ else:
+ v_video, v_audio = model_forward(model, fk, step)
+ # The model outputs are inference tensors. Keep their disposable
+ # fp32 velocity updates in the same context so ``out=velocity``
+ # can reuse the output storage without an extra clone.
+ mv_video_t = v_video.float()
+ mv_audio_t = v_audio[audio_target_slice].float()
+
+ video_target = video_rows[video_target_slice]
+ _minimax_h3_update_target_rows_(
+ video_target,
+ mv_video_t,
+ sigma_t=video_sigma_t[step],
+ sigma_curr=s_v,
+ sigma_ratio=video_sigma_ratios[step],
+ one_minus_sigma_ratio=video_one_minus_sigma_ratios[step],
+ denoised_scratch=video_denoised_scratch,
+ )
+
+ audio_target = audio_rows[audio_target_slice]
+ _minimax_h3_update_target_rows_(
+ audio_target,
+ mv_audio_t,
+ sigma_t=audio_sigma_t[step],
+ sigma_curr=s_a,
+ sigma_ratio=audio_sigma_ratios[step],
+ one_minus_sigma_ratio=audio_one_minus_sigma_ratios[step],
+ denoised_scratch=audio_denoised_scratch,
+ )
+ if on_step is not None:
+ on_step(step, video_rows, audio_rows)
+
+ return video_rows, audio_rows
+
+
+__all__ = [
+ "MINIMAX_H3_AUDIO_REF_COND_TIMESTEP",
+ "MINIMAX_H3_AUDIO_ROW_WIDTH",
+ "MINIMAX_H3_IMGVID_COND_TIMESTEP",
+ "MINIMAX_H3_VIDEO_ROW_WIDTH",
+ "MiniMaxH3DenoiseBranch",
+ "minimax_h3_denoise_loop",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/keyframe_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/keyframe_encoding.py
new file mode 100644
index 000000000..e59564715
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/keyframe_encoding.py
@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 keyframe (imgvid) condition encoding.
+
+Condition anchor row recipe:
+
+- ``video_vae.encode_images(PIL, use_fp16_latent=True)`` under a
+ scoped seed-42 RNG fork — the DiagonalGaussian is SAMPLED
+ (use_mean=False) with seed 42, so the seed is part of
+ the contract, not a convenience
+- normalize ``(z - latents_mean) / latents_std`` with the loader-injected
+ ``MiniMaxH3VideoVAEArchConfig`` values
+- patchify [1, 2, 2] into packed cond rows, fp32
+"""
+
+from __future__ import annotations
+
+import contextlib
+import functools
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
+ MiniMaxH3VideoVAEArchConfig,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
+ minimax_h3_patchify_video_latent,
+)
+
+MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
+MINIMAX_H3_KEYFRAME_PATCH_SIZE = (1, 2, 2)
+
+
+@contextlib.contextmanager
+def minimax_h3_scoped_encode_rng(seed: int, device: torch.device | None = None):
+ """Seed torch RNGs for a deterministic sampled VAE encode without leaking state.
+
+ The encode recipes seed the default torch generators right before a
+ posterior-sampled VAE encode. Forking restores the process-global CPU and
+ CUDA generators after the encode while preserving the exact sampled result.
+ """
+ devices: list[torch.device] = []
+ if device is not None and device.type == "cuda" and torch.cuda.is_available():
+ devices = [device]
+ with torch.random.fork_rng(devices=devices):
+ torch.default_generator.manual_seed(int(seed))
+ for forked_device in devices:
+ with torch.cuda.device(forked_device):
+ torch.cuda.manual_seed(int(seed))
+ yield
+
+
+@contextlib.contextmanager
+def minimax_h3_scoped_encode_fp32(video_vae: Any):
+ """Scope the video VAE to fp32 for one or more keyframe/reference encodes.
+
+ encode_keyframe_cond_rows and encode_reference_video_rows each also guard
+ their own cast (skipping it when already fp32), so nesting this around a
+ caller that does more than one encode -- FL2VA's two keyframes, ref2va's
+ image reference plus video reference -- turns their per-call casts into
+ no-ops instead of toggling the whole VAE's dtype once per encode.
+ """
+ parameter = next(video_vae.parameters())
+ prev_dtype = parameter.dtype
+ if prev_dtype != torch.float32:
+ video_vae.to(torch.float32)
+ try:
+ yield
+ finally:
+ if prev_dtype != torch.float32:
+ video_vae.to(prev_dtype)
+
+
+@functools.lru_cache(maxsize=None)
+def _cached_latent_mean_std(
+ mean_values: tuple[float, ...],
+ std_values: tuple[float, ...],
+ view_shape: tuple[int, ...],
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """CPU mean/std tensors for a fixed (values, shape) triple, built once.
+
+ arch_config.latents_mean/std are static config values fixed for the
+ life of a loaded VAE, so every call with the same values reconstructs
+ an identical tensor; cache it instead of rebuilding on every encode.
+ """
+ mean = torch.tensor(mean_values).view(view_shape)
+ std = torch.tensor(std_values).view(view_shape)
+ return mean, std
+
+
+@torch.inference_mode()
+def minimax_h3_encode_keyframe_cond_rows(
+ video_vae: Any,
+ image: Any,
+ arch_config: MiniMaxH3VideoVAEArchConfig,
+) -> torch.Tensor:
+ """Encode a target-canvas PIL image into packed imgvid cond rows.
+
+ Returns [n_rows, 24 * patch_h * patch_w] fp32 on CPU.
+ """
+ seed = MINIMAX_H3_KEYFRAME_ENCODE_SEED
+ # The encode recipe runs on fp32 weights. Normal H3 residency already keeps
+ # the shared video VAE in fp32; retain the scoped cast for standalone use.
+ parameter = next(video_vae.parameters())
+ prev_dtype = parameter.dtype
+ if prev_dtype != torch.float32:
+ video_vae.to(torch.float32)
+ try:
+ with minimax_h3_scoped_encode_rng(seed, parameter.device):
+ z = video_vae.encode_images(image, use_fp16_latent=True)[0]
+ finally:
+ if prev_dtype != torch.float32:
+ video_vae.to(prev_dtype)
+ z = z.cpu().float()
+ if z.dim() == 4:
+ z = z[None]
+ latent_channels = arch_config.latent_channels
+ if z.dim() != 5 or int(z.shape[1]) != latent_channels:
+ raise ValueError(f"unexpected imgvid latent shape {list(z.shape)}")
+ mean, std = _cached_latent_mean_std(
+ tuple(arch_config.latents_mean),
+ tuple(arch_config.latents_std),
+ (1, latent_channels, 1, 1, 1),
+ )
+ z.sub_(mean).div_(std)
+ rows = minimax_h3_patchify_video_latent(
+ z, patch_size=list(MINIMAX_H3_KEYFRAME_PATCH_SIZE)
+ )
+ return rows.to(torch.float32)
+
+
+__all__ = [
+ "MINIMAX_H3_KEYFRAME_ENCODE_SEED",
+ "MINIMAX_H3_KEYFRAME_PATCH_SIZE",
+ "_cached_latent_mean_std",
+ "minimax_h3_encode_keyframe_cond_rows",
+ "minimax_h3_scoped_encode_rng",
+ "minimax_h3_scoped_encode_fp32",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/material_io.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/material_io.py
new file mode 100644
index 000000000..a6ae54f12
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/material_io.py
@@ -0,0 +1,913 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Request-owned material URI localization for the MiniMax H3 pipeline.
+
+The canonical MiniMax H3 contract intentionally carries semantic URIs rather
+than worker-local paths. Direct media consumers (Pillow, ffmpeg and
+torchaudio) cannot consume every URI scheme in that contract, so localization
+belongs at the model-specific material boundary. Materialized sources and
+derived work directories are registered on ``Req.extra`` and explicitly
+released by the encoder stages.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import math
+import shutil
+import subprocess
+import tempfile
+import urllib.parse
+import urllib.request
+from pathlib import Path
+from typing import Any, Iterable
+
+MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY = "minimax_h3_material_localization"
+MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY = "minimax_h3_material_probe_facts"
+MINIMAX_H3_TEMP_DIRS_EXTRA_KEY = "minimax_h3_request_temp_dirs"
+MINIMAX_H3_HTTP_READ_CHUNK_BYTES = 1024 * 1024
+MINIMAX_H3_BASE64_DECODE_CHUNK_CHARS = 1024 * 1024
+MINIMAX_H3_BASE64_HEADER_MAX_CHARS = 4 * 1024
+MINIMAX_H3_TAR_HEADER_MAX_ENCODED_CHARS = 64 * 1024
+
+_BASE64_ALPHABET = frozenset(
+ b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=_-"
+)
+
+_DEFAULT_SUFFIX_BY_TYPE = {
+ "image": ".png",
+ "video": ".mp4",
+ "video_audio": ".mp4",
+ "audio": ".wav",
+}
+_SUFFIX_BY_MEDIA_TYPE = {
+ "image/jpeg": ".jpg",
+ "image/png": ".png",
+ "image/webp": ".webp",
+ "video/mp4": ".mp4",
+ "video/quicktime": ".mov",
+ "audio/mpeg": ".mp3",
+ "audio/mp4": ".m4a",
+ "audio/wav": ".wav",
+ "audio/x-wav": ".wav",
+ "audio/flac": ".flac",
+}
+
+
+def minimax_h3_register_temp_dir(batch: Any, path: str, *, owner: str) -> str:
+ """Register one request-owned directory and return *path* unchanged."""
+
+ registry = batch.extra.setdefault(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY, {})
+ paths = registry.setdefault(str(owner), [])
+ normalized = str(path)
+ if normalized not in paths:
+ paths.append(normalized)
+ return normalized
+
+
+def minimax_h3_cleanup_temp_dirs(
+ batch: Any, *, owners: Iterable[str] | None = None
+) -> None:
+ """Remove registered request directories, tolerating repeated cleanup."""
+
+ registry = batch.extra.get(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY)
+ selected = (
+ list(registry)
+ if owners is None and isinstance(registry, dict)
+ else [str(owner) for owner in (owners or ())]
+ )
+ if isinstance(registry, dict):
+ for owner in selected:
+ paths = registry.pop(owner, [])
+ if isinstance(paths, (list, tuple)):
+ for path in paths:
+ shutil.rmtree(str(path), ignore_errors=True)
+ if not registry:
+ batch.extra.pop(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY, None)
+ if owners is None or "material" in selected:
+ batch.extra.pop(MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY, None)
+ batch.extra.pop(MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY, None)
+
+
+def _base64_uri_payload_start(uri: str) -> tuple[int, str | None]:
+ media_type = None
+ if uri.startswith("data:"):
+ separator = uri.find(",")
+ if separator < 0:
+ raise ValueError("data URI must contain a comma separator")
+ if separator > MINIMAX_H3_BASE64_HEADER_MAX_CHARS:
+ raise ValueError("data URI header is too large")
+ header = uri[:separator]
+ if ";base64" not in header:
+ raise ValueError("data URI must use ;base64 encoding")
+ media_type = header[5:].split(";", 1)[0].lower() or None
+ payload_start = separator + 1
+ elif uri.startswith("base64://"):
+ payload_start = len("base64://")
+ separator = uri.find(",", payload_start)
+ if separator >= 0:
+ if separator - payload_start > MINIMAX_H3_BASE64_HEADER_MAX_CHARS:
+ raise ValueError("base64 URI header is too large")
+ header = uri[payload_start:separator]
+ media_type = header.split(";", 1)[0].lower() or None
+ payload_start = separator + 1
+ else: # pragma: no cover - guarded by the caller
+ raise ValueError("not a base64 material URI")
+ return payload_start, media_type
+
+
+def _iter_base64_payload_bytes(uri: str, payload_start: int):
+ """Yield validated, unquoted base64 bytes without copying the payload."""
+
+ index = payload_start
+ while index < len(uri):
+ character = uri[index]
+ if character == "%":
+ if index + 2 >= len(uri):
+ raise ValueError("material URI has an invalid percent escape")
+ try:
+ value = int(uri[index + 1 : index + 3], 16)
+ except ValueError as exc:
+ raise ValueError("material URI has an invalid percent escape") from exc
+ index += 3
+ character = chr(value)
+ else:
+ index += 1
+
+ if character.isspace():
+ continue
+ if len(character) != 1 or ord(character) > 127:
+ raise ValueError("material URI base64 payload must be ASCII")
+ value = ord(character)
+ if value not in _BASE64_ALPHABET:
+ raise ValueError(
+ f"material URI has an invalid base64 character {character!r}"
+ )
+ yield value
+
+
+def _parse_tar_member_uri(uri: str) -> tuple[Path, int, int, str | None]:
+ if uri.startswith("tar+offset://"):
+ prefix = "tar+offset://"
+ elif uri.startswith("tar+b64header://"):
+ prefix = "tar+b64header://"
+ else:
+ raise ValueError("unsupported tar material URI")
+ try:
+ tar_path, encoded_header = uri[len(prefix) :].rsplit(":", 1)
+ except ValueError as exc:
+ raise ValueError(
+ "tar material URI must contain ':'"
+ ) from exc
+ if len(encoded_header) > MINIMAX_H3_TAR_HEADER_MAX_ENCODED_CHARS:
+ raise ValueError("tar material URI encoded header is too large")
+ padded = encoded_header + "=" * (-len(encoded_header) % 4)
+ try:
+ header = json.loads(
+ base64.b64decode(
+ padded.encode("ascii"), altchars=b"-_", validate=True
+ ).decode("utf-8")
+ )
+ except Exception as exc:
+ raise ValueError("tar material URI has an invalid encoded header") from exc
+ if not isinstance(header, dict):
+ raise ValueError("tar material URI header must be a JSON object")
+ if header.get("schema") != "sglang.tar_member_ref/v1":
+ raise ValueError(
+ f"unsupported tar material header schema: {header.get('schema')!r}"
+ )
+ try:
+ offset = int(header["offset_data"])
+ size = int(header["size"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError(
+ "tar material header requires integer offset_data and size"
+ ) from exc
+ if offset < 0 or size < 0:
+ raise ValueError("tar material offset_data and size must be non-negative")
+ return (
+ Path(tar_path).expanduser(),
+ offset,
+ size,
+ str(header.get("member") or "") or None,
+ )
+
+
+def _safe_suffix(value: str | None) -> str | None:
+ if not value:
+ return None
+ suffix = Path(urllib.parse.urlsplit(value).path).suffix.lower()
+ if suffix and len(suffix) <= 10 and suffix[1:].isalnum():
+ return suffix
+ return None
+
+
+def _checked_material_file(path: Path, *, label: str) -> str:
+ """Validate that a localized source exists and is non-empty."""
+
+ if not path.is_file():
+ raise FileNotFoundError(f"{label} does not exist or is not a file: {path}")
+ if path.stat().st_size <= 0:
+ raise ValueError(f"{label} is empty: {path}")
+ return str(path)
+
+
+def _parse_frame_rate(value: Any) -> float:
+ if value in {None, "", "N/A", "0/0"}:
+ return 0.0
+ raw = str(value)
+ try:
+ if "/" in raw:
+ numerator, denominator = raw.split("/", 1)
+ denominator_value = float(denominator)
+ parsed = float(numerator) / denominator_value if denominator_value else 0.0
+ else:
+ parsed = float(raw)
+ return parsed if math.isfinite(parsed) else 0.0
+ except (TypeError, ValueError, ZeroDivisionError):
+ return 0.0
+
+
+def _parse_display_ratio(value: Any) -> float:
+ if value in {None, "", "N/A", "0:1", "0/1"}:
+ return 0.0
+ raw = str(value).strip()
+ separator = ":" if ":" in raw else "/" if "/" in raw else None
+ try:
+ if separator is None:
+ ratio = float(raw)
+ else:
+ numerator, denominator = raw.split(separator, 1)
+ ratio = float(numerator) / float(denominator)
+ except (TypeError, ValueError, ZeroDivisionError):
+ return 0.0
+ return ratio if math.isfinite(ratio) and ratio > 0 else 0.0
+
+
+def _stream_rotation_degrees(stream: dict[str, Any]) -> float:
+ values: list[Any] = []
+ side_data = stream.get("side_data_list")
+ if isinstance(side_data, list):
+ values.extend(
+ item.get("rotation") for item in side_data if isinstance(item, dict)
+ )
+ tags = stream.get("tags")
+ if isinstance(tags, dict):
+ values.append(tags.get("rotate"))
+ for value in values:
+ if value in {None, "", "N/A"}:
+ continue
+ try:
+ rotation = float(value)
+ except (TypeError, ValueError):
+ continue
+ if math.isfinite(rotation):
+ return rotation % 360.0
+ return 0.0
+
+
+def _display_geometry(stream: dict[str, Any]) -> tuple[float, float, float, float]:
+ """Return square-pixel display width/height, SAR, and rotation."""
+
+ coded_width = int(stream.get("width") or 0)
+ coded_height = int(stream.get("height") or 0)
+ sar = _parse_display_ratio(stream.get("sample_aspect_ratio")) or 1.0
+ dar = _parse_display_ratio(stream.get("display_aspect_ratio"))
+ physical_height = float(coded_height)
+ physical_width = dar * physical_height if dar > 0.0 else float(coded_width) * sar
+ rotation = _stream_rotation_degrees(stream)
+ quarter_turns = round(rotation / 90.0)
+ if abs(rotation - quarter_turns * 90.0) <= 1e-6:
+ if quarter_turns % 2:
+ return physical_height, physical_width, sar, rotation
+ return physical_width, physical_height, sar, rotation
+ radians = math.radians(rotation)
+ cosine = abs(math.cos(radians))
+ sine = abs(math.sin(radians))
+ display_width = physical_width * cosine + physical_height * sine
+ display_height = physical_width * sine + physical_height * cosine
+ return display_width, display_height, sar, rotation
+
+
+_FFPROBE_STREAM_ENTRIES = (
+ "stream=codec_type,width,height,duration,sample_rate,channels,"
+ "avg_frame_rate,r_frame_rate,nb_frames,sample_aspect_ratio,display_aspect_ratio"
+ ":stream_tags=rotate"
+)
+
+# ffprobe gained the per-stream "stream_side_data" section in 6.0 and rejects
+# the whole -show_entries spec without it. Older builds (e.g. Ubuntu 22.04's
+# 4.4) report rotation through the "rotate" stream tag, which
+# _stream_rotation_degrees already reads, so drop the section on fallback.
+_FFPROBE_ENTRY_VARIANTS = (
+ f"{_FFPROBE_STREAM_ENTRIES}:stream_side_data=rotation:format=format_name,duration",
+ f"{_FFPROBE_STREAM_ENTRIES}:format=format_name,duration",
+)
+_ffprobe_entries: str | None = None
+
+
+def _ffprobe_media(path: str) -> dict[str, Any]:
+ global _ffprobe_entries
+
+ variants = (
+ (_ffprobe_entries,) if _ffprobe_entries is not None else _FFPROBE_ENTRY_VARIANTS
+ )
+
+ last_error: Exception | None = None
+ for entries in variants:
+ try:
+ result = subprocess.run(
+ [
+ "ffprobe",
+ "-v",
+ "error",
+ "-protocol_whitelist",
+ "file",
+ "-format_whitelist",
+ "mov,mp4,m4a,3gp,3g2,mj2,matroska,webm,wav,mp3,flac,ogg",
+ "-show_entries",
+ entries,
+ "-of",
+ "json",
+ "-i",
+ path,
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ except subprocess.CalledProcessError as exc:
+ # Only an unknown section name is worth retrying; a genuinely bad
+ # input must fail on the first variant.
+ if "No match for section" not in (exc.stderr or ""):
+ raise
+ last_error = exc
+ continue
+ _ffprobe_entries = entries
+ return json.loads(result.stdout)
+
+ raise last_error # type: ignore[misc]
+
+
+def _validate_localized_media(
+ path: str,
+ *,
+ condition_type: str,
+) -> dict[str, Any]:
+ """Probe one localized source and return facts used by MiniMax H3 admission.
+
+ This is deliberately a model-facing validity check: it verifies that the
+ source is non-empty, parseable and contains the stream type requested by
+ the condition. Generic transport/resource ceilings do not belong here;
+ the model's shape and temporal contracts are resolved separately.
+ """
+
+ if condition_type == "image":
+ try:
+ from PIL import Image, ImageOps
+
+ with Image.open(path) as image:
+ coded_width, coded_height = image.size
+ image_format = str(image.format or "").upper()
+ if coded_width <= 0 or coded_height <= 0:
+ raise ValueError("image has no positive dimensions")
+ display_image = ImageOps.exif_transpose(image)
+ width, height = display_image.size
+ except Exception as exc:
+ raise ValueError("MiniMax H3 image material is invalid") from exc
+ if image_format not in {"JPEG", "PNG", "WEBP"}:
+ raise ValueError("MiniMax H3 image material uses an unsupported format")
+ if width <= 0 or height <= 0:
+ raise ValueError(
+ "MiniMax H3 image material has no positive display geometry"
+ )
+ return {
+ "condition_type": "image",
+ "coded_width": int(coded_width),
+ "coded_height": int(coded_height),
+ "display_width": int(width),
+ "display_height": int(height),
+ "image_format": image_format,
+ "exif_transposed": (coded_width, coded_height) != (width, height),
+ }
+
+ if condition_type not in {"audio", "video", "video_audio"}:
+ raise ValueError(f"unsupported MiniMax H3 condition type {condition_type!r}")
+ try:
+ payload = _ffprobe_media(path)
+ except Exception as exc:
+ raise ValueError("MiniMax H3 media material is invalid") from exc
+
+ streams = payload.get("streams") or []
+ format_names = set(
+ str((payload.get("format") or {}).get("format_name") or "").split(",")
+ )
+ allowed_formats = {
+ "mov",
+ "mp4",
+ "m4a",
+ "3gp",
+ "3g2",
+ "mj2",
+ "matroska",
+ "webm",
+ "wav",
+ "mp3",
+ "flac",
+ "ogg",
+ }
+ if not format_names or not format_names.issubset(allowed_formats):
+ raise ValueError("MiniMax H3 media container format is not allowed")
+ video_streams = [s for s in streams if s.get("codec_type") == "video"]
+ audio_streams = [s for s in streams if s.get("codec_type") == "audio"]
+ if condition_type in {"video", "video_audio"} and not video_streams:
+ raise ValueError("MiniMax H3 video material has no video stream")
+ if condition_type in {"audio", "video_audio"} and not audio_streams:
+ raise ValueError("MiniMax H3 audio material has no audio stream")
+
+ primary_video: dict[str, Any] | None = None
+ for stream in video_streams:
+ try:
+ width = int(stream.get("width") or 0)
+ height = int(stream.get("height") or 0)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(
+ "MiniMax H3 video material has invalid dimensions"
+ ) from exc
+ if width <= 0 or height <= 0:
+ raise ValueError("MiniMax H3 video material has no positive dimensions")
+ fps = _parse_frame_rate(stream.get("avg_frame_rate")) or _parse_frame_rate(
+ stream.get("r_frame_rate")
+ )
+ if fps <= 0:
+ raise ValueError("MiniMax H3 video material has no usable frame rate")
+ if primary_video is None:
+ primary_video = stream
+
+ for stream in audio_streams:
+ try:
+ sample_rate = int(stream.get("sample_rate") or 0)
+ channels = int(stream.get("channels") or 0)
+ except (TypeError, ValueError) as exc:
+ raise ValueError("MiniMax H3 audio material has invalid metadata") from exc
+ if sample_rate <= 0:
+ raise ValueError("MiniMax H3 audio material has no usable sample rate")
+ if channels <= 0:
+ raise ValueError("MiniMax H3 audio material has no usable channel count")
+
+ durations: list[float] = []
+ for value in [
+ (payload.get("format") or {}).get("duration"),
+ *(stream.get("duration") for stream in streams),
+ ]:
+ if value in {None, "", "N/A"}:
+ continue
+ try:
+ duration = float(value)
+ except (TypeError, ValueError):
+ continue
+ if math.isfinite(duration) and duration > 0:
+ durations.append(duration)
+ if not durations:
+ raise ValueError("MiniMax H3 media material has no positive duration")
+
+ duration_seconds = max(durations)
+ facts: dict[str, Any] = {
+ "condition_type": condition_type,
+ "duration_seconds": duration_seconds,
+ "has_audio": bool(audio_streams),
+ }
+ if primary_video is not None:
+ coded_width = int(primary_video.get("width") or 0)
+ coded_height = int(primary_video.get("height") or 0)
+ display_width, display_height, sar, rotation = _display_geometry(primary_video)
+ fps = _parse_frame_rate(primary_video.get("avg_frame_rate"))
+ if fps <= 0:
+ fps = _parse_frame_rate(primary_video.get("r_frame_rate"))
+ raw_count = primary_video.get("nb_frames")
+ try:
+ frame_count = int(raw_count)
+ except (TypeError, ValueError):
+ frame_count = max(1, int(round(duration_seconds * fps)))
+ if frame_count <= 0:
+ frame_count = max(1, int(round(duration_seconds * fps)))
+ try:
+ video_duration_seconds = float(primary_video.get("duration"))
+ except (TypeError, ValueError):
+ video_duration_seconds = 0.0
+ if not math.isfinite(video_duration_seconds) or video_duration_seconds <= 0:
+ video_duration_seconds = frame_count / fps
+ facts.update(
+ {
+ "coded_width": coded_width,
+ "coded_height": coded_height,
+ "display_width": display_width,
+ "display_height": display_height,
+ "sample_aspect_ratio": str(
+ primary_video.get("sample_aspect_ratio") or "1:1"
+ ),
+ "sample_aspect_ratio_value": sar,
+ "display_aspect_ratio": display_width / display_height,
+ "rotation_degrees": rotation,
+ "fps": fps,
+ "frame_count": frame_count,
+ "video_duration_seconds": video_duration_seconds,
+ }
+ )
+ if audio_streams:
+ facts["audio_sample_rate"] = int(audio_streams[0].get("sample_rate") or 0)
+ facts["audio_channels"] = int(audio_streams[0].get("channels") or 0)
+ try:
+ audio_duration_seconds = float(audio_streams[0].get("duration"))
+ except (TypeError, ValueError):
+ audio_duration_seconds = 0.0
+ facts["audio_duration_seconds"] = (
+ audio_duration_seconds
+ if math.isfinite(audio_duration_seconds) and audio_duration_seconds > 0
+ else duration_seconds
+ )
+ return facts
+
+
+def _validate_material_once(
+ batch: Any,
+ uri: str,
+ path: str,
+ *,
+ condition_type: str,
+) -> dict[str, Any]:
+ probe_facts = batch.extra.setdefault(MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY, {})
+ key = (uri, condition_type)
+ cached = probe_facts.get(key)
+ if isinstance(cached, dict):
+ return cached
+ facts = _validate_localized_media(path, condition_type=condition_type)
+ if not isinstance(facts, dict):
+ raise ValueError("MiniMax H3 material probe did not return facts")
+ probe_facts[key] = facts
+ return facts
+
+
+def _material_output_paths(
+ batch: Any,
+ *,
+ condition_type: str,
+ condition_index: int,
+ source_name: str | None = None,
+ media_type: str | None = None,
+) -> tuple[Path, Path]:
+ suffix = (
+ _safe_suffix(source_name)
+ or _SUFFIX_BY_MEDIA_TYPE.get(str(media_type or "").lower())
+ or _DEFAULT_SUFFIX_BY_TYPE.get(condition_type, ".bin")
+ )
+ output_path = (
+ Path(_material_workdir(batch)) / f"condition_{int(condition_index):04d}{suffix}"
+ )
+ return output_path, output_path.with_name(output_path.name + ".partial")
+
+
+def _material_workdir(batch: Any) -> str:
+ registry = batch.extra.setdefault(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY, {})
+ material_dirs = registry.get("material")
+ if isinstance(material_dirs, list) and material_dirs:
+ return str(material_dirs[0])
+ return minimax_h3_register_temp_dir(
+ batch,
+ tempfile.mkdtemp(prefix="minimax_h3_material_"),
+ owner="material",
+ )
+
+
+def _decode_base64_chunk(encoded: bytes | bytearray) -> bytes:
+ try:
+ return base64.b64decode(encoded, altchars=b"-_", validate=True)
+ except Exception as exc:
+ raise ValueError("material URI has an invalid base64 payload") from exc
+
+
+def _stream_base64_material(
+ batch: Any,
+ uri: str,
+ *,
+ condition_type: str,
+ condition_index: int,
+) -> str:
+ payload_start, media_type = _base64_uri_payload_start(uri)
+
+ output_path, partial_path = _material_output_paths(
+ batch,
+ condition_type=condition_type,
+ condition_index=condition_index,
+ media_type=media_type,
+ )
+ total = 0
+ encoded_size = 0
+ padding = 0
+ saw_padding = False
+ encoded_chunk = bytearray()
+ try:
+ with partial_path.open("wb") as output:
+ for value in _iter_base64_payload_bytes(uri, payload_start):
+ encoded_size += 1
+ if value == ord("="):
+ saw_padding = True
+ padding += 1
+ if padding > 2:
+ raise ValueError("material URI has invalid base64 padding")
+ elif saw_padding:
+ raise ValueError("material URI has data after base64 padding")
+ encoded_chunk.append(value)
+ if len(encoded_chunk) == MINIMAX_H3_BASE64_DECODE_CHUNK_CHARS:
+ decoded_chunk = _decode_base64_chunk(encoded_chunk)
+ total += len(decoded_chunk)
+ output.write(decoded_chunk)
+ encoded_chunk.clear()
+ if encoded_size == 0:
+ raise ValueError("material URI base64 payload is empty")
+ if encoded_size % 4 == 1 or (padding and encoded_size % 4):
+ raise ValueError("material URI has an invalid base64 payload length")
+ if encoded_chunk:
+ encoded_chunk.extend(b"=" * (-len(encoded_chunk) % 4))
+ decoded_chunk = _decode_base64_chunk(encoded_chunk)
+ total += len(decoded_chunk)
+ output.write(decoded_chunk)
+ decoded_size = (encoded_size * 3) // 4 - padding
+ if decoded_size <= 0:
+ raise ValueError("material URI decoded payload is empty")
+ if total != decoded_size:
+ raise ValueError(
+ f"MiniMax H3 base64 decoded size {total} != expected {decoded_size}"
+ )
+ partial_path.replace(output_path)
+ except Exception:
+ partial_path.unlink(missing_ok=True)
+ output_path.unlink(missing_ok=True)
+ raise
+ return str(output_path)
+
+
+def _stream_tar_member_material(
+ batch: Any,
+ uri: str,
+ *,
+ condition_type: str,
+ condition_index: int,
+) -> str:
+ source_path, offset, size, member = _parse_tar_member_uri(uri)
+ if size <= 0:
+ raise ValueError("tar material payload is empty")
+ if not source_path.is_file():
+ raise FileNotFoundError(
+ f"tar material source does not exist or is not a file: {source_path}"
+ )
+ source_size = source_path.stat().st_size
+ if offset > source_size or size > source_size - offset:
+ available = max(0, source_size - offset)
+ raise ValueError(
+ f"tar material payload is truncated: expected {size} bytes, "
+ f"only {available} available"
+ )
+
+ output_path, partial_path = _material_output_paths(
+ batch,
+ condition_type=condition_type,
+ condition_index=condition_index,
+ source_name=member,
+ )
+ remaining = size
+ try:
+ with source_path.open("rb") as source, partial_path.open("wb") as output:
+ source.seek(offset)
+ while remaining:
+ chunk = source.read(min(MINIMAX_H3_HTTP_READ_CHUNK_BYTES, remaining))
+ if not chunk:
+ raise ValueError(
+ f"tar material payload is truncated with {remaining} bytes left"
+ )
+ if len(chunk) > remaining:
+ raise ValueError("tar material reader returned too many bytes")
+ output.write(chunk)
+ remaining -= len(chunk)
+ partial_path.replace(output_path)
+ except Exception:
+ partial_path.unlink(missing_ok=True)
+ output_path.unlink(missing_ok=True)
+ raise
+ return str(output_path)
+
+
+def _http_media_type(response: Any) -> str | None:
+ headers = getattr(response, "headers", None)
+ if headers is None:
+ return None
+ get_content_type = getattr(headers, "get_content_type", None)
+ if callable(get_content_type):
+ value = get_content_type()
+ else:
+ value = headers.get("Content-Type") or headers.get("content-type")
+ if isinstance(value, str):
+ value = value.split(";", 1)[0].strip()
+ return str(value).lower() if value else None
+
+
+def _stream_http_material(
+ batch: Any,
+ uri: str,
+ *,
+ condition_type: str,
+ condition_index: int,
+ timeout_s: float,
+) -> str:
+ # Use the repository's legacy urllib behavior. Model-specific material
+ # localization intentionally does not perform the shared public SSRF
+ # policy or a cumulative request deadline.
+ with urllib.request.urlopen(uri, timeout=timeout_s) as response:
+ media_type = _http_media_type(response)
+ suffix = (
+ _safe_suffix(uri)
+ or _SUFFIX_BY_MEDIA_TYPE.get(str(media_type or "").lower())
+ or _DEFAULT_SUFFIX_BY_TYPE.get(condition_type, ".bin")
+ )
+ output_path = (
+ Path(_material_workdir(batch))
+ / f"condition_{int(condition_index):04d}{suffix}"
+ )
+ partial_path = output_path.with_name(output_path.name + ".partial")
+ total = 0
+ try:
+ with partial_path.open("wb") as output:
+ while True:
+ chunk = response.read(MINIMAX_H3_HTTP_READ_CHUNK_BYTES)
+ if not chunk:
+ break
+ if not isinstance(chunk, bytes):
+ raise TypeError(
+ "HTTP material response.read() must return bytes, got "
+ f"{type(chunk).__name__}"
+ )
+ total += len(chunk)
+ output.write(chunk)
+ if total == 0:
+ raise ValueError(f"HTTP material body is empty: {uri}")
+ partial_path.replace(output_path)
+ except Exception:
+ partial_path.unlink(missing_ok=True)
+ output_path.unlink(missing_ok=True)
+ raise
+ return str(output_path)
+
+
+def minimax_h3_localize_material_uri(
+ batch: Any,
+ uri: str,
+ *,
+ condition_type: str,
+ condition_index: int,
+ timeout_s: float = 120.0,
+) -> str:
+ """Return a local path for a canonical condition URI.
+
+ Local paths and local ``file://`` URIs are validated and returned without
+ copying. HTTP(S), base64/data and direct tar-member URIs are materialized
+ once per request and cached for all MiniMax H3 consumers.
+ """
+
+ if not isinstance(uri, str) or not uri:
+ raise ValueError("condition URI must be a non-empty string")
+ parsed = None
+ for special_scheme in ("data", "base64", "tar+offset", "tar+b64header"):
+ if uri.startswith(special_scheme + ":"):
+ scheme = special_scheme
+ break
+ else:
+ parsed = urllib.parse.urlsplit(uri)
+ scheme = parsed.scheme
+
+ if scheme == "file":
+ assert parsed is not None
+ if parsed.netloc not in {"", "localhost"}:
+ raise ValueError(f"file URI host must be local, got {parsed.netloc!r}")
+ output_path = _checked_material_file(
+ Path(urllib.parse.unquote(parsed.path)),
+ label="MiniMax H3 material source",
+ )
+ _validate_material_once(batch, uri, output_path, condition_type=condition_type)
+ return output_path
+ if not scheme:
+ output_path = _checked_material_file(
+ Path(uri).expanduser(),
+ label="MiniMax H3 material source",
+ )
+ _validate_material_once(batch, uri, output_path, condition_type=condition_type)
+ return output_path
+ if scheme == "s3":
+ raise NotImplementedError(
+ "MiniMax H3 s3:// material URIs require a configured artifact resolver"
+ )
+
+ cache = batch.extra.setdefault(MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY, {})
+ cached = cache.get(uri)
+ if isinstance(cached, str):
+ cached_path = Path(cached)
+ if cached_path.exists():
+ output_path = _checked_material_file(
+ cached_path,
+ label="cached MiniMax H3 material",
+ )
+ _validate_material_once(
+ batch, uri, output_path, condition_type=condition_type
+ )
+ return output_path
+ cache.pop(uri, None)
+
+ if scheme in {"http", "https"}:
+ output_path = _stream_http_material(
+ batch,
+ uri,
+ condition_type=condition_type,
+ condition_index=condition_index,
+ timeout_s=timeout_s,
+ )
+ try:
+ _validate_material_once(
+ batch, uri, output_path, condition_type=condition_type
+ )
+ except Exception:
+ Path(output_path).unlink(missing_ok=True)
+ raise
+ cache[uri] = output_path
+ return output_path
+
+ if scheme in {"data", "base64"}:
+ output_path = _stream_base64_material(
+ batch,
+ uri,
+ condition_type=condition_type,
+ condition_index=condition_index,
+ )
+ try:
+ _validate_material_once(
+ batch, uri, output_path, condition_type=condition_type
+ )
+ except Exception:
+ Path(output_path).unlink(missing_ok=True)
+ raise
+ cache[uri] = output_path
+ return output_path
+
+ if scheme in {"tar+offset", "tar+b64header"}:
+ output_path = _stream_tar_member_material(
+ batch,
+ uri,
+ condition_type=condition_type,
+ condition_index=condition_index,
+ )
+ try:
+ _validate_material_once(
+ batch, uri, output_path, condition_type=condition_type
+ )
+ except Exception:
+ Path(output_path).unlink(missing_ok=True)
+ raise
+ cache[uri] = output_path
+ return output_path
+
+ raise NotImplementedError(
+ f"MiniMax H3 material localization does not support URI scheme {scheme!r}"
+ )
+
+
+def minimax_h3_probe_material(
+ batch: Any,
+ uri: str,
+ *,
+ condition_type: str,
+ condition_index: int,
+) -> dict[str, Any]:
+ """Localize and return cached display-geometry facts for one condition."""
+
+ path = minimax_h3_localize_material_uri(
+ batch,
+ uri,
+ condition_type=condition_type,
+ condition_index=condition_index,
+ )
+ key = (uri, condition_type)
+ facts = batch.extra.get(MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY, {}).get(key)
+ if not isinstance(facts, dict) or not facts:
+ raise RuntimeError(
+ "MiniMax H3 material localization completed without cached probe facts"
+ )
+ return {"local_path": path, **facts}
+
+
+__all__ = [
+ "MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY",
+ "MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY",
+ "MINIMAX_H3_TEMP_DIRS_EXTRA_KEY",
+ "minimax_h3_cleanup_temp_dirs",
+ "minimax_h3_localize_material_uri",
+ "minimax_h3_probe_material",
+ "minimax_h3_register_temp_dir",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py
new file mode 100644
index 000000000..d6f1de66f
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py
@@ -0,0 +1,502 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 packed-sequence materialization from the validated workspace
+builder, covering fl2va and t2va layouts.
+
+Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P].
+Builder rules:
+- block-derived position infos, update masks, token tags, and cu_seqlens
+- img_position_ids fp64 grid: text rows (row_idx,0,0); video/cond t counter
+ continues text_len with temporal interp spans (frame_rescale 5/3 x
+ frame_per_token (1,4,4,4,4)); each spatial sqrt_area axis uses evenly spaced
+ coordinates excluding the right endpoint, then scales them by INTERP;
+ audio channel-major blocks pinned to the w-grid extremes.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+ MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+)
+
+_INTERP = 32
+_T_GROUP = 5
+_FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
+_FRAME_RESCALE = 5.0 / 3.0
+_PATCH_H = 2
+_PATCH_W = 2
+
+
+def _keyframe_cond_frame_indices(
+ *,
+ include_keyframe_cond: bool,
+ keyframe_frame_indices: list[int] | tuple[int, ...] | None,
+) -> list[int]:
+ if not include_keyframe_cond:
+ if keyframe_frame_indices is not None:
+ raise ValueError(
+ "keyframe_frame_indices must be omitted when keyframe cond is not included"
+ )
+ return []
+ if keyframe_frame_indices is None:
+ raise ValueError("strict fl2va packed layout requires keyframe_frame_indices")
+ if any(
+ isinstance(value, bool) or not isinstance(value, int)
+ for value in keyframe_frame_indices
+ ):
+ raise ValueError(
+ "strict fl2va packed layout requires integer keyframe_frame_indices"
+ )
+ out = list(keyframe_frame_indices)
+ if tuple(out) not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+ raise ValueError(
+ "strict fl2va packed layout requires keyframe_frame_indices in "
+ f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {out!r}"
+ )
+ return out
+
+
+def _resolve_keyframe_frame_indices(
+ frame_indices: Sequence[int],
+ *,
+ frame_count: int | None,
+) -> list[int]:
+ if frame_indices and frame_count is None:
+ raise ValueError(
+ "frame_count is required when keyframe_frame_indices are provided"
+ )
+ if frame_count is None:
+ return []
+ if frame_count <= 0:
+ raise ValueError("frame_count must be positive")
+ seen: dict[int, int] = {}
+ resolved: list[int] = []
+ for block_index, semantic_index in enumerate(frame_indices):
+ if semantic_index == -1:
+ resolved_index = frame_count - 1
+ elif 0 <= semantic_index < frame_count:
+ resolved_index = semantic_index
+ else:
+ raise ValueError(
+ f"keyframe frame index {semantic_index} must be -1 or in "
+ f"[0, {frame_count})"
+ )
+ previous = seen.get(resolved_index)
+ if previous is not None:
+ raise ValueError(
+ f"keyframe frame index at block {block_index} resolves to "
+ f"{resolved_index}, already bound by block {previous}"
+ )
+ seen[resolved_index] = block_index
+ resolved.append(resolved_index)
+ return resolved
+
+
+def _temporal_position_span(temporal_length: int) -> float:
+ """Temporal position span for patch_t=1, in fp64.
+
+ NOTE: intentionally NOT merged with ``_video_t_span``. This variant sums
+ via numpy (pairwise summation), matching the fl2va anchor computation,
+ while ``_video_t_span`` sums sequentially, matching the ref2va
+ t-origin accumulation. The two orders diverge in the last ulp
+ from n=16 onward, so each path must keep its own summation order.
+ """
+ spans = np.ones(int(temporal_length), dtype=np.float64) * _FRAME_RESCALE
+ for token_index in range(_T_GROUP):
+ spans[token_index::_T_GROUP] *= _FRAME_PER_TOKEN[token_index]
+ return float(spans.sum())
+
+
+def minimax_h3_packed_sequence(
+ *,
+ text_len: int,
+ latent_t: int,
+ latent_h: int,
+ latent_w: int,
+ audio_t: int,
+ audio_channel: int = 2,
+ include_keyframe_cond: bool,
+ keyframe_frame_indices: list[int] | tuple[int, ...] | None = None,
+ frame_count: int | None = None,
+) -> dict[str, Any]:
+ """Build the packed-sequence structural fields for one CFG branch.
+
+ The used length is padded up to a multiple of 64.
+ """
+ ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
+ frame_rows = ph * pw
+ cond_frame_indices = _keyframe_cond_frame_indices(
+ include_keyframe_cond=include_keyframe_cond,
+ keyframe_frame_indices=keyframe_frame_indices,
+ )
+ resolved_cond_frame_indices = _resolve_keyframe_frame_indices(
+ cond_frame_indices,
+ frame_count=frame_count,
+ )
+ cond_rows = len(cond_frame_indices) * frame_rows
+ video_rows = latent_t * frame_rows
+ audio_rows = audio_t * audio_channel
+ used = text_len + cond_rows + audio_rows + video_rows
+ seq_len = (
+ (used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1)
+ // MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
+ * MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
+ )
+
+ text_sl = slice(0, text_len)
+ cond_sl = slice(text_len, text_len + cond_rows)
+ audio_sl = slice(cond_sl.stop, cond_sl.stop + audio_rows)
+ video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
+ target_img_pos = torch.arange(video_sl.start, video_sl.stop)
+ img_pos = (
+ torch.cat([torch.arange(cond_sl.start, cond_sl.stop), target_img_pos])
+ if cond_rows
+ else target_img_pos
+ )
+ update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
+ update_mask[cond_rows:] = True
+ audio_pos = torch.arange(audio_sl.start, audio_sl.stop)
+ text_pos = torch.arange(0, text_len)
+
+ g = torch.zeros(seq_len, 3, dtype=torch.float64)
+ g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)
+
+ t_grid = _video_t_grid(latent_t, float(text_len))
+ sqrt_area = np.sqrt(latent_h * latent_w)
+ h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, sqrt_area)
+ w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, sqrt_area)
+ hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
+ frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
+ video_g = g[video_sl].view(latent_t, frame_rows, 3)
+ video_g[:, :, 0] = t_grid[:, None]
+ video_g[:, :, 1:] = frame[None]
+ for block_index, pixel_index in enumerate(resolved_cond_frame_indices):
+ sl = slice(
+ cond_sl.start + block_index * frame_rows,
+ cond_sl.start + (block_index + 1) * frame_rows,
+ )
+ if pixel_index == 0:
+ cond_t = float(text_len)
+ elif frame_count is not None and pixel_index == frame_count - 1:
+ cond_t = (
+ float(text_len) + _temporal_position_span(latent_t) - _FRAME_RESCALE
+ )
+ else:
+ raise ValueError(
+ "fl2va packed layout only supports first/last keyframe anchors, "
+ f"got resolved frame index {pixel_index}"
+ )
+ g[sl, 0] = cond_t
+ g[sl, 1:] = frame
+ audio_t_grid = float(text_len) + torch.arange(audio_t, dtype=torch.float64)
+ g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
+ g[audio_sl.start : audio_sl.start + audio_t, 2] = float(w_grid[0])
+ g[audio_sl.start + audio_t : audio_sl.stop, 2] = float(w_grid[-1])
+
+ token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING
+ token_tags[text_sl] = 1 # TEXT (fl2va image-segment override happens upstream)
+ token_tags[audio_sl] = 2 # AUDIO
+ token_tags[img_pos] = 0 # VIDEO
+
+ cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
+ return {
+ "seq_len": seq_len,
+ "img_pos": img_pos,
+ "audio_pos": audio_pos,
+ "text_pos": text_pos,
+ "update_mask": update_mask,
+ "img_position_ids": g,
+ "token_tags": token_tags,
+ "cu_seqlens": cu,
+ }
+
+
+def _positive_int(
+ block: Mapping[str, object],
+ key: str,
+ path: str,
+ *,
+ allow_zero: bool = False,
+) -> int:
+ value = block.get(key)
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise ValueError(f"{path}.{key} must be an integer")
+ if value < 0 or (value == 0 and not allow_zero):
+ predicate = "non-negative" if allow_zero else "positive"
+ raise ValueError(f"{path}.{key} must be {predicate}")
+ return int(value)
+
+
+def _axis_from_sqrt_area(dim: int, patch: int, sqrt_area: float) -> torch.Tensor:
+ ratio = dim / sqrt_area
+ left = (1.0 - ratio) * 1.0 / 2.0
+ right = left + ratio * 1.0
+ grid = np.linspace(left, right, dim // patch, endpoint=False) * _INTERP
+ return torch.from_numpy(grid).to(torch.float64)
+
+
+def _video_t_grid(n: int, origin: float) -> torch.Tensor:
+ spans = torch.tensor(
+ [_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)],
+ dtype=torch.float64,
+ )
+ return origin + torch.cat(
+ [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)]
+ )
+
+
+def _video_t_span(n: int) -> float:
+ # Sequential fp64 summation on purpose — see _temporal_position_span for
+ # why the two span implementations must not be unified.
+ return sum(_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n))
+
+
+def _range_for_slice(sl: slice) -> torch.Tensor:
+ return torch.arange(sl.start, sl.stop, dtype=torch.long)
+
+
+def _cat_ranges(parts: list[torch.Tensor]) -> torch.Tensor:
+ if len(parts) == 1:
+ return parts[0]
+ if parts:
+ return torch.cat(parts)
+ return torch.empty(0, dtype=torch.long)
+
+
+def minimax_h3_packed_sequence_ref2va_blocks(
+ *,
+ text_len: int,
+ latent_t: int,
+ latent_h: int,
+ latent_w: int,
+ audio_t: int,
+ ref_blocks: Sequence[Mapping[str, object]],
+ audio_channel: int = 2,
+ seq_len: int | None = None,
+) -> dict[str, Any]:
+ """General ref2va-family packed layout.
+
+ ``ref_blocks`` are consumed in request/plan order:
+ - ``{"kind": "image", "latent_h": H, "latent_w": W}``
+ - ``{"kind": "audio", "ref_audio_t": T}``
+ - ``{"kind": "video"|"video_audio", "ref_audio_t": T,
+ "latent_t": RT, "latent_h": RH, "latent_w": RW}``
+
+ Video-bearing blocks pack their audio rows immediately before their video
+ rows; both share the same temporal origin and advance by the longer of the
+ audio and video spans. Standalone audio advances the target origin by its
+ own T, and image blocks advance it by one integer slot.
+ """
+ if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)):
+ raise ValueError("ref_blocks must be a sequence")
+
+ parsed: list[dict[str, object]] = []
+ ref_visual_rows = 0
+ ref_audio_rows = 0
+ for index, raw in enumerate(ref_blocks):
+ path = f"ref_blocks[{index}]"
+ if not isinstance(raw, Mapping):
+ raise ValueError(f"{path} must be an object")
+ kind = raw.get("kind", raw.get("type"))
+ if not isinstance(kind, str) or not kind:
+ raise ValueError(f"{path}.kind must be a non-empty string")
+ if kind == "image":
+ rh = _positive_int(raw, "latent_h", path)
+ rw = _positive_int(raw, "latent_w", path)
+ rows = (rh // _PATCH_H) * (rw // _PATCH_W)
+ item = {"kind": kind, "latent_h": rh, "latent_w": rw, "rows": rows}
+ ref_visual_rows += rows
+ elif kind == "audio":
+ rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
+ rows = rt * audio_channel
+ item = {"kind": kind, "ref_audio_t": rt, "audio_rows": rows}
+ ref_audio_rows += rows
+ elif kind in ("video", "video_audio"):
+ rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
+ vt = _positive_int(raw, "latent_t", path)
+ vh = _positive_int(raw, "latent_h", path)
+ vw = _positive_int(raw, "latent_w", path)
+ frame_rows = (vh // _PATCH_H) * (vw // _PATCH_W)
+ audio_rows = rt * audio_channel
+ video_rows = vt * frame_rows
+ item = {
+ "kind": kind,
+ "ref_audio_t": rt,
+ "latent_t": vt,
+ "latent_h": vh,
+ "latent_w": vw,
+ "frame_rows": frame_rows,
+ "audio_rows": audio_rows,
+ "video_rows": video_rows,
+ }
+ ref_audio_rows += audio_rows
+ ref_visual_rows += video_rows
+ else:
+ raise ValueError(f"{path}.kind unsupported for ref2va: {kind!r}")
+ parsed.append(item)
+
+ ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
+ frame_rows = ph * pw
+ video_rows = latent_t * frame_rows
+ audio_rows = audio_t * audio_channel
+ ref_rows = ref_visual_rows + ref_audio_rows
+ used = text_len + ref_rows + audio_rows + video_rows
+ if seq_len is None:
+ seq_len = (
+ (used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1)
+ // MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
+ * MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
+ )
+ if seq_len < used:
+ raise ValueError(f"seq_len {seq_len} < used rows {used}")
+
+ text_sl = slice(0, text_len)
+ cursor = text_len
+ block_slices: list[dict[str, object]] = []
+ for item in parsed:
+ kind = str(item["kind"])
+ if kind == "image":
+ rows = int(item["rows"])
+ visual_sl = slice(cursor, cursor + rows)
+ cursor = visual_sl.stop
+ block_slices.append({**item, "visual_sl": visual_sl})
+ elif kind == "audio":
+ rows = int(item["audio_rows"])
+ audio_sl = slice(cursor, cursor + rows)
+ cursor = audio_sl.stop
+ block_slices.append({**item, "audio_sl": audio_sl})
+ else:
+ a_rows = int(item["audio_rows"])
+ v_rows = int(item["video_rows"])
+ audio_sl = slice(cursor, cursor + a_rows)
+ visual_sl = slice(audio_sl.stop, audio_sl.stop + v_rows)
+ cursor = visual_sl.stop
+ block_slices.append({**item, "audio_sl": audio_sl, "visual_sl": visual_sl})
+
+ audio_sl = slice(cursor, cursor + audio_rows)
+ video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
+ ref_img_pos_parts: list[torch.Tensor] = []
+ ref_audio_pos_parts: list[torch.Tensor] = []
+ g = torch.zeros(seq_len, 3, dtype=torch.float64)
+ g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)
+
+ target_area = np.sqrt(latent_h * latent_w)
+ h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, target_area)
+ w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, target_area)
+ hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
+ target_frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
+
+ t_cursor = float(text_len)
+ for item in block_slices:
+ kind = str(item["kind"])
+ if kind == "image":
+ visual_sl = item["visual_sl"]
+ assert isinstance(visual_sl, slice)
+ ref_img_pos_parts.append(_range_for_slice(visual_sl))
+ rh = int(item["latent_h"])
+ rw = int(item["latent_w"])
+ area = np.sqrt(rh * rw)
+ ref_hh, ref_ww = torch.meshgrid(
+ _axis_from_sqrt_area(rh, _PATCH_H, area),
+ _axis_from_sqrt_area(rw, _PATCH_W, area),
+ indexing="ij",
+ )
+ g[visual_sl, 0] = t_cursor
+ g[visual_sl, 1] = ref_hh.reshape(-1)
+ g[visual_sl, 2] = ref_ww.reshape(-1)
+ t_cursor += 1.0
+ elif kind == "audio":
+ audio_ref_sl = item["audio_sl"]
+ assert isinstance(audio_ref_sl, slice)
+ ref_t = int(item["ref_audio_t"])
+ ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
+ ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
+ g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
+ if ref_t:
+ g[audio_ref_sl.start : audio_ref_sl.start + ref_t, 2] = float(w_grid[0])
+ g[audio_ref_sl.start + ref_t : audio_ref_sl.stop, 2] = float(w_grid[-1])
+ t_cursor += float(ref_t)
+ else:
+ audio_ref_sl = item["audio_sl"]
+ visual_sl = item["visual_sl"]
+ assert isinstance(audio_ref_sl, slice)
+ assert isinstance(visual_sl, slice)
+ ref_t = int(item["ref_audio_t"])
+ vt = int(item["latent_t"])
+ vh = int(item["latent_h"])
+ vw = int(item["latent_w"])
+ ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
+ ref_img_pos_parts.append(_range_for_slice(visual_sl))
+
+ ref_area = np.sqrt(vh * vw)
+ rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area)
+ rv_w_grid = _axis_from_sqrt_area(vw, _PATCH_W, ref_area)
+ rv_hh, rv_ww = torch.meshgrid(rv_h_grid, rv_w_grid, indexing="ij")
+
+ ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
+ g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
+ if ref_t:
+ g[audio_ref_sl.start : audio_ref_sl.start + ref_t, 2] = float(
+ rv_w_grid[0]
+ )
+ g[audio_ref_sl.start + ref_t : audio_ref_sl.stop, 2] = float(
+ rv_w_grid[-1]
+ )
+
+ rv_frame = torch.stack([rv_hh.reshape(-1), rv_ww.reshape(-1)], dim=-1)
+ rv_g = g[visual_sl].view(vt, int(item["frame_rows"]), 3)
+ rv_g[:, :, 0] = _video_t_grid(vt, t_cursor)[:, None]
+ rv_g[:, :, 1:] = rv_frame[None]
+ t_cursor += max(float(ref_t), _video_t_span(vt))
+
+ audio_t_grid = t_cursor + torch.arange(audio_t, dtype=torch.float64)
+ g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
+ g[audio_sl.start : audio_sl.start + audio_t, 2] = float(w_grid[0])
+ g[audio_sl.start + audio_t : audio_sl.stop, 2] = float(w_grid[-1])
+
+ video_g = g[video_sl].view(latent_t, frame_rows, 3)
+ video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None]
+ video_g[:, :, 1:] = target_frame[None]
+
+ target_img_pos = _range_for_slice(video_sl)
+ target_audio_pos = _range_for_slice(audio_sl)
+ img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos])
+ audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos])
+
+ update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
+ update_mask[ref_visual_rows:] = True
+ audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool)
+ audio_update_mask[ref_audio_rows:] = True
+ text_pos = torch.arange(0, text_len)
+
+ token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING
+ token_tags[text_sl] = 1 # TEXT
+ token_tags[audio_pos] = 2 # AUDIO (refs + target)
+ token_tags[img_pos] = 0 # VIDEO (refs + target)
+
+ cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
+ return {
+ "seq_len": seq_len,
+ "img_pos": img_pos,
+ "audio_pos": audio_pos,
+ "text_pos": text_pos,
+ "update_mask": update_mask,
+ "audio_update_mask": audio_update_mask,
+ "img_position_ids": g,
+ "token_tags": token_tags,
+ "cu_seqlens": cu,
+ }
+
+
+__all__ = [
+ "minimax_h3_packed_sequence",
+ "minimax_h3_packed_sequence_ref2va_blocks",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_tokens.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_tokens.py
new file mode 100644
index 000000000..50b84ee52
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_tokens.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def _int_tuple(value: Sequence[int], name: str, length: int) -> tuple[int, ...]:
+ if len(value) != length:
+ raise ValueError(f"{name} must have length {length}, got {list(value)!r}")
+ out = tuple(int(item) for item in value)
+ if any(item <= 0 for item in out):
+ raise ValueError(f"{name} values must be positive, got {list(value)!r}")
+ return out
+
+
+def _rank(tensor: torch.Tensor, name: str, rank: int) -> None:
+ if tensor.ndim != rank:
+ raise ValueError(f"{name} must be rank {rank}, got shape={list(tensor.shape)}")
+
+
+def minimax_h3_patchify_video_latent(
+ latent: torch.Tensor,
+ *,
+ patch_size: Sequence[int],
+) -> torch.Tensor:
+ """Pack SGLang video latent [B,C,T,H,W] into DiT token rows."""
+
+ _rank(latent, "video latent", 5)
+ pt, ph, pw = _int_tuple(patch_size, "patch_size", 3)
+ batch, channel, full_t, full_h, full_w = (int(dim) for dim in latent.shape)
+ if full_t % pt or full_h % ph or full_w % pw:
+ raise ValueError(
+ "video latent spatial/time dims must be divisible by patch_size: "
+ f"shape={list(latent.shape)}, patch_size={[pt, ph, pw]}"
+ )
+ t, h, w = full_t // pt, full_h // ph, full_w // pw
+ packed = latent.reshape(batch, channel, t, pt, h, ph, w, pw)
+ packed = torch.einsum("nctrhpwq->nthwcrpq", packed)
+ return packed.reshape(batch * t * h * w, channel * pt * ph * pw).contiguous()
+
+
+def minimax_h3_unpatchify_video_tokens(
+ rows: torch.Tensor,
+ *,
+ latent_shape: Sequence[int],
+ patch_size: Sequence[int],
+) -> torch.Tensor:
+ """Unpack DiT video token rows into SGLang latent [B,C,T,H,W]."""
+
+ _rank(rows, "video token rows", 2)
+ t, h, w, channel = _int_tuple(latent_shape, "latent_shape", 4)
+ pt, ph, pw = _int_tuple(patch_size, "patch_size", 3)
+ expected_dim = pt * ph * pw * channel
+ if int(rows.shape[-1]) != expected_dim:
+ raise ValueError(
+ f"video token dim {int(rows.shape[-1])} != patch volume * channel "
+ f"{expected_dim} for latent_shape={list(latent_shape)}, "
+ f"patch_size={[pt, ph, pw]}"
+ )
+ rows_per_sample = t * h * w
+ if int(rows.shape[0]) % rows_per_sample:
+ raise ValueError(
+ f"video rows {int(rows.shape[0])} must be divisible by t*h*w "
+ f"{rows_per_sample} for latent_shape={list(latent_shape)}"
+ )
+ packed = rows.reshape(-1, t, h, w, channel, pt, ph, pw)
+ latent = torch.einsum("nthwcrpq->nctrhpwq", packed)
+ return latent.reshape(-1, channel, t * pt, h * ph, w * pw).contiguous()
+
+
+def minimax_h3_unpack_audio_tokens(
+ rows: torch.Tensor,
+ *,
+ audio_t: int,
+ audio_channel: int,
+) -> torch.Tensor:
+ """Unpack DiT audio token rows into SGLang audio VAE latent [C,latent_dim,T]."""
+
+ _rank(rows, "audio token rows", 2)
+ audio_t = int(audio_t)
+ audio_channel = int(audio_channel)
+ if audio_t <= 0 or audio_channel <= 0:
+ raise ValueError(
+ f"audio_t and audio_channel must be positive, got {audio_t=} "
+ f"{audio_channel=}"
+ )
+ if int(rows.shape[0]) != audio_t:
+ raise ValueError(f"audio rows {int(rows.shape[0])} != audio_t {audio_t}")
+ if audio_t % audio_channel:
+ raise ValueError(
+ f"audio_t must be divisible by audio_channel, got {audio_t=} "
+ f"{audio_channel=}"
+ )
+ native = rows.reshape(audio_channel, audio_t // audio_channel, int(rows.shape[-1]))
+ return native.permute(0, 2, 1).contiguous()
+
+
+__all__ = [
+ "minimax_h3_patchify_video_latent",
+ "minimax_h3_unpack_audio_tokens",
+ "minimax_h3_unpatchify_video_tokens",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/prequeue.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/prequeue.py
new file mode 100644
index 000000000..5355b9136
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/prequeue.py
@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 probe -> resolve-once admission hook.
+
+This module is intentionally data/CPU only. It localizes condition media,
+caches display-geometry facts, freezes every target/material canvas, and
+resolves the real aligned workload before a video job is published or sent to
+the scheduler.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_MAX_DURATION_SECONDS,
+ MINIMAX_H3_MIN_DURATION_SECONDS,
+ MINIMAX_H3_SUPPORTED_FPS,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ MINIMAX_H3_TEMP_DIRS_EXTRA_KEY,
+ minimax_h3_cleanup_temp_dirs,
+ minimax_h3_probe_material,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ MINIMAX_H3_BASE_SHORT_EDGE,
+ MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY,
+ MiniMaxH3ResolvedPlan,
+ minimax_h3_plan_from_batch,
+ minimax_h3_resolve_spatial_shape,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+ minimax_h3_align_frame_count,
+ minimax_h3_audio_latent_t,
+ minimax_h3_video_latent_t,
+)
+
+MINIMAX_H3_PROBE_FACTS_EXTRA_KEY = "minimax_h3_probe_facts_by_condition"
+MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY = "minimax_h3_resolved_material_shapes"
+
+
+def _replace_plan_shape(
+ plan: MiniMaxH3ResolvedPlan, shape: dict[str, Any]
+) -> MiniMaxH3ResolvedPlan:
+ return MiniMaxH3ResolvedPlan(
+ task=plan.task,
+ prompt=plan.prompt,
+ seed=plan.seed,
+ materials=plan.materials,
+ encoders=plan.encoders,
+ branches=plan.branches,
+ default_flow_shift=plan.default_flow_shift,
+ default_audio_flow_shift=plan.default_audio_flow_shift,
+ flow_shift=plan.flow_shift,
+ audio_flow_shift=plan.audio_flow_shift,
+ shape=shape,
+ condition_mask=plan.condition_mask,
+ )
+
+
+def _display_shape(facts: dict[str, Any], *, label: str) -> tuple[float, float]:
+ try:
+ width = float(facts["display_width"])
+ height = float(facts["display_height"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError(f"{label} has no usable display geometry") from exc
+ if width <= 0 or height <= 0:
+ raise ValueError(f"{label} has no usable display geometry")
+ return width, height
+
+
+def _resolve_deferred_spatial_shape(
+ plan: MiniMaxH3ResolvedPlan,
+ shape: dict[str, Any],
+ probe_facts: dict[int, dict[str, Any]],
+) -> None:
+ if str(shape.get("geometry")) != "deferred":
+ return
+ if plan.task == "fl2va":
+ candidates = [
+ material
+ for material in plan.materials
+ if material.material_chain == "image.target_canvas"
+ ]
+ if len(candidates) not in {1, 2}:
+ raise ValueError(
+ f"fl2va requires one or two keyframe materials, got {len(candidates)}"
+ )
+ for material in candidates:
+ if material.frame_index not in {0, -1}:
+ raise ValueError(
+ "fl2va deferred geometry requires semantic frame_index 0 or "
+ f"-1, got {material.frame_index!r} for "
+ f"conditions[{material.condition_index}]"
+ )
+ # Select by semantic time, not request/material iteration order. The
+ # last-frame sentinel sorts after the first-frame anchor.
+ source = min(
+ candidates,
+ key=lambda material: (
+ material.frame_index == -1,
+ int(material.frame_index),
+ int(material.condition_index),
+ ),
+ )
+ geometry_source = (
+ "first_keyframe" if source.frame_index == 0 else "last_keyframe"
+ )
+ else:
+ raise ValueError(f"task {plan.task!r} has unsupported deferred target geometry")
+ width, height = _display_shape(
+ probe_facts[int(source.condition_index)],
+ label=f"conditions[{source.condition_index}]",
+ )
+ shape.update(
+ minimax_h3_resolve_spatial_shape(
+ width=width,
+ height=height,
+ base_short_edge=int(shape["base_short_edge"]),
+ )
+ )
+ shape["geometry_source"] = geometry_source
+ shape["geometry_source_condition_index"] = int(source.condition_index)
+ if source.frame_index is not None:
+ shape["geometry_source_frame_index"] = int(source.frame_index)
+
+
+def _resolve_deferred_temporal_shape(
+ plan: MiniMaxH3ResolvedPlan,
+ shape: dict[str, Any],
+ probe_facts: dict[int, dict[str, Any]],
+) -> None:
+ if str(shape.get("temporal")) != "deferred_from_audio_reference":
+ return
+ sources = [
+ material
+ for material in plan.materials
+ if material.condition_type in {"audio", "video", "video_audio"}
+ and bool(probe_facts[int(material.condition_index)].get("has_audio"))
+ ]
+ if len(sources) != 1:
+ raise ValueError(
+ "audio-derived target duration requires exactly one probed "
+ f"condition with an audio stream, got {len(sources)}"
+ )
+ source = sources[0]
+ facts = probe_facts[int(source.condition_index)]
+ try:
+ duration_seconds = float(facts["audio_duration_seconds"]) - float(
+ source.start_time_seconds
+ )
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError("audio reference has no positive probed duration") from exc
+ if duration_seconds <= 0:
+ raise ValueError("audio reference has no positive probed duration")
+ if not (
+ MINIMAX_H3_MIN_DURATION_SECONDS
+ <= duration_seconds
+ <= MINIMAX_H3_MAX_DURATION_SECONDS
+ ):
+ raise ValueError(
+ "audio reference duration must be in "
+ f"[{MINIMAX_H3_MIN_DURATION_SECONDS:g}, "
+ f"{MINIMAX_H3_MAX_DURATION_SECONDS:g}] seconds, got {duration_seconds:g}"
+ )
+ fps = MINIMAX_H3_SUPPORTED_FPS
+ frame_count = minimax_h3_align_frame_count(int(round(duration_seconds * fps)))
+ aligned_duration = frame_count / fps
+ shape.update(
+ {
+ "temporal": "resolved_from_audio_reference",
+ "duration_seconds": aligned_duration,
+ "frame_count": frame_count,
+ "video_latent_t": minimax_h3_video_latent_t(frame_count),
+ "audio_latent_t": minimax_h3_audio_latent_t(aligned_duration),
+ }
+ )
+
+
+def _validate_reference_start_times(
+ plan: MiniMaxH3ResolvedPlan,
+ probe_facts: dict[int, dict[str, Any]],
+) -> None:
+ for material in plan.materials:
+ start_time_seconds = float(material.start_time_seconds)
+ if start_time_seconds == 0:
+ continue
+ condition_index = int(material.condition_index)
+ facts = probe_facts[condition_index]
+ try:
+ video_duration_seconds = float(facts["video_duration_seconds"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError(
+ f"conditions[{condition_index}].start_time_seconds requires "
+ "a video with a positive probed duration"
+ ) from exc
+ if start_time_seconds >= video_duration_seconds:
+ raise ValueError(
+ f"conditions[{condition_index}].start_time_seconds must be less "
+ f"than the video duration {video_duration_seconds:g}, got "
+ f"{start_time_seconds:g}"
+ )
+ if bool(facts.get("has_audio")):
+ try:
+ audio_duration_seconds = float(facts["audio_duration_seconds"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError(
+ f"conditions[{condition_index}] has no usable audio duration"
+ ) from exc
+ if start_time_seconds >= audio_duration_seconds:
+ raise ValueError(
+ f"conditions[{condition_index}].start_time_seconds must be "
+ f"less than the soundtrack duration {audio_duration_seconds:g}, "
+ f"got {start_time_seconds:g}"
+ )
+
+
+def _resolved_work_frame_count(
+ plan: MiniMaxH3ResolvedPlan,
+ shape: dict[str, Any],
+ probe_facts: dict[int, dict[str, Any]],
+) -> int:
+ if shape.get("frame_count") is not None:
+ return int(shape["frame_count"])
+ raise ValueError("MiniMax H3 target frame count remained unresolved before queue")
+
+
+def _preserve_prequeue_material_dirs(batch: Any) -> None:
+ temp_registry = batch.extra.get(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY)
+ if isinstance(temp_registry, dict) and "material" in temp_registry:
+ # Multi-output dispatch shallow-copies request extras. Keep the
+ # single localized source closure owned by the API request until
+ # every output finishes; encoder-stage "material" cleanup must not
+ # delete it after the first expanded output.
+ paths = temp_registry.pop("material")
+ prequeue_paths = temp_registry.setdefault("prequeue_material", [])
+ for path in paths if isinstance(paths, list) else ():
+ if path not in prequeue_paths:
+ prequeue_paths.append(path)
+
+
+def minimax_h3_prepare_for_queue(batch: Any) -> MiniMaxH3ResolvedPlan:
+ """Freeze MiniMax H3 media/shape facts before queue admission."""
+
+ try:
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ raise ValueError(
+ "MiniMax H3 pre-queue validation requires a canonical request"
+ )
+
+ probe_facts: dict[int, dict[str, Any]] = {}
+ for material in plan.materials:
+ probe_facts[int(material.condition_index)] = minimax_h3_probe_material(
+ batch,
+ material.uri,
+ condition_type=material.condition_type,
+ condition_index=int(material.condition_index),
+ )
+ batch.extra[MINIMAX_H3_PROBE_FACTS_EXTRA_KEY] = probe_facts
+
+ shape = dict(plan.shape)
+ _validate_reference_start_times(plan, probe_facts)
+ _resolve_deferred_spatial_shape(plan, shape, probe_facts)
+ _resolve_deferred_temporal_shape(plan, shape, probe_facts)
+ if str(shape.get("geometry")) != "resolved_v2":
+ raise ValueError(
+ "MiniMax H3 target geometry remained unresolved before queue"
+ )
+
+ material_shapes: dict[int, dict[str, Any]] = {}
+ for material in plan.materials:
+ condition_index = int(material.condition_index)
+ if material.material_chain == "image.target_canvas":
+ resolved = {
+ key: shape[key]
+ for key in (
+ "geometry",
+ "shape_policy_version",
+ "base_short_edge",
+ "effective_short_edge",
+ "size_mode",
+ "max_pixels",
+ "multiple",
+ "rounding",
+ "width",
+ "height",
+ )
+ if key in shape
+ }
+ elif material.material_chain in {
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+ }:
+ width, height = _display_shape(
+ probe_facts[condition_index],
+ label=f"conditions[{condition_index}]",
+ )
+ resolved = minimax_h3_resolve_spatial_shape(
+ width=width,
+ height=height,
+ base_short_edge=MINIMAX_H3_BASE_SHORT_EDGE,
+ )
+ elif material.material_chain == "image.reference_preserve":
+ width, height = _display_shape(
+ probe_facts[condition_index],
+ label=f"conditions[{condition_index}]",
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.reference_encoding import (
+ minimax_h3_resolve_reference_image_shape,
+ )
+
+ resolved = minimax_h3_resolve_reference_image_shape(
+ width=width,
+ height=height,
+ )
+ else:
+ continue
+ resolved = dict(resolved)
+ resolved["condition_index"] = condition_index
+ material_shapes[condition_index] = resolved
+ batch.extra[MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY] = material_shapes
+
+ work_frames = _resolved_work_frame_count(plan, shape, probe_facts)
+ resolved_plan = _replace_plan_shape(plan, shape)
+ batch.extra[MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY] = resolved_plan
+ # Req delegates these fields to SamplingParams. Freezing them here makes
+ # queue metadata and dynamic-batch signatures use the same final shape
+ # that the MiniMax H3 stages consume.
+ batch.width = int(shape["width"])
+ batch.height = int(shape["height"])
+ batch.fps = MINIMAX_H3_SUPPORTED_FPS
+ batch.num_frames = int(work_frames)
+ _preserve_prequeue_material_dirs(batch)
+ return resolved_plan
+ except Exception:
+ minimax_h3_cleanup_temp_dirs(batch)
+ batch.extra.pop(MINIMAX_H3_PROBE_FACTS_EXTRA_KEY, None)
+ batch.extra.pop(MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY, None)
+ raise
+
+
+__all__ = [
+ "MINIMAX_H3_PROBE_FACTS_EXTRA_KEY",
+ "MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY",
+ "minimax_h3_prepare_for_queue",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py
new file mode 100644
index 000000000..7473d6390
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py
@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 Qwen presentation building.
+
+Builds the positive presentation token stream:
+- fl2va: ': ' label + vision block (<|vision_start|> +
+ N*<|image_pad|> + <|vision_end|>) + prompt text.
+- t2va: prompt text only (no vision block).
+Prompt text passes through verbatim (no stripping or rewriting).
+
+All presentation variants are emitted through the shared ``_Presentation``
+accumulator so ids and AdaLN token tags cannot drift apart.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+
+VISION_START = "<|vision_start|>"
+VISION_END = "<|vision_end|>"
+IMAGE_PAD = "<|image_pad|>"
+VIDEO_PAD = "<|video_pad|>"
+
+_TEXT_TAG = 1
+_VIDEO_TAG = 0
+
+
+def _text_ids(tokenizer: Any, text: str) -> list[int]:
+ return list(tokenizer(text, add_special_tokens=False)["input_ids"])
+
+
+def _vision_block_ids(tokenizer: Any, pad_token: str, count: int) -> list[int]:
+ return (
+ [tokenizer.convert_tokens_to_ids(VISION_START)]
+ + [tokenizer.convert_tokens_to_ids(pad_token)] * int(count)
+ + [tokenizer.convert_tokens_to_ids(VISION_END)]
+ )
+
+
+class _Presentation:
+ """Accumulates aligned (ids, token_tags) presentation segments."""
+
+ def __init__(self) -> None:
+ self.ids: list[int] = []
+ self.tags: list[int] = []
+
+ def text(self, token_ids: list[int]) -> None:
+ self.ids += token_ids
+ self.tags += [_TEXT_TAG] * len(token_ids)
+
+ def vision(self, token_ids: list[int]) -> None:
+ self.ids += token_ids
+ self.tags += [_VIDEO_TAG] * len(token_ids)
+
+ def build(self) -> tuple[torch.Tensor, torch.Tensor]:
+ return (
+ torch.tensor(self.ids, dtype=torch.long),
+ torch.tensor(self.tags, dtype=torch.long),
+ )
+
+
+def _timestamped_video_blocks(
+ presentation: _Presentation,
+ tokenizer: Any,
+ *,
+ counts: Sequence[int],
+ timestamps: Sequence[float],
+ context: str,
+) -> None:
+ """Emit per-temporal-block ``<{t:.1f} seconds>`` text + VIDEO vision."""
+
+ counts = [int(value) for value in counts]
+ timestamps = [float(value) for value in timestamps]
+ if not counts or len(counts) != len(timestamps):
+ raise ValueError(f"{context}video block token counts and timestamps must align")
+ for count, timestamp in zip(counts, timestamps):
+ if count <= 0:
+ raise ValueError(f"{context}video block token count must be positive")
+ presentation.text(_text_ids(tokenizer, f"<{timestamp:.1f} seconds>"))
+ presentation.vision(_vision_block_ids(tokenizer, VIDEO_PAD, count))
+
+
+def minimax_h3_text_only_ids(tokenizer: Any, prompt: str) -> torch.Tensor:
+ """t2va presentation: verbatim prompt, no special tokens."""
+ if not prompt:
+ raise ValueError("prompt must be non-empty")
+ return torch.tensor(_text_ids(tokenizer, prompt), dtype=torch.long)
+
+
+def minimax_h3_multi_image_presentation(
+ tokenizer: Any,
+ *,
+ prompt: str,
+ image_token_counts: list[int],
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if not image_token_counts:
+ raise ValueError("image_token_counts must be non-empty")
+ presentation = _Presentation()
+ for index, count in enumerate(image_token_counts, start=1):
+ if int(count) <= 0:
+ raise ValueError("image_token_count must be positive")
+ presentation.text(_text_ids(tokenizer, f": "))
+ presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
+ presentation.text(_text_ids(tokenizer, prompt))
+ return presentation.build()
+
+
+def minimax_h3_ref2va_presentation(
+ tokenizer: Any,
+ *,
+ prompt: str,
+ condition_labels: list[tuple[str, int]],
+ image_token_count: int | list[int] | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """ref2va positive presentation:
+
+ per condition in request order — image i: ``: `` label followed
+ by the vision block; audio j: ``: `` label only (audio content
+ never enters Qwen) — then the verbatim prompt. Returns (ids, token_tags)
+ with the vision block tagged VIDEO(0) and everything else TEXT(1).
+
+ condition_labels: [("image", 1), ("audio", 1), ...] with 1-based ordinals
+ per type.
+ """
+ return minimax_h3_ref2va_video_presentation(
+ tokenizer,
+ prompt=prompt,
+ condition_labels=condition_labels,
+ image_token_count=image_token_count,
+ video_block_token_counts=None,
+ video_block_timestamps=None,
+ )
+
+
+def _as_int_list(value: int | Sequence[int] | None, *, name: str) -> list[int]:
+ if value is None:
+ return []
+ if isinstance(value, int):
+ return [int(value)]
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+ raise ValueError(f"{name} must be an int or a sequence of ints")
+ return [int(item) for item in value]
+
+
+def _as_nested_int_list(
+ value: Sequence[int] | Sequence[Sequence[int]] | None,
+ *,
+ name: str,
+) -> list[list[int]]:
+ if value is None:
+ return []
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+ raise ValueError(f"{name} must be a sequence")
+ if len(value) == 0:
+ return []
+ first = value[0]
+ if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
+ out: list[list[int]] = []
+ for group in value:
+ if not isinstance(group, Sequence) or isinstance(group, (str, bytes)):
+ raise ValueError(f"{name} must not mix nested and flat entries")
+ out.append([int(item) for item in group])
+ return out
+ return [[int(item) for item in value]]
+
+
+def _as_nested_float_list(
+ value: Sequence[float] | Sequence[Sequence[float]] | None,
+ *,
+ name: str,
+) -> list[list[float]]:
+ if value is None:
+ return []
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+ raise ValueError(f"{name} must be a sequence")
+ if len(value) == 0:
+ return []
+ first = value[0]
+ if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
+ out: list[list[float]] = []
+ for group in value:
+ if not isinstance(group, Sequence) or isinstance(group, (str, bytes)):
+ raise ValueError(f"{name} must not mix nested and flat entries")
+ out.append([float(item) for item in group])
+ return out
+ return [[float(item) for item in value]]
+
+
+def minimax_h3_ref2va_video_presentation(
+ tokenizer: Any,
+ *,
+ prompt: str,
+ condition_labels: list[tuple[str, int]],
+ image_token_count: int | list[int] | None,
+ video_block_token_counts: list[int] | list[list[int]] | None,
+ video_block_timestamps: list[float] | list[list[float]] | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """ref2va (optionally with video refs) positive presentation:
+
+ per condition in request order —
+ - image i: ``: `` label + one image vision block;
+ - audio j: ``: `` label only (audio content never enters Qwen);
+ - video k: ``: `` label, then per temporal block a timestamp
+ text ``<{t:.1f} seconds>`` followed by a VIDEO vision block
+ (<|vision_start|> + <|video_pad|> x n + <|vision_end|>). Timestamps are
+ the mean of each merged frame pair (Qwen3VL temporal merge 2; odd frame
+ counts repeat the last frame), emitting the
+ ``<0.2 seconds>`` ..
+ ``<4.0 seconds>`` sequence — note Python bankers-rounding at .1f.
+ then the verbatim prompt. Vision blocks are tagged VIDEO(0), everything
+ else TEXT(1).
+ """
+ if not prompt:
+ raise ValueError("prompt must be non-empty")
+ presentation = _Presentation()
+ image_token_counts = _as_int_list(image_token_count, name="image_token_count")
+ video_counts_by_ref = _as_nested_int_list(
+ video_block_token_counts,
+ name="video_block_token_counts",
+ )
+ video_timestamps_by_ref = _as_nested_float_list(
+ video_block_timestamps,
+ name="video_block_timestamps",
+ )
+ if len(video_counts_by_ref) != len(video_timestamps_by_ref):
+ raise ValueError("video block token counts and timestamps must align")
+ image_seen = 0
+ video_seen = 0
+ for cond_type, ordinal in condition_labels:
+ if cond_type == "image":
+ image_seen += 1
+ if image_seen > len(image_token_counts):
+ raise ValueError("image_token_count required for an image reference")
+ count = int(image_token_counts[image_seen - 1])
+ if count <= 0:
+ raise ValueError("image_token_count required for an image reference")
+ presentation.text(_text_ids(tokenizer, f": "))
+ presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
+ elif cond_type == "audio":
+ presentation.text(_text_ids(tokenizer, f": "))
+ elif cond_type == "video":
+ video_seen += 1
+ if video_seen > len(video_counts_by_ref):
+ raise ValueError(
+ "video reference requires block token counts and timestamps"
+ )
+ counts = video_counts_by_ref[video_seen - 1]
+ timestamps = video_timestamps_by_ref[video_seen - 1]
+ if not counts or not timestamps:
+ raise ValueError(
+ "video reference requires block token counts and timestamps"
+ )
+ presentation.text(_text_ids(tokenizer, f": "))
+ _timestamped_video_blocks(
+ presentation,
+ tokenizer,
+ counts=counts,
+ timestamps=timestamps,
+ context="",
+ )
+ else:
+ raise ValueError(f"unsupported ref2va condition type {cond_type!r}")
+ if image_seen != len(image_token_counts):
+ raise ValueError("unused image_token_count entries")
+ if video_seen != len(video_counts_by_ref):
+ raise ValueError("unused video block token count entries")
+ presentation.text(_text_ids(tokenizer, prompt))
+ return presentation.build()
+
+
+__all__ = [
+ "minimax_h3_multi_image_presentation",
+ "minimax_h3_ref2va_presentation",
+ "minimax_h3_ref2va_video_presentation",
+ "minimax_h3_text_only_ids",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py
new file mode 100644
index 000000000..311581398
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py
@@ -0,0 +1,741 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 ref2va reference-material encoding.
+
+Encoding recipes for user-provided reference materials:
+
+- image reference: independent 2048px short-edge resize with upscale enabled,
+ LANCZOS, and nearest-32 dimensions, then the SAME keyframe tokenizer recipe
+ as fl2va (seed-42 sampled encode, normalize, [1,2,2] patchify);
+- audio reference: the audio material chain (pure
+ audio is losslessly normalized
+ to stereo; video soundtracks are extracted as 44.1 kHz stereo), then a
+ single resample to 32 kHz,
+ audio VAE posterior MEAN (encoder -> optional pre_block -> mean_proj; no
+ sampling), canonical [2, T, 32], normalize with loader-injected audio stats,
+ channel-major rows.
+"""
+
+from __future__ import annotations
+
+import functools
+import math
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
+ MiniMaxH3AudioVAEArchConfig,
+)
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
+ MiniMaxH3VideoVAEArchConfig,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_SUPPORTED_FPS,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.keyframe_encoding import (
+ _cached_latent_mean_std,
+ minimax_h3_scoped_encode_rng,
+)
+
+MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048
+MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32
+MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
+MINIMAX_H3_AUDIO_CHANNELS = 2
+
+
+class _AudioVAEDeterminismContext:
+ """Scoped determinism config for the audio encode.
+
+ Disables TF32, forces deterministic algorithms, DISABLES cuDNN entirely
+ for the encode (convs run on the fallback kernels), and pins SDP to the
+ math backend. This configuration is required for a deterministic encode.
+ Everything is restored on exit
+ so the decode path keeps its own configuration.
+
+ Reentrant via a shared depth counter: a caller encoding several reference
+ materials in one request (audio_encoding.py's per-material loop) can wrap
+ the whole loop in one of these: only the outermost enter/exit actually
+ touches torch.backends, and each per-material call's own nested
+ with-block becomes a no-op increment/decrement instead of redundantly
+ saving and restoring the same flags per material.
+ """
+
+ _depth = 0
+ _saved: tuple | None = None
+
+ def __enter__(self):
+ if _AudioVAEDeterminismContext._depth == 0:
+ b = torch.backends
+ _AudioVAEDeterminismContext._saved = (
+ b.cuda.matmul.allow_tf32,
+ b.cudnn.allow_tf32,
+ b.cudnn.benchmark,
+ b.cudnn.deterministic,
+ b.cudnn.enabled,
+ b.cuda.flash_sdp_enabled(),
+ b.cuda.mem_efficient_sdp_enabled(),
+ b.cuda.math_sdp_enabled(),
+ )
+ b.cuda.matmul.allow_tf32 = False
+ b.cudnn.allow_tf32 = False
+ b.cudnn.benchmark = False
+ b.cudnn.deterministic = True
+ b.cudnn.enabled = False
+ b.cuda.enable_flash_sdp(False)
+ b.cuda.enable_mem_efficient_sdp(False)
+ b.cuda.enable_math_sdp(True)
+ _AudioVAEDeterminismContext._depth += 1
+ return self
+
+ def __exit__(self, exc_type, exc, tb):
+ _AudioVAEDeterminismContext._depth -= 1
+ if _AudioVAEDeterminismContext._depth == 0:
+ b = torch.backends
+ (
+ b.cuda.matmul.allow_tf32,
+ b.cudnn.allow_tf32,
+ b.cudnn.benchmark,
+ b.cudnn.deterministic,
+ b.cudnn.enabled,
+ flash,
+ mem_eff,
+ math_sdp,
+ ) = _AudioVAEDeterminismContext._saved
+ b.cuda.enable_flash_sdp(flash)
+ b.cuda.enable_mem_efficient_sdp(mem_eff)
+ b.cuda.enable_math_sdp(math_sdp)
+ _AudioVAEDeterminismContext._saved = None
+
+
+def _nearest_multiple(value: float, multiple: int) -> int:
+ return max(multiple, int(round(float(value) / multiple)) * multiple)
+
+
+def minimax_h3_resolve_reference_image_shape(
+ *,
+ width: int | float,
+ height: int | float,
+) -> dict[str, Any]:
+ """Resolve a ref2va image independently from the target canvas.
+
+ The image keeps its display ratio, always targets a 2048px short edge (even
+ when that requires upscaling), and rounds both dimensions independently to
+ the nearest 32px grid. Unlike target/video ``adapt_shape_v1``, reference
+ images have no area-cap branch.
+ """
+
+ try:
+ source_width = float(width)
+ source_height = float(height)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(
+ "reference image width and height must be positive finite numbers"
+ ) from exc
+ if (
+ not math.isfinite(source_width)
+ or not math.isfinite(source_height)
+ or source_width <= 0.0
+ or source_height <= 0.0
+ ):
+ raise ValueError(
+ "reference image width and height must be positive finite numbers"
+ )
+ if source_width > 4.0 * source_height or source_height > 4.0 * source_width:
+ raise ValueError(
+ "reference image ratio must be within the inclusive range "
+ f"1:4 to 4:1, got {source_width:g}x{source_height:g}"
+ )
+
+ scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(source_width, source_height)
+ target_width = _nearest_multiple(
+ source_width * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE
+ )
+ target_height = _nearest_multiple(
+ source_height * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE
+ )
+ return {
+ "geometry": "reference_image_resolved",
+ "shape_policy_version": "reference_image_short_edge_v1",
+ "base_short_edge": MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE,
+ "effective_short_edge": min(target_width, target_height),
+ "size_mode": "short_edge",
+ "multiple": MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE,
+ "rounding": "nearest",
+ "allow_upscale": True,
+ "width": target_width,
+ "height": target_height,
+ }
+
+
+def minimax_h3_resize_reference_image(
+ image: Any,
+ *,
+ target_width: int,
+ target_height: int,
+) -> Any:
+ """Resize a reference image to the shape fixed by pre-queue admission."""
+
+ from PIL import Image
+
+ if target_width <= 0 or target_height <= 0:
+ raise ValueError("reference image target dimensions must be positive")
+ if (
+ target_width % MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE
+ or target_height % MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE
+ ):
+ raise ValueError(
+ "reference image target dimensions must be aligned to "
+ f"{MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE}"
+ )
+ image = image.convert("RGB")
+ if (target_width, target_height) == image.size:
+ return image
+ return image.resize((target_width, target_height), Image.Resampling.LANCZOS)
+
+
+def _load_waveform(
+ path: str,
+ *,
+ material_chain: str = "audio",
+ max_duration_seconds: float | None = None,
+ start_time_seconds: float = 0.0,
+ source_sample_rate: int | None = None,
+) -> tuple[torch.Tensor, int]:
+ """Apply the audio material chain.
+
+ Pure-audio references preserve their source rate while normalizing to
+ stereo. Video-bearing references first extract 44.1 kHz stereo PCM. The
+ audio VAE boundary then performs the single 32 kHz resample below. ffmpeg
+ writes bounded interleaved float PCM directly to stdout, avoiding a
+ temporary lossless file plus a second decode.
+ """
+
+ import subprocess
+
+ import numpy as np
+
+ if max_duration_seconds is not None:
+ max_duration_seconds = float(max_duration_seconds)
+ if not math.isfinite(max_duration_seconds) or max_duration_seconds <= 0:
+ raise ValueError("reference audio duration bound must be positive")
+ start_time_seconds = float(start_time_seconds)
+ if not math.isfinite(start_time_seconds) or start_time_seconds < 0:
+ raise ValueError("reference audio start time must be non-negative")
+
+ if material_chain == "audio":
+ if source_sample_rate is None or int(source_sample_rate) <= 0:
+ raise ValueError("reference audio sample rate must be positive")
+ source_rate = int(source_sample_rate)
+ elif material_chain in {
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+ }:
+ source_rate = 44100
+ else:
+ raise ValueError(
+ f"unsupported MiniMax H3 audio material chain {material_chain!r}"
+ )
+
+ command = [
+ "ffmpeg",
+ "-v",
+ "error",
+ ]
+ if start_time_seconds > 0:
+ command += ["-ss", f"{start_time_seconds:.9g}"]
+ command += [
+ "-i",
+ str(path),
+ "-map",
+ "0:a:0",
+ "-vn",
+ "-ac",
+ str(MINIMAX_H3_AUDIO_CHANNELS),
+ ]
+ if material_chain != "audio":
+ command += ["-ar", str(source_rate)]
+ if max_duration_seconds is not None:
+ command += ["-t", f"{max_duration_seconds:.9g}"]
+ command += ["-f", "f32le", "pipe:1"]
+ decoded = subprocess.run(command, check=True, capture_output=True)
+ payload = decoded.stdout
+ if not isinstance(payload, bytes):
+ raise TypeError("ffmpeg float PCM output must be bytes")
+ frame_bytes = MINIMAX_H3_AUDIO_CHANNELS * torch.float32.itemsize
+ if len(payload) % frame_bytes:
+ raise ValueError(
+ "ffmpeg returned a partial reference-audio sample frame: "
+ f"{len(payload)} bytes"
+ )
+ waveform = torch.from_numpy(
+ np.frombuffer(payload, dtype=np.float32)
+ .reshape(-1, MINIMAX_H3_AUDIO_CHANNELS)
+ .T.copy()
+ )
+ return waveform, source_rate
+
+
+@functools.lru_cache(maxsize=8)
+def _audio_resampler(source_rate: int):
+ import torchaudio
+
+ return torchaudio.transforms.Resample(source_rate, MINIMAX_H3_AUDIO_SAMPLE_RATE)
+
+
+@torch.inference_mode()
+def minimax_h3_encode_reference_audio_rows(
+ audio_vae: Any,
+ audio_path: str,
+ arch_config: MiniMaxH3AudioVAEArchConfig,
+ *,
+ material_chain: str = "audio",
+ max_duration_seconds: float | None = None,
+ start_time_seconds: float = 0.0,
+ source_sample_rate: int | None = None,
+) -> dict[str, Any]:
+ """Encode a reference audio file into normalized channel-major rows.
+
+ Returns {"rows": [2*T, 32] fp32 cpu, "ref_audio_t": T,
+ "duration_seconds": float}.
+ """
+ model = audio_vae
+ device = next(model.parameters()).device
+ waveform, source_rate = _load_waveform(
+ audio_path,
+ material_chain=material_chain,
+ max_duration_seconds=max_duration_seconds,
+ start_time_seconds=start_time_seconds,
+ source_sample_rate=source_sample_rate,
+ )
+ if waveform.numel() == 0:
+ raise ValueError(f"reference audio is empty: {audio_path}")
+ if int(source_rate) != MINIMAX_H3_AUDIO_SAMPLE_RATE:
+ waveform = _audio_resampler(int(source_rate))(waveform)
+ waveform = waveform.to(device)
+
+ with _AudioVAEDeterminismContext():
+ audio_data = model.preprocess(
+ waveform.unsqueeze(1), MINIMAX_H3_AUDIO_SAMPLE_RATE
+ )
+ z = model.encoder(audio_data)
+ if bool(getattr(model, "attn_proj", False)):
+ z = model.pre_block(z.transpose(1, 2)).transpose(1, 2)
+ if not hasattr(model, "mean_proj"):
+ raise AttributeError(
+ "audio VAE model must expose mean_proj for deterministic mean encoding"
+ )
+ latent = model.mean_proj(z).float() # [2, 32, T] or [2, T, 32]
+ if latent.ndim != 3:
+ raise ValueError(f"expected 3D audio latent, got {list(latent.shape)}")
+ latent_channels = arch_config.latent_channels
+ if int(latent.shape[-1]) != latent_channels:
+ if int(latent.shape[1]) != latent_channels:
+ raise ValueError(f"cannot canonicalize audio latent {list(latent.shape)}")
+ latent = latent.transpose(1, 2).contiguous() # -> [2, T, 32]
+ latent = latent.cpu()
+
+ mean, std = _cached_latent_mean_std(
+ tuple(arch_config.latents_mean),
+ tuple(arch_config.latents_std),
+ (1, 1, latent_channels),
+ )
+ latent.sub_(mean).div_(std)
+ rows = latent.reshape(-1, latent_channels).to(torch.float32).contiguous()
+ ref_audio_t = int(latent.shape[1])
+ return {
+ "rows": rows,
+ "ref_audio_t": ref_audio_t,
+ "duration_seconds": float(waveform.shape[-1])
+ / float(MINIMAX_H3_AUDIO_SAMPLE_RATE),
+ }
+
+
+MINIMAX_H3_PREPARED_REFERENCE_IMAGE_EXTRA_KEY = "minimax_h3_prepared_reference_image"
+
+
+def minimax_h3_decode_reference_video_frames(
+ video_path: str,
+ *,
+ target_width: int,
+ target_height: int,
+ target_frame_count: int,
+ fps: float = MINIMAX_H3_SUPPORTED_FPS,
+ start_time_seconds: float = 0.0,
+) -> Any:
+ """Decode, transform, and truncate a reference video in one ffmpeg pass.
+
+ ffmpeg applies display rotation, CFR sampling, direct Lanczos scaling, and
+ square-pixel normalization before writing bounded RGB24 frames to stdout.
+ The returned array is shared by Qwen and the visual VAE, so conditioning
+ never passes through a lossy x264 intermediate or a second video decode.
+ """
+ import subprocess
+
+ import numpy as np
+
+ if target_frame_count <= 0:
+ raise ValueError("target_frame_count must be positive")
+ if target_width <= 0 or target_height <= 0:
+ raise ValueError("target reference-video dimensions must be positive")
+ if not math.isfinite(float(fps)) or float(fps) <= 0:
+ raise ValueError("reference-video fps must be positive")
+ start_time_seconds = float(start_time_seconds)
+ if not math.isfinite(start_time_seconds) or start_time_seconds < 0:
+ raise ValueError("reference-video start time must be non-negative")
+
+ filters = (
+ f"fps={float(fps):g},"
+ f"scale={target_width}:{target_height}:flags=lanczos,"
+ "setsar=1"
+ )
+ command = ["ffmpeg", "-v", "error"]
+ if start_time_seconds > 0:
+ # Input seeking remains accurate while transcoding and avoids decoding
+ # the unused prefix of a long reference into RGB frames.
+ command += ["-ss", f"{start_time_seconds:.9g}"]
+ command += [
+ "-i",
+ str(video_path),
+ "-map",
+ "0:v:0",
+ "-an",
+ "-vf",
+ filters,
+ "-frames:v",
+ str(target_frame_count),
+ "-f",
+ "rawvideo",
+ "-pix_fmt",
+ "rgb24",
+ "pipe:1",
+ ]
+ decoded = subprocess.run(
+ command,
+ check=True,
+ capture_output=True,
+ )
+ payload = decoded.stdout
+ if not isinstance(payload, bytes):
+ raise TypeError("ffmpeg RGB24 output must be bytes")
+ frame_bytes = target_width * target_height * 3
+ if len(payload) % frame_bytes:
+ raise ValueError(
+ "ffmpeg returned a partial reference-video frame: "
+ f"{len(payload)} bytes for {target_width}x{target_height} RGB24"
+ )
+ frame_count = len(payload) // frame_bytes
+ if frame_count <= 0:
+ raise ValueError(f"reference video has no frames: {video_path}")
+ return np.frombuffer(payload, dtype=np.uint8).reshape(
+ frame_count, target_height, target_width, 3
+ )
+
+
+MINIMAX_H3_REFERENCE_VIDEO_ENCODE_SEED = 42
+MINIMAX_H3_REFERENCE_VIDEO_PATCH_SIZE = (1, 2, 2)
+
+
+@torch.inference_mode()
+def minimax_h3_encode_reference_video_rows(
+ video_vae: Any,
+ frames: Any,
+ arch_config: MiniMaxH3VideoVAEArchConfig,
+) -> tuple[torch.Tensor, int, int, int]:
+ """Encode transformed reference-video frames into packed imgvid cond rows.
+
+ Frames come from the request's single ffmpeg transformation pass, then use
+ the SAME ``encode_videos`` recipe as the fl2va keyframe sink (fp32 weights,
+ configured complete-tile parallelism, torch seed pinned at 42 because the
+ encode SAMPLES the DiagonalGaussian, fp16 latent), then normalize
+ and [1,2,2]-patchify. The VAE's clip_length=17 / token_drop=3 give the
+ 17-frames-per-5-latents temporal grouping (107 frames -> 32 latents).
+
+ Returns (rows [n, 96] fp32 cpu, latent_t, latent_h, latent_w).
+ """
+ import numpy as np
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
+ minimax_h3_patchify_video_latent,
+ )
+
+ frames = np.asarray(frames)
+ if (
+ frames.ndim != 4
+ or int(frames.shape[-1]) != 3
+ or frames.dtype != np.uint8
+ or int(frames.shape[0]) <= 0
+ ):
+ raise ValueError(
+ "reference-video frames must be non-empty [T,H,W,3] uint8, got "
+ f"shape={list(frames.shape)}, dtype={frames.dtype}"
+ )
+
+ parameter = next(video_vae.parameters())
+ prev_dtype = parameter.dtype
+ if prev_dtype != torch.float32:
+ video_vae.to(torch.float32)
+ try:
+ with minimax_h3_scoped_encode_rng(
+ MINIMAX_H3_REFERENCE_VIDEO_ENCODE_SEED, parameter.device
+ ):
+ z = video_vae.encode_videos(frames, use_fp16_latent=True)[0]
+ finally:
+ if prev_dtype != torch.float32:
+ video_vae.to(prev_dtype)
+ z = z.cpu().float()
+ if z.dim() == 4:
+ z = z[None]
+ latent_channels = arch_config.latent_channels
+ if z.dim() != 5 or int(z.shape[1]) != latent_channels:
+ raise ValueError(f"unexpected reference video latent shape {list(z.shape)}")
+ latent_t, latent_h, latent_w = int(z.shape[2]), int(z.shape[3]), int(z.shape[4])
+ mean, std = _cached_latent_mean_std(
+ tuple(arch_config.latents_mean),
+ tuple(arch_config.latents_std),
+ (1, latent_channels, 1, 1, 1),
+ )
+ z.sub_(mean).div_(std)
+ rows = minimax_h3_patchify_video_latent(
+ z, patch_size=list(MINIMAX_H3_REFERENCE_VIDEO_PATCH_SIZE)
+ )
+ return rows.to(torch.float32), latent_t, latent_h, latent_w
+
+
+MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS = 2.0
+MINIMAX_H3_QWEN_TEMPORAL_PATCH = 2
+
+
+def minimax_h3_sample_reference_video_frames(
+ frames: Any,
+) -> dict[str, Any]:
+ """Sample Qwen frames from the shared transformed RGB array.
+
+ Frame-sampling recipe (24 FPS -> 2 FPS strided view) plus the qwen3
+ timestamp rule (indices padded to the
+ temporal patch size with the last frame, block ts = mean of the pair at
+ sample fps; text is rendered later with ``f"<{ts:.1f} seconds>"``).
+
+ Returns {"frames": np.ndarray TxHxWx3 u8, "block_timestamps": [float]}.
+ """
+ import numpy as np
+
+ frames = np.asarray(frames)
+ if frames.ndim != 4 or int(frames.shape[0]) <= 0:
+ raise ValueError(
+ "Qwen reference-video sampling requires non-empty [T,H,W,C] frames"
+ )
+ sample_stride = int(MINIMAX_H3_SUPPORTED_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS)
+ sampled_frames = frames[::sample_stride]
+ ts = [
+ i / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS
+ for i in range(int(sampled_frames.shape[0]))
+ ]
+ pad = (-len(ts)) % MINIMAX_H3_QWEN_TEMPORAL_PATCH
+ ts = ts + [ts[-1]] * pad
+ block_timestamps = [
+ (ts[i] + ts[i + MINIMAX_H3_QWEN_TEMPORAL_PATCH - 1]) / 2
+ for i in range(0, len(ts), MINIMAX_H3_QWEN_TEMPORAL_PATCH)
+ ]
+ return {"frames": sampled_frames, "block_timestamps": block_timestamps}
+
+
+def _reference_video_materials(plan: Any) -> list[Any]:
+ return [
+ m
+ for m in plan.materials
+ if m.material_chain
+ in {"video.reference_preserve", "video_audio.reference_preserve"}
+ ]
+
+
+def _reference_video_target_frame_count(
+ *,
+ plan: Any,
+) -> int:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+ minimax_h3_align_frame_count,
+ minimax_h3_frame_count_from_video_latent_t,
+ )
+
+ shape = plan.shape
+ fps = int(shape["fps"])
+ duration = shape.get("duration_seconds")
+ if duration is not None:
+ return minimax_h3_align_frame_count(int(round(float(duration) * fps)))
+ if shape.get("video_latent_t") is not None:
+ return minimax_h3_frame_count_from_video_latent_t(int(shape["video_latent_t"]))
+ raise ValueError(
+ "reference-video preparation requires pre-queue resolved temporal dimensions"
+ )
+
+
+def minimax_h3_prepared_reference_videos(batch: Any, plan: Any) -> dict[str, Any]:
+ """Decode the bounded reference-video RGB frames once per request.
+
+ BOTH the visual-condition tokenizer and Qwen consume the same transformed
+ array. Its frame cap comes from the resolved target duration (17n+5 rule).
+ The original path travels alongside for direct soundtrack decoding.
+ """
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY,
+ )
+
+ cached = batch.extra.get(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY)
+ if cached is not None:
+ return cached
+ videos = _reference_video_materials(plan)
+ if not videos:
+ raise NotImplementedError(
+ "ref2va video preparation requires a video or video_audio reference"
+ )
+
+ prepared_videos = []
+ for material in videos:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_localize_material_uri,
+ )
+
+ video_path = minimax_h3_localize_material_uri(
+ batch,
+ material.uri,
+ condition_type=material.condition_type,
+ condition_index=int(material.condition_index),
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.prequeue import (
+ MINIMAX_H3_PROBE_FACTS_EXTRA_KEY,
+ MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY,
+ )
+
+ condition_index = int(material.condition_index)
+ source_facts = batch.extra.get(MINIMAX_H3_PROBE_FACTS_EXTRA_KEY, {}).get(
+ condition_index
+ )
+ resolved_material_shape = batch.extra.get(
+ MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY, {}
+ ).get(condition_index)
+ if not isinstance(source_facts, dict) or not isinstance(
+ resolved_material_shape, dict
+ ):
+ raise ValueError(
+ "reference-video preparation requires cached pre-queue probe "
+ f"and shape facts for conditions[{condition_index}]"
+ )
+ input_has_audio = bool(source_facts.get("has_audio"))
+ target_frames = _reference_video_target_frame_count(plan=plan)
+ frames = minimax_h3_decode_reference_video_frames(
+ video_path,
+ target_width=int(resolved_material_shape["width"]),
+ target_height=int(resolved_material_shape["height"]),
+ target_frame_count=target_frames,
+ fps=float(plan.shape["fps"]),
+ start_time_seconds=float(material.start_time_seconds),
+ )
+ prepared_videos.append(
+ {
+ "frames": frames,
+ "original_path": video_path,
+ "target_frame_count": target_frames,
+ "frame_count": int(frames.shape[0]),
+ "condition_index": int(material.condition_index),
+ "material_chain": str(material.material_chain),
+ "start_time_seconds": float(material.start_time_seconds),
+ "input_has_audio": input_has_audio,
+ "width": int(resolved_material_shape["width"]),
+ "height": int(resolved_material_shape["height"]),
+ }
+ )
+ prepared = {
+ key: value for key, value in prepared_videos[0].items() if key != "frames"
+ }
+ prepared["videos"] = prepared_videos
+ batch.extra[MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY] = prepared
+ return prepared
+
+
+def minimax_h3_prepared_reference_image(batch: Any, plan: Any) -> dict[str, Any]:
+ """Resize ref2va image references to their pre-queue-resolved shapes.
+
+ Qwen (pixel_values) and the visual-condition tokenizer consume the identical
+ prepared image. The runtime never recomputes geometry from ``plan.shape``;
+ it consumes the per-material width/height admitted before queueing.
+ """
+ cached = batch.extra.get(MINIMAX_H3_PREPARED_REFERENCE_IMAGE_EXTRA_KEY)
+ if cached is not None:
+ return cached
+ images = [
+ m for m in plan.materials if m.material_chain == "image.reference_preserve"
+ ]
+ if not images:
+ raise ValueError("ref2va requires at least one image reference")
+ from PIL import Image, ImageOps
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_localize_material_uri,
+ )
+
+ prepared_images = []
+ for material in images:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.prequeue import (
+ MINIMAX_H3_PROBE_FACTS_EXTRA_KEY,
+ MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY,
+ )
+
+ condition_index = int(material.condition_index)
+ source_facts = batch.extra.get(MINIMAX_H3_PROBE_FACTS_EXTRA_KEY, {}).get(
+ condition_index
+ )
+ resolved_shape = batch.extra.get(
+ MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY, {}
+ ).get(condition_index)
+ if not isinstance(source_facts, dict) or not isinstance(resolved_shape, dict):
+ raise ValueError(
+ "reference-image preparation requires cached pre-queue probe "
+ f"and shape facts for conditions[{condition_index}]"
+ )
+ image_path = minimax_h3_localize_material_uri(
+ batch,
+ material.uri,
+ condition_type=material.condition_type,
+ condition_index=condition_index,
+ )
+ with Image.open(image_path) as source_image:
+ image = ImageOps.exif_transpose(source_image).convert("RGB")
+ expected_size = (
+ int(resolved_shape["width"]),
+ int(resolved_shape["height"]),
+ )
+ prepared_image = minimax_h3_resize_reference_image(
+ image,
+ target_width=expected_size[0],
+ target_height=expected_size[1],
+ )
+ if prepared_image.size != expected_size:
+ raise ValueError(
+ "reference image preparation disagrees with pre-queue shape: "
+ f"expected={expected_size}, actual={prepared_image.size}"
+ )
+ prepared_images.append(
+ {
+ "image": prepared_image,
+ "condition_index": condition_index,
+ }
+ )
+ prepared = {
+ # single-image consumers keep the existing keys
+ "image": prepared_images[0]["image"],
+ "condition_index": prepared_images[0]["condition_index"],
+ "images": prepared_images,
+ }
+ batch.extra[MINIMAX_H3_PREPARED_REFERENCE_IMAGE_EXTRA_KEY] = prepared
+ return prepared
+
+
+__all__ = [
+ "minimax_h3_decode_reference_video_frames",
+ "minimax_h3_encode_reference_audio_rows",
+ "minimax_h3_encode_reference_video_rows",
+ "minimax_h3_prepared_reference_image",
+ "minimax_h3_prepared_reference_videos",
+ "minimax_h3_resolve_reference_image_shape",
+ "minimax_h3_sample_reference_video_frames",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py
new file mode 100644
index 000000000..14e375a3d
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py
@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Public MiniMax H3 model-index admission contract."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_QUALITY_PROFILES,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ canonical_minimax_h3_task,
+ partition_for_task,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+_MINIMAX_H3_QUALITY_WORKLOAD = {
+ "task": "t2va",
+ "width": 1344,
+ "height": 768,
+ "fps": 24,
+ "frame_count": 124,
+ "num_inference_steps": 50,
+ "flow_shift": 12.0,
+ "audio_flow_shift": 3.0,
+}
+
+
+def _string_list(value: Any, path: str) -> tuple[str, ...]:
+ if not isinstance(value, list) or not value:
+ raise ValueError(f"{path} must be a non-empty list")
+ values = tuple(value)
+ if any(not isinstance(item, str) or not item for item in values):
+ raise ValueError(f"{path} must contain non-empty strings")
+ if len(set(values)) != len(values):
+ raise ValueError(f"{path} must not contain duplicates")
+ return values
+
+
+@dataclass(frozen=True)
+class MiniMaxH3ReleaseMetadata:
+ schema_version: int
+ partition: str
+ tasks: tuple[str, ...]
+ task_aliases: Mapping[str, str]
+ video_sigma_shift: float
+ audio_sigma_shift: float
+
+ @classmethod
+ def from_model_index(
+ cls, model_index: Mapping[str, Any]
+ ) -> MiniMaxH3ReleaseMetadata:
+ raw = model_index.get("_minimax_h3")
+ if not isinstance(raw, Mapping):
+ raise ValueError("model_index.json._minimax_h3 must be an object")
+ if raw.get("schema_version") != 1:
+ raise ValueError("model_index.json._minimax_h3.schema_version must be 1")
+ partition = raw.get("partition")
+ if partition not in {"fl2va", "ref2va"}:
+ raise ValueError(
+ "model_index.json._minimax_h3.partition must be one of " "fl2va, ref2va"
+ )
+ tasks = _string_list(raw.get("tasks"), "model_index.json._minimax_h3.tasks")
+ aliases = raw.get("task_aliases", {})
+ if not isinstance(aliases, Mapping) or any(
+ not isinstance(key, str)
+ or not key
+ or not isinstance(value, str)
+ or not value
+ for key, value in aliases.items()
+ ):
+ raise ValueError(
+ "model_index.json._minimax_h3.task_aliases must map strings to strings"
+ )
+ scales = raw.get("sigma_shift_scales")
+ if not isinstance(scales, Mapping):
+ raise ValueError(
+ "model_index.json._minimax_h3.sigma_shift_scales must be an object"
+ )
+ try:
+ video_sigma = float(scales["video"])
+ audio_sigma = float(scales["audio"])
+ except (KeyError, TypeError, ValueError) as exc:
+ raise ValueError(
+ "model_index.json._minimax_h3.sigma_shift_scales requires numeric "
+ "video and audio values"
+ ) from exc
+ metadata = cls(
+ schema_version=1,
+ partition=partition,
+ tasks=tasks,
+ task_aliases=dict(aliases),
+ video_sigma_shift=video_sigma,
+ audio_sigma_shift=audio_sigma,
+ )
+ for task in metadata.tasks:
+ if canonical_minimax_h3_task(task) != task:
+ raise ValueError(
+ f"tasks must contain canonical task names, got {task!r}"
+ )
+ if partition_for_task(task) != partition:
+ raise ValueError(
+ f"task {task!r} does not belong to partition {partition!r}"
+ )
+ for alias, target in metadata.task_aliases.items():
+ if target not in metadata.tasks:
+ raise ValueError(
+ f"task alias {alias!r} targets undeclared task {target!r}"
+ )
+ if canonical_minimax_h3_task(alias) != target:
+ raise ValueError(
+ f"unsupported task alias mapping {alias!r} -> {target!r}"
+ )
+ return metadata
+
+ @property
+ def sigma_shift_scales(self) -> dict[str, float]:
+ return {"video": self.video_sigma_shift, "audio": self.audio_sigma_shift}
+
+ def canonical_task(self, task: str) -> str:
+ normalized = task.strip().lower()
+ canonical = self.task_aliases.get(normalized, normalized)
+ if canonical not in self.tasks:
+ raise ValueError(
+ f"task {task!r} is not served by MiniMax H3 partition {self.partition!r}; "
+ f"supported tasks: {list(self.tasks)!r}"
+ )
+ if partition_for_task(canonical) != self.partition:
+ raise ValueError(
+ f"task {task!r} resolves outside partition {self.partition!r}"
+ )
+ return canonical
+
+
+class MiniMaxH3PartitionAdmissionStage(PipelineStage):
+ def __init__(self, metadata: MiniMaxH3ReleaseMetadata) -> None:
+ super().__init__()
+ self.metadata = metadata
+
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ task = None if batch.sampling_params is None else batch.sampling_params.task
+ if not isinstance(task, str) or not task.strip():
+ raise ValueError("MiniMax H3 request task must be a non-empty string")
+ self.metadata.canonical_task(task)
+ quality = getattr(batch.sampling_params, "quality", "lossless")
+ if quality not in MINIMAX_H3_QUALITY_PROFILES:
+ raise ValueError(
+ f"unsupported MiniMax-H3 quality profile {quality!r}; supported: "
+ f"{list(MINIMAX_H3_QUALITY_PROFILES)}"
+ )
+ approximate = quality != "lossless"
+ attention_backend = str(server_args.attention_backend or "").strip().lower()
+ if attention_backend == "sage_attn" and not batch.is_warmup:
+ raise ValueError(
+ "MiniMax-H3 does not support SageAttention: the current packed "
+ "varlen path does not preserve model output"
+ )
+ if approximate and not batch.is_warmup:
+ server_args.pipeline_config.validate_quality_deployment(server_args)
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ raise ValueError(
+ "MiniMax-H3 approximate quality profiles require a resolved "
+ "request plan"
+ )
+ shape = plan.shape
+ actual = {
+ "task": plan.task,
+ "width": int(shape["width"]),
+ "height": int(shape["height"]),
+ "fps": int(shape["fps"]),
+ "frame_count": int(shape["frame_count"]),
+ "num_inference_steps": int(batch.num_inference_steps),
+ "flow_shift": float(
+ plan.flow_shift
+ if plan.flow_shift is not None
+ else plan.default_flow_shift
+ ),
+ "audio_flow_shift": float(
+ plan.audio_flow_shift
+ if plan.audio_flow_shift is not None
+ else plan.default_audio_flow_shift
+ ),
+ }
+ exact_fields = (
+ "task",
+ "width",
+ "height",
+ "fps",
+ "frame_count",
+ "num_inference_steps",
+ )
+ exact = all(
+ actual[name] == _MINIMAX_H3_QUALITY_WORKLOAD[name]
+ for name in exact_fields
+ )
+ shifts = math.isclose(
+ actual["flow_shift"],
+ _MINIMAX_H3_QUALITY_WORKLOAD["flow_shift"],
+ abs_tol=1e-9,
+ ) and math.isclose(
+ actual["audio_flow_shift"],
+ _MINIMAX_H3_QUALITY_WORKLOAD["audio_flow_shift"],
+ abs_tol=1e-9,
+ )
+ if not exact or not shifts:
+ raise ValueError(
+ "MiniMax-H3 approximate quality profiles are validated only for "
+ f"{_MINIMAX_H3_QUALITY_WORKLOAD}; got {actual}"
+ )
+ return batch
+
+
+__all__ = ["MiniMaxH3PartitionAdmissionStage", "MiniMaxH3ReleaseMetadata"]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py
new file mode 100644
index 000000000..84a05979e
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/request_validation.py
@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 canonical request validation.
+
+Entry fail-fast for `minimax_h3.request/v1`: every violation raises ValueError
+with the offending field path. Output is a normalized canonical dict (frame
+indices validated but semantic -1 preserved, nothing else rewritten — prompt passes through verbatim and
+conditions order is semantic, never reordered).
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_MAX_DURATION_SECONDS,
+ MINIMAX_H3_MIN_DURATION_SECONDS,
+ MINIMAX_H3_SUPPORTED_FPS,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_CONDITION_ROLE_KEYFRAME,
+ MINIMAX_H3_CONDITION_ROLE_REFERENCE,
+ MINIMAX_H3_FINITE_ASPECT_RATIOS,
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+ MINIMAX_H3_TASK_FL2VA,
+ MINIMAX_H3_TASK_REF2VA,
+ MINIMAX_H3_TASK_T2VA,
+ MiniMaxH3TaskProfile,
+ canonical_minimax_h3_task,
+ minimax_h3_task_profile,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+ minimax_h3_align_frame_count,
+)
+
+MINIMAX_H3_REQUEST_SCHEMA = "minimax_h3.request/v1"
+MINIMAX_H3_MAX_SIGNED_SEED = (1 << 63) - 1
+_ALLOWED_CONDITION_KEYS = frozenset(
+ {"type", "uri", "role", "frame_index", "start_time_seconds"}
+)
+
+
+def _require_str(value: Any, path: str) -> str:
+ if not isinstance(value, str) or value == "":
+ raise ValueError(f"{path} must be a non-empty string")
+ return value
+
+
+def _require_int(value: Any, path: str) -> int:
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise ValueError(f"{path} must be an integer")
+ return value
+
+
+def _optional_positive_finite_float(value: Any, path: str) -> float | None:
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{path} must be a number")
+ normalized = float(value)
+ if not math.isfinite(normalized) or normalized <= 0.0:
+ raise ValueError(f"{path} must be a positive finite number")
+ return normalized
+
+
+def _optional_nonnegative_finite_float(value: Any, path: str) -> float | None:
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{path} must be a number")
+ normalized = float(value)
+ if not math.isfinite(normalized) or normalized < 0.0:
+ raise ValueError(f"{path} must be a non-negative finite number")
+ return normalized
+
+
+def _validate_target(target: Any, *, profile: MiniMaxH3TaskProfile) -> dict[str, Any]:
+ path = "target"
+ if not isinstance(target, Mapping):
+ raise ValueError(f"{path} is required and must be an object")
+ # The canonical target has a deliberately small projection. Transport
+ # compatibility keys are ignored; only these three declared values are
+ # validated and emitted below.
+ short_edge = _require_int(target.get("short_edge"), f"{path}.short_edge")
+ if short_edge != 768:
+ raise ValueError(
+ f"{path}.short_edge must be 768 for minimax_h3, got {short_edge}"
+ )
+ aspect_ratio = _require_str(target.get("aspect_ratio"), f"{path}.aspect_ratio")
+ if profile.aspect_ratio_forced_auto and aspect_ratio != "auto":
+ raise ValueError(
+ f'{path}.aspect_ratio must be "auto" for task {profile.task!r}, '
+ f"got {aspect_ratio!r}"
+ )
+ has_duration = target.get("duration_seconds") is not None
+ if (
+ profile.task in {MINIMAX_H3_TASK_T2VA, MINIMAX_H3_TASK_REF2VA}
+ and aspect_ratio != "auto"
+ and aspect_ratio not in MINIMAX_H3_FINITE_ASPECT_RATIOS
+ ):
+ raise ValueError(
+ f"{path}.aspect_ratio for task {profile.task!r} must be 'auto' or "
+ f"one of {list(MINIMAX_H3_FINITE_ASPECT_RATIOS)!r}, got "
+ f"{aspect_ratio!r}"
+ )
+ if not has_duration:
+ if not profile.duration_from_audio_reference:
+ raise ValueError(f"{path}.duration_seconds is required")
+ # ref2va: duration may derive from a reference audio; the
+ # audio-condition presence is enforced after conditions validate.
+ out: dict[str, Any] = {
+ "short_edge": short_edge,
+ "aspect_ratio": aspect_ratio,
+ }
+ if has_duration:
+ duration = target["duration_seconds"]
+ if isinstance(duration, bool) or not isinstance(duration, (int, float)):
+ raise ValueError(f"{path}.duration_seconds must be a number")
+ if duration <= 0:
+ raise ValueError(f"{path}.duration_seconds must be positive")
+ if not (
+ MINIMAX_H3_MIN_DURATION_SECONDS
+ <= float(duration)
+ <= MINIMAX_H3_MAX_DURATION_SECONDS
+ ):
+ raise ValueError(
+ f"{path}.duration_seconds must be in "
+ f"[{MINIMAX_H3_MIN_DURATION_SECONDS:g}, "
+ f"{MINIMAX_H3_MAX_DURATION_SECONDS:g}], got {duration}"
+ )
+ out["duration_seconds"] = float(duration)
+ return out
+
+
+def _validate_conditions(
+ conditions: Any,
+ *,
+ profile: MiniMaxH3TaskProfile,
+ frame_count: int | None,
+) -> list[dict[str, Any]]:
+ path = "conditions"
+ if conditions is None:
+ conditions = []
+ if not isinstance(conditions, Sequence) or isinstance(conditions, (str, bytes)):
+ raise ValueError(f"{path} must be a list")
+
+ if not profile.conditions_required:
+ if len(conditions) > 0:
+ raise ValueError(
+ f"{path} must be empty for task {profile.task!r} "
+ f"(got {len(conditions)} entries)"
+ )
+ return []
+ if len(conditions) == 0:
+ raise ValueError(
+ f"{path} requires at least one entry for task {profile.task!r}"
+ )
+ if (
+ profile.min_condition_count is not None
+ and len(conditions) < profile.min_condition_count
+ ):
+ raise ValueError(
+ f"{path} requires at least {profile.min_condition_count} entries "
+ f"for task {profile.task!r}, got {len(conditions)}"
+ )
+ if (
+ profile.max_condition_count is not None
+ and len(conditions) > profile.max_condition_count
+ ):
+ raise ValueError(
+ f"{path} allows at most {profile.max_condition_count} entries "
+ f"for task {profile.task!r}, got {len(conditions)}"
+ )
+
+ aligned_frame_count = (
+ minimax_h3_align_frame_count(frame_count) if frame_count is not None else None
+ )
+ normalized: list[dict[str, Any]] = []
+ seen_frame_indices: dict[int, int] = {}
+ for index, cond in enumerate(conditions):
+ cpath = f"{path}[{index}]"
+ if not isinstance(cond, Mapping):
+ raise ValueError(f"{cpath} must be an object")
+ unknown = set(cond) - _ALLOWED_CONDITION_KEYS
+ if unknown:
+ raise ValueError(f"{cpath} has unknown fields: {sorted(unknown)}")
+ role = _require_str(cond.get("role"), f"{cpath}.role")
+ if role not in (
+ MINIMAX_H3_CONDITION_ROLE_KEYFRAME,
+ MINIMAX_H3_CONDITION_ROLE_REFERENCE,
+ ):
+ raise ValueError(
+ f"{cpath}.role must be keyframe or reference, " f"got {role!r}"
+ )
+ cond_type = _require_str(cond.get("type"), f"{cpath}.type")
+ try:
+ rule = profile.rule_for(role=role, condition_type=cond_type)
+ except ValueError as exc:
+ raise ValueError(f"{cpath}: {exc}") from exc
+ uri = _require_str(cond.get("uri"), f"{cpath}.uri")
+
+ entry: dict[str, Any] = {"type": cond_type, "uri": uri, "role": role}
+ if rule.requires_frame_index:
+ frame_index = _require_int(cond.get("frame_index"), f"{cpath}.frame_index")
+ if aligned_frame_count is None:
+ raise ValueError(
+ f"{cpath}.frame_index requires a resolved target duration"
+ )
+ if frame_index == -1:
+ resolved = aligned_frame_count - 1
+ elif 0 <= frame_index < aligned_frame_count:
+ resolved = frame_index
+ else:
+ raise ValueError(
+ f"{cpath}.frame_index must be -1 or in "
+ f"[0, {aligned_frame_count}) after 17n+5 frame alignment, "
+ f"got {frame_index}"
+ )
+ if resolved in seen_frame_indices:
+ raise ValueError(
+ f"{cpath}.frame_index resolves to {resolved}, already "
+ f"bound by conditions[{seen_frame_indices[resolved]}]"
+ )
+ seen_frame_indices[resolved] = index
+ # Preserve the request-level semantic index. In particular, -1 is
+ # the canonical last-frame sentinel; the resolved pixel frame is
+ # carried separately by MiniMaxH3ResolvedPlan.
+ entry["frame_index"] = frame_index
+ elif cond.get("frame_index") is not None:
+ raise ValueError(f"{cpath}.frame_index is not allowed for role={role!r}")
+ start_time_seconds = _optional_nonnegative_finite_float(
+ cond.get("start_time_seconds"), f"{cpath}.start_time_seconds"
+ )
+ if start_time_seconds is not None:
+ if cond_type not in {"video", "video_audio"}:
+ raise ValueError(
+ f"{cpath}.start_time_seconds is only allowed for video "
+ "or video_audio references"
+ )
+ entry["start_time_seconds"] = start_time_seconds
+ normalized.append(entry)
+ return normalized
+
+
+def _validate_fl2va_conditions(conditions: Sequence[Mapping[str, Any]]) -> None:
+ """Enforce the public FL contract after per-entry schema validation."""
+
+ frame_indices = tuple(condition.get("frame_index") for condition in conditions)
+ if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+ raise ValueError(
+ "conditions for task 'fl2va' must be one or two ordered "
+ "image/keyframe entries with frame_index [0], [-1], or [0, -1], "
+ f"got {list(frame_indices)!r}"
+ )
+
+
+def minimax_h3_validate_canonical_request(
+ *,
+ task: Any,
+ prompt: Any,
+ conditions: Any,
+ target: Any,
+ flow_shift: Any = None,
+ audio_flow_shift: Any = None,
+ seed: Any = None,
+ **_extra_kwargs: Any,
+) -> dict[str, Any]:
+ """Validate and normalize a `minimax_h3.request/v1` canonical request.
+
+ Returns the normalized canonical dict; raises ValueError with a field
+ path on any violation. Conditions order is preserved (it is semantic:
+ prompt ordinal labels reference it). seed=0 is a legal value.
+ """
+ # Accept transport wrappers and compatibility kwargs at this boundary, but
+ # never copy them into the canonical request.
+ del _extra_kwargs
+ # Normalize the task name before profile lookup so offline callers match
+ # the adapter behaviour.
+ task_name = canonical_minimax_h3_task(_require_str(task, "task"))
+ profile = minimax_h3_task_profile(task_name)
+ prompt_text = _require_str(prompt, "prompt")
+
+ normalized_target = _validate_target(target, profile=profile)
+ requested_frame_count = None
+ if normalized_target.get("duration_seconds") is not None:
+ requested_frame_count = int(
+ round(
+ float(normalized_target["duration_seconds"]) * MINIMAX_H3_SUPPORTED_FPS
+ )
+ )
+ normalized_conditions = _validate_conditions(
+ conditions,
+ profile=profile,
+ frame_count=requested_frame_count,
+ )
+ if profile.task == MINIMAX_H3_TASK_FL2VA:
+ _validate_fl2va_conditions(normalized_conditions)
+ # ref2va accepts ordered material streams containing any mix of
+ # image/audio/video/video_audio references. Type admission is handled by
+ # the task profile; temporal ambiguity is validated later when target
+ # duration is omitted.
+ if not profile.video_reference_supported:
+ for index, cond in enumerate(normalized_conditions):
+ if cond["type"] in ("video", "video_audio"):
+ raise ValueError(
+ f"conditions[{index}]: video references are not supported "
+ f"in v1 for task {profile.task!r} (image/audio only)"
+ )
+ if normalized_target.get("duration_seconds") is None:
+ # Only reachable for duration_from_audio_reference profiles.
+ duration_sources = [
+ cond
+ for cond in normalized_conditions
+ if cond["type"] in ("audio", "video", "video_audio")
+ ]
+ if not duration_sources:
+ raise ValueError(
+ "target.duration_seconds is required, or exactly one "
+ "audio reference to derive duration from (including "
+ f"video/video_audio soundtracks; task {profile.task!r})"
+ )
+ if len(duration_sources) > 1:
+ raise ValueError(
+ "target.duration_seconds is required when multiple "
+ "audio-bearing references are provided"
+ )
+
+ canonical: dict[str, Any] = {
+ "schema": MINIMAX_H3_REQUEST_SCHEMA,
+ "task": task_name,
+ "prompt": prompt_text,
+ "conditions": normalized_conditions,
+ "target": normalized_target,
+ }
+ normalized_flow_shift = _optional_positive_finite_float(flow_shift, "flow_shift")
+ normalized_audio_flow_shift = _optional_positive_finite_float(
+ audio_flow_shift, "audio_flow_shift"
+ )
+ if normalized_flow_shift is not None:
+ canonical["flow_shift"] = normalized_flow_shift
+ if normalized_audio_flow_shift is not None:
+ canonical["audio_flow_shift"] = normalized_audio_flow_shift
+ if seed is not None:
+ normalized_seed = _require_int(seed, "seed")
+ if normalized_seed < 0:
+ raise ValueError(f"seed must be non-negative, got {normalized_seed}")
+ if normalized_seed > MINIMAX_H3_MAX_SIGNED_SEED:
+ raise ValueError(
+ f"seed must not exceed the signed int64 maximum, got {normalized_seed}"
+ )
+ canonical["seed"] = normalized_seed
+ return canonical
+
+
+__all__ = [
+ "MINIMAX_H3_REQUEST_SCHEMA",
+ "MINIMAX_H3_MAX_SIGNED_SEED",
+ "MINIMAX_H3_SUPPORTED_FPS",
+ "minimax_h3_validate_canonical_request",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/resolved_plan.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/resolved_plan.py
new file mode 100644
index 000000000..e77ec94b3
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/resolved_plan.py
@@ -0,0 +1,447 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 ResolvedPlan: the data-only per-request execution plan.
+
+`minimax_h3_resolve_plan` turns a validated canonical request (see
+request_validation.py) into the data-only plan consumed by stages 1-8.
+Stages never branch on task names; skips must be explicit in the plan.
+
+Scope notes (adapt_shape_v1):
+- all target and material-derived ratios use the single adaptive spatial
+ resolver exported by this module. It starts from a 768px nominal short edge,
+ applies the 768x1344 soft area cap, then rounds both axes independently to
+ the nearest 32px grid.
+- ``auto`` uses the task profile: t2va/ref2va resolve to the 16:9 policy
+ default, while fl2va defers geometry until material probe facts are
+ available. Consumers must fail fast if required evidence is missing.
+- per-modality request overrides and task defaults are retained separately so
+ the timestep stage can apply request > model config > task default priority.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from typing import Any
+
+import msgspec
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_SUPPORTED_FPS,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+ minimax_h3_task_profile,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+ minimax_h3_align_frame_count,
+ minimax_h3_audio_latent_t,
+ minimax_h3_video_latent_t,
+)
+
+MINIMAX_H3_SHAPE_POLICY_VERSION = "adapt_shape_v1"
+MINIMAX_H3_BASE_SHORT_EDGE = 768
+MINIMAX_H3_MAX_PIXELS = MINIMAX_H3_BASE_SHORT_EDGE * 1344
+MINIMAX_H3_CANVAS_MULTIPLE = 32
+MINIMAX_H3_MIN_ASPECT_RATIO = 1.0 / 4.0
+MINIMAX_H3_MAX_ASPECT_RATIO = 4.0
+
+
+class MiniMaxH3MaterialPlanItem(msgspec.Struct, frozen=True):
+ condition_index: int
+ role: str
+ condition_type: str
+ uri: str
+ material_chain: str
+ # Request-level semantic frame index. -1 remains the last-frame sentinel.
+ frame_index: int | None = None
+ # Concrete pixel-frame index after target 17n+5 alignment.
+ resolved_frame_index: int | None = None
+ # Per-reference seek applied identically to the visual and audio streams.
+ start_time_seconds: float = 0.0
+
+
+class MiniMaxH3ResolvedPlan(msgspec.Struct, frozen=True):
+ task: str
+ prompt: str
+ seed: int | None
+ materials: tuple[MiniMaxH3MaterialPlanItem, ...]
+ encoders: dict
+ branches: tuple[dict, ...]
+ default_flow_shift: float
+ default_audio_flow_shift: float
+ flow_shift: float | None
+ audio_flow_shift: float | None
+ shape: dict
+ condition_mask: dict
+
+
+def _parse_aspect_ratio(value: str) -> tuple[int, int]:
+ parts = value.split(":")
+ if len(parts) != 2:
+ raise ValueError(f"target.aspect_ratio must be 'W:H' or 'auto', got {value!r}")
+ try:
+ w, h = int(parts[0]), int(parts[1])
+ except ValueError as exc:
+ raise ValueError(
+ f"target.aspect_ratio must be integer 'W:H', got {value!r}"
+ ) from exc
+ if w <= 0 or h <= 0:
+ raise ValueError(
+ f"target.aspect_ratio components must be positive, got {value!r}"
+ )
+ return w, h
+
+
+def _nearest_multiple(value: float, multiple: int) -> int:
+ return max(multiple, int(round(float(value) / multiple)) * multiple)
+
+
+def _validate_base_short_edge(value: Any) -> int:
+ try:
+ short_edge = int(value)
+ except (TypeError, ValueError) as exc:
+ raise ValueError("target.short_edge must be 768") from exc
+ if short_edge != MINIMAX_H3_BASE_SHORT_EDGE or value != short_edge:
+ raise ValueError(
+ f"target.short_edge must be 768 for MiniMax H3 shape policy v2, got {value!r}"
+ )
+ return short_edge
+
+
+def minimax_h3_resolve_spatial_shape(
+ *,
+ width: int | float,
+ height: int | float,
+ base_short_edge: int = MINIMAX_H3_BASE_SHORT_EDGE,
+) -> dict[str, Any]:
+ """Resolve one display ratio with the ``adapt_shape_v1`` math.
+
+ This is the only implementation of adaptive target geometry. Callers may
+ pass an explicit aspect-ratio pair or probed display dimensions; only the
+ ratio is significant. The supported ratio range is inclusive 1:4 to 4:1.
+ The returned dimensions are always 32px aligned; nearest-grid rounding may
+ leave the final area slightly above the pre-round soft pixel budget.
+ """
+ base_short_edge = _validate_base_short_edge(base_short_edge)
+ try:
+ source_width = float(width)
+ source_height = float(height)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(
+ "shape width and height must be positive finite numbers"
+ ) from exc
+ if (
+ not math.isfinite(source_width)
+ or not math.isfinite(source_height)
+ or source_width <= 0.0
+ or source_height <= 0.0
+ ):
+ raise ValueError("shape width and height must be positive finite numbers")
+
+ ratio = source_width / source_height
+ if not math.isfinite(ratio) or ratio <= 0.0:
+ raise ValueError("shape ratio must be a positive finite number")
+ if not MINIMAX_H3_MIN_ASPECT_RATIO <= ratio <= MINIMAX_H3_MAX_ASPECT_RATIO:
+ raise ValueError(
+ "adapt_shape_v1 ratio must be within the inclusive range "
+ f"1:4 to 4:1, got {source_width:g}:{source_height:g}"
+ )
+
+ if ratio >= 1.0:
+ nominal_width = float(base_short_edge) * ratio
+ nominal_height = float(base_short_edge)
+ else:
+ nominal_width = float(base_short_edge)
+ nominal_height = float(base_short_edge) / ratio
+ nominal_area = nominal_width * nominal_height
+ if nominal_area > MINIMAX_H3_MAX_PIXELS:
+ size_mode = "area"
+ scale = math.sqrt(float(MINIMAX_H3_MAX_PIXELS) / nominal_area)
+ nominal_width *= scale
+ nominal_height *= scale
+ else:
+ size_mode = "short_edge"
+
+ resolved_width = _nearest_multiple(nominal_width, MINIMAX_H3_CANVAS_MULTIPLE)
+ resolved_height = _nearest_multiple(nominal_height, MINIMAX_H3_CANVAS_MULTIPLE)
+
+ return {
+ "geometry": "resolved_v2",
+ "shape_policy_version": MINIMAX_H3_SHAPE_POLICY_VERSION,
+ "base_short_edge": base_short_edge,
+ "effective_short_edge": min(resolved_width, resolved_height),
+ "size_mode": size_mode,
+ "max_pixels": MINIMAX_H3_MAX_PIXELS,
+ "multiple": MINIMAX_H3_CANVAS_MULTIPLE,
+ "rounding": "nearest",
+ "width": resolved_width,
+ "height": resolved_height,
+ }
+
+
+def _resolve_shape(
+ target: Mapping[str, Any],
+ *,
+ geometry_source: str,
+ auto_aspect_ratio: str | None = None,
+ auto_geometry_source: str | None = None,
+) -> dict[str, Any]:
+ fps = MINIMAX_H3_SUPPORTED_FPS
+ if "duration_seconds" not in target:
+ # ref2va duration_from_audio_reference: temporal shape resolves at
+ # material time from the reference audio probe. Validation
+ # guarantees an audio condition exists.
+ shape: dict[str, Any] = {
+ "fps": fps,
+ "temporal": "deferred_from_audio_reference",
+ "geometry_source": geometry_source,
+ }
+ return _resolve_spatial(
+ shape,
+ target,
+ auto_aspect_ratio=auto_aspect_ratio,
+ auto_geometry_source=auto_geometry_source,
+ )
+ frame_count = minimax_h3_align_frame_count(
+ int(round(float(target["duration_seconds"]) * fps))
+ )
+ duration_seconds = frame_count / fps
+ shape = {
+ "fps": fps,
+ "frame_count": frame_count,
+ "video_latent_t": minimax_h3_video_latent_t(frame_count),
+ "audio_latent_t": minimax_h3_audio_latent_t(duration_seconds),
+ "geometry_source": geometry_source,
+ }
+ return _resolve_spatial(
+ shape,
+ target,
+ auto_aspect_ratio=auto_aspect_ratio,
+ auto_geometry_source=auto_geometry_source,
+ )
+
+
+def _resolve_spatial(
+ shape: dict[str, Any],
+ target: Mapping[str, Any],
+ *,
+ auto_aspect_ratio: str | None,
+ auto_geometry_source: str | None,
+) -> dict[str, Any]:
+ aspect_ratio = str(target["aspect_ratio"])
+ base_short_edge = _validate_base_short_edge(target.get("short_edge"))
+ if aspect_ratio == "auto":
+ if auto_aspect_ratio is None:
+ # Deferred: canvas comes from material/model geometry at prepare time.
+ shape["geometry"] = "deferred"
+ shape["geometry_source"] = auto_geometry_source or shape["geometry_source"]
+ shape["shape_policy_version"] = MINIMAX_H3_SHAPE_POLICY_VERSION
+ shape["base_short_edge"] = base_short_edge
+ shape["size_mode"] = "deferred"
+ return shape
+ aspect_ratio = auto_aspect_ratio
+ shape["geometry_source"] = auto_geometry_source or "policy_default"
+ ar_w, ar_h = _parse_aspect_ratio(aspect_ratio)
+ shape.update(
+ minimax_h3_resolve_spatial_shape(
+ width=ar_w,
+ height=ar_h,
+ base_short_edge=base_short_edge,
+ )
+ )
+ return shape
+
+
+def minimax_h3_resolve_plan(canonical: Mapping[str, Any]) -> MiniMaxH3ResolvedPlan:
+ """Canonical request (already validated) -> ResolvedPlan."""
+ if not isinstance(canonical, Mapping):
+ raise ValueError("canonical request must be a mapping")
+ allowed_keys = {
+ "schema",
+ "task",
+ "prompt",
+ "conditions",
+ "target",
+ "seed",
+ "flow_shift",
+ "audio_flow_shift",
+ }
+ unknown = set(canonical) - allowed_keys
+ if unknown:
+ raise ValueError(f"canonical request has unknown fields: {sorted(unknown)}")
+ for key in ("schema", "task", "prompt", "conditions", "target"):
+ if key not in canonical:
+ raise ValueError(f"canonical request missing {key!r}")
+ profile = minimax_h3_task_profile(str(canonical["task"]))
+ if profile.task == "fl2va":
+ conditions = canonical["conditions"]
+ signatures = (
+ [
+ (
+ condition.get("type"),
+ condition.get("role"),
+ condition.get("frame_index"),
+ )
+ for condition in conditions
+ ]
+ if isinstance(conditions, (list, tuple))
+ and all(isinstance(condition, Mapping) for condition in conditions)
+ else []
+ )
+ frame_signature = tuple(signature[2] for signature in signatures)
+ if (
+ not signatures
+ or any(signature[:2] != ("image", "keyframe") for signature in signatures)
+ or frame_signature not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES
+ ):
+ raise ValueError(
+ "fl2va ResolvedPlan requires one or two ordered image/keyframe "
+ "conditions with frame_index [0], [-1], or [0, -1], got "
+ f"{signatures!r}"
+ )
+ shape = _resolve_shape(
+ canonical["target"],
+ geometry_source=profile.geometry_source,
+ auto_aspect_ratio=profile.auto_aspect_ratio,
+ auto_geometry_source=profile.auto_geometry_source,
+ )
+
+ materials: list[MiniMaxH3MaterialPlanItem] = []
+ visual_encode: list[int] = []
+ audio_encode: list[int] = []
+ keyframe_semantic_indices: list[int] = []
+ keyframe_pixel_indices: list[int] = []
+ seen_keyframe_pixel_indices: dict[int, int] = {}
+ for index, cond in enumerate(canonical["conditions"]):
+ rule = profile.rule_for(
+ role=str(cond["role"]), condition_type=str(cond["type"])
+ )
+ frame_index = cond.get("frame_index")
+ resolved_frame_index = None
+ if rule.requires_frame_index:
+ if frame_index is None:
+ raise ValueError(f"conditions[{index}].frame_index is required")
+ semantic_frame_index = int(frame_index)
+ frame_count = int(shape["frame_count"])
+ if semantic_frame_index == -1:
+ resolved_frame_index = frame_count - 1
+ elif 0 <= semantic_frame_index < frame_count:
+ resolved_frame_index = semantic_frame_index
+ else:
+ raise ValueError(
+ f"conditions[{index}].frame_index must be -1 or in "
+ f"[0, {frame_count}) after 17n+5 frame alignment, got "
+ f"{semantic_frame_index}"
+ )
+ previous = seen_keyframe_pixel_indices.get(resolved_frame_index)
+ if previous is not None:
+ raise ValueError(
+ f"conditions[{index}].frame_index resolves to "
+ f"{resolved_frame_index}, already bound by "
+ f"conditions[{previous}]"
+ )
+ seen_keyframe_pixel_indices[resolved_frame_index] = index
+ keyframe_semantic_indices.append(semantic_frame_index)
+ keyframe_pixel_indices.append(resolved_frame_index)
+ materials.append(
+ MiniMaxH3MaterialPlanItem(
+ condition_index=index,
+ role=str(cond["role"]),
+ condition_type=str(cond["type"]),
+ uri=str(cond["uri"]),
+ material_chain=rule.material_chain,
+ frame_index=frame_index,
+ resolved_frame_index=resolved_frame_index,
+ start_time_seconds=float(cond.get("start_time_seconds", 0.0)),
+ )
+ )
+ if rule.visual_tokenizer_encode:
+ visual_encode.append(index)
+ if rule.audio_tokenizer_encode:
+ audio_encode.append(index)
+
+ encoders = {
+ "qwen": {
+ "prompt": canonical["prompt"],
+ "ordered_condition_indices": list(range(len(canonical["conditions"]))),
+ },
+ "visual": visual_encode,
+ "audio": audio_encode,
+ }
+
+ condition_mask: dict[str, Any] = {}
+ if keyframe_pixel_indices:
+ condition_mask = {
+ # Both arrays are request-ordered. Semantic indices feed Qwen and
+ # the RoPE rule; resolved
+ # indices are concrete output frames.
+ "semantic_frame_indices": keyframe_semantic_indices,
+ "pixel_frame_indices": keyframe_pixel_indices,
+ }
+
+ return MiniMaxH3ResolvedPlan(
+ task=profile.task,
+ prompt=str(canonical["prompt"]),
+ seed=canonical.get("seed"),
+ materials=tuple(materials),
+ encoders=encoders,
+ branches=profile.branches,
+ default_flow_shift=float(profile.default_flow_shift),
+ default_audio_flow_shift=float(profile.default_audio_flow_shift),
+ flow_shift=(
+ float(canonical["flow_shift"])
+ if canonical.get("flow_shift") is not None
+ else None
+ ),
+ audio_flow_shift=(
+ float(canonical["audio_flow_shift"])
+ if canonical.get("audio_flow_shift") is not None
+ else None
+ ),
+ shape=shape,
+ condition_mask=condition_mask,
+ )
+
+
+MINIMAX_H3_CANONICAL_REQUEST_EXTRA_KEY = "minimax_h3_canonical_request"
+MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY = "minimax_h3_resolved_plan"
+
+
+def minimax_h3_plan_from_batch(batch: Any) -> MiniMaxH3ResolvedPlan | None:
+ """Resolve (once) and cache the plan for a Req carrying a canonical request.
+
+ Returns None when the request predates the canonical schema (such
+ requests keep their existing behavior).
+ """
+ extra = getattr(batch, "extra", None)
+ if not isinstance(extra, Mapping):
+ return None
+ cached = extra.get(MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY)
+ if cached is not None:
+ if not isinstance(cached, MiniMaxH3ResolvedPlan):
+ raise ValueError(
+ f"batch.extra[{MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY!r}] must be a "
+ "MiniMaxH3ResolvedPlan"
+ )
+ canonical = extra.get(MINIMAX_H3_CANONICAL_REQUEST_EXTRA_KEY)
+ if cached is not None:
+ return cached
+ if canonical is None:
+ return None
+ plan = minimax_h3_resolve_plan(canonical)
+ if isinstance(extra, dict):
+ extra[MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY] = plan
+ return plan
+
+
+__all__ = [
+ "MINIMAX_H3_BASE_SHORT_EDGE",
+ "MINIMAX_H3_CANVAS_MULTIPLE",
+ "MINIMAX_H3_CANONICAL_REQUEST_EXTRA_KEY",
+ "MINIMAX_H3_MAX_PIXELS",
+ "MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY",
+ "MiniMaxH3ResolvedPlan",
+ "minimax_h3_plan_from_batch",
+ "minimax_h3_resolve_plan",
+ "minimax_h3_resolve_spatial_shape",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/__init__.py
new file mode 100644
index 000000000..f65515057
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pipeline lifecycle stages for the native MiniMax H3 implementation."""
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/audio_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/audio_encoding.py
new file mode 100644
index 000000000..d348b4b50
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/audio_encoding.py
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+ ComponentUse,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.condition_encoding import (
+ ConditionEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class MiniMaxH3AudioEncodingStage(ConditionEncodingStage):
+ deduplicated_extra_output_keys = (MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY,)
+
+ def __init__(self, audio_vae, vae_arch_config) -> None:
+ super().__init__()
+ self.audio_vae = audio_vae
+ self.vae_arch_config = vae_arch_config
+
+ @property
+ def role_affinity(self) -> RoleType:
+ return RoleType.ENCODER
+
+ def component_uses(
+ self, server_args: ServerArgs, stage_name: str | None = None
+ ) -> list[ComponentUse]:
+ stage_name = self._component_stage_name(stage_name)
+ return [ComponentUse(stage_name, "audio_vae")]
+
+ def build_dedup_fingerprint(self, batch: Req, server_args: ServerArgs):
+ parent_request_id = batch.extra.get("parent_request_id")
+ return (
+ ("expanded_outputs", parent_request_id)
+ if parent_request_id is not None
+ else id(batch)
+ )
+
+ @torch.no_grad()
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_cleanup_temp_dirs,
+ )
+
+ try:
+ return self._forward(batch, server_args)
+ finally:
+ # Audio encoding is the final material consumer in the MiniMax H3
+ # encoder pipeline, including requests with no routed audio.
+ minimax_h3_cleanup_temp_dirs(batch, owners=("material",))
+
+ def _forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is not None:
+ routed = plan.encoders.get("audio")
+ if not routed:
+ return batch
+ self._encode_references_from_plan(batch, plan, routed)
+ return batch
+ if (
+ batch.sampling_params is not None
+ and batch.sampling_params.audio_path is not None
+ ):
+ raise NotImplementedError(
+ "MiniMaxH3AudioEncodingStage direct audio tokenizer encode "
+ "requires a canonical minimax_h3 request (resolved plan); "
+ "legacy audio_path-only requests are unsupported."
+ )
+ return batch
+
+ def _encode_references_from_plan(self, batch: Req, plan, routed) -> None:
+ """Direct reference-audio encode: audio VAE posterior mean ->
+ normalized channel-major rows in batch.extra."""
+ routed_set = set(routed)
+ routed_materials = [
+ material
+ for material in plan.materials
+ if material.condition_index in routed_set
+ ]
+ from .replica_broadcast import (
+ minimax_h3_replica_broadcast_error,
+ minimax_h3_replica_broadcast_extra,
+ minimax_h3_replica_ctx,
+ )
+
+ _, replica_rank = minimax_h3_replica_ctx()
+ owner_exception = None
+ owner_error = None
+ if (
+ replica_rank == 0
+ and MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY not in batch.extra
+ ):
+ try:
+ with self.use_declared_component(
+ component_name="audio_vae",
+ module=self.audio_vae,
+ ) as audio_vae:
+ assert audio_vae is not None
+ self.audio_vae = audio_vae
+ batch.extra[MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY] = (
+ self._encode_reference_payload(
+ batch,
+ plan,
+ routed_materials,
+ )
+ )
+ except Exception as exc:
+ owner_exception = exc
+ owner_error = f"{type(exc).__name__}: {exc}"
+ owner_error = minimax_h3_replica_broadcast_error(owner_error)
+ if owner_error is not None:
+ if owner_exception is not None:
+ raise owner_exception
+ raise RuntimeError(
+ f"MiniMax H3 audio encode failed on rank 0: {owner_error}"
+ )
+ minimax_h3_replica_broadcast_extra(
+ batch, MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY
+ )
+
+ def _encode_reference_payload(self, batch: Req, plan, materials) -> dict:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_localize_material_uri,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.prequeue import (
+ MINIMAX_H3_PROBE_FACTS_EXTRA_KEY,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.reference_encoding import (
+ _AudioVAEDeterminismContext,
+ minimax_h3_encode_reference_audio_rows,
+ )
+
+ if not materials:
+ raise ValueError("ref2va audio routing selected no reference materials")
+ entries = []
+ max_duration_seconds = (
+ float(plan.shape["frame_count"]) / float(plan.shape["fps"])
+ if plan.shape.get("frame_count") is not None
+ and plan.shape.get("fps") is not None
+ else None
+ )
+ # One determinism-flag toggle for the whole routed set, not one per
+ # material: _AudioVAEDeterminismContext is reentrant, so each
+ # material's own nested context (inside
+ # minimax_h3_encode_reference_audio_rows) becomes a no-op depth
+ # increment/decrement under this outer scope.
+ with _AudioVAEDeterminismContext():
+ for material in materials:
+ audio_path = minimax_h3_localize_material_uri(
+ batch,
+ material.uri,
+ condition_type=material.condition_type,
+ condition_index=int(material.condition_index),
+ )
+ material_chain = str(material.material_chain)
+ source_facts = batch.extra.get(
+ MINIMAX_H3_PROBE_FACTS_EXTRA_KEY, {}
+ ).get(int(material.condition_index))
+ if not isinstance(source_facts, dict):
+ raise ValueError(
+ "reference-audio encoding requires cached pre-queue probe "
+ f"facts for conditions[{int(material.condition_index)}]"
+ )
+ input_has_audio = bool(source_facts.get("has_audio", True))
+ if material_chain == "video.reference_preserve" and not input_has_audio:
+ # Keep the visual reference block in request order while
+ # representing the absent soundtrack as a zero-length
+ # audio condition.
+ out = {
+ "rows": torch.empty((0, 32), dtype=torch.float32),
+ "ref_audio_t": 0,
+ "duration_seconds": 0.0,
+ }
+ else:
+ out = minimax_h3_encode_reference_audio_rows(
+ self.audio_vae,
+ audio_path,
+ self.vae_arch_config,
+ material_chain=material_chain,
+ max_duration_seconds=max_duration_seconds,
+ start_time_seconds=float(material.start_time_seconds),
+ source_sample_rate=(
+ int(source_facts["audio_sample_rate"])
+ if material_chain == "audio"
+ else None
+ ),
+ )
+ entries.append(
+ {
+ **out,
+ "condition_index": int(material.condition_index),
+ "material_chain": material_chain,
+ }
+ )
+ payload = dict(entries[0]) if len(entries) == 1 else {}
+ payload["audios"] = entries
+ return payload
+
+
+__all__ = ["MiniMaxH3AudioEncodingStage"]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py
new file mode 100644
index 000000000..77d8e074a
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py
@@ -0,0 +1,435 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import functools
+from collections.abc import Mapping
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.distributed import (
+ get_world_group,
+ model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+ ComponentUse,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+ StageParallelismType,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ VerificationResult,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.precision import (
+ autocast_enabled,
+ resolve_decode_precision,
+ resolve_precision,
+)
+from sglang.multimodal_gen.runtime.utils.torch_compile import (
+ ActiveTargetCompiledCallable,
+)
+
+
+def _required_tensor(value, path: str) -> torch.Tensor:
+ if not isinstance(value, torch.Tensor):
+ raise ValueError(f"{path} must be a torch.Tensor")
+ return value
+
+
+@functools.lru_cache(maxsize=None)
+def _cached_decode_mean_std(
+ mean_values: tuple[float, ...],
+ std_values: tuple[float, ...],
+ device: torch.device,
+ dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Device/dtype-keyed mean/std tensors, built once per distinct combination.
+
+ mean_values/std_values come from the loaded arch_config and are fixed
+ for the process lifetime, so the same (values, device, dtype) always
+ reconstructs an identical tensor; cache it instead of rebuilding it on
+ every decode call.
+ """
+ mean = torch.as_tensor(mean_values, device=device, dtype=dtype)
+ std = torch.as_tensor(std_values, device=device, dtype=dtype)
+ return mean, std
+
+
+def _reverse_normalize_latents_(
+ latents: torch.Tensor,
+ *,
+ mean_values,
+ std_values,
+ name: str,
+) -> torch.Tensor:
+ mean, std = _cached_decode_mean_std(
+ tuple(mean_values), tuple(std_values), latents.device, latents.dtype
+ )
+ if mean.ndim != 1:
+ raise ValueError(f"{name}.latents_mean must be 1-D, got {tuple(mean.shape)}")
+ if std.ndim != 1:
+ raise ValueError(f"{name}.latents_std must be 1-D, got {tuple(std.shape)}")
+ if mean.shape != std.shape:
+ raise ValueError(
+ f"{name} latent normalization shape mismatch: "
+ f"mean={tuple(mean.shape)} std={tuple(std.shape)}"
+ )
+ if latents.ndim < 2:
+ raise ValueError(f"{name} latents must have a channel dimension")
+ if int(latents.shape[1]) != int(mean.shape[0]):
+ raise ValueError(
+ f"{name} latent normalization channel mismatch: "
+ f"latents.shape[1]={int(latents.shape[1])} mean_len={int(mean.shape[0])}"
+ )
+ view_shape = [1] * latents.ndim
+ view_shape[1] = int(mean.shape[0])
+ return latents.mul_(std.view(*view_shape)).add_(mean.view(*view_shape))
+
+
+def _crop_to_target_canvas(batch: Req, frames: torch.Tensor) -> torch.Tensor:
+ """Crop decoded frames [B,C,T,H,W] back to the target canvas.
+
+ The visual VAE pads the latent grid to its tile multiples (padding lands
+ at the bottom/right), so a non-tile-aligned geometry decodes larger than the
+ requested canvas (e.g. 1344x768 for a 1280x704 target). Target dims come
+ from the direct-mode denoise state (latent_h/w * 16); requests without
+ that state keep the raw decode.
+ """
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
+ )
+
+ state = batch.extra.get(MINIMAX_H3_DENOISE_STATE_EXTRA_KEY)
+ if state is None:
+ return frames
+ target_h = int(state["latent_h"]) * 16
+ target_w = int(state["latent_w"]) * 16
+ h, w = int(frames.shape[-2]), int(frames.shape[-1])
+ if h < target_h or w < target_w:
+ raise ValueError(
+ f"decoded frames {h}x{w} smaller than target canvas {target_h}x{target_w}"
+ )
+ if h == target_h and w == target_w:
+ return frames
+ return frames[..., :target_h, :target_w]
+
+
+def _canonical_visual_video_frames(
+ frames: torch.Tensor, *, batch_size: int
+) -> torch.Tensor:
+ if frames.ndim == 4:
+ if int(frames.shape[0]) % batch_size != 0:
+ raise ValueError(
+ f"Decoded visual video shape {tuple(frames.shape)} is incompatible "
+ f"with batch_size={batch_size}"
+ )
+ frames = frames.reshape(
+ batch_size, int(frames.shape[0]) // batch_size, *frames.shape[1:]
+ )
+ frames = frames.transpose(1, 2)
+ elif frames.ndim == 5:
+ if int(frames.shape[0]) != batch_size:
+ raise ValueError(
+ f"Decoded visual video batch mismatch: frames.shape[0]={int(frames.shape[0])} "
+ f"batch_size={batch_size}"
+ )
+ else:
+ raise ValueError(
+ f"Decoded visual video shape {tuple(frames.shape)} is not supported"
+ )
+ return frames
+
+
+def _canonical_output_audio_waveform(
+ audio_waveform: torch.Tensor, *, batch_size: int
+) -> torch.Tensor:
+ """Project audio-VAE-native ``[C, 1, L]`` audio to output ``[1, C, L]``.
+
+ The audio VAE treats stereo channels as its decoder batch and returns
+ ``[2, 1, samples]`` for MiniMax H3's one generated sample. The generic output
+ path instead selects generated samples along dimension zero. Keep the audio VAE
+ tensor unchanged for decoder artifacts, then make the singleton generated-
+ sample dimension explicit only at the ``OutputBatch`` boundary.
+ """
+ if audio_waveform.ndim != 3:
+ raise ValueError(
+ "Decoded audio VAE waveform must be [C, 1, L], got "
+ f"{tuple(audio_waveform.shape)}"
+ )
+ if batch_size != 1:
+ raise ValueError(
+ "MiniMax H3 audio VAE output only supports one generated sample, "
+ f"got visual batch_size={batch_size}"
+ )
+ if int(audio_waveform.shape[1]) != 1:
+ raise ValueError(
+ "Decoded audio VAE waveform must have shape [C, 1, L], got "
+ f"{tuple(audio_waveform.shape)}"
+ )
+ return audio_waveform.permute(1, 0, 2).contiguous()
+
+
+_MINIMAX_H3_DECODER_TASKS = frozenset({"t2va", "fl2va", "ref2va"})
+_MINIMAX_H3_CANONICAL_REQUEST_EXTRA_KEY = "minimax_h3_canonical_request"
+_MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY = "minimax_h3_resolved_plan"
+
+
+def _minimax_h3_decoder_task(batch: Req) -> str | None:
+ """Return the validated request task used for output-decoder routing.
+
+ Debug requests have no canonical task and retain the
+ generic decoder.
+ """
+
+ extra = getattr(batch, "extra", None)
+ if not isinstance(extra, Mapping):
+ return None
+ canonical = extra.get(_MINIMAX_H3_CANONICAL_REQUEST_EXTRA_KEY)
+ if canonical is not None and not isinstance(canonical, Mapping):
+ raise ValueError("minimax_h3_canonical_request must be a mapping")
+ canonical_task = canonical.get("task") if isinstance(canonical, Mapping) else None
+ resolved = extra.get(_MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY)
+ resolved_task = getattr(resolved, "task", None) if resolved is not None else None
+ if canonical_task is not None and resolved_task is not None:
+ if str(canonical_task) != str(resolved_task):
+ raise ValueError(
+ "MiniMax H3 decoder task mismatch between canonical request and "
+ "resolved plan"
+ )
+ task_value = resolved_task if resolved_task is not None else canonical_task
+ if task_value is None:
+ return None
+ if not isinstance(task_value, str) or task_value not in _MINIMAX_H3_DECODER_TASKS:
+ raise ValueError(f"unsupported MiniMax H3 decoder task {task_value!r}")
+ return task_value
+
+
+class MiniMaxH3DecodingStage(DecodingStage):
+ def __init__(self, video_vae, audio_vae) -> None:
+ super().__init__(vae=video_vae, component_name="video_vae")
+ self.video_vae = video_vae
+ self.audio_vae = audio_vae
+ self._compiled_audio_vae_decode = ActiveTargetCompiledCallable()
+
+ @property
+ def role_affinity(self) -> RoleType:
+ return RoleType.DECODER
+
+ @property
+ def parallelism_type(self) -> StageParallelismType:
+ # Every decode-group rank owns a subset of visual VAE tiles. The GPU
+ # worker only materializes/saves the final OutputBatch on world rank 0.
+ return StageParallelismType.REPLICATED
+
+ def component_uses(
+ self, server_args: ServerArgs, stage_name: str | None = None
+ ) -> list[ComponentUse]:
+ stage_name = self._component_stage_name(stage_name)
+ video_vae_dtype = resolve_precision(
+ server_args, "video_vae", precision_attr="vae_precision"
+ )
+ audio_vae_dtype = resolve_precision(
+ server_args, "audio_vae", precision_attr="audio_vae_precision"
+ )
+ uses = [
+ ComponentUse(stage_name, "video_vae", target_dtype=video_vae_dtype),
+ ]
+ uses.append(ComponentUse(stage_name, "audio_vae", target_dtype=audio_vae_dtype))
+ return uses
+
+ def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+ result.add_check(
+ "audio_latents",
+ batch.audio_latents,
+ [V.is_tensor, V.with_dims(3)],
+ )
+ return result
+
+ def verify_output(
+ self, batch: OutputBatch, server_args: ServerArgs
+ ) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
+ result.add_check("audio", batch.audio, [V.is_tensor, V.with_dims(3)])
+ result.add_check("audio_sample_rate", batch.audio_sample_rate, V.positive_int)
+ return result
+
+ def _decode_audio(
+ self,
+ audio_latent: torch.Tensor,
+ server_args: ServerArgs,
+ ) -> dict:
+ with self.use_declared_component(
+ component_name="audio_vae",
+ module=self.audio_vae,
+ ) as audio_vae:
+ assert audio_vae is not None
+ self.audio_vae = audio_vae
+ if audio_vae.training:
+ audio_vae.eval()
+ audio_arch_config = server_args.pipeline_config.audio_vae_config.arch_config
+ audio_decode_latent = _reverse_normalize_latents_(
+ audio_latent,
+ mean_values=audio_arch_config.latents_mean,
+ std_values=audio_arch_config.latents_std,
+ name="audio_vae",
+ )
+ audio_vae_dtype = resolve_precision(
+ server_args, "audio_vae", precision_attr="audio_vae_precision"
+ )
+ audio_autocast_enabled = (
+ audio_latent.device.type == "cuda"
+ and autocast_enabled(audio_vae_dtype, server_args.disable_autocast)
+ )
+ with torch.autocast(
+ device_type=audio_latent.device.type,
+ dtype=audio_vae_dtype,
+ enabled=audio_autocast_enabled,
+ ):
+ audio_decode = self._get_vae_decode_fn(
+ audio_vae,
+ server_args,
+ decode_fn=audio_vae.decode,
+ compiled_callable=self._compiled_audio_vae_decode,
+ )
+ waveform = _required_tensor(
+ audio_decode(audio_decode_latent), "audio_vae.decode"
+ )
+ return {
+ "waveform": waveform,
+ "sample_rate": int(audio_vae.sample_rate),
+ }
+
+ @torch.no_grad()
+ def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
+ _minimax_h3_decoder_task(batch)
+ visual_latent = _required_tensor(batch.latents, "batch.latents")
+ audio_latent = _required_tensor(batch.audio_latents, "batch.audio_latents")
+ if visual_latent.ndim != 5:
+ raise ValueError("batch.latents must be [B, C, T, H, W]")
+ if audio_latent.ndim != 3:
+ raise ValueError(
+ "batch.audio_latents must be [audio_channel, latent_dim, T]"
+ )
+
+ if self.video_vae is None:
+ raise RuntimeError("MiniMax H3 tasks require the video_vae output decoder")
+ with self.use_declared_component(
+ component_name="video_vae",
+ module=self.video_vae,
+ ) as selected_video_vae:
+ if selected_video_vae is None:
+ raise RuntimeError("video_vae became unavailable during decode")
+ self.video_vae = selected_video_vae
+ if selected_video_vae.training:
+ selected_video_vae.eval()
+ visual_arch_config = server_args.pipeline_config.vae_config.arch_config
+ visual_decode_latent = _reverse_normalize_latents_(
+ visual_latent,
+ mean_values=visual_arch_config.latents_mean,
+ std_values=visual_arch_config.latents_std,
+ name="video_vae",
+ )
+ video_vae_dtype = resolve_decode_precision(server_args, "video_vae")
+ visual_autocast_enabled = (
+ visual_latent.device.type == "cuda"
+ and autocast_enabled(video_vae_dtype, server_args.disable_autocast)
+ )
+ if visual_autocast_enabled:
+ selected_video_vae.prepare_decoder_autocast_weights(video_vae_dtype)
+ with torch.autocast(
+ device_type=visual_latent.device.type,
+ dtype=video_vae_dtype,
+ enabled=visual_autocast_enabled,
+ ):
+ video_decode = self._get_vae_decode_fn(
+ selected_video_vae,
+ server_args,
+ decode_fn=selected_video_vae.decode_base,
+ )
+ visual_frames = video_decode(visual_decode_latent)
+ visual_frames = selected_video_vae.processor.revert_tensor(
+ visual_frames
+ )
+ visual_frames = _required_tensor(
+ visual_frames,
+ "video_vae.processor.revert_tensor",
+ )
+ visual_frames = _canonical_visual_video_frames(
+ visual_frames, batch_size=int(visual_latent.shape[0])
+ )
+ visual_frames = _crop_to_target_canvas(batch, visual_frames)
+ if (
+ visual_frames.dtype != torch.float32
+ or not visual_frames.is_contiguous()
+ ):
+ canonical_frames = torch.empty_like(
+ visual_frames,
+ dtype=torch.float32,
+ memory_format=torch.contiguous_format,
+ )
+ canonical_frames.copy_(visual_frames)
+ visual_frames = canonical_frames
+
+ # DP is currently rejected by ServerArgs validation, so the world group
+ # is one request replica (TP/CFG/SP ranks), not a collection of
+ # independent requests. Decode the non-sharded audio VAE once per
+ # request and distribute its output to the ranks that decoded video.
+ world_group = get_world_group() if model_parallel_is_initialized() else None
+ is_audio_owner = world_group is None or world_group.rank_in_group == 0
+ owner_exception = None
+ owner_error = None
+ audio_payload = None
+ if is_audio_owner:
+ try:
+ audio_payload = self._decode_audio(audio_latent, server_args)
+ except Exception as exc:
+ owner_exception = exc
+ owner_error = f"{type(exc).__name__}: {exc}"
+ if world_group is not None:
+ owner_error = world_group.broadcast_object(owner_error, src=0)
+ if owner_error is not None:
+ if owner_exception is not None:
+ raise owner_exception
+ raise RuntimeError(
+ f"MiniMax H3 audio decode failed on rank 0: {owner_error}"
+ )
+ if world_group is not None:
+ audio_payload = world_group.broadcast_tensor_dict(audio_payload, src=0)
+ if not isinstance(audio_payload, dict):
+ raise RuntimeError("MiniMax H3 audio decode produced no output payload")
+ audio_waveform = _required_tensor(
+ audio_payload.get("waveform"), "audio_vae.decode"
+ )
+ audio_sample_rate = int(audio_payload["sample_rate"])
+
+ visual_frames = server_args.pipeline_config.post_decoding(
+ visual_frames, server_args
+ )
+ output_audio_waveform = _canonical_output_audio_waveform(
+ audio_waveform, batch_size=int(visual_frames.shape[0])
+ )
+ return OutputBatch(
+ output=visual_frames,
+ audio=output_audio_waveform,
+ audio_sample_rate=audio_sample_rate,
+ trajectory_timesteps=batch.trajectory_timesteps,
+ trajectory_latents=batch.trajectory_latents,
+ rollout_trajectory_data=batch.rollout_trajectory_data,
+ trajectory_decoded=None,
+ metrics=batch.metrics,
+ noise_pred=None,
+ )
+
+
+__all__ = [
+ "MiniMaxH3DecodingStage",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
new file mode 100644
index 000000000..802388c6f
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
@@ -0,0 +1,963 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 denoise sink for packed-token DiT stepping, CFG-distilled
+single-branch execution, and payload validation.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from contextlib import contextmanager
+from functools import partial
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
+ CacheDitConfig,
+ disable_cache_on_transformer,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
+ is_fsdp_managed_module,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
+ DenoisingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_QUALITY_PROFILES,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ VerificationResult,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import maybe_nvtx_range
+from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
+
+logger = init_logger(__name__)
+
+_REF2VA_VIDEO_CHAINS = {
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+}
+
+
+def minimax_h3_condition_noise_aug(sampling: Any) -> tuple[float, float]:
+ """Resolve condition timesteps using the model defaults."""
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import (
+ MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
+ MINIMAX_H3_IMGVID_COND_TIMESTEP,
+ )
+
+ imgvid_noise_aug = getattr(
+ sampling,
+ "imgvid_cond_noise_aug_for_inference",
+ None,
+ )
+ if imgvid_noise_aug is None:
+ # The model default uses imgvid noise aug 0.999. A request selecting
+ # imgvid noise aug 1.0 still overrides this.
+ imgvid_noise_aug = MINIMAX_H3_IMGVID_COND_TIMESTEP
+ audio_noise_aug = getattr(
+ sampling,
+ "audio_cond_noise_aug_for_inference",
+ None,
+ )
+ if audio_noise_aug is None:
+ audio_noise_aug = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP
+ return float(imgvid_noise_aug), float(audio_noise_aug)
+
+
+def _validate_fl2va_keyframe_payload(plan: Any, keyframe: Any) -> None:
+ """Reject stale/middle/reordered keyframe payloads at the DiT sink."""
+
+ task = None if plan is None else str(plan.task)
+ if task != "fl2va":
+ if keyframe is not None:
+ raise ValueError(
+ "keyframe condition rows are only valid for plan.task='fl2va'"
+ )
+ return
+ if not isinstance(keyframe, Mapping):
+ raise ValueError("fl2va denoising requires encoded keyframe condition rows")
+
+ semantic_indices = tuple(keyframe.get("semantic_frame_indices") or ())
+ if semantic_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+ raise ValueError(
+ "fl2va denoising requires semantic_frame_indices in "
+ f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, "
+ f"got {semantic_indices!r}"
+ )
+ frame_count = keyframe.get("frame_count")
+ if isinstance(frame_count, bool) or not isinstance(frame_count, int):
+ raise ValueError("fl2va keyframe payload requires an integer frame_count")
+ if frame_count <= 1:
+ raise ValueError("fl2va keyframe payload frame_count must be greater than one")
+ pixel_indices = keyframe.get("pixel_frame_indices")
+ expected_pixel_indices = [
+ frame_count - 1 if index == -1 else index for index in semantic_indices
+ ]
+ if pixel_indices != expected_pixel_indices:
+ raise ValueError(
+ "fl2va denoising requires pixel_frame_indices resolved from the "
+ "semantic anchors, "
+ f"got {pixel_indices!r} for frame_count={frame_count}"
+ )
+
+ entries = keyframe.get("keyframes")
+ if (
+ not isinstance(entries, list)
+ or len(entries) != len(semantic_indices)
+ or any(not isinstance(entry, Mapping) for entry in entries)
+ ):
+ raise ValueError(
+ "fl2va denoising requires one encoded keyframe per semantic anchor"
+ )
+ if [entry.get("frame_index") for entry in entries] != list(semantic_indices):
+ raise ValueError("fl2va encoded keyframes must remain in semantic anchor order")
+ if [
+ entry.get("resolved_frame_index") for entry in entries
+ ] != expected_pixel_indices:
+ raise ValueError(
+ "fl2va encoded keyframes must carry matching resolved_frame_index values"
+ )
+
+ latent_h = int(keyframe.get("latent_h") or 0)
+ latent_w = int(keyframe.get("latent_w") or 0)
+ rows = keyframe.get("rows")
+ expected_rows = len(semantic_indices) * (latent_h // 2) * (latent_w // 2)
+ if (
+ latent_h <= 0
+ or latent_w <= 0
+ or not isinstance(rows, torch.Tensor)
+ or int(rows.shape[0]) != expected_rows
+ ):
+ actual_rows = None if not isinstance(rows, torch.Tensor) else int(rows.shape[0])
+ raise ValueError(
+ "fl2va encoded keyframe rows do not match target-canvas blocks: "
+ f"expected={expected_rows}, actual={actual_rows}"
+ )
+
+
+def _imgvid_condition_shapes(
+ *,
+ ref2va_blocks: list[dict[str, int | str]] | None,
+ keyframe: Any,
+ is_ref2va: bool,
+) -> list[tuple[int, int, int]]:
+ """Return visual-condition ``(T,H,W)`` in packed anchor-row order."""
+
+ if ref2va_blocks is not None:
+ shapes = []
+ for block in ref2va_blocks:
+ kind = str(block["kind"])
+ if kind == "image":
+ shapes.append((1, int(block["latent_h"]), int(block["latent_w"])))
+ elif kind in {"video", "video_audio"}:
+ shapes.append(
+ (
+ int(block["latent_t"]),
+ int(block["latent_h"]),
+ int(block["latent_w"]),
+ )
+ )
+ return shapes
+
+ if is_ref2va:
+ # ref2va always carries a resolved plan, so ordered blocks are
+ # supplied above; reaching here indicates an upstream bug.
+ raise ValueError("ref2va visual-condition shapes require ordered blocks")
+
+ if not isinstance(keyframe, Mapping):
+ return []
+ entries = keyframe.get("keyframes")
+ if isinstance(entries, list) and entries:
+ return [
+ (1, int(entry["latent_h"]), int(entry["latent_w"])) for entry in entries
+ ]
+
+ latent_h = int(keyframe["latent_h"])
+ latent_w = int(keyframe["latent_w"])
+ frame_rows = (latent_h // 2) * (latent_w // 2)
+ rows = keyframe["rows"]
+ if frame_rows <= 0 or int(rows.shape[0]) % frame_rows:
+ raise ValueError(
+ "legacy keyframe rows cannot be split into visual-condition frames"
+ )
+ return [(1, latent_h, latent_w)] * (int(rows.shape[0]) // frame_rows)
+
+
+def _ref2va_payload_entry(
+ payload: Any,
+ *,
+ list_key: str,
+ condition_index: int,
+ path: str,
+) -> Mapping[str, Any]:
+ if not isinstance(payload, Mapping):
+ raise ValueError(f"{path} is required for ref2va condition rows")
+ entries = payload.get(list_key)
+ if isinstance(entries, list):
+ for entry in entries:
+ if (
+ isinstance(entry, Mapping)
+ and entry.get("condition_index") is not None
+ and int(entry["condition_index"]) == int(condition_index)
+ ):
+ return entry
+ if len(entries) == 1 and isinstance(entries[0], Mapping):
+ return entries[0]
+ raise ValueError(
+ f"{path}.{list_key} missing entry for condition_index={condition_index}"
+ )
+ if payload.get("condition_index") is None or int(payload["condition_index"]) == int(
+ condition_index
+ ):
+ return payload
+ raise ValueError(
+ f"{path}.{list_key} missing entry for condition_index={condition_index}"
+ )
+
+
+def _cat_optional(rows: list[torch.Tensor]) -> torch.Tensor | None:
+ if not rows:
+ return None
+ return rows[0] if len(rows) == 1 else torch.cat(rows, dim=0)
+
+
+def _ref2va_ordered_blocks_and_rows(
+ *,
+ plan: Any,
+ ref_image: Any,
+ ref_audio: Any,
+ ref_video: Any,
+) -> tuple[list[dict[str, int | str]], torch.Tensor | None, torch.Tensor | None]:
+ blocks: list[dict[str, int | str]] = []
+ visual_rows: list[torch.Tensor] = []
+ audio_rows: list[torch.Tensor] = []
+ for material in plan.materials:
+ chain = str(material.material_chain)
+ condition_index = int(material.condition_index)
+ if chain == "image.reference_preserve":
+ entry = _ref2va_payload_entry(
+ ref_image,
+ list_key="images",
+ condition_index=condition_index,
+ path="batch.extra.minimax_h3_reference_image_rows",
+ )
+ blocks.append(
+ {
+ "kind": "image",
+ "latent_h": int(entry["latent_h"]),
+ "latent_w": int(entry["latent_w"]),
+ }
+ )
+ visual_rows.append(entry["rows"])
+ elif chain == "audio":
+ entry = _ref2va_payload_entry(
+ ref_audio,
+ list_key="audios",
+ condition_index=condition_index,
+ path="batch.extra.minimax_h3_reference_audio_rows",
+ )
+ ref_audio_t = int(entry["ref_audio_t"])
+ blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t})
+ if ref_audio_t > 0:
+ audio_rows.append(entry["rows"])
+ elif chain in _REF2VA_VIDEO_CHAINS:
+ video_entry = _ref2va_payload_entry(
+ ref_video,
+ list_key="videos",
+ condition_index=condition_index,
+ path="batch.extra.minimax_h3_reference_video_rows",
+ )
+ audio_entry = _ref2va_payload_entry(
+ ref_audio,
+ list_key="audios",
+ condition_index=condition_index,
+ path="batch.extra.minimax_h3_reference_audio_rows",
+ )
+ ref_audio_t = int(audio_entry["ref_audio_t"])
+ blocks.append(
+ {
+ "kind": (
+ "video_audio"
+ if chain == "video_audio.reference_preserve"
+ else "video"
+ ),
+ "ref_audio_t": ref_audio_t,
+ "latent_t": int(video_entry["latent_t"]),
+ "latent_h": int(video_entry["latent_h"]),
+ "latent_w": int(video_entry["latent_w"]),
+ }
+ )
+ visual_rows.append(video_entry["rows"])
+ if ref_audio_t > 0:
+ audio_rows.append(audio_entry["rows"])
+ else:
+ raise ValueError(f"unsupported ref2va material chain {chain!r}")
+ return blocks, _cat_optional(visual_rows), _cat_optional(audio_rows)
+
+
+def _resolve_denoise_model(
+ transformer: Any,
+ device: torch.device,
+ *,
+ placement_managed: bool = False,
+) -> Any:
+ """Resolve the DiT module and place it for denoise.
+
+ ComponentManager and FSDP own their device placement; move only unmanaged
+ plain modules.
+ """
+ model = getattr(transformer, "model", transformer)
+ if placement_managed or is_fsdp_managed_module(model):
+ if model.training:
+ model.eval()
+ return model
+ return model.to(device).eval()
+
+
+def _precompute_refined_prompt_embeds(
+ model: Any,
+ positive: Any,
+ *,
+ device: torch.device,
+) -> bool:
+ """Move request-static text refinement out of the denoise hot loop."""
+ refine = getattr(model, "refine_prompt_embeds", None)
+ if not callable(refine):
+ return False
+
+ static_kwargs = positive.static_kwargs
+ prompt_embeds = static_kwargs["prompt_embeds"]
+ refiner_params = static_kwargs["refiner_packed_seq_params"]
+ if isinstance(refiner_params, dict):
+ refiner_cu = refiner_params["cu_seqlens_q"]
+ else:
+ refiner_cu = refiner_params.cu_seqlens_q
+ with torch.inference_mode():
+ refined = refine(
+ prompt_embeds,
+ refiner_cu,
+ device=device,
+ )
+ if not torch.is_tensor(refined):
+ raise TypeError("MiniMax H3 refine_prompt_embeds must return a torch.Tensor")
+ if int(refined.shape[0]) != int(prompt_embeds.shape[0]):
+ raise ValueError(
+ "MiniMax H3 refined prompt row count changed: "
+ f"{int(prompt_embeds.shape[0])} -> {int(refined.shape[0])}"
+ )
+ static_kwargs["prompt_embeds"] = refined
+ static_kwargs["refined_prompt_embeds_length"] = int(refined.shape[0])
+ return True
+
+
+def _precompute_rope_cache(
+ model: Any,
+ positive: Any,
+ *,
+ device: torch.device,
+) -> bool:
+ """Move request-static RoPE construction out of the denoise hot loop."""
+ build = getattr(model, "build_rope_cache", None)
+ if not callable(build):
+ return False
+ static_kwargs = positive.static_kwargs
+ with torch.inference_mode():
+ static_kwargs["rope_cache"] = build(
+ static_kwargs["img_position_ids"],
+ device=device,
+ )
+ return True
+
+
+class MiniMaxH3DenoisingStage(DenoisingStage):
+ def __init__(self, transformer, pipeline=None) -> None:
+ super().__init__(
+ transformer=transformer,
+ scheduler=None,
+ pipeline=pipeline,
+ )
+ self._minimax_h3_quality_profile = "lossless"
+ self._minimax_h3_cache_mode: str | None = None
+
+ def _owns_compile_warmup_lifecycle(self) -> bool:
+ return True
+
+ def _cache_dit_requested(self) -> bool:
+ return (
+ getattr(self, "_minimax_h3_quality_profile", "lossless") != "lossless"
+ or super()._cache_dit_requested()
+ )
+
+ def _maybe_enable_cache_dit(
+ self, num_inference_steps: int | tuple[int, int], batch: Req
+ ) -> None:
+ quality = getattr(batch.sampling_params, "quality", "lossless")
+ if quality not in MINIMAX_H3_QUALITY_PROFILES:
+ raise ValueError(f"unsupported MiniMax-H3 quality profile {quality!r}")
+ explicit_fields = getattr(batch.sampling_params, "_explicit_fields", ())
+ generic_requested = (
+ super()._cache_dit_requested() and "quality" not in explicit_fields
+ )
+ desired_mode = (
+ quality
+ if quality != "lossless"
+ else ("generic" if generic_requested else None)
+ )
+ current_mode = getattr(self, "_minimax_h3_cache_mode", None)
+ self._minimax_h3_quality_profile = quality
+
+ # H3 is monolithic-only, and the scheduler executes one worker batch at
+ # a time. Combined with `quality` in the dynamic-batch signature, this
+ # makes the process-wide hook transition safe at this batch boundary.
+ if self._cache_dit_enabled and current_mode != desired_mode:
+ self.transformer = disable_cache_on_transformer(self.transformer)
+ self._cache_dit_enabled = False
+ self._cached_num_steps = None
+ self._minimax_h3_cache_mode = None
+
+ if desired_mode is None:
+ return
+ super()._maybe_enable_cache_dit(num_inference_steps, batch)
+ if self._cache_dit_enabled:
+ self._minimax_h3_cache_mode = desired_mode
+
+ def _cache_dit_scm_masks(
+ self, primary_num_steps: int, secondary_num_steps: int | None = None
+ ) -> tuple[str, str, list[int] | None, list[int] | None]:
+ if getattr(self, "_minimax_h3_quality_profile", "lossless") != "lossless":
+ return "none", "dynamic", None, None
+ return super()._cache_dit_scm_masks(primary_num_steps, secondary_num_steps)
+
+ def _build_cache_dit_config(
+ self,
+ num_inference_steps: int,
+ steps_computation_mask: list[int] | None,
+ scm_policy: str,
+ *,
+ secondary: bool = False,
+ ) -> CacheDitConfig:
+ quality = getattr(self, "_minimax_h3_quality_profile", "lossless")
+ profile = MINIMAX_H3_QUALITY_PROFILES[quality]
+ if profile is None or secondary:
+ return super()._build_cache_dit_config(
+ num_inference_steps,
+ steps_computation_mask,
+ scm_policy,
+ secondary=secondary,
+ )
+ warmup, threshold, max_cached = profile
+ return CacheDitConfig(
+ enabled=True,
+ Fn_compute_blocks=1,
+ Bn_compute_blocks=0,
+ max_warmup_steps=warmup,
+ residual_diff_threshold=threshold,
+ max_continuous_cached_steps=max_cached,
+ enable_taylorseer=False,
+ taylorseer_order=1,
+ num_inference_steps=num_inference_steps,
+ steps_computation_mask=steps_computation_mask,
+ steps_computation_policy=scm_policy,
+ )
+
+ @torch.no_grad()
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ with self._offload_for_torch_compile_warmup(batch):
+ return self._forward_native(batch, server_args)
+
+ def _forward_native(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
+ MINIMAX_H3_SIGMAS_EXTRA_KEY,
+ MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY,
+ )
+
+ if MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY in batch.extra:
+ for required in (
+ MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
+ MINIMAX_H3_SIGMAS_EXTRA_KEY,
+ ):
+ if required not in batch.extra:
+ raise ValueError(
+ f"direct full-loop denoise requires batch.extra[{required!r}]"
+ )
+ self._run_full_loop(batch, server_args)
+ return batch
+ if (
+ batch.latents is not None
+ or batch.audio_latents is not None
+ or batch.timestep is not None
+ or batch.timesteps is not None
+ ):
+ raise NotImplementedError(
+ "MiniMaxH3DenoisingStage requires the canonical direct pipeline "
+ "to populate text embeddings, denoise state, and sigma schedules."
+ )
+ return batch
+
+ def _run_full_loop(self, batch: Req, server_args: ServerArgs) -> None:
+ """Assemble the cfg-distilled positive input and run the full loop.
+
+ The heavy lifting is decomposed into per-phase helpers:
+ context/payload resolution, condition-row assembly, packed-layout
+ construction, condition noise augmentation, initial-row expansion,
+ the denoise loop itself, and output publication.
+ """
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import (
+ MiniMaxH3DenoiseBranch,
+ minimax_h3_denoise_loop,
+ )
+
+ ctx = _resolve_full_loop_context(batch)
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA")
+ device = torch.device("cuda")
+ sigmas_video = [float(v) for v in ctx.sigmas["video"]]
+ self._maybe_enable_cache_dit_and_torch_compile(
+ len(sigmas_video) - 1,
+ batch,
+ )
+
+ _assemble_condition_rows(ctx)
+
+ emb = ctx.embeddings["positive"]
+ packed = _build_packed_layout(ctx, emb)
+ tags = packed["token_tags"]
+ tags[packed["text_pos"].view(-1)] = (
+ emb["text_token_tags"].view(-1).to(torch.long)
+ )
+
+ sampling = batch.sampling_params
+ imgvid_noise_aug, audio_noise_aug = minimax_h3_condition_noise_aug(sampling)
+ _apply_condition_noise_aug(
+ ctx,
+ sampling=sampling,
+ imgvid_noise_aug=imgvid_noise_aug,
+ audio_noise_aug=audio_noise_aug,
+ )
+
+ placement_managed = self._component_residency_manager is not None
+ if placement_managed:
+ self._manage_dit_use_site(self.transformer, "transformer", batch)
+ try:
+ model = _resolve_denoise_model(
+ self.transformer,
+ device,
+ placement_managed=placement_managed,
+ )
+ positive = MiniMaxH3DenoiseBranch(
+ packed=packed,
+ text_embeddings=emb["hidden_states"],
+ token_tags=tags,
+ device=device,
+ )
+ _precompute_refined_prompt_embeds(
+ model,
+ positive,
+ device=device,
+ )
+ _precompute_rope_cache(
+ model,
+ positive,
+ device=device,
+ )
+ initial_video, initial_audio = _expand_initial_rows(ctx, positive)
+ with (
+ maybe_nvtx_range("denoising_loop", self.current_use_nvtx),
+ self.progress_bar(
+ total=len(sigmas_video) - 1,
+ batch=batch,
+ desc="minimax_h3 denoise",
+ ) as progress_bar,
+ ):
+
+ def on_step(_step, _video_rows, _audio_rows):
+ progress_bar.update()
+ if not batch.is_warmup:
+ self.step_profile()
+
+ video_rows, audio_rows = minimax_h3_denoise_loop(
+ model=model,
+ model_forward=partial(self._forward_dit, batch=batch),
+ positive=positive,
+ initial_video_rows=initial_video,
+ initial_audio_rows=initial_audio,
+ keyframe_cond_rows=ctx.cond_rows,
+ audio_ref_rows=ctx.audio_ref_rows,
+ sigmas_video=sigmas_video,
+ sigmas_audio=[float(v) for v in ctx.sigmas["audio"]],
+ device=device,
+ imgvid_cond_noise_aug_for_inference=float(imgvid_noise_aug),
+ audio_cond_noise_aug_for_inference=float(audio_noise_aug),
+ on_step=on_step,
+ step_profiler=partial(
+ self._profile_denoising_step,
+ batch=batch,
+ ),
+ )
+ finally:
+ self._finish_active_component_use()
+ _publish_full_loop_outputs(
+ ctx,
+ batch=batch,
+ positive=positive,
+ video_rows=video_rows,
+ audio_rows=audio_rows,
+ )
+
+ @contextmanager
+ def _profile_denoising_step(self, step_index: int, *, batch: Req):
+ with (
+ maybe_nvtx_range(
+ f"denoising_step_{step_index}",
+ self.current_use_nvtx,
+ ),
+ StageProfiler(
+ f"denoising_step_{step_index}",
+ logger=logger,
+ metrics=batch.metrics,
+ perf_dump_path_provided=batch.perf_dump_path is not None,
+ record_as_step=True,
+ ),
+ ):
+ yield
+
+ def _forward_dit(
+ self,
+ model: Any,
+ call_kwargs: dict[str, Any],
+ step_index: int,
+ *,
+ batch: Req,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Route the custom full loop through the native denoising runner."""
+
+ from sglang.multimodal_gen.runtime.managers.forward_context import (
+ set_forward_context,
+ )
+
+ with set_forward_context(
+ current_timestep=step_index,
+ attn_metadata=None,
+ forward_batch=batch,
+ ):
+ runner = self._maybe_get_bcg_runner(model)
+ if runner is None:
+ return model(**call_kwargs)
+ return self._bcg_run(runner, call_kwargs, model)
+
+ def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+ result.add_check(
+ "audio_latents",
+ batch.audio_latents,
+ [V.is_tensor, V.with_dims(3)],
+ )
+ return result
+
+
+class _FullLoopContext:
+ """Mutable per-request state threaded through the full-loop phases."""
+
+ __slots__ = (
+ "plan",
+ "keyframe",
+ "ref_image",
+ "ref_audio",
+ "ref_video",
+ "is_ref2va",
+ "embeddings",
+ "state",
+ "sigmas",
+ "latent_t",
+ "latent_h",
+ "latent_w",
+ "audio_t",
+ "ref2va_positive_blocks",
+ "cond_rows",
+ "audio_ref_rows",
+ "include_cond",
+ "keyframe_frame_indices",
+ "keyframe_frame_count",
+ )
+
+ def __init__(self) -> None:
+ for name in self.__slots__:
+ setattr(self, name, None)
+ self.is_ref2va = False
+ self.include_cond = False
+
+
+def _resolve_full_loop_context(batch: Req) -> _FullLoopContext:
+ """Read/validate extras and the denoise state into a loop context.
+
+ Enforces the task-payload exclusivity rules (keyframe vs reference
+ exclusivity) and cross-checks the resolved latent dims.
+ """
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
+ MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY,
+ MINIMAX_H3_SIGMAS_EXTRA_KEY,
+ MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ ctx = _FullLoopContext()
+ ctx.embeddings = batch.extra[MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY]
+ ctx.state = batch.extra[MINIMAX_H3_DENOISE_STATE_EXTRA_KEY]
+ ctx.sigmas = batch.extra[MINIMAX_H3_SIGMAS_EXTRA_KEY]
+ ctx.plan = minimax_h3_plan_from_batch(batch)
+ ctx.keyframe = batch.extra.get(MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY)
+ ctx.ref_image = batch.extra.get(MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY)
+ ctx.ref_audio = batch.extra.get(MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY)
+ ctx.ref_video = batch.extra.get(MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY)
+ ctx.is_ref2va = (
+ ctx.ref_image is not None
+ or ctx.ref_audio is not None
+ or ctx.ref_video is not None
+ )
+ if ctx.is_ref2va and ctx.keyframe is not None:
+ raise ValueError("keyframe and reference extras are mutually exclusive")
+ _validate_fl2va_keyframe_payload(ctx.plan, ctx.keyframe)
+
+ ctx.latent_t = int(ctx.state["latent_t"])
+ ctx.latent_h = int(ctx.state["latent_h"])
+ ctx.latent_w = int(ctx.state["latent_w"])
+ ctx.audio_t = int(ctx.state["audio_t"])
+ return ctx
+
+
+def _assemble_condition_rows(ctx: _FullLoopContext) -> None:
+ """Populate cond/audio reference rows and keyframe metadata per task."""
+
+ ctx.include_cond = (
+ ctx.keyframe is not None
+ or ctx.ref_image is not None
+ or ctx.ref_video is not None
+ )
+ if ctx.is_ref2va:
+ if ctx.plan is None:
+ raise ValueError(
+ "ref2va reference extras require a resolved plan; "
+ "plan-less ref2va requests are unsupported"
+ )
+ ctx.ref2va_positive_blocks, ctx.cond_rows, ctx.audio_ref_rows = (
+ _ref2va_ordered_blocks_and_rows(
+ plan=ctx.plan,
+ ref_image=ctx.ref_image,
+ ref_audio=ctx.ref_audio,
+ ref_video=ctx.ref_video,
+ )
+ )
+ ctx.include_cond = ctx.cond_rows is not None
+ else:
+ ctx.cond_rows = ctx.keyframe["rows"] if ctx.include_cond else None
+ ctx.audio_ref_rows = (
+ ctx.ref_audio["rows"] if ctx.ref_audio is not None else None
+ )
+ if ctx.keyframe is not None:
+ raw_indices = ctx.keyframe.get("semantic_frame_indices")
+ ctx.keyframe_frame_indices = [int(v) for v in raw_indices]
+ ctx.keyframe_frame_count = int(ctx.keyframe["frame_count"])
+
+
+def _build_packed_layout(
+ ctx: _FullLoopContext,
+ emb: Mapping[str, Any],
+) -> dict[str, torch.Tensor]:
+ """Build the per-task packed layout for the positive branch."""
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import (
+ minimax_h3_packed_sequence,
+ minimax_h3_packed_sequence_ref2va_blocks,
+ )
+
+ if ctx.is_ref2va:
+ if ctx.ref2va_positive_blocks is None:
+ raise ValueError("ref2va ordered reference blocks missing")
+ packed = minimax_h3_packed_sequence_ref2va_blocks(
+ text_len=int(emb["text_len"]),
+ latent_t=ctx.latent_t,
+ latent_h=ctx.latent_h,
+ latent_w=ctx.latent_w,
+ audio_t=ctx.audio_t,
+ ref_blocks=ctx.ref2va_positive_blocks,
+ )
+ else:
+ packed = minimax_h3_packed_sequence(
+ text_len=int(emb["text_len"]),
+ latent_t=ctx.latent_t,
+ latent_h=ctx.latent_h,
+ latent_w=ctx.latent_w,
+ audio_t=ctx.audio_t,
+ include_keyframe_cond=ctx.include_cond,
+ keyframe_frame_indices=(
+ ctx.keyframe_frame_indices if ctx.include_cond else None
+ ),
+ frame_count=ctx.keyframe_frame_count,
+ )
+ return packed
+
+
+def _condition_audio_lengths(ctx: _FullLoopContext) -> list[int]:
+ """Per-task condition audio T list for the noise-aug recipe."""
+
+ condition_audio_t: list[int] = []
+ if ctx.ref2va_positive_blocks is not None:
+ for block in ctx.ref2va_positive_blocks:
+ if str(block["kind"]) in {"audio", "video", "video_audio"}:
+ ref_audio_t = int(block["ref_audio_t"])
+ if ref_audio_t > 0:
+ condition_audio_t.append(ref_audio_t)
+ elif isinstance(ctx.ref_audio, Mapping):
+ entries = ctx.ref_audio.get("audios")
+ if isinstance(entries, list):
+ condition_audio_t.extend(
+ int(entry["ref_audio_t"])
+ for entry in entries
+ if int(entry["ref_audio_t"]) > 0
+ )
+ elif ctx.ref_audio.get("ref_audio_t") is not None:
+ ref_audio_t = int(ctx.ref_audio["ref_audio_t"])
+ if ref_audio_t > 0:
+ condition_audio_t.append(ref_audio_t)
+ return condition_audio_t
+
+
+def _apply_condition_noise_aug(
+ ctx: _FullLoopContext,
+ *,
+ sampling: Any,
+ imgvid_noise_aug: float,
+ audio_noise_aug: float,
+) -> None:
+ """Apply condition noise augmentation to cond rows."""
+
+ noise_visual_conditions = (
+ ctx.cond_rows is not None and float(imgvid_noise_aug) < 1.0
+ )
+ noise_audio_conditions = (
+ ctx.audio_ref_rows is not None and float(audio_noise_aug) < 1.0
+ )
+ if not (noise_visual_conditions or noise_audio_conditions):
+ return
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.condition_noise import (
+ minimax_h3_audio_cond_noise_aug_rows,
+ minimax_h3_imgvid_cond_noise_aug_rows,
+ )
+
+ noise_seed = getattr(ctx.plan, "seed", None) if ctx.plan is not None else None
+ if noise_seed is None:
+ noise_seed = getattr(sampling, "seed", None)
+ if noise_seed is None:
+ noise_seed = 42
+
+ if noise_visual_conditions:
+ condition_shapes = _imgvid_condition_shapes(
+ ref2va_blocks=ctx.ref2va_positive_blocks,
+ keyframe=ctx.keyframe,
+ is_ref2va=ctx.is_ref2va,
+ )
+ imgvid_cond_num_frames = len(condition_shapes)
+ if not condition_shapes:
+ raise ValueError("imgvid condition rows are missing shape metadata")
+ # Imgvid conditions (ref2va blocks / keyframes) contribute one frame
+ # count per condition entry.
+ ctx.cond_rows = minimax_h3_imgvid_cond_noise_aug_rows(
+ ctx.cond_rows,
+ condition_shapes=condition_shapes,
+ target_latent_t=ctx.latent_t,
+ imgvid_cond_num_frames=imgvid_cond_num_frames,
+ seed=int(noise_seed),
+ noise_aug=float(imgvid_noise_aug),
+ )
+
+ if noise_audio_conditions:
+ condition_audio_t = _condition_audio_lengths(ctx)
+ if not condition_audio_t:
+ raise ValueError("audio condition rows are missing length metadata")
+ ctx.audio_ref_rows = minimax_h3_audio_cond_noise_aug_rows(
+ ctx.audio_ref_rows,
+ condition_audio_t=condition_audio_t,
+ seed=int(noise_seed),
+ noise_aug=float(audio_noise_aug),
+ )
+
+
+def _expand_initial_rows(
+ ctx: _FullLoopContext,
+ positive: Any,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Scatter target-row noise into the full packed layout when cond rows exist."""
+
+ initial_video = ctx.state["initial_video_rows"]
+ if ctx.include_cond:
+ # layout target rows = noise; cond anchors appended by the loop
+ n_video = positive.img_pos.shape[0]
+ full = torch.zeros(int(n_video), initial_video.shape[1], dtype=torch.float32)
+ full[positive.update_mask] = initial_video
+ initial_video = full
+
+ initial_audio = ctx.state["initial_audio_rows"]
+ if ctx.audio_ref_rows is not None:
+ n_audio = positive.audio_pos.shape[0]
+ full_audio = torch.zeros(
+ int(n_audio), initial_audio.shape[1], dtype=torch.float32
+ )
+ full_audio[positive.audio_update_mask] = initial_audio
+ initial_audio = full_audio
+ return initial_video, initial_audio
+
+
+def _publish_full_loop_outputs(
+ ctx: _FullLoopContext,
+ *,
+ batch: Req,
+ positive: Any,
+ video_rows: torch.Tensor,
+ audio_rows: torch.Tensor,
+) -> None:
+ """Compose the generated target latents onto the batch."""
+
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
+ minimax_h3_unpack_audio_tokens,
+ minimax_h3_unpatchify_video_tokens,
+ )
+
+ target_rows = video_rows[positive.video_target_slice]
+ # Keep latents on CUDA so decode can reuse them without a device round-trip;
+ # decode autocast is enabled only for CUDA inputs.
+ batch.latents = minimax_h3_unpatchify_video_tokens(
+ target_rows,
+ latent_shape=[ctx.latent_t, ctx.latent_h // 2, ctx.latent_w // 2, 24],
+ patch_size=[1, 2, 2],
+ )
+ audio_target_rows = audio_rows[positive.audio_target_slice]
+ batch.audio_latents = minimax_h3_unpack_audio_tokens(
+ audio_target_rows, audio_t=ctx.audio_t * 2, audio_channel=2
+ )
+
+
+__all__ = [
+ "MiniMaxH3DenoisingStage",
+ "minimax_h3_condition_noise_aug",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py
new file mode 100644
index 000000000..86adbe91b
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/latent_preparation.py
@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ VerificationResult,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class MiniMaxH3LatentPreparationStage(PipelineStage):
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ raise NotImplementedError(
+ f"{self.__class__.__name__} is a MiniMax H3 contract stage "
+ "and has no implementation yet."
+ )
+ self._prepare_denoise_state_from_plan(batch, plan)
+ self._publish_native_latent_state(batch)
+ return batch
+
+ def run_grouped_requests(
+ self,
+ batches: list[Req],
+ server_args: ServerArgs,
+ ) -> list[Req]:
+ """Preserve H3's independent per-modality RNG streams per request."""
+ return [self(batch, server_args) for batch in batches]
+
+ @staticmethod
+ def _publish_native_latent_state(batch: Req) -> None:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
+ )
+
+ state = batch.extra.get(MINIMAX_H3_DENOISE_STATE_EXTRA_KEY)
+ if not isinstance(state, dict):
+ raise ValueError("MiniMax H3 denoise state must be a mapping")
+ video_rows = state.get("initial_video_rows")
+ audio_rows = state.get("initial_audio_rows")
+ if not isinstance(video_rows, torch.Tensor) or video_rows.ndim != 2:
+ raise ValueError("MiniMax H3 initial_video_rows must be a rank-2 tensor")
+ if not isinstance(audio_rows, torch.Tensor) or audio_rows.ndim != 2:
+ raise ValueError("MiniMax H3 initial_audio_rows must be a rank-2 tensor")
+
+ latent_t = int(state["latent_t"])
+ latent_h = int(state["latent_h"])
+ latent_w = int(state["latent_w"])
+ audio_t = int(state["audio_t"])
+ batch.latents = video_rows
+ batch.audio_latents = audio_rows
+ batch.raw_latent_shape = (1, 24, latent_t, latent_h, latent_w)
+ batch.raw_audio_latent_shape = (2, 32, audio_t)
+
+ def _prepare_denoise_state_from_plan(self, batch: Req, plan) -> None:
+ """Direct initial-noise materialization (t2va recipe):
+ torch.Generator().manual_seed(seed); video rows drawn first,
+ then audio rows, CPU fp32. Every task consumes the final latent grid
+ frozen by the pre-queue shape resolver."""
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_DENOISE_STATE_EXTRA_KEY,
+ )
+
+ if MINIMAX_H3_DENOISE_STATE_EXTRA_KEY in batch.extra:
+ return
+ shape = plan.shape
+ geometry = str(shape["geometry"])
+ if geometry != "resolved_v2":
+ raise ValueError(
+ "MiniMax H3 latent preparation requires pre-queue resolved_v2 "
+ f"geometry, got {geometry!r}"
+ )
+ latent_h = int(shape["height"]) // 16
+ latent_w = int(shape["width"]) // 16
+ if shape.get("video_latent_t") is None or shape.get("audio_latent_t") is None:
+ raise ValueError(
+ "MiniMax H3 latent preparation requires pre-queue resolved "
+ "temporal dimensions"
+ )
+ latent_t = int(shape["video_latent_t"])
+ audio_t = int(shape["audio_latent_t"])
+
+ seed = plan.seed
+ if seed is None:
+ seed = 42 # pinned default seed
+ video_rows_n = latent_t * (latent_h // 2) * (latent_w // 2)
+ audio_rows_n = audio_t * 2
+ # Noise semantics:
+ # - video noise is drawn on the RAW latent tensor
+ # [1, 24, T, H_lat, W_lat] in tensor layout, then patchified
+ # into packed row order;
+ # - audio uses an INDEPENDENT generator re-seeded with the same
+ # seed (each modality re-seeds its own generator);
+ # - no extra cond-frame noise is drawn for image-conditioned
+ # requests.
+ # The same seed always reproduces the same noise.
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
+ minimax_h3_patchify_video_latent,
+ )
+
+ gen_v = torch.Generator().manual_seed(int(seed))
+ video_tensor = torch.randn(
+ 1,
+ 24,
+ latent_t,
+ latent_h,
+ latent_w,
+ generator=gen_v,
+ dtype=torch.float32,
+ )
+ video_noise = minimax_h3_patchify_video_latent(
+ video_tensor, patch_size=[1, 2, 2]
+ ).to(torch.float32)
+ gen_a = torch.Generator().manual_seed(int(seed))
+ audio_noise = torch.randn(
+ audio_rows_n, 32, generator=gen_a, dtype=torch.float32
+ )
+ if list(video_noise.shape) != [video_rows_n, 96]:
+ raise ValueError(
+ f"aligned video noise shape {list(video_noise.shape)} != "
+ f"[{video_rows_n}, 96]"
+ )
+ batch.extra[MINIMAX_H3_DENOISE_STATE_EXTRA_KEY] = {
+ "initial_video_rows": video_noise,
+ "initial_audio_rows": audio_noise,
+ "latent_t": latent_t,
+ "latent_h": latent_h,
+ "latent_w": latent_w,
+ "audio_t": audio_t,
+ }
+
+ def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check(
+ "prompt_or_embeds",
+ None,
+ lambda _: V.string_or_list_strings(batch.prompt)
+ or V.list_not_empty(batch.prompt_embeds),
+ )
+ result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors)
+ result.add_check(
+ "num_videos_per_prompt", batch.num_outputs_per_prompt, V.positive_int
+ )
+ result.add_check("generator", batch.generator, V.generator_or_list_generators)
+ result.add_check("num_frames", batch.num_frames, V.positive_int)
+ result.add_check("height", batch.height, V.positive_int)
+ result.add_check("width", batch.width, V.positive_int)
+ result.add_check("latents", batch.latents, V.none_or_tensor)
+ return result
+
+ def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(2)])
+ result.add_check(
+ "audio_latents", batch.audio_latents, [V.is_tensor, V.with_dims(2)]
+ )
+ result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple)
+ result.add_check(
+ "raw_audio_latent_shape", batch.raw_audio_latent_shape, V.is_tuple
+ )
+ return result
+
+
+__all__ = [
+ "MiniMaxH3LatentPreparationStage",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py
new file mode 100644
index 000000000..d3499642e
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/replica_broadcast.py
@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Request-replica broadcast helpers for MiniMax H3 encoding stages."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def minimax_h3_replica_ctx() -> tuple[int, int]:
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_world_group,
+ model_parallel_is_initialized,
+ )
+
+ if not model_parallel_is_initialized():
+ return 1, 0
+ # ServerArgs currently rejects DP>1 and H3 rejects disaggregation, so the
+ # world group contains exactly one request replica (TP/CFG/SP ranks).
+ group = get_world_group()
+ return int(group.world_size), int(group.rank_in_group)
+
+
+def minimax_h3_replica_broadcast_extra(batch: Any, key: str) -> None:
+ world, rank = minimax_h3_replica_ctx()
+ if world <= 1:
+ return
+
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_world_group,
+ )
+
+ group = get_world_group()
+ payload = {"value": batch.extra.get(key)} if rank == 0 else None
+ payload = group.broadcast_tensor_dict(payload, src=0)
+ value = payload.get("value") if isinstance(payload, dict) else None
+ if value is None:
+ raise RuntimeError(f"replica broadcast of batch.extra[{key!r}] got None")
+ if rank != 0:
+ batch.extra[key] = value
+
+
+def minimax_h3_replica_broadcast_error(error: str | None) -> str | None:
+ world, rank = minimax_h3_replica_ctx()
+ if world <= 1:
+ return error
+
+ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ get_world_group,
+ )
+
+ group = get_world_group()
+ return group.broadcast_object(error if rank == 0 else None, src=0)
+
+
+__all__ = [
+ "minimax_h3_replica_broadcast_error",
+ "minimax_h3_replica_broadcast_extra",
+ "minimax_h3_replica_ctx",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py
new file mode 100644
index 000000000..007d9cc53
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py
@@ -0,0 +1,524 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
+ TextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class MiniMaxH3TextEncodingStage(TextEncodingStage):
+ deduplicated_output_fields = ("prompt_embeds", "prompt_seq_lens")
+ deduplicated_extra_output_keys = (MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY,)
+
+ def __init__(self, text_encoder, tokenizer, processor) -> None:
+ super().__init__(
+ text_encoders=[text_encoder],
+ tokenizers=[tokenizer],
+ )
+ self.text_encoder = text_encoder
+ self.tokenizer = tokenizer
+ if processor is None:
+ raise ValueError(
+ "MiniMaxH3TextEncodingStage requires the pipeline processor "
+ "component (model_index.json: processor)"
+ )
+ self.processor = processor
+
+ @torch.no_grad()
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is not None:
+ try:
+ self._encode_from_plan(batch, plan)
+ self._publish_native_text_conditioning(batch)
+ except Exception:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_cleanup_temp_dirs,
+ )
+
+ minimax_h3_cleanup_temp_dirs(batch)
+ raise
+ return batch
+ if batch.sampling_params is not None and (
+ batch.sampling_params.prompt is not None
+ or batch.sampling_params.prompt_path is not None
+ ):
+ raise NotImplementedError(
+ "MiniMaxH3TextEncodingStage direct Qwen3VL encoder forward requires "
+ "a canonical minimax_h3 request (resolved plan); legacy prompt-only "
+ "requests are unsupported."
+ )
+ return batch
+
+ def build_dedup_fingerprint(self, batch: Req, server_args: ServerArgs):
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ return super().build_dedup_fingerprint(batch, server_args)
+ materials = tuple(
+ (
+ item.condition_index,
+ item.role,
+ item.condition_type,
+ item.uri,
+ item.material_chain,
+ item.frame_index,
+ item.resolved_frame_index,
+ item.start_time_seconds,
+ )
+ for item in plan.materials
+ )
+ return (
+ plan.task,
+ plan.prompt,
+ materials,
+ self.freeze_for_dedup(plan.shape),
+ )
+
+ def run_grouped_requests(
+ self,
+ batches: list[Req],
+ server_args: ServerArgs,
+ ) -> list[Req]:
+ """Distribute independent H3 presentations over replicated encoders.
+
+ H3 presentations have variable multimodal layouts, so they cannot be
+ stacked into the generic text batch without changing padding/kernels.
+ Assigning one complete request to a rank preserves the exact
+ single-request encoder path, then broadcasts that request's native
+ payload to the other ranks.
+ """
+ grouped = self._group_requests_by_fingerprint(
+ batches,
+ lambda batch: self.build_dedup_fingerprint(batch, server_args),
+ )
+ if not grouped:
+ return []
+
+ encoder_config = server_args.pipeline_config.text_encoder_configs[0]
+ dp_group = self._text_encode_dp_group(
+ server_args,
+ encoder_config,
+ len(grouped),
+ self.text_encoder,
+ )
+ if dp_group is None:
+ return super().run_grouped_requests(batches, server_args)
+
+ results: list[Req | None] = [None] * len(batches)
+ for group_index, (_, equivalent) in enumerate(grouped):
+ owner = group_index % dp_group.world_size
+ first_index, first_batch = equivalent[0]
+ owner_exception = None
+ owner_error = None
+ payload = None
+ if dp_group.rank_in_group == owner:
+ try:
+ first_result = self(first_batch, server_args)
+ payload = first_result.extra.get(
+ MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY
+ )
+ if not isinstance(payload, dict):
+ raise ValueError(
+ "MiniMax H3 text encode produced no native payload"
+ )
+ except Exception as exc:
+ owner_exception = exc
+ owner_error = f"{type(exc).__name__}: {exc}"
+
+ owner_error = dp_group.broadcast_object(owner_error, src=owner)
+ if owner_error is not None:
+ if owner_exception is not None:
+ raise owner_exception
+ raise RuntimeError(
+ f"MiniMax H3 text encode failed on rank {owner}: {owner_error}"
+ )
+ payload = dp_group.broadcast_tensor_dict(payload, src=owner)
+ if not isinstance(payload, dict):
+ raise RuntimeError("MiniMax H3 text payload broadcast failed")
+
+ if dp_group.rank_in_group != owner:
+ first_batch.extra[MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY] = payload
+ self._publish_native_text_conditioning(first_batch)
+ first_result = first_batch
+ results[first_index] = first_result
+
+ for index, batch in equivalent[1:]:
+ self.copy_deduplicated_outputs(first_result, batch)
+ results[index] = batch
+
+ return [result for result in results if result is not None]
+
+ def _log_dp_choice(self, batch_size: int, world_size: int) -> None:
+ if self._dp_choice_logged:
+ return
+ self._dp_choice_logged = True
+ logger.info(
+ "encoder_parallel: distributing %d independent MiniMax H3 "
+ "presentations over %d replicated encoder ranks",
+ batch_size,
+ world_size,
+ )
+
+ @staticmethod
+ def _publish_native_text_conditioning(batch: Req) -> None:
+ """Mirror H3's rich payload onto the native text-stage fields.
+
+ H3 keeps token tags and presentation metadata in ``Req.extra``, but
+ the shared TextEncodingStage contract still owns ``prompt_embeds``.
+ Publishing the same tensor there preserves native verification,
+ grouped-request deduplication, and downstream memory accounting
+ without duplicating the embedding storage.
+ """
+ payload = batch.extra.get(MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY)
+ positive = payload.get("positive") if isinstance(payload, dict) else None
+ hidden_states = (
+ positive.get("hidden_states") if isinstance(positive, dict) else None
+ )
+ text_len = positive.get("text_len") if isinstance(positive, dict) else None
+ if not isinstance(hidden_states, torch.Tensor) or hidden_states.ndim < 2:
+ raise ValueError(
+ "MiniMax H3 text payload must contain positive.hidden_states "
+ "with at least two dimensions"
+ )
+ if not isinstance(text_len, int) or text_len != int(hidden_states.shape[0]):
+ raise ValueError(
+ "MiniMax H3 text payload positive.text_len must match the "
+ "hidden-state sequence dimension"
+ )
+ batch.prompt_embeds = [hidden_states]
+ batch.prompt_seq_lens = [[text_len]]
+
+ def _encode_from_plan(self, batch: Req, plan) -> None:
+ """Encode the positive Qwen3VL presentation into layer-50 states.
+
+ MiniMax H3 only supports the CFG-distilled model path, so every task
+ emits exactly one positive embedding payload. ComponentManager owns
+ residency/offload, while every folded-TP rank enters the encoder
+ collectives and receives the same replicated hidden states.
+ """
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.presentation import (
+ minimax_h3_text_only_ids,
+ )
+
+ prompt = plan.prompt
+ keyframes = [
+ m for m in plan.materials if m.material_chain == "image.target_canvas"
+ ]
+ if plan.task == "fl2va":
+ frame_indices = tuple(material.frame_index for material in keyframes)
+ if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+ raise ValueError(
+ "fl2va text encoding requires an ordered keyframe signature "
+ f"in {MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got "
+ f"{frame_indices!r}"
+ )
+ elif keyframes:
+ raise ValueError(
+ f"task {plan.task!r} cannot carry image.target_canvas materials"
+ )
+ if MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY in batch.extra:
+ return
+ if self.text_encoder is None:
+ raise ValueError(
+ "MiniMaxH3TextEncodingStage direct encode requires a text_encoder "
+ "component"
+ )
+ encode_ids = getattr(self.text_encoder, "encode_ids", None)
+ if not callable(encode_ids):
+ raise TypeError(
+ "MiniMax H3 text_encoder component must expose callable "
+ "encode_ids(...) for direct encode (MiniMaxH3Qwen3VLEncoder)"
+ )
+ if self.tokenizer is None:
+ raise ValueError(
+ "MiniMaxH3TextEncodingStage direct encode requires a tokenizer component"
+ )
+ self._manage_text_encoder_use(0)
+ with set_forward_context(current_timestep=0, attn_metadata=None):
+ if plan.task == "ref2va":
+ embeddings = self._encode_ref2va(batch, plan, encode_ids)
+ elif keyframes:
+ embeddings = self._encode_fl2va_keyframes(
+ batch,
+ plan,
+ encode_ids,
+ prompt=prompt,
+ )
+ else:
+ positive_ids = minimax_h3_text_only_ids(self.tokenizer, prompt)
+ embeddings = {
+ "positive": {
+ "hidden_states": encode_ids(positive_ids),
+ "text_len": int(positive_ids.shape[0]),
+ "text_token_tags": torch.ones(
+ int(positive_ids.shape[0]), dtype=torch.long
+ ),
+ }
+ }
+ batch.extra[MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY] = embeddings
+
+ def _encode_fl2va_keyframes(
+ self,
+ batch: Req,
+ plan,
+ encode_ids,
+ *,
+ prompt: str,
+ ) -> dict:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.canvas import (
+ minimax_h3_prepared_keyframes,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.presentation import (
+ minimax_h3_multi_image_presentation,
+ )
+
+ # The SAME prepared target-canvas images feed
+ # Qwen and the visual-condition tokenizer; preparation is cached per request.
+ prepared = minimax_h3_prepared_keyframes(batch, plan)
+ images = [item["image"] for item in prepared["images"]]
+ frame_indices = tuple(prepared.get("semantic_frame_indices") or ())
+ if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES or len(
+ images
+ ) != len(frame_indices):
+ raise ValueError(
+ "fl2va Qwen preparation requires one or two ordered images with "
+ "a supported semantic_frame_indices signature, got "
+ f"{frame_indices!r}"
+ )
+ processor = self.processor
+ vision = processor.image_processor(images=images, return_tensors="pt")
+ pixel_values = vision["pixel_values"]
+ image_grid_thw = vision["image_grid_thw"]
+ if int(image_grid_thw.shape[0]) != len(images):
+ raise ValueError(
+ f"expected {len(images)} image grids, got {list(image_grid_thw.shape)}"
+ )
+ merge = int(processor.image_processor.merge_size) ** 2
+ image_token_counts = [
+ int(image_grid_thw[i].prod().item()) // merge for i in range(len(images))
+ ]
+ pos_ids, pos_tags = minimax_h3_multi_image_presentation(
+ self.tokenizer,
+ prompt=prompt,
+ image_token_counts=image_token_counts,
+ )
+ pos_hidden = encode_ids(
+ pos_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ )
+ return {
+ "positive": {
+ "hidden_states": pos_hidden,
+ "text_len": int(pos_ids.shape[0]),
+ "text_token_tags": pos_tags,
+ },
+ }
+
+ def _encode_ref2va(self, batch: Req, plan, encode_ids) -> dict:
+ """Encode the positive ref2va presentation.
+
+ Per condition in order — image i: ': ' label +
+ vision block (prepared reference image); audio j: ': ' label
+ only — then the verbatim prompt.
+ """
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.presentation import (
+ minimax_h3_ref2va_presentation,
+ minimax_h3_ref2va_video_presentation,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.reference_encoding import (
+ minimax_h3_prepared_reference_image,
+ minimax_h3_prepared_reference_videos,
+ minimax_h3_sample_reference_video_frames,
+ )
+
+ prepared_videos = None
+ if any(
+ material.material_chain
+ in ("video.reference_preserve", "video_audio.reference_preserve")
+ for material in plan.materials
+ ):
+ prepared_videos = minimax_h3_prepared_reference_videos(batch, plan)
+ video_has_audio: dict[int, bool] = {}
+ for video_index, item in enumerate((prepared_videos or {}).get("videos") or []):
+ if item.get("condition_index") is None:
+ continue
+ if "input_has_audio" not in item:
+ raise ValueError(
+ f"prepared reference video {video_index} is missing "
+ "'input_has_audio'; the canonical minimax_h3 producer must "
+ "supply the audio probe for every video condition"
+ )
+ video_has_audio[int(item["condition_index"])] = bool(
+ item["input_has_audio"]
+ )
+ condition_labels: list[tuple[str, int]] = []
+ counters = {"image": 0, "audio": 0, "video": 0}
+ has_image = False
+ has_video = False
+ for material in plan.materials:
+ if material.material_chain == "image.reference_preserve":
+ counters["image"] += 1
+ condition_labels.append(("image", counters["image"]))
+ has_image = True
+ elif material.material_chain == "audio":
+ counters["audio"] += 1
+ condition_labels.append(("audio", counters["audio"]))
+ elif material.material_chain in (
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+ ):
+ # A plain video contributes an Audio label only when its
+ # probed source actually has a soundtrack. ``video_audio`` is
+ # an explicit caller promise and remains fail-closed in the
+ # audio stage if its stream is missing.
+ if material.material_chain == "video_audio.reference_preserve":
+ contributes_audio = True
+ else:
+ condition_index = int(material.condition_index)
+ if condition_index not in video_has_audio:
+ raise KeyError(
+ "prepared reference videos carry no "
+ f"'input_has_audio' probe for condition "
+ f"{condition_index}; the canonical minimax_h3 "
+ "producer must supply it"
+ )
+ contributes_audio = video_has_audio[condition_index]
+ if contributes_audio:
+ counters["audio"] += 1
+ condition_labels.append(("audio", counters["audio"]))
+ counters["video"] += 1
+ condition_labels.append(("video", counters["video"]))
+ has_video = True
+ else:
+ raise NotImplementedError(
+ f"ref2va does not support chain {material.material_chain!r}"
+ )
+
+ pixel_values = None
+ image_grid_thw = None
+ n_image_tokens = None
+ processor = self.processor
+
+ if has_image:
+ prepared = minimax_h3_prepared_reference_image(batch, plan)
+ images = [item["image"] for item in prepared["images"]]
+ proc = processor
+ vision = proc.image_processor(images=images, return_tensors="pt")
+ pixel_values = vision["pixel_values"]
+ image_grid_thw = vision["image_grid_thw"]
+ if int(image_grid_thw.shape[0]) != len(images):
+ raise ValueError(
+ f"expected {len(images)} image grids, got "
+ f"{list(image_grid_thw.shape)}"
+ )
+ merge = int(proc.image_processor.merge_size) ** 2
+ counts = [
+ int(image_grid_thw[i].prod().item()) // merge
+ for i in range(len(images))
+ ]
+ # presentation takes an int for one image, a list for several
+ n_image_tokens = counts[0] if len(counts) == 1 else counts
+
+ pixel_values_videos = None
+ video_grid_thw = None
+ video_block_token_counts = None
+ video_block_timestamps = None
+ if has_video:
+ prepared = prepared_videos or minimax_h3_prepared_reference_videos(
+ batch, plan
+ )
+ videos = []
+ sampled_videos = []
+ for item in prepared["videos"]:
+ sampled = minimax_h3_sample_reference_video_frames(item["frames"])
+ videos.append(torch.from_numpy(sampled["frames"]).permute(0, 3, 1, 2))
+ sampled_videos.append(sampled)
+ proc = processor
+ vout = proc.video_processor(
+ videos=videos,
+ do_sample_frames=False,
+ input_data_format="channels_first",
+ return_tensors="pt",
+ )
+ pixel_values_videos = vout["pixel_values_videos"]
+ video_grid_thw = vout["video_grid_thw"]
+ if int(video_grid_thw.shape[0]) != len(videos):
+ raise ValueError(
+ f"expected {len(videos)} video grids, got "
+ f"{list(video_grid_thw.shape)}"
+ )
+ merge = int(proc.image_processor.merge_size) ** 2
+ video_block_token_counts = []
+ video_block_timestamps = []
+ for index, sampled in enumerate(sampled_videos):
+ n_blocks = int(video_grid_thw[index, 0])
+ per_block = (
+ int(video_grid_thw[index, 1])
+ * int(video_grid_thw[index, 2])
+ // merge
+ )
+ timestamps = [float(ts) for ts in sampled["block_timestamps"]]
+ if len(timestamps) != n_blocks:
+ raise ValueError(
+ f"video block count mismatch: processor {n_blocks} vs "
+ f"timestamps {len(timestamps)} for video {index}"
+ )
+ video_block_token_counts.append([per_block] * n_blocks)
+ video_block_timestamps.append(timestamps)
+
+ if has_video:
+ pos_ids, pos_tags = minimax_h3_ref2va_video_presentation(
+ self.tokenizer,
+ prompt=plan.prompt,
+ condition_labels=condition_labels,
+ image_token_count=n_image_tokens,
+ video_block_token_counts=video_block_token_counts,
+ video_block_timestamps=video_block_timestamps,
+ )
+ else:
+ pos_ids, pos_tags = minimax_h3_ref2va_presentation(
+ self.tokenizer,
+ prompt=plan.prompt,
+ condition_labels=condition_labels,
+ image_token_count=n_image_tokens,
+ )
+ pos_hidden = encode_ids(
+ pos_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ pixel_values_videos=pixel_values_videos,
+ video_grid_thw=video_grid_thw,
+ )
+ return {
+ "positive": {
+ "hidden_states": pos_hidden,
+ "text_len": int(pos_ids.shape[0]),
+ "text_token_tags": pos_tags,
+ },
+ }
+
+
+__all__ = ["MiniMaxH3TextEncodingStage"]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py
new file mode 100644
index 000000000..512fb5aa9
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/timestep_preparation.py
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+ VerificationResult,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+from ..constants import MINIMAX_H3_SIGMAS_EXTRA_KEY
+
+
+class MiniMaxH3TimestepPreparationStage(PipelineStage):
+ deduplicated_tensor_tree_output_fields = ("timesteps", "sigmas")
+ deduplicated_extra_tensor_tree_output_keys = (MINIMAX_H3_SIGMAS_EXTRA_KEY,)
+
+ def __init__(self, sigma_shift_scales=None) -> None:
+ super().__init__()
+ # Per-model sigma shift override (model_index.json "_minimax_h3" release
+ # block, sigma_shift_scales): the schedule constants are a MODEL
+ # serving contract — fl2va and ref2va use video 12 / audio 3 by default.
+ self.sigma_shift_scales = sigma_shift_scales
+
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ raise NotImplementedError(
+ f"{self.__class__.__name__} is a MiniMax H3 contract stage "
+ "and has no implementation yet."
+ )
+ self._generate_sigmas_from_plan(batch, plan)
+ self._publish_native_timestep_state(batch)
+ return batch
+
+ def build_dedup_fingerprint(self, batch: Req, server_args: ServerArgs):
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ raise NotImplementedError(
+ f"{self.__class__.__name__} is a MiniMax H3 contract stage "
+ "and has no implementation yet."
+ )
+ return (
+ batch.num_inference_steps,
+ plan.flow_shift,
+ plan.audio_flow_shift,
+ plan.default_flow_shift,
+ plan.default_audio_flow_shift,
+ self.freeze_for_dedup(self.sigma_shift_scales),
+ )
+
+ @staticmethod
+ def _publish_native_timestep_state(batch: Req) -> None:
+ sigmas = batch.extra.get(MINIMAX_H3_SIGMAS_EXTRA_KEY)
+ if not isinstance(sigmas, dict):
+ raise ValueError("MiniMax H3 sigma schedules must be a mapping")
+ video_sigmas = sigmas.get("video")
+ audio_sigmas = sigmas.get("audio")
+ if (
+ not isinstance(video_sigmas, list)
+ or not isinstance(audio_sigmas, list)
+ or len(video_sigmas) != len(audio_sigmas)
+ or len(video_sigmas) < 2
+ ):
+ raise ValueError(
+ "MiniMax H3 video/audio sigma schedules must be equal-length lists"
+ )
+ batch.sigmas = list(video_sigmas)
+ batch.timesteps = torch.tensor(
+ [1.0 - float(sigma) for sigma in video_sigmas[:-1]],
+ dtype=torch.float32,
+ )
+
+ def _generate_sigmas_from_plan(self, batch: Req, plan) -> None:
+ """Generate the fixed per-modality float32 time-shift schedules."""
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+ minimax_h3_time_shift_sigmas,
+ )
+
+ if MINIMAX_H3_SIGMAS_EXTRA_KEY in batch.extra:
+ return
+ requested_num_steps = getattr(batch, "num_inference_steps", None)
+ if requested_num_steps is None:
+ sampling = getattr(batch, "sampling_params", None)
+ requested_num_steps = getattr(sampling, "num_inference_steps", None)
+ if requested_num_steps is None:
+ requested_num_steps = 50
+ if (
+ isinstance(requested_num_steps, bool)
+ or not isinstance(requested_num_steps, int)
+ or requested_num_steps <= 0
+ ):
+ raise ValueError(
+ "num_inference_steps must be a positive integer, got "
+ f"{requested_num_steps!r}"
+ )
+
+ model_scales = self.sigma_shift_scales
+ if model_scales is not None and not isinstance(model_scales, Mapping):
+ raise ValueError("model sigma_shift_scales must be an object")
+
+ def resolved_scale(
+ *, modality: str, request_value, task_default: float
+ ) -> float:
+ value = request_value
+ source = (
+ "request "
+ f"{'flow_shift' if modality == 'video' else 'audio_flow_shift'}"
+ )
+ if value is None and model_scales is not None:
+ value = model_scales.get(modality)
+ source = f"model sigma_shift_scales.{modality}"
+ if value is None:
+ value = task_default
+ source = (
+ "task default "
+ f"{'flow_shift' if modality == 'video' else 'audio_flow_shift'}"
+ )
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError(f"{source} must be a positive finite number")
+ scale = float(value)
+ if not math.isfinite(scale) or scale <= 0.0:
+ raise ValueError(f"{source} must be a positive finite number")
+ return scale
+
+ scales = {
+ "video": resolved_scale(
+ modality="video",
+ request_value=plan.flow_shift,
+ task_default=plan.default_flow_shift,
+ ),
+ "audio": resolved_scale(
+ modality="audio",
+ request_value=plan.audio_flow_shift,
+ task_default=plan.default_audio_flow_shift,
+ ),
+ }
+ sigmas: dict[str, list[float]] = {}
+ for modality in ("video", "audio"):
+ sigmas[modality] = minimax_h3_time_shift_sigmas(
+ num_steps=requested_num_steps,
+ shift_scale=scales[modality],
+ )
+ batch.extra[MINIMAX_H3_SIGMAS_EXTRA_KEY] = sigmas
+
+ def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check(
+ "num_inference_steps", batch.num_inference_steps, V.positive_int
+ )
+ result.add_check("timesteps", batch.timesteps, V.none_or_tensor)
+ result.add_check("sigmas", batch.sigmas, V.none_or_list)
+ return result
+
+ def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+ result = VerificationResult()
+ result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.with_dims(1)])
+ result.add_check("sigmas", batch.sigmas, V.list_not_empty)
+ result.add_check(
+ MINIMAX_H3_SIGMAS_EXTRA_KEY,
+ batch.extra.get(MINIMAX_H3_SIGMAS_EXTRA_KEY),
+ lambda value: isinstance(value, Mapping),
+ )
+ return result
+
+
+__all__ = ["MiniMaxH3TimestepPreparationStage"]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py
new file mode 100644
index 000000000..d94f39979
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/visual_encoding.py
@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+ ComponentUse,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.condition_encoding import (
+ ConditionEncodingStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class MiniMaxH3VisualEncodingStage(ConditionEncodingStage):
+ deduplicated_extra_output_keys = (
+ MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY,
+ )
+
+ def __init__(
+ self,
+ video_vae,
+ vae_arch_config,
+ ) -> None:
+ super().__init__()
+ self.video_vae = video_vae
+ self.vae_arch_config = vae_arch_config
+
+ @property
+ def role_affinity(self) -> RoleType:
+ return RoleType.ENCODER
+
+ def component_uses(
+ self, server_args: ServerArgs, stage_name: str | None = None
+ ) -> list[ComponentUse]:
+ stage_name = self._component_stage_name(stage_name)
+ return [ComponentUse(stage_name, "video_vae")]
+
+ def build_dedup_fingerprint(self, batch: Req, server_args: ServerArgs):
+ parent_request_id = batch.extra.get("parent_request_id")
+ return (
+ ("expanded_outputs", parent_request_id)
+ if parent_request_id is not None
+ else id(batch)
+ )
+
+ @torch.no_grad()
+ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_cleanup_temp_dirs,
+ )
+
+ try:
+ return self._forward(batch, server_args)
+ except Exception:
+ # No later stage will run after an encoder failure.
+ minimax_h3_cleanup_temp_dirs(batch)
+ raise
+
+ def _forward(self, batch: Req, server_args: ServerArgs) -> Req:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is not None:
+ routed = plan.encoders.get("visual")
+ if not routed:
+ return batch
+ from .replica_broadcast import (
+ minimax_h3_replica_broadcast_error,
+ minimax_h3_replica_broadcast_extra,
+ minimax_h3_replica_ctx,
+ )
+
+ output_keys = self._visual_output_keys(plan, routed)
+ replica_world, replica_rank = minimax_h3_replica_ctx()
+ parallel_encode = replica_world > 1 and bool(self.video_vae.parallel_tiling)
+ owner_exception = None
+ owner_error = None
+ if (parallel_encode or replica_rank == 0) and any(
+ key not in batch.extra for key in output_keys
+ ):
+ try:
+ with self.use_declared_component(
+ component_name="video_vae",
+ module=self.video_vae,
+ ) as video_vae:
+ assert video_vae is not None
+ self.video_vae = video_vae
+ self._encode_keyframes_from_plan(batch, plan, routed)
+ except Exception as exc:
+ owner_exception = exc
+ owner_error = f"{type(exc).__name__}: {exc}"
+ if parallel_encode:
+ if owner_exception is not None:
+ raise owner_exception
+ else:
+ owner_error = minimax_h3_replica_broadcast_error(owner_error)
+ if owner_error is not None:
+ if owner_exception is not None:
+ raise owner_exception
+ raise RuntimeError(
+ f"MiniMax H3 visual encode failed on rank 0: {owner_error}"
+ )
+ for key in output_keys:
+ minimax_h3_replica_broadcast_extra(batch, key)
+ return batch
+ if (
+ batch.sampling_params is not None
+ and batch.sampling_params.image_path is not None
+ ):
+ raise NotImplementedError(
+ "MiniMaxH3VisualEncodingStage direct visual tokenizer encode "
+ "requires a canonical minimax_h3 request (resolved plan); "
+ "legacy image_path-only requests are unsupported."
+ )
+ return batch
+
+ @staticmethod
+ def _visual_output_keys(plan, routed) -> tuple[str, ...]:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY,
+ MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY,
+ )
+
+ routed_set = set(routed)
+ chains = {
+ material.material_chain
+ for material in plan.materials
+ if material.condition_index in routed_set
+ }
+ keys = []
+ if "image.target_canvas" in chains:
+ keys.append(MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY)
+ if "image.reference_preserve" in chains:
+ keys.append(MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY)
+ if chains & {
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+ }:
+ keys.append(MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY)
+ if not keys:
+ raise ValueError("MiniMax H3 visual routing selected no visual materials")
+ return tuple(keys)
+
+ def _encode_keyframes_from_plan(self, batch: Req, plan, routed) -> None:
+ """Direct keyframe encode: seeded sampled encode_images ->
+ normalized [n,96] cond rows in batch.extra."""
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.keyframe_encoding import (
+ minimax_h3_encode_keyframe_cond_rows,
+ minimax_h3_scoped_encode_fp32,
+ )
+
+ materials = [m for m in plan.materials if m.condition_index in set(routed)]
+ chains = {m.material_chain for m in materials}
+ keyframe_materials = [
+ material
+ for material in materials
+ if material.material_chain == "image.target_canvas"
+ ]
+ if str(plan.task) == "fl2va":
+ frame_indices = tuple(
+ material.frame_index for material in keyframe_materials
+ )
+ if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
+ raise ValueError(
+ "fl2va visual encoding requires an ordered keyframe signature "
+ f"in {MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got "
+ f"{frame_indices!r}"
+ )
+ elif keyframe_materials:
+ raise ValueError(
+ f"task {plan.task!r} cannot carry image.target_canvas materials"
+ )
+ if MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY in batch.extra:
+ return
+ if chains == {"image.reference_preserve"}:
+ with minimax_h3_scoped_encode_fp32(self.video_vae):
+ self._encode_reference_image(batch, plan)
+ return
+ video_chains = {
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+ }
+ if chains and chains <= {"image.reference_preserve", *video_chains}:
+ # One VAE dtype toggle for both encodes below, not one each.
+ with minimax_h3_scoped_encode_fp32(self.video_vae):
+ if "image.reference_preserve" in chains:
+ self._encode_reference_image(batch, plan)
+ if chains & video_chains:
+ self._encode_reference_video(batch, plan)
+ return
+ unsupported = [
+ m.material_chain
+ for m in materials
+ if m.material_chain != "image.target_canvas"
+ ]
+ if unsupported:
+ raise NotImplementedError(
+ "MiniMaxH3VisualEncodingStage direct encode only supports "
+ f"image.target_canvas / image.reference_preserve, got {unsupported}"
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.canvas import (
+ minimax_h3_prepared_keyframes,
+ )
+
+ # Parallel tiling gives each replicated rank complete tiles, then gathers
+ # them before the seeded posterior sample.
+ prepared = minimax_h3_prepared_keyframes(batch, plan)
+ prepared_indices = tuple(prepared.get("semantic_frame_indices") or ())
+ if prepared_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES or len(
+ prepared.get("images") or ()
+ ) != len(prepared_indices):
+ raise ValueError(
+ "fl2va visual preparation requires one or two ordered images "
+ "with a supported semantic_frame_indices signature"
+ )
+ encoded = []
+ rows_list = []
+ # One VAE dtype toggle for the whole signature (up to two keyframes),
+ # not one per keyframe.
+ with minimax_h3_scoped_encode_fp32(self.video_vae):
+ for item in prepared["images"]:
+ image = item["image"]
+ width, height = item["canvas_width"], item["canvas_height"]
+ # The encode sampling seed is pinned at 42 (the VAE sample
+ # seed is part of the contract), independent of the
+ # request seed.
+ rows = minimax_h3_encode_keyframe_cond_rows(
+ self.video_vae,
+ image,
+ self.vae_arch_config,
+ )
+ encoded.append(
+ {
+ "rows": rows,
+ "latent_h": height // 16,
+ "latent_w": width // 16,
+ "canvas_height": height,
+ "canvas_width": width,
+ "frame_index": item.get("frame_index"),
+ "resolved_frame_index": item.get("resolved_frame_index"),
+ "condition_index": item.get("condition_index"),
+ }
+ )
+ rows_list.append(rows)
+ rows = rows_list[0] if len(rows_list) == 1 else torch.cat(rows_list, dim=0)
+ first = encoded[0]
+ batch.extra[MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY] = {
+ "rows": rows,
+ "latent_h": first["latent_h"],
+ "latent_w": first["latent_w"],
+ "canvas_height": first["canvas_height"],
+ "canvas_width": first["canvas_width"],
+ "keyframes": encoded,
+ "semantic_frame_indices": prepared.get("semantic_frame_indices"),
+ "pixel_frame_indices": prepared.get("pixel_frame_indices"),
+ "frame_count": prepared.get("frame_count"),
+ }
+
+ def _encode_reference_video(self, batch: Req, plan) -> None:
+ """ref2va video/video_audio encode from shared transformed frames."""
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.reference_encoding import (
+ minimax_h3_encode_reference_video_rows,
+ minimax_h3_prepared_reference_videos,
+ )
+
+ if MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY in batch.extra:
+ return
+ prepared = minimax_h3_prepared_reference_videos(batch, plan)
+ videos = prepared.get("videos")
+ if not isinstance(videos, list) or not videos:
+ raise ValueError(
+ "prepared reference videos payload must carry a non-empty "
+ f"'videos' list, got {videos!r}"
+ )
+ entries = []
+ for item in videos:
+ rows, latent_t, latent_h, latent_w = minimax_h3_encode_reference_video_rows(
+ self.video_vae,
+ item["frames"],
+ self.vae_arch_config,
+ )
+ entries.append(
+ {
+ "rows": rows,
+ "latent_t": latent_t,
+ "latent_h": latent_h,
+ "latent_w": latent_w,
+ "condition_index": int(item["condition_index"]),
+ "material_chain": str(item["material_chain"]),
+ }
+ )
+ for item in videos:
+ # Text and visual encoding are the only RGB-frame consumers. Drop
+ # the large request-local arrays as soon as both have completed.
+ item.pop("frames", None)
+ payload = dict(entries[0])
+ payload["videos"] = entries
+ batch.extra[MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY] = payload
+
+ def _encode_reference_image(self, batch: Req, plan) -> None:
+ """ref2va reference image encode: cap_resize (intrinsic
+ geometry) + the verified keyframe recipe; rows use the image's OWN
+ latent grid, not the target canvas."""
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.keyframe_encoding import (
+ minimax_h3_encode_keyframe_cond_rows,
+ )
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.reference_encoding import (
+ minimax_h3_prepared_reference_image,
+ )
+
+ if MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY in batch.extra:
+ return
+ prepared = minimax_h3_prepared_reference_image(batch, plan)
+ entries = []
+ for item in prepared["images"]:
+ image = item["image"]
+ rows = minimax_h3_encode_keyframe_cond_rows(
+ self.video_vae,
+ image,
+ self.vae_arch_config,
+ )
+ width, height = image.size
+ entries.append(
+ {
+ "rows": rows,
+ "latent_h": height // 16,
+ "latent_w": width // 16,
+ "condition_index": int(item["condition_index"]),
+ "material_chain": "image.reference_preserve",
+ }
+ )
+ payload = dict(entries[0]) # single-image consumers keep the keys
+ payload["images"] = entries
+ batch.extra[MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY] = payload
+
+
+__all__ = ["MiniMaxH3VisualEncodingStage"]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/task_profiles.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/task_profiles.py
new file mode 100644
index 000000000..aa1cd2b51
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/task_profiles.py
@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax H3 task profiles.
+
+Data table driving request validation and plan resolution for the three v1
+tasks (t2va / fl2va / ref2va). One row per task; stages and the request
+projector consume rows instead of branching on task names.
+
+Design summary: keyframes bind target geometry; references remain independent.
+"""
+
+from __future__ import annotations
+
+from types import MappingProxyType
+
+import msgspec
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_DEFAULT_BRANCHES,
+)
+
+MINIMAX_H3_CONDITION_ROLE_KEYFRAME = "keyframe"
+MINIMAX_H3_CONDITION_ROLE_REFERENCE = "reference"
+
+MINIMAX_H3_TASK_T2VA = "t2va"
+MINIMAX_H3_TASK_FL2VA = "fl2va"
+MINIMAX_H3_TASK_REF2VA = "ref2va"
+
+# Public dispatcher contract. SGLang only needs the stable task mapping.
+MINIMAX_H3_TASK_PARTITIONS = MappingProxyType(
+ {
+ "t2va": "fl2va",
+ "fl2va": "fl2va",
+ "ref2va": "ref2va",
+ }
+)
+
+
+def canonical_minimax_h3_task(task: str) -> str:
+ """Normalize a public task name (no aliases are currently defined)."""
+
+ return task
+
+
+def partition_for_task(task: str) -> str:
+ """Return the deployment partition serving *task*."""
+
+ if not isinstance(task, str) or not task.strip():
+ raise ValueError("MiniMax H3 task must be a non-empty string")
+ normalized = task.strip().lower()
+ try:
+ return MINIMAX_H3_TASK_PARTITIONS[normalized]
+ except KeyError as exc:
+ raise ValueError(f"unsupported MiniMax H3 task {task!r}") from exc
+
+
+MINIMAX_H3_FINITE_ASPECT_RATIOS = (
+ "21:9",
+ "16:9",
+ "4:3",
+ "1:1",
+ "3:4",
+ "9:16",
+)
+
+# ``fl2va`` remains the single public task name for all keyframe variants:
+# first-frame-only, last-frame-only, and first+last. Keep the accepted
+# semantic signatures centralized so validation and every downstream sink can
+# reject middle/reversed/stale payloads consistently.
+MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: tuple[tuple[int, ...], ...] = (
+ (0,),
+ (-1,),
+ (0, -1),
+)
+
+# Canonical material-chain names are part of the direct runtime plan.
+# Profile validation keeps only the wire-level allowlist here.
+MINIMAX_H3_CANONICAL_MATERIAL_CHAINS = frozenset(
+ {
+ "image.target_canvas",
+ "image.reference_preserve",
+ "video.reference_preserve",
+ "video_audio.reference_preserve",
+ "audio",
+ }
+)
+
+
+class MiniMaxH3ConditionRule(msgspec.Struct, frozen=True):
+ """Per-(role, type) admission + routing rule."""
+
+ role: str
+ condition_type: str
+ material_chain: str
+ requires_frame_index: bool = False
+ visual_tokenizer_encode: bool = False
+ audio_tokenizer_encode: bool = False
+
+
+class MiniMaxH3TaskProfile(msgspec.Struct, frozen=True):
+ """One row of the task table."""
+
+ task: str
+ conditions_required: bool
+ condition_rules: tuple[MiniMaxH3ConditionRule, ...]
+ branches: tuple[dict, ...]
+ default_flow_shift: float
+ default_audio_flow_shift: float
+ required_components: tuple[str, ...]
+ aspect_ratio_forced_auto: bool
+ geometry_source: str # "explicit_target" | "first_keyframe" | "model_default"
+ # For ref2va video, target may omit duration_seconds
+ # when an audio reference is present — duration derives from the
+ # reference audio probe and is then checked against the shared 4-15s range.
+ duration_from_audio_reference: bool = False
+ # Some deployments keep video references disabled while retaining the
+ # routing rule for projection/contract compatibility.
+ video_reference_supported: bool = True
+ min_condition_count: int | None = None
+ max_condition_count: int | None = None
+ # Resolution policy for target.aspect_ratio="auto". A concrete ratio
+ # resolves immediately; None keeps geometry deferred to the named source.
+ auto_aspect_ratio: str | None = None
+ auto_geometry_source: str | None = None
+
+ def rule_for(self, *, role: str, condition_type: str) -> MiniMaxH3ConditionRule:
+ for rule in self.condition_rules:
+ if rule.role == role and rule.condition_type == condition_type:
+ return rule
+ raise ValueError(
+ f"task {self.task!r} does not allow condition "
+ f"role={role!r} type={condition_type!r}"
+ )
+
+
+_BASE_COMPONENTS = (
+ "processor",
+ "text_encoder",
+ "tokenizer",
+ "transformer",
+ "video_vae",
+ "audio_vae",
+)
+
+_BRANCHES = tuple(dict(branch) for branch in MINIMAX_H3_DEFAULT_BRANCHES)
+MINIMAX_H3_TASK_PROFILES: dict[str, MiniMaxH3TaskProfile] = {
+ MINIMAX_H3_TASK_T2VA: MiniMaxH3TaskProfile(
+ task=MINIMAX_H3_TASK_T2VA,
+ conditions_required=False,
+ condition_rules=(),
+ branches=_BRANCHES,
+ default_flow_shift=12.0,
+ default_audio_flow_shift=3.0,
+ required_components=_BASE_COMPONENTS,
+ aspect_ratio_forced_auto=False,
+ geometry_source="explicit_target",
+ auto_aspect_ratio="16:9",
+ auto_geometry_source="policy_default",
+ ),
+ MINIMAX_H3_TASK_FL2VA: MiniMaxH3TaskProfile(
+ task=MINIMAX_H3_TASK_FL2VA,
+ conditions_required=True,
+ condition_rules=(
+ MiniMaxH3ConditionRule(
+ role=MINIMAX_H3_CONDITION_ROLE_KEYFRAME,
+ condition_type="image",
+ material_chain="image.target_canvas",
+ requires_frame_index=True,
+ visual_tokenizer_encode=True,
+ ),
+ ),
+ branches=_BRANCHES,
+ default_flow_shift=12.0,
+ default_audio_flow_shift=3.0,
+ required_components=_BASE_COMPONENTS,
+ # FL follows the resolved target geometry projected by the caller.
+ # Its first/last keyframes are prepared against that target canvas.
+ aspect_ratio_forced_auto=False,
+ geometry_source="explicit_target",
+ min_condition_count=1,
+ max_condition_count=2,
+ auto_geometry_source="first_keyframe",
+ ),
+ MINIMAX_H3_TASK_REF2VA: MiniMaxH3TaskProfile(
+ task=MINIMAX_H3_TASK_REF2VA,
+ conditions_required=True,
+ condition_rules=(
+ MiniMaxH3ConditionRule(
+ role=MINIMAX_H3_CONDITION_ROLE_REFERENCE,
+ condition_type="image",
+ material_chain="image.reference_preserve",
+ visual_tokenizer_encode=True,
+ ),
+ MiniMaxH3ConditionRule(
+ role=MINIMAX_H3_CONDITION_ROLE_REFERENCE,
+ condition_type="video",
+ material_chain="video.reference_preserve",
+ visual_tokenizer_encode=True,
+ # ref2va video: the reference video's ORIGINAL soundtrack is the
+ # audio reference; -17 rows cover the full pre-truncation track.
+ # Route the same material into the audio encoder.
+ audio_tokenizer_encode=True,
+ ),
+ MiniMaxH3ConditionRule(
+ role=MINIMAX_H3_CONDITION_ROLE_REFERENCE,
+ condition_type="video_audio",
+ material_chain="video_audio.reference_preserve",
+ visual_tokenizer_encode=True,
+ audio_tokenizer_encode=True,
+ ),
+ MiniMaxH3ConditionRule(
+ role=MINIMAX_H3_CONDITION_ROLE_REFERENCE,
+ condition_type="audio",
+ material_chain="audio",
+ audio_tokenizer_encode=True,
+ ),
+ ),
+ branches=_BRANCHES,
+ default_flow_shift=12.0,
+ default_audio_flow_shift=3.0,
+ required_components=_BASE_COMPONENTS,
+ # ref2va video may use an explicit aspect ratio such as "7:4":
+ # ref2va allows explicit aspect ratios; "auto" falls back to the
+ # policy default (references never bind target geometry).
+ aspect_ratio_forced_auto=False,
+ geometry_source="explicit_target",
+ duration_from_audio_reference=True,
+ auto_aspect_ratio="16:9",
+ auto_geometry_source="policy_default",
+ # ref2va video uses a subject image plus a reference video with soundtrack.
+ video_reference_supported=True,
+ ),
+}
+
+
+def minimax_h3_task_profile(task: str) -> MiniMaxH3TaskProfile:
+ profile = MINIMAX_H3_TASK_PROFILES.get(task)
+ if profile is None:
+ raise ValueError(
+ f"unknown minimax_h3 task {task!r}; supported: "
+ f"{sorted(MINIMAX_H3_TASK_PROFILES)}"
+ )
+ return profile
+
+
+def _validate_profiles() -> None:
+ for task, profile in MINIMAX_H3_TASK_PROFILES.items():
+ if profile.task != task:
+ raise ValueError(f"profile key/task mismatch: {task} vs {profile.task}")
+ roles = {rule.role for rule in profile.condition_rules}
+ if len(roles) > 1:
+ raise ValueError(
+ f"task {task}: condition roles must not mix, got {sorted(roles)}"
+ )
+ if profile.min_condition_count is not None and profile.min_condition_count <= 0:
+ raise ValueError(f"task {task}: min_condition_count must be positive")
+ if profile.max_condition_count is not None and profile.max_condition_count <= 0:
+ raise ValueError(f"task {task}: max_condition_count must be positive")
+ if (
+ profile.min_condition_count is not None
+ and profile.max_condition_count is not None
+ and profile.min_condition_count > profile.max_condition_count
+ ):
+ raise ValueError(
+ f"task {task}: min_condition_count exceeds max_condition_count"
+ )
+ for rule in profile.condition_rules:
+ if rule.material_chain not in MINIMAX_H3_CANONICAL_MATERIAL_CHAINS:
+ raise ValueError(
+ f"task {task}: unknown material chain {rule.material_chain!r}"
+ )
+
+
+_validate_profiles()
+
+__all__ = [
+ "MINIMAX_H3_CONDITION_ROLE_KEYFRAME",
+ "MINIMAX_H3_FINITE_ASPECT_RATIOS",
+ "MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES",
+ "MINIMAX_H3_CONDITION_ROLE_REFERENCE",
+ "MINIMAX_H3_CANONICAL_MATERIAL_CHAINS",
+ "MINIMAX_H3_TASK_FL2VA",
+ "MINIMAX_H3_TASK_PROFILES",
+ "MINIMAX_H3_TASK_REF2VA",
+ "MINIMAX_H3_TASK_T2VA",
+ "MiniMaxH3TaskProfile",
+ "canonical_minimax_h3_task",
+ "minimax_h3_task_profile",
+ "partition_for_task",
+]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/time_request.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/time_request.py
new file mode 100644
index 000000000..e6fb5f392
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/time_request.py
@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+
+def minimax_h3_align_frame_count(frame_count: int) -> int:
+ """Snap ``frame_count`` up to the MiniMax H3 17n+5 frame boundary."""
+ if frame_count <= 0:
+ return 1
+ current = int(frame_count)
+ return current + (5 - current) % 17
+
+
+def minimax_h3_video_latent_t(frame_count: int) -> int:
+ if frame_count <= 5:
+ return 2
+ return ((int(frame_count) - 5) // 17) * 5 + 2
+
+
+def minimax_h3_frame_count_from_video_latent_t(out_t: int) -> int:
+ if out_t == 1:
+ return 1
+ if out_t < 2 or (out_t - 2) % 5 != 0:
+ raise ValueError("MiniMax H3 video latent T must be 1 or match 5n+2")
+ return 17 * ((int(out_t) - 2) // 5) + 5
+
+
+def minimax_h3_audio_latent_t(duration_seconds: float) -> int:
+ # Rounding happens at the 40 Hz audio latent boundary.
+ return int(round(float(duration_seconds) * 40.0))
+
+
+def minimax_h3_time_shift_sigmas(
+ *,
+ num_steps: int = 50,
+ shift_scale: float = 6.0,
+) -> list[float]:
+ if shift_scale <= 0:
+ raise ValueError("MiniMax H3 shift_scale must be > 0")
+ if num_steps <= 0:
+ raise ValueError("MiniMax H3 num_steps must be > 0")
+
+ import torch
+
+ # The rectified-flow sigma range is fixed at [1.0, 0.0].
+ base = torch.linspace(
+ 1.0,
+ 0.0,
+ int(num_steps),
+ device="cpu",
+ dtype=torch.float32,
+ )
+ shifted = float(shift_scale) * base / (1 + (float(shift_scale) - 1) * base)
+ shifted, _ = torch.unique_consecutive(shifted, return_counts=True)
+ # A one-point request is still exactly one point. Normal serving uses
+ # multiple points, but preserving the requested cardinality keeps
+ # ``num_inference_steps`` the sole schedule-size control.
+ if num_steps > 1 and shifted[-1].item() > 0.0:
+ shifted = torch.cat([shifted, torch.tensor([0.0], dtype=shifted.dtype)])
+ return [float(value) for value in shifted.tolist()]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py
new file mode 100644
index 000000000..7ab71a609
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py
@@ -0,0 +1,545 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Video API lowering and strict delivery hooks for MiniMax H3."""
+
+from __future__ import annotations
+
+import json
+import math
+import subprocess
+from typing import TYPE_CHECKING, Any
+
+from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
+ MiniMaxH3PipelineConfig,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+ VideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
+ MINIMAX_H3_SUPPORTED_FPS,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ canonical_minimax_h3_task,
+)
+
+if TYPE_CHECKING:
+ from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+
+
+def _extra_value(request: VideoGenerationsRequest, name: str) -> Any:
+ return (request.model_extra or {}).get(name)
+
+
+def _parse_extra_value(value: Any) -> Any:
+ if not isinstance(value, str):
+ return value
+ try:
+ return json.loads(value)
+ except (json.JSONDecodeError, TypeError, ValueError):
+ return value
+
+
+def _format_video_seconds(value: float) -> str:
+ rounded = round(float(value))
+ if abs(float(value) - rounded) < 1e-9:
+ return str(int(rounded))
+ return f"{float(value):.6f}".rstrip("0").rstrip(".")
+
+
+class MiniMaxH3VideoModelAdapter:
+ """Canonical request and fail-closed AV delivery hooks for MiniMax H3."""
+
+ strict_file_delivery = True
+ model_specific_fields = frozenset(
+ {
+ "task",
+ "conditions",
+ "target",
+ "audio_flow_shift",
+ "audio_guidance_scale",
+ "quality",
+ "output_mode",
+ "imgvid_cond_noise_aug_for_inference",
+ "audio_cond_noise_aug_for_inference",
+ }
+ )
+ supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+
+ def validate_task_gate(self, task: Any, *, provided: bool) -> None:
+ if not provided or task is None:
+ raise ValueError(
+ "task is required for MiniMax H3; supported tasks: "
+ "fl2va, ref2va, t2va"
+ )
+ if not isinstance(task, str):
+ raise ValueError("task must be a non-empty string for MiniMax H3")
+ if not task.strip():
+ raise ValueError("task must be a non-empty string for MiniMax H3")
+ if task not in self.supported_tasks:
+ raise ValueError(
+ f"unsupported MiniMax H3 task {task!r}; supported tasks: "
+ "fl2va, ref2va, t2va"
+ )
+
+ @staticmethod
+ def _positive_finite_extra(
+ request: VideoGenerationsRequest,
+ name: str,
+ ) -> float | None:
+ value = _parse_extra_value(_extra_value(request, name))
+ if value is None:
+ return None
+ if isinstance(value, bool):
+ raise ValueError(f"{name} must be a positive finite float")
+ try:
+ parsed = float(value)
+ except (TypeError, ValueError) as exc:
+ raise ValueError(f"{name} must be a positive finite float") from exc
+ if not math.isfinite(parsed) or parsed <= 0:
+ raise ValueError(f"{name} must be a positive finite float")
+ return parsed
+
+ @staticmethod
+ def _quality_extra(
+ request: VideoGenerationsRequest,
+ name: str,
+ ) -> str | None:
+ value = _parse_extra_value(_extra_value(request, name))
+ if value is None:
+ return None
+ if not isinstance(value, str) or not value.strip():
+ raise ValueError(f"{name} must be a non-empty string")
+ return value.strip().lower()
+
+ @staticmethod
+ def _reject_retired_cfg_fields(kwargs: dict[str, Any]) -> None:
+ """Fail fast on retired CFG knobs.
+
+ MiniMax H3 serves only the CFG-distilled single-positive-branch
+ checkpoint, so these fields cannot influence generation. Any
+ provided value is an error; only absent (``None``) is accepted.
+ """
+
+ for field_name in (
+ "guidance_scale",
+ "guidance_scale_2",
+ "true_cfg_scale",
+ "negative_prompt",
+ ):
+ if kwargs.pop(field_name, None) is not None:
+ raise ValueError(
+ f"{field_name} is not supported: MiniMax H3 serves only the "
+ "CFG-distilled single-positive-branch checkpoint"
+ )
+
+ @staticmethod
+ def _reject_transport_timing_fields(request: VideoGenerationsRequest) -> None:
+ """Fail fast on explicit transport timing fields.
+
+ Canonical MiniMax H3 timing is resolved exclusively from
+ ``target.duration_seconds`` (or the single ref2va audio source), so
+ explicit ``num_frames``/``fps`` cannot influence generation. Only
+ absent (``None``, the protocol default) is accepted.
+ """
+
+ for field_name in ("num_frames", "fps"):
+ if getattr(request, field_name) is not None:
+ raise ValueError(
+ f"{field_name} is not supported: MiniMax H3 derives the "
+ "temporal shape from target.duration_seconds"
+ )
+
+ def lower_video_request_kwargs(
+ self,
+ request: VideoGenerationsRequest,
+ kwargs: dict[str, Any],
+ ) -> dict[str, Any]:
+ self.validate_transport_options(request, model_path=None)
+ kwargs = dict(kwargs)
+ self._reject_transport_timing_fields(request)
+ # Drop the transport-synthesized timing defaults; canonical MiniMax H3
+ # timing never flows through num_frames/fps.
+ kwargs.pop("num_frames", None)
+ kwargs.pop("fps", None)
+ self._reject_retired_cfg_fields(kwargs)
+ quality = self._quality_extra(request, "quality")
+ kwargs.update(
+ {
+ "audio_flow_shift": self._positive_finite_extra(
+ request, "audio_flow_shift"
+ ),
+ "task": canonical_minimax_h3_task(
+ _parse_extra_value(_extra_value(request, "task"))
+ ),
+ "conditions": _parse_extra_value(_extra_value(request, "conditions")),
+ "target": _parse_extra_value(_extra_value(request, "target")),
+ "imgvid_cond_noise_aug_for_inference": _parse_extra_value(
+ _extra_value(request, "imgvid_cond_noise_aug_for_inference")
+ ),
+ "audio_cond_noise_aug_for_inference": _parse_extra_value(
+ _extra_value(request, "audio_cond_noise_aug_for_inference")
+ ),
+ }
+ )
+ if quality is not None:
+ kwargs["quality"] = quality
+ return kwargs
+
+ def validate_transport_options(
+ self,
+ request: VideoGenerationsRequest,
+ *,
+ model_path: str | None,
+ ) -> None:
+ extras = request.model_extra or {}
+ self.validate_task_gate(extras.get("task"), provided="task" in extras)
+ del model_path
+ self._positive_finite_extra(request, "audio_flow_shift")
+ if _extra_value(request, "audio_guidance_scale") is not None:
+ raise ValueError(
+ "audio_guidance_scale is not supported: MiniMax H3 serves only "
+ "the CFG-distilled single-positive-branch checkpoint"
+ )
+ output_mode = _extra_value(request, "output_mode")
+ if output_mode not in (None, "decoded_files"):
+ raise ValueError(
+ "MiniMax H3 SGLang backend only supports "
+ f"output_mode='decoded_files', got {output_mode!r}"
+ )
+ if request.enable_frame_interpolation:
+ raise ValueError(
+ "MiniMax H3 does not support enable_frame_interpolation: the "
+ "accepted delivery contract is the canonical 24 fps output"
+ )
+ if request.enable_upscaling:
+ raise ValueError(
+ "MiniMax H3 does not support enable_upscaling: the accepted "
+ "delivery contract is the resolved target canvas"
+ )
+
+ def validate_sampling_params(self, sampling_params: SamplingParams) -> None:
+ """Apply the HTTP task/delivery gate to offline requests as well."""
+
+ task = getattr(sampling_params, "task", None)
+ self.validate_task_gate(task, provided=task is not None)
+ sampling_params.task = canonical_minimax_h3_task(task)
+ if getattr(sampling_params, "enable_frame_interpolation", False):
+ raise ValueError(
+ "MiniMax H3 does not support enable_frame_interpolation: the "
+ "accepted delivery contract is the canonical 24 fps output"
+ )
+ if getattr(sampling_params, "enable_upscaling", False):
+ raise ValueError(
+ "MiniMax H3 does not support enable_upscaling: the accepted "
+ "delivery contract is the resolved target canvas"
+ )
+ if not bool(getattr(sampling_params, "save_output", False)) or not getattr(
+ sampling_params, "output_path", None
+ ):
+ raise ValueError(
+ "MiniMax H3 DiffGenerator requires save_output=True and a non-empty "
+ "output_path for validated file delivery"
+ )
+ output_mode = getattr(sampling_params, "output_mode", None)
+ if output_mode not in (None, "decoded_files"):
+ raise ValueError(
+ "MiniMax H3 SGLang backend only supports "
+ f"output_mode='decoded_files', got {output_mode!r}"
+ )
+
+ def prepare_for_queue_sync(self, batch: Req) -> None:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.prequeue import (
+ minimax_h3_prepare_for_queue,
+ )
+
+ minimax_h3_prepare_for_queue(batch)
+
+ def cleanup_request_sync(self, batch: Req) -> None:
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
+ minimax_h3_cleanup_temp_dirs,
+ )
+
+ minimax_h3_cleanup_temp_dirs(batch)
+
+ @staticmethod
+ def _resolved_shape(batch: Req) -> dict[str, Any] | None:
+ canonical = getattr(batch, "extra", {}).get("minimax_h3_canonical_request")
+ if not isinstance(canonical, dict) or not all(
+ key in canonical
+ for key in ("schema", "task", "prompt", "conditions", "target")
+ ):
+ return None
+ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_plan_from_batch,
+ )
+
+ plan = minimax_h3_plan_from_batch(batch)
+ if plan is None:
+ return None
+ shape = plan.shape
+ if str(shape.get("geometry") or "") != "resolved_v2":
+ raise ValueError(
+ "queued MiniMax H3 jobs require pre-queue resolved_v2 geometry"
+ )
+ if shape.get("frame_count") is None:
+ raise ValueError(
+ "queued MiniMax H3 jobs require pre-queue resolved temporal dimensions"
+ )
+ return shape
+
+ def project_queued_job_fields(self, batch: Req) -> dict[str, str]:
+ shape = self._resolved_shape(batch)
+ if shape is None:
+ return {}
+ fields: dict[str, str] = {}
+ if shape.get("width") is not None and shape.get("height") is not None:
+ fields["size"] = f"{int(shape['width'])}x{int(shape['height'])}"
+ queued_frame_count = shape.get("frame_count")
+ if queued_frame_count is not None:
+ fields["seconds"] = _format_video_seconds(
+ int(queued_frame_count) / float(shape["fps"])
+ )
+ quality = getattr(batch.sampling_params, "quality", None)
+ explicit_fields = getattr(batch.sampling_params, "_explicit_fields", ())
+ if quality and "quality" in explicit_fields:
+ fields["quality"] = str(quality)
+ return fields
+
+ def validate_final_outputs_sync(
+ self,
+ output_paths: list[str],
+ batch: Req,
+ ) -> dict[str, str]:
+ expected_outputs = int(getattr(batch, "num_outputs_per_prompt", 1))
+ if len(output_paths) != expected_outputs:
+ raise RuntimeError(
+ "MiniMax H3 video generation produced "
+ f"{len(output_paths)} output files, expected {expected_outputs}"
+ )
+
+ expected_frame_count = None
+ expected_size = None
+ shape = self._resolved_shape(batch)
+ if shape is not None:
+ if shape.get("frame_count") is not None:
+ expected_frame_count = int(shape["frame_count"])
+ if shape.get("width") is not None and shape.get("height") is not None:
+ expected_size = (int(shape["width"]), int(shape["height"]))
+
+ final_media_fields: dict[str, str] = {}
+ for output_index, output_path in enumerate(output_paths):
+ media_fields = _probe_minimax_h3_output_fields(
+ output_path,
+ expected_frame_count=expected_frame_count,
+ expected_size=expected_size,
+ )
+ if output_index == 0:
+ final_media_fields = media_fields
+ elif media_fields != final_media_fields:
+ raise RuntimeError(
+ "generated MiniMax H3 outputs have inconsistent media metadata: "
+ f"output 0={final_media_fields}, output "
+ f"{output_index}={media_fields}"
+ )
+ return final_media_fields
+
+
+def _probe_minimax_h3_output_fields(
+ path: str,
+ *,
+ expected_frame_count: int | None = None,
+ expected_size: tuple[int, int] | None = None,
+) -> dict[str, str]:
+ """Validate one final MiniMax H3 AV file and derive truthful metadata."""
+
+ try:
+ probe = subprocess.run(
+ [
+ "ffprobe",
+ "-v",
+ "error",
+ "-show_entries",
+ "stream=codec_type,codec_name,pix_fmt,width,height,avg_frame_rate,nb_frames,duration,sample_rate,channels:format=format_name,duration",
+ "-of",
+ "json",
+ str(path),
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ timeout=30,
+ )
+ except subprocess.TimeoutExpired as exc:
+ raise RuntimeError(
+ "MiniMax H3 final output ffprobe timed out after 30 seconds"
+ ) from exc
+ except FileNotFoundError as exc:
+ raise RuntimeError(
+ "ffprobe is required to validate final MiniMax H3 output"
+ ) from exc
+ except subprocess.CalledProcessError as exc:
+ detail = (exc.stderr or "").strip()
+ suffix = f": {detail}" if detail else ""
+ raise RuntimeError(
+ f"ffprobe failed for final MiniMax H3 output{suffix}"
+ ) from exc
+
+ try:
+ payload = json.loads(probe.stdout)
+ except (TypeError, ValueError) as exc:
+ raise RuntimeError(
+ "ffprobe returned invalid JSON for MiniMax H3 output"
+ ) from exc
+ if not isinstance(payload, dict):
+ raise RuntimeError("ffprobe returned invalid JSON for MiniMax H3 output")
+ streams = payload.get("streams") or []
+ if not isinstance(streams, list):
+ raise RuntimeError("ffprobe returned invalid stream metadata")
+ video_streams = [
+ stream
+ for stream in streams
+ if isinstance(stream, dict) and stream.get("codec_type") == "video"
+ ]
+ audio_streams = [
+ stream
+ for stream in streams
+ if isinstance(stream, dict) and stream.get("codec_type") == "audio"
+ ]
+ if len(streams) != 2 or len(video_streams) != 1 or len(audio_streams) != 1:
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 must contain exactly one video stream and "
+ f"one audio stream, got {len(video_streams)} video, "
+ f"{len(audio_streams)} audio, and {len(streams)} total streams"
+ )
+ video_stream = video_streams[0]
+ audio_stream = audio_streams[0]
+ width = int(video_stream.get("width") or 0)
+ height = int(video_stream.get("height") or 0)
+ if width <= 0 or height <= 0:
+ raise RuntimeError(
+ f"generated MiniMax H3 MP4 has invalid size {width}x{height}"
+ )
+ if expected_size is not None and (width, height) != expected_size:
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 size does not match the resolved request: "
+ f"expected {expected_size[0]}x{expected_size[1]}, got {width}x{height}"
+ )
+ rate_raw = str(video_stream.get("avg_frame_rate") or "")
+ try:
+ if "/" in rate_raw:
+ numerator, denominator = rate_raw.split("/", 1)
+ fps = float(numerator) / float(denominator)
+ else:
+ fps = float(rate_raw)
+ except (TypeError, ValueError, ZeroDivisionError):
+ fps = 0.0
+ if abs(fps - float(MINIMAX_H3_SUPPORTED_FPS)) > 1e-6:
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 frame rate must be "
+ f"{MINIMAX_H3_SUPPORTED_FPS} fps, got {rate_raw!r}"
+ )
+ try:
+ audio_sample_rate = int(audio_stream.get("sample_rate") or 0)
+ except (TypeError, ValueError):
+ audio_sample_rate = 0
+ expected_sample_rate = MiniMaxH3PipelineConfig.output_audio_sample_rate
+ if expected_sample_rate is not None and audio_sample_rate != expected_sample_rate:
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 audio sample rate must be "
+ f"{expected_sample_rate} Hz, "
+ f"got {audio_sample_rate}"
+ )
+ try:
+ audio_channels = int(audio_stream.get("channels") or 0)
+ except (TypeError, ValueError):
+ audio_channels = 0
+ expected_channels = MiniMaxH3PipelineConfig.output_audio_channels
+ if expected_channels is not None and audio_channels != expected_channels:
+ channel_contract = (
+ "stereo (2 channels)"
+ if expected_channels == 2
+ else f"{expected_channels} channels"
+ )
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 audio must be "
+ f"{channel_contract}, got {audio_channels}"
+ )
+ try:
+ duration = float(video_stream.get("duration") or 0.0)
+ except (TypeError, ValueError):
+ duration = 0.0
+ frames_raw = video_stream.get("nb_frames")
+ frame_count = None
+ if frames_raw not in {None, "", "N/A"}:
+ try:
+ frame_count = int(frames_raw)
+ except (TypeError, ValueError):
+ frame_count = None
+ if expected_frame_count is not None and frame_count != expected_frame_count:
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 frame count does not match the resolved "
+ f"request: expected {expected_frame_count}, got {frame_count}"
+ )
+ if frame_count is not None and "/" in rate_raw and fps > 0:
+ duration = frame_count / fps
+ if duration <= 0:
+ try:
+ format_metadata = payload.get("format")
+ if not isinstance(format_metadata, dict):
+ format_metadata = {}
+ duration = float(format_metadata.get("duration") or 0.0)
+ except (TypeError, ValueError):
+ duration = 0.0
+ if duration <= 0:
+ raise RuntimeError("generated MiniMax H3 MP4 has no positive video duration")
+ try:
+ audio_duration = float(audio_stream.get("duration") or 0.0)
+ except (TypeError, ValueError):
+ audio_duration = 0.0
+ if audio_duration <= 0:
+ raise RuntimeError("generated MiniMax H3 MP4 has no positive audio duration")
+ drift_tolerance_s = MiniMaxH3PipelineConfig.output_av_drift_tolerance_s
+ if (
+ drift_tolerance_s is not None
+ and abs(audio_duration - duration) > drift_tolerance_s
+ ):
+ raise RuntimeError(
+ "generated MiniMax H3 MP4 audio/video duration drift exceeds "
+ f"{drift_tolerance_s:g}s: "
+ f"video={duration:.6f}s audio={audio_duration:.6f}s"
+ )
+ format_metadata = payload.get("format")
+ if not isinstance(format_metadata, dict):
+ format_metadata = {}
+ format_names = {
+ token.strip().lower()
+ for token in str(format_metadata.get("format_name") or "").split(",")
+ if token.strip()
+ }
+ if "mp4" not in format_names:
+ raise RuntimeError(
+ "generated MiniMax H3 output container must be MP4-family, got "
+ f"{format_metadata.get('format_name')!r}"
+ )
+ video_codec = str(video_stream.get("codec_name") or "").lower()
+ if video_codec != "h264":
+ raise RuntimeError(
+ f"generated MiniMax H3 MP4 video codec must be h264, got {video_codec!r}"
+ )
+ audio_codec = str(audio_stream.get("codec_name") or "").lower()
+ if audio_codec != "aac":
+ raise RuntimeError(
+ f"generated MiniMax H3 MP4 audio codec must be aac, got {audio_codec!r}"
+ )
+ pixel_format = str(video_stream.get("pix_fmt") or "").lower()
+ if pixel_format != "yuv420p":
+ raise RuntimeError(
+ f"generated MiniMax H3 MP4 pixel format must be yuv420p, got {pixel_format!r}"
+ )
+ return {
+ "size": f"{width}x{height}",
+ "seconds": _format_video_seconds(duration),
+ }
+
+
+__all__ = ["MiniMaxH3VideoModelAdapter"]
diff --git a/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/python/sglang/multimodal_gen/runtime/platforms/__init__.py
index 5ee20138e..91d7c6a9c 100644
--- a/python/sglang/multimodal_gen/runtime/platforms/__init__.py
+++ b/python/sglang/multimodal_gen/runtime/platforms/__init__.py
@@ -50,6 +50,27 @@ def cuda_platform_plugin() -> str | None:
if cuda_is_jetson():
is_cuda = True
+ else:
+ # NVML is NVIDIA-specific. CUDA-compatible stacks (e.g. Iluvatar
+ # CoreX) expose devices through torch's CUDA API without shipping
+ # libnvidia-ml. Only fall back when NVML itself is unavailable;
+ # a successful NVML init that reports zero devices must keep the
+ # CPU-build-on-GPU-machine edge case above.
+ try:
+ import torch
+
+ # ROCm also exposes devices through torch.cuda, so keep this
+ # non-NVML fallback limited to non-HIP runtimes.
+ is_cuda = (
+ getattr(torch.version, "hip", None) is None
+ and torch.cuda.is_available()
+ and torch.cuda.device_count() > 0
+ )
+ if is_cuda:
+ logger.debug("CUDA detected via torch (NVML unavailable)")
+ except Exception as exc:
+ logger.debug("torch CUDA detection failed: %s", exc)
+
if is_cuda:
logger.debug("CUDA is available")
diff --git a/python/sglang/multimodal_gen/runtime/platforms/aiter.py b/python/sglang/multimodal_gen/runtime/platforms/aiter.py
index d15240964..c06ef820a 100644
--- a/python/sglang/multimodal_gen/runtime/platforms/aiter.py
+++ b/python/sglang/multimodal_gen/runtime/platforms/aiter.py
@@ -3,8 +3,10 @@
from sglang.srt.utils import (
get_bool_env_var,
is_gfx95_supported,
+ is_gfx942_supported,
is_hip,
)
USE_AITER = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
+USE_AITER_GFX942 = USE_AITER and is_gfx942_supported()
USE_AITER_GFX95 = USE_AITER and is_gfx95_supported()
diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py
index 99504537a..2458fa7dc 100644
--- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py
+++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py
@@ -73,17 +73,26 @@ class ServerArgsAutoTuner:
if args.performance_mode == "speed":
logger.info("Applying performance_mode=speed")
- if not args.enable_torch_compile and not args.is_arg_explicitly_set(
- "enable_torch_compile"
+ if (
+ self._deployment_config().speed_mode_enable_torch_compile_by_default
+ and not args.enable_torch_compile
+ and not args.is_arg_explicitly_set("enable_torch_compile")
):
# speed means fastest: compile by default. An explicit
# --enable-torch-compile false still wins (e.g. models where
- # compile measures slower, like short-step Z-Image runs).
+ # compile is slower or changes the numerical contract).
args.enable_torch_compile = True
logger.info(
"performance_mode=speed enables torch.compile "
"(pass --enable-torch-compile false to opt out)"
)
+ elif not args.enable_torch_compile and not args.is_arg_explicitly_set(
+ "enable_torch_compile"
+ ):
+ logger.info(
+ "performance_mode=speed keeps torch.compile disabled for "
+ "this model (pass --enable-torch-compile true to opt in)"
+ )
if args.num_gpus >= 2 and self._can_apply_fsdp_policy(
require_memory_headroom=False
):
diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py
index 7e4b1b274..b27f445a0 100644
--- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py
+++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py
@@ -139,6 +139,8 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
"ideogram-v4-instant",
"ideogram-ai/ideogram-4-fp8",
"ideogram-ai/ideogram-4-nf4",
+ "minimax-h3",
+ "minimaxai/minimax-h3",
"qwen/qwen-image",
"qwen/qwen-image-2512",
"qwen-image",
@@ -155,6 +157,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset(
{
"GlmImagePipelineConfig",
"Ideogram4PipelineConfig",
+ "MiniMaxH3PipelineConfig",
"QwenImagePipelineConfig",
"ZImagePipelineConfig",
}
@@ -179,6 +182,8 @@ def _normalized_bcg_model_refs(model_ref: str | None) -> set[str]:
class ServerArgs(DisaggServerArgsMixin):
# Model and path configuration (for convenience)
model_path: str
+ model_subfolder: str | None = None
+ model_variant: str | None = None
# explicit model ID override (e.g. "Qwen-Image")
model_id: str | None = None
@@ -503,6 +508,7 @@ class ServerArgs(DisaggServerArgsMixin):
self._validate_cfg_parallel()
self._validate_batching()
self._validate_breakable_cuda_graph()
+ self.pipeline_config.validate_server_args(self)
def resolved_bcg_text_buckets(self) -> tuple[int, ...]:
"""Sorted, de-duplicated, positive BCG text buckets.
@@ -550,9 +556,10 @@ class ServerArgs(DisaggServerArgsMixin):
return
logger.warning(
- "[Diffusion BCG] disabled for %s: only Ideogram-4, Qwen/Qwen-Image, "
- "Qwen/Qwen-Image-2512, Tongyi-MAI/Z-Image/Z-Image-Turbo, "
- "and zai-org/GLM-Image are currently supported.",
+ "[Diffusion BCG] disabled for %s: only Ideogram-4, MiniMax-H3, "
+ "Qwen/Qwen-Image, Qwen/Qwen-Image-2512, "
+ "Tongyi-MAI/Z-Image/Z-Image-Turbo, and zai-org/GLM-Image are "
+ "currently supported.",
pipeline_config_name,
)
self.enable_breakable_cuda_graph = False
@@ -1287,6 +1294,7 @@ class ServerArgs(DisaggServerArgsMixin):
# Convert string disagg_role to enum (from CLI/config)
if isinstance(self.disagg_role, str):
self.disagg_role = RoleType.from_string(self.disagg_role)
+ self._validate_disagg_capability()
self.gpu_ids = normalize_gpu_ids(self.gpu_ids)
# 1. adjust parameters
@@ -1311,6 +1319,26 @@ class ServerArgs(DisaggServerArgsMixin):
type=str,
help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
)
+ parser.add_argument(
+ "--model-subfolder",
+ type=str,
+ default=ServerArgs.model_subfolder,
+ help=(
+ "Advanced override for a Diffusers pipeline subfolder inside the "
+ "model repository. Prefer --model-variant when a model exposes "
+ "semantic variant-to-weights routing."
+ ),
+ )
+ parser.add_argument(
+ "--model-variant",
+ type=str,
+ default=ServerArgs.model_variant,
+ help=(
+ "Semantic checkpoint variant to serve. Models with partitioned "
+ "checkpoints use this value to select the compatible weights "
+ "without exposing repository subfolder layout."
+ ),
+ )
parser.add_argument(
"--model-id",
type=str,
@@ -1397,7 +1425,6 @@ class ServerArgs(DisaggServerArgsMixin):
"Explicit offload/FSDP/parallelism flags take precedence."
),
)
-
# Parallelism
parser.add_argument(
"--num-gpus",
@@ -1454,13 +1481,11 @@ class ServerArgs(DisaggServerArgsMixin):
help=(
"Text/image encoder parallelism across a multi-rank replica. "
"`auto` folds encoders wide enough to benefit (best "
- "single-request latency) and data-parallels the rest at "
- "batch>1; `fold` always tensor-parallels the encoder weights; "
- "`dp` never folds and splits the batch across ranks (best "
- "batched throughput; also raises --batching-max-size to the "
- "replica size unless set explicitly); `replicate` disables "
- "both. `sglang serve` defaults to `dp`; other entrypoints to "
- "`auto`."
+ "single-request latency) and data-parallels eligible native "
+ "text encoders at batch>1; `fold` always tensor-parallels the "
+ "encoder weights; `dp` never folds and splits the batch across "
+ "ranks (best batched throughput; requires TP=1 and DP=1); "
+ "`replicate` disables both. The default is `auto`."
),
)
parser.add_argument(
@@ -2293,6 +2318,20 @@ class ServerArgs(DisaggServerArgsMixin):
raise ValueError("pipeline_config is not set in ServerArgs")
self.pipeline_config.check_pipeline_config()
+ self._validate_disagg_capability()
+
+ def _validate_disagg_capability(self) -> None:
+ if self.pipeline_config is None:
+ return
+ if (
+ self.disagg_role != RoleType.MONOLITHIC
+ and not self.pipeline_config.supports_disaggregation()
+ ):
+ raise ValueError(
+ f"{type(self.pipeline_config).__name__} only supports monolithic "
+ f"deployment; disaggregation role {self.disagg_role.value!r} "
+ "is not supported"
+ )
def _validate_offload(self):
# validate dit_offload_prefetch_size
@@ -2443,7 +2482,14 @@ class ServerArgs(DisaggServerArgsMixin):
)
def _validate_cfg_parallel(self):
- if self.enable_cfg_parallel and self.num_gpus == 1:
+ if not self.enable_cfg_parallel:
+ return
+ deployment_config = self.pipeline_config.get_model_deployment_config()
+ if not deployment_config.supports_cfg_parallel:
+ raise ValueError(
+ f"{type(self.pipeline_config).__name__} does not support CFG parallelism"
+ )
+ if self.num_gpus == 1:
raise ValueError(
"CFG Parallelism is enabled via `--enable-cfg-parallel`, but num_gpus == 1"
)
@@ -2455,6 +2501,10 @@ class ServerArgs(DisaggServerArgsMixin):
raise ValueError("batching_max_size must be >= 1")
if self.batching_delay_ms < 0:
raise ValueError("batching_delay_ms must be >= 0")
+ if self.encoder_parallel == "dp" and (
+ (self.tp_size or 1) != 1 or (self.dp_size or 1) != 1
+ ):
+ raise ValueError("encoder_parallel=dp requires tp_size=1 and dp_size=1")
def _set_default_attention_backend(self) -> None:
"""Configure ROCm defaults when users do not specify an attention backend."""
diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
index 749dbe2e7..9f642ae0c 100644
--- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
+++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
@@ -32,6 +32,7 @@ from diffusers.loaders.lora_base import (
_best_guess_weight_name, # watch out for potetential removal from diffusers
)
from huggingface_hub.errors import (
+ EntryNotFoundError,
LocalEntryNotFoundError,
RepositoryNotFoundError,
RevisionNotFoundError,
@@ -72,6 +73,43 @@ _WEIGHT_FILE_PATTERNS = (
)
+def _model_hub_name() -> str:
+ return "ModelScope" if envs.SGLANG_USE_MODELSCOPE.get() else "Hugging Face Hub"
+
+
+def _snapshot_has_files(
+ local_path: str,
+ allow_patterns: Optional[Union[list[str], str]],
+) -> bool:
+ patterns = (
+ [allow_patterns]
+ if isinstance(allow_patterns, str)
+ else allow_patterns or ["**/*"]
+ )
+ return any(
+ os.path.isfile(candidate)
+ for pattern in patterns
+ for candidate in glob.iglob(
+ os.path.join(local_path, pattern),
+ recursive=True,
+ )
+ )
+
+
+def _is_modelscope_not_found_error(error: BaseException) -> bool:
+ if not envs.SGLANG_USE_MODELSCOPE.get():
+ return False
+
+ from modelscope.hub.errors import NotExistError
+
+ current: Optional[BaseException] = error
+ while current is not None:
+ if isinstance(current, NotExistError):
+ return True
+ current = current.__cause__
+ return False
+
+
def _is_diffusers_component_entry(value: Any) -> bool:
return (
isinstance(value, (list, tuple))
@@ -633,11 +671,9 @@ def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
def _resolve_remote_repo_model_index_path(model_name_or_path: str) -> str:
"""Return a local path to a remote repo's ``model_index.json``"""
- from huggingface_hub.errors import EntryNotFoundError
-
try:
- # Cache-aware: no local_dir, so HF reuses the cache and revalidates the
- # ETag against the Hub, re-downloading only when the remote changed.
+ # Cache-aware: no local_dir, so the selected Hub reuses its cache and
+ # revalidates the remote file when online.
return hf_hub_download(repo_id=model_name_or_path, filename="model_index.json")
except EntryNotFoundError:
# Repo exists but has no model_index.json (single-model repo); let the
@@ -677,8 +713,6 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
Returns:
The parsed model_index.json as a dictionary
"""
- from huggingface_hub.errors import EntryNotFoundError
-
overlay_config = maybe_load_overlay_model_index(
model_name_or_path,
snapshot_download_fn=snapshot_download,
@@ -825,7 +859,9 @@ def maybe_download_model(
# Try to read from HF cache without network access
try:
logger.info(
- "Checking for cached model in HF Hub cache for %s...", model_name_or_path
+ "Checking for cached model in %s cache for %s...",
+ _model_hub_name(),
+ model_name_or_path,
)
local_path = snapshot_download(
repo_id=model_name_or_path,
@@ -863,14 +899,17 @@ def maybe_download_model(
f"Model {model_name_or_path} found in cache but is incomplete and download=False."
)
logger.info(
- "Model found in cache but incomplete, will download from HF Hub"
+ "Model found in cache but incomplete, will download from %s",
+ _model_hub_name(),
)
except LocalEntryNotFoundError:
if not download:
raise ValueError(
f"Model {model_name_or_path} not found in local cache and download=False."
)
- logger.info("Model not found in cache, will download from HF Hub")
+ logger.info(
+ "Model not found in cache, will download from %s", _model_hub_name()
+ )
except Exception as e:
logger.warning(
"Unexpected error while checking cache for %s: %s, will attempt download",
@@ -887,7 +926,8 @@ def maybe_download_model(
for attempt in range(MAX_RETRIES):
try:
logger.info(
- "Downloading model snapshot from HF Hub for %s (attempt %d/%d)...",
+ "Downloading model snapshot from %s for %s (attempt %d/%d)...",
+ _model_hub_name(),
model_name_or_path,
attempt + 1,
MAX_RETRIES,
@@ -942,10 +982,15 @@ def maybe_download_model(
f"Model or revision not found at {model_name_or_path}. "
f"Please check the model ID or ensure you have access to the repository. Error: {e}"
) from e
- except (RequestException, RequestsConnectionError) as e:
+ except (RequestException, RequestsConnectionError, ConnectionError) as e:
+ if _is_modelscope_not_found_error(e):
+ raise ValueError(
+ f"Model or revision not found at {model_name_or_path}. "
+ "Please check the model ID or ensure you have access to the repository."
+ ) from e
if attempt == MAX_RETRIES - 1:
raise ValueError(
- f"Could not find model at {model_name_or_path} and failed to download from HF Hub "
+ f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()} "
f"after {MAX_RETRIES} attempts due to network error: {e}"
) from e
wait_time = 2**attempt
@@ -960,7 +1005,7 @@ def maybe_download_model(
time.sleep(wait_time)
except Exception as e:
raise ValueError(
- f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}"
+ f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()}: {e}"
) from e
@@ -974,13 +1019,19 @@ def hf_hub_download(
"""Unified hf_hub_download that supports both Hugging Face Hub and ModelScope."""
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import model_file_download
+ from modelscope.hub.errors import NotExistError
- return model_file_download(
- model_id=repo_id,
- file_path=filename,
- cache_dir=local_dir,
- **kwargs,
- )
+ try:
+ return model_file_download(
+ model_id=repo_id,
+ file_path=filename,
+ local_dir=str(local_dir) if local_dir is not None else None,
+ **kwargs,
+ )
+ except NotExistError as exc:
+ # Keep the Hugging Face-compatible exception contract used by
+ # maybe_download_model_index for repositories without model_index.json.
+ raise EntryNotFoundError(str(exc)) from exc
else:
from huggingface_hub import hf_hub_download as _hf_hub_download
@@ -1005,16 +1056,25 @@ def snapshot_download(
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import snapshot_download as _ms_snapshot_download
+ # ModelScope validates cached files on every online snapshot request and
+ # has no force_download argument. Dropping it preserves the caller's
+ # intended online revalidation without leaking Hub-specific kwargs.
+ kwargs.pop("force_download", None)
ms_kwargs = {
"model_id": repo_id,
- "local_dir": local_dir,
+ "local_dir": str(local_dir) if local_dir is not None else None,
"ignore_patterns": ignore_patterns,
"allow_patterns": allow_patterns,
"local_files_only": local_files_only,
"max_workers": max_workers,
}
ms_kwargs.update(kwargs)
- return _ms_snapshot_download(**ms_kwargs)
+ local_path = _ms_snapshot_download(**ms_kwargs)
+ if local_files_only and not _snapshot_has_files(local_path, allow_patterns):
+ raise LocalEntryNotFoundError(
+ f"No cached files for {repo_id} match {allow_patterns or '**/*'}"
+ )
+ return local_path
else:
from huggingface_hub import snapshot_download as _hf_snapshot_download
diff --git a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py
index 44bd6ae73..9ca30e927 100644
--- a/python/sglang/multimodal_gen/runtime/warmup_request_builder.py
+++ b/python/sglang/multimodal_gen/runtime/warmup_request_builder.py
@@ -398,6 +398,9 @@ def build_warmup_reqs(
req.suppress_logs = True
req.metrics.suppress_stage_breakdown = True
req.extra["server_internal_prewarm"] = True
+ req.sampling_params.prepare_synthetic_warmup_request_for_queue(
+ req, server_args
+ )
if return_warmup_result:
req.extra["return_warmup_result"] = True
if server_based_warmup:
diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py
index 5a0ae5b20..f79e7c54b 100644
--- a/python/sglang/multimodal_gen/test/server/gpu_cases.py
+++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py
@@ -592,6 +592,71 @@ else:
ONE_GPU_B200_CASES = ONE_GPU_MODELOPT_NVFP4_CASES
+MINIMAX_H3_FOUR_GPU_H100_CASES = [
+ DiffusionTestCase(
+ "minimax_h3_fl2va_first_frame_4gpu_h100",
+ DiffusionServerArgs(
+ model_path="MiniMaxAI/MiniMax-H3",
+ modality="video",
+ num_gpus=4,
+ tp_size=2,
+ ulysses_degree=2,
+ extras=[
+ "--model-variant",
+ "fl2va",
+ "--performance-mode",
+ "speed",
+ "--enable-torch-compile",
+ "false",
+ ],
+ ),
+ DiffusionSamplingParams(
+ prompt=(
+ "A static night view of a narrow London alley in soft rain, wet "
+ "pavement reflecting a yellow streetlamp, the blue K. West sign "
+ "glowing above a doorway, cardboard boxes near the wall, a pale "
+ "parked car in the distance, and a slender glam-rock figure "
+ "holding a guitar under the lamp; preserve the album-cover "
+ "composition, brick storefronts, muted teal and amber colors, "
+ "subtle rain shimmer only."
+ ),
+ output_size="1344x768",
+ seconds=5,
+ output_format="mp4",
+ num_outputs_per_prompt=1,
+ extras={
+ "task": "fl2va",
+ "conditions": [
+ {
+ "type": "image",
+ "uri": (
+ "https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/"
+ "5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/"
+ "825646124206.jpg/600x600bb.jpg"
+ ),
+ "role": "keyframe",
+ "frame_index": 0,
+ }
+ ],
+ "target": {
+ "short_edge": 768,
+ "aspect_ratio": "16:9",
+ "duration_seconds": 5.0,
+ },
+ "num_inference_steps": 2,
+ "flow_shift": 12.0,
+ "audio_flow_shift": 3.0,
+ "seed": 42,
+ },
+ ),
+ run_perf_check=False,
+ run_consistency_check=False,
+ run_component_accuracy_check=False,
+ run_models_api_check=False,
+ run_t2v_input_reference_check=False,
+ )
+]
+
TWO_GPU_CASES = [
DiffusionTestCase(
"flux2_modelopt_fp8_tp2_t2i",
@@ -968,6 +1033,9 @@ FILE_SUITES = {
"1-gpu-b200": [
"test_server_b200.py",
],
+ "4-gpu-h100": [
+ "test_server_4_gpu_h100.py",
+ ],
}
PARAMETRIZED_CASE_GROUPS = {
diff --git a/python/sglang/multimodal_gen/test/server/test_server_4_gpu_h100.py b/python/sglang/multimodal_gen/test/server/test_server_4_gpu_h100.py
new file mode 100644
index 000000000..031b89ae9
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/server/test_server_4_gpu_h100.py
@@ -0,0 +1,18 @@
+"""PR-blocking diffusion smoke tests that require four H100 GPUs."""
+
+from __future__ import annotations
+
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
+ diffusion_case_fixture,
+)
+from sglang.multimodal_gen.test.server.gpu_cases import (
+ MINIMAX_H3_FOUR_GPU_H100_CASES,
+)
+from sglang.multimodal_gen.test.server.test_server_common import ( # noqa: F401
+ DiffusionServerBase,
+ diffusion_server,
+)
+
+
+class TestDiffusionServerFourGpuH100(DiffusionServerBase):
+ case = diffusion_case_fixture(MINIMAX_H3_FOUR_GPU_H100_CASES)
diff --git a/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py b/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py
index c1b9aecf1..eeabcec06 100644
--- a/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py
+++ b/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py
@@ -8,6 +8,9 @@ from unittest.mock import patch
class _FakeDBCacheConfig:
+ def __init__(self, **kwargs):
+ self.kwargs = kwargs
+
def reset(self, **kwargs):
return kwargs
@@ -21,9 +24,17 @@ class _FakeForwardPattern:
def _install_cache_dit_stub():
cache_dit = types.ModuleType("cache_dit")
+ cache_dit.enable_calls = []
+ cache_dit.disable_calls = []
cache_dit.refresh_calls = []
cache_dit.steps_mask_calls = []
+ def enable_cache(target, **kwargs):
+ cache_dit.enable_calls.append({"target": target, **kwargs})
+
+ def disable_cache(target):
+ cache_dit.disable_calls.append(target)
+
def refresh_context(transformer, cache_config, verbose=False):
cache_dit.refresh_calls.append(
{
@@ -39,6 +50,8 @@ def _install_cache_dit_stub():
)
return [1] * total_steps
+ cache_dit.enable_cache = enable_cache
+ cache_dit.disable_cache = disable_cache
cache_dit.refresh_context = refresh_context
cache_dit.steps_mask = steps_mask
cache_dit.BlockAdapter = types.SimpleNamespace
@@ -50,9 +63,11 @@ def _install_cache_dit_stub():
block_adapters = types.ModuleType("cache_dit.caching.block_adapters")
class _FakeBlockAdapterRegister:
- @staticmethod
- def is_supported(_transformer):
- return True
+ supported = True
+
+ @classmethod
+ def is_supported(cls, _transformer):
+ return cls.supported
block_adapters.BlockAdapterRegister = _FakeBlockAdapterRegister
@@ -279,6 +294,35 @@ class TestBuildCustomBlockAdapter(unittest.TestCase):
)
self.assertFalse(adapter_turbo.has_separate_cfg)
+ def test_minimax_h3_uses_main_blocks_with_hidden_state_pattern(self):
+ module = _import_module_with_stub()
+ blocks = ["block_0", "block_1"]
+ transformer = _make_transformer("MiniMaxH3DiTModel")
+ transformer.blocks = blocks
+
+ adapter = module._build_custom_block_adapter(transformer)
+
+ self.assertEqual(adapter.blocks, blocks)
+ self.assertEqual(adapter.forward_pattern, "Pattern_3")
+ self.assertFalse(adapter.has_separate_cfg)
+
+ def test_custom_adapter_is_retained_until_disable(self):
+ module = _import_module_with_stub()
+ module.BlockAdapterRegister.supported = False
+ transformer = _make_transformer("MiniMaxH3DiTModel")
+ transformer.blocks = ["block_0"]
+ config = module.CacheDitConfig(enabled=True, num_inference_steps=50)
+
+ returned = module.enable_cache_on_transformer(transformer, config)
+
+ self.assertIs(returned, transformer)
+ adapter = transformer._sglang_cache_dit_adapter
+ self.assertIs(module.cache_dit.enable_calls[0]["target"], adapter)
+
+ self.assertIs(module.disable_cache_on_transformer(transformer), transformer)
+ self.assertEqual(module.cache_dit.disable_calls, [adapter])
+ self.assertFalse(hasattr(transformer, "_sglang_cache_dit_adapter"))
+
if __name__ == "__main__":
unittest.main()
diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py
index f83f18197..c0086492b 100644
--- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py
+++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py
@@ -37,6 +37,10 @@ class Ideogram4Transformer2DModel(torch.nn.Module):
pass
+class MiniMaxH3DiTModel(torch.nn.Module):
+ pass
+
+
class ZImageTransformer2DModel(torch.nn.Module):
def rotary_emb(self, pos_ids):
return torch.zeros(pos_ids.shape[0], 8, device=pos_ids.device)
@@ -47,6 +51,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
self.stage = DenoisingStage.__new__(DenoisingStage)
self.qwen_model = QwenImageTransformer2DModel()
self.ideogram_model = Ideogram4Transformer2DModel()
+ self.minimax_h3_model = MiniMaxH3DiTModel()
self.zimage_model = ZImageTransformer2DModel()
self.other_model = OtherTransformer2DModel()
@@ -192,6 +197,93 @@ class TestDiffusionBCGPadding(unittest.TestCase):
self.assertEqual(short["freqs_cis"][0].shape, (64, 8))
self.assertEqual(_signature_kwargs(short), _signature_kwargs(longer))
+ def _minimax_h3_kwargs(self, text_seq: int):
+ image_seq = 4
+ audio_seq = 2
+ media_seq = image_seq + audio_seq
+ used = text_seq + media_seq
+ packed_seq = 64
+ audio_pos = torch.arange(text_seq, text_seq + audio_seq)
+ img_pos = torch.arange(text_seq + audio_seq, used)
+ return {
+ "x": torch.zeros(1, packed_seq, 96),
+ "audio_x": torch.zeros(1, packed_seq, 32),
+ "img_position_ids": torch.zeros(1, packed_seq, 3),
+ "unique_timesteps": torch.zeros(1),
+ "inverse_indices": torch.zeros(packed_seq, dtype=torch.long),
+ "update_mask": torch.ones(image_seq, dtype=torch.bool),
+ "token_tags": torch.full((packed_seq,), -1, dtype=torch.long),
+ "skip_mask_out_condition": False,
+ "prompt_embeds": torch.ones(text_seq, 8),
+ "refined_prompt_embeds_length": text_seq,
+ "img_pos_info": {"position_ids": img_pos},
+ "audio_pos_info": {"position_ids": audio_pos},
+ "text_pos_info": {"position_ids": torch.arange(text_seq, dtype=torch.long)},
+ "img_pos_for_infer_output_info": {"position_ids": img_pos},
+ "local_embedding_layout": {
+ "text_source_ids": torch.arange(text_seq),
+ "text_row_ids": torch.arange(text_seq),
+ },
+ "packed_seq_params": {
+ "cu_seqlens_q": torch.tensor([0, used, packed_seq], dtype=torch.int32),
+ "max_seqlen_q": used,
+ },
+ "refiner_packed_seq_params": {
+ "cu_seqlens_q": torch.tensor(
+ [0, text_seq, text_seq], dtype=torch.int32
+ ),
+ "max_seqlen_q": text_seq,
+ },
+ }
+
+ def test_minimax_h3_prompt_lengths_share_bucket_signature(self):
+ with self._patch_buckets(64, 128):
+ short = self.stage._bcg_pad_prompt_kwargs(
+ self._minimax_h3_kwargs(19), current_model=self.minimax_h3_model
+ )
+ longer = self.stage._bcg_pad_prompt_kwargs(
+ self._minimax_h3_kwargs(47), current_model=self.minimax_h3_model
+ )
+
+ # H3 keeps the original aligned main sequence: growing it changes SP
+ # row partitions and loses bitwise eager equivalence.
+ self.assertEqual(short["x"].shape, (1, 64, 96))
+ self.assertEqual(longer["x"].shape, (1, 64, 96))
+ self.assertEqual(short["audio_x"].shape, (1, 64, 32))
+ self.assertEqual(short["prompt_embeds"].shape, (64, 8))
+ self.assertEqual(longer["prompt_embeds"].shape, (64, 8))
+ self.assertNotIn("local_embedding_layout", short)
+ self.assertNotIn("local_embedding_layout", longer)
+ self.assertEqual(
+ short["refined_prompt_embeds_length"].shape,
+ torch.Size([]),
+ )
+ self.assertEqual(short["refined_prompt_embeds_length"].item(), 19)
+ self.assertEqual(longer["refined_prompt_embeds_length"].item(), 47)
+ self.assertEqual(
+ short["packed_seq_params"]["cu_seqlens_q"].tolist(), [0, 25, 64]
+ )
+ self.assertEqual(
+ longer["packed_seq_params"]["cu_seqlens_q"].tolist(), [0, 53, 64]
+ )
+ self.assertEqual(short["packed_seq_params"]["max_seqlen_q"], 64)
+ self.assertEqual(short["refiner_packed_seq_params"]["max_seqlen_q"], 64)
+ self.assertEqual(
+ short["refiner_packed_seq_params"]["cu_seqlens_q"].tolist(),
+ [0, 19, 64],
+ )
+ # Dummy prompt rows occupy the independent packed-attention segment;
+ # real text/media remain in the first segment.
+ self.assertEqual(
+ short["text_pos_info"]["position_ids"][19:].tolist(),
+ list(range(25, 70)),
+ )
+ self.assertEqual(
+ longer["text_pos_info"]["position_ids"][47:].tolist(),
+ list(range(53, 70)),
+ )
+ self.assertEqual(_signature_kwargs(short), _signature_kwargs(longer))
+
def _ideogram_kwargs(self, text_seq: int, *, image_seq: int = 4):
total_seq = text_seq + image_seq
indicator = torch.zeros(1, total_seq, dtype=torch.long)
diff --git a/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py b/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py
index cc55ef344..441ceaf67 100644
--- a/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py
+++ b/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py
@@ -221,7 +221,7 @@ def _finalize(
policy,
mode="world",
group_size=2,
- batched=False,
+ prefer_dp=False,
measured=True,
):
monkeypatch.setattr(
@@ -234,7 +234,7 @@ def _finalize(
)
enc = _enc(hidden, heads, inter)
enc.parallel_folding_mode = mode
- finalize_encoder_folding(enc, policy, batched=batched)
+ finalize_encoder_folding(enc, policy, prefer_dp=prefer_dp)
return enc.parallel_folding_mode
@@ -249,11 +249,13 @@ def test_finalize_auto_keeps_wide_clears_narrow(monkeypatch):
assert _finalize(monkeypatch, 2560, 32, 9728, "auto") is None # below threshold
-def test_finalize_auto_leaves_dp_capable_unsharded_when_batched(monkeypatch):
+def test_finalize_auto_leaves_dp_capable_unsharded_when_dp_preferred(monkeypatch):
# with a batch, dp (one all_gather) beats folding (an all_reduce per layer)
- assert _finalize(monkeypatch, 4096, 64, 10240, "auto", batched=True) is None
+ assert _finalize(monkeypatch, 4096, 64, 10240, "auto", prefer_dp=True) is None
+ # TP or DiT-DP makes encoder DP ineligible, so keep the useful fold
+ assert _finalize(monkeypatch, 4096, 64, 10240, "auto") == "world"
# CLIP-L cannot dp either, so folding remains the only question
- assert _finalize(monkeypatch, 768, 12, 3072, "auto", batched=True) is None
+ assert _finalize(monkeypatch, 768, 12, 3072, "auto", prefer_dp=True) is None
def test_finalize_auto_needs_a_measured_topology(monkeypatch):
diff --git a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py
new file mode 100644
index 000000000..bf60422e6
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py
@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+from unittest.mock import patch
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.runtime.loader import fsdp_load
+
+
+class _UniformDtypeModel(nn.Module):
+ param_names_mapping = {}
+
+ def __init__(self) -> None:
+ super().__init__()
+
+ def post_load_weights(self) -> None:
+ pass
+
+
+class _MixedDtypeModel(_UniformDtypeModel):
+ _fsdp_mixed_dtype_params = True
+
+
+class TestFSDPMixedPrecisionPolicy(unittest.TestCase):
+ def _load_and_capture_policy(
+ self,
+ model_cls: type[nn.Module],
+ *,
+ fsdp_inference: bool,
+ ):
+ with (
+ patch.object(fsdp_load.current_platform, "is_mps", return_value=False),
+ patch.object(fsdp_load, "init_device_mesh", return_value=object()),
+ patch.object(fsdp_load, "shard_model") as shard_model,
+ patch.object(
+ fsdp_load,
+ "safetensors_weights_iterator",
+ return_value=iter(()),
+ ),
+ patch.object(fsdp_load, "load_model_from_full_model_state_dict"),
+ patch.object(fsdp_load, "set_mixed_precision_policy") as set_policy,
+ ):
+ fsdp_load.maybe_load_fsdp_model(
+ model_cls=model_cls,
+ init_params={},
+ weight_dir_list=[],
+ device=torch.device("cpu"),
+ hsdp_replicate_dim=1,
+ hsdp_shard_dim=1,
+ param_dtype=torch.bfloat16,
+ reduce_dtype=torch.float32,
+ fsdp_inference=fsdp_inference,
+ )
+
+ state_kwargs = set_policy.call_args.kwargs
+ return state_kwargs["mp_policy"], state_kwargs, shard_model
+
+ def test_uniform_dtype_model_uses_requested_fsdp_param_dtype(self):
+ policy, state_kwargs, shard_model = self._load_and_capture_policy(
+ _UniformDtypeModel,
+ fsdp_inference=True,
+ )
+
+ self.assertEqual(policy.param_dtype, torch.bfloat16)
+ self.assertEqual(policy.reduce_dtype, torch.float32)
+ self.assertIsNone(policy.output_dtype)
+ self.assertFalse(policy.cast_forward_inputs)
+ self.assertEqual(state_kwargs["param_dtype"], torch.bfloat16)
+ self.assertIs(shard_model.call_args.kwargs["mp_policy"], policy)
+
+ def test_mixed_dtype_model_preserves_original_fsdp_param_dtypes(self):
+ policy, state_kwargs, shard_model = self._load_and_capture_policy(
+ _MixedDtypeModel,
+ fsdp_inference=True,
+ )
+
+ self.assertIsNone(policy.param_dtype)
+ self.assertEqual(policy.reduce_dtype, torch.float32)
+ self.assertIsNone(policy.output_dtype)
+ self.assertFalse(policy.cast_forward_inputs)
+ self.assertEqual(state_kwargs["param_dtype"], torch.bfloat16)
+ self.assertIs(shard_model.call_args.kwargs["mp_policy"], policy)
+
+ def test_mixed_dtype_opt_in_does_not_change_non_fsdp_policy(self):
+ policy, state_kwargs, shard_model = self._load_and_capture_policy(
+ _MixedDtypeModel,
+ fsdp_inference=False,
+ )
+
+ self.assertEqual(policy.param_dtype, torch.bfloat16)
+ self.assertEqual(state_kwargs["param_dtype"], torch.bfloat16)
+ shard_model.assert_not_called()
diff --git a/python/sglang/multimodal_gen/test/unit/test_gpu_worker_cpu_threads.py b/python/sglang/multimodal_gen/test/unit/test_gpu_worker_cpu_threads.py
new file mode 100644
index 000000000..b772db821
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_gpu_worker_cpu_threads.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Contract for the per-worker CPU intra-op thread budget."""
+
+import unittest
+from unittest.mock import patch
+
+from sglang.multimodal_gen.runtime.managers.gpu_worker import (
+ _worker_cpu_intra_op_threads,
+)
+
+
+class TestWorkerCpuIntraOpThreads(unittest.TestCase):
+ def test_divides_host_cores_across_colocated_workers(self):
+ with (
+ patch.dict("os.environ", {}, clear=False),
+ patch("os.cpu_count", return_value=128),
+ ):
+ import os
+
+ os.environ.pop("OMP_NUM_THREADS", None)
+ self.assertEqual(_worker_cpu_intra_op_threads(8), 16)
+ self.assertEqual(_worker_cpu_intra_op_threads(4), 16) # capped
+ self.assertEqual(_worker_cpu_intra_op_threads(128), 1)
+ self.assertEqual(_worker_cpu_intra_op_threads(256), 1) # floor
+
+ def test_single_gpu_keeps_cap(self):
+ with (
+ patch.dict("os.environ", {}, clear=False),
+ patch("os.cpu_count", return_value=8),
+ ):
+ import os
+
+ os.environ.pop("OMP_NUM_THREADS", None)
+ self.assertEqual(_worker_cpu_intra_op_threads(1), 8)
+
+ def test_explicit_omp_setting_wins(self):
+ with patch.dict("os.environ", {"OMP_NUM_THREADS": "32"}):
+ self.assertIsNone(_worker_cpu_intra_op_threads(8))
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py b/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py
index 614a5134b..f396de93d 100644
--- a/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py
+++ b/python/sglang/multimodal_gen/test/unit/test_hf_diffusers_utils.py
@@ -1,9 +1,15 @@
import json
+import modelscope
+import pytest
+from huggingface_hub.errors import LocalEntryNotFoundError
+
+from sglang.multimodal_gen.runtime.utils import hf_diffusers_utils
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
_check_index_files_for_missing_shards,
_verify_diffusers_model_complete,
)
+from sglang.srt.environ import envs
def _write_model_index(root):
@@ -70,3 +76,91 @@ def test_diffusers_cache_validation_checks_declared_component_shards(tmp_path):
assert not is_valid
assert "transformer/missing.safetensors" in missing_files
assert "transformer" in checked_subdirs
+
+
+def test_modelscope_file_download_preserves_local_dir(monkeypatch, tmp_path):
+ calls = []
+
+ def model_file_download(**kwargs):
+ calls.append(kwargs)
+ return str(tmp_path / kwargs["file_path"])
+
+ monkeypatch.setattr(modelscope, "model_file_download", model_file_download)
+
+ with envs.SGLANG_USE_MODELSCOPE.override(True):
+ result = hf_diffusers_utils.hf_hub_download(
+ "MiniMax/MiniMax-H3",
+ "FL2VA/model_index.json",
+ local_dir=tmp_path,
+ revision="master",
+ )
+
+ assert result == str(tmp_path / "FL2VA/model_index.json")
+ assert calls == [
+ {
+ "model_id": "MiniMax/MiniMax-H3",
+ "file_path": "FL2VA/model_index.json",
+ "local_dir": str(tmp_path),
+ "revision": "master",
+ }
+ ]
+
+
+def test_modelscope_snapshot_download_selects_h3_partition(monkeypatch, tmp_path):
+ calls = []
+
+ def snapshot_download(**kwargs):
+ calls.append(kwargs)
+ return str(tmp_path)
+
+ monkeypatch.setattr(modelscope, "snapshot_download", snapshot_download)
+
+ with envs.SGLANG_USE_MODELSCOPE.override(True):
+ result = hf_diffusers_utils.snapshot_download(
+ "MiniMax/MiniMax-H3",
+ local_dir=tmp_path,
+ allow_patterns=["Ref2VA/**"],
+ force_download=True,
+ )
+
+ assert result == str(tmp_path)
+ assert calls == [
+ {
+ "model_id": "MiniMax/MiniMax-H3",
+ "local_dir": str(tmp_path),
+ "ignore_patterns": None,
+ "allow_patterns": ["Ref2VA/**"],
+ "local_files_only": False,
+ "max_workers": 8,
+ }
+ ]
+
+
+def test_modelscope_empty_selected_partition_is_a_cache_miss(monkeypatch, tmp_path):
+ monkeypatch.setattr(modelscope, "snapshot_download", lambda **_: str(tmp_path))
+
+ with (
+ envs.SGLANG_USE_MODELSCOPE.override(True),
+ pytest.raises(LocalEntryNotFoundError, match="Ref2VA"),
+ ):
+ hf_diffusers_utils.snapshot_download(
+ "MiniMax/MiniMax-H3",
+ allow_patterns=["Ref2VA/**"],
+ local_files_only=True,
+ )
+
+
+def test_modelscope_selected_partition_cache_hit_requires_a_file(monkeypatch, tmp_path):
+ model_index = tmp_path / "FL2VA" / "model_index.json"
+ model_index.parent.mkdir()
+ model_index.write_text("{}")
+ monkeypatch.setattr(modelscope, "snapshot_download", lambda **_: str(tmp_path))
+
+ with envs.SGLANG_USE_MODELSCOPE.override(True):
+ result = hf_diffusers_utils.snapshot_download(
+ "MiniMax/MiniMax-H3",
+ allow_patterns=["FL2VA/**"],
+ local_files_only=True,
+ )
+
+ assert result == str(tmp_path)
diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
index d73487817..542e7a77f 100644
--- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
+++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
@@ -91,6 +91,13 @@ class _NestedDummyModel(torch.nn.Module, LayerwiseOffloadableModuleMixin):
self.encoder = _DummyModel()
+class _NestedSameNamedBlocksModel(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.token_refiner = _DummyModel()
+ self.blocks = torch.nn.ModuleList([_DummyBlock()])
+
+
class _SharedBuffer(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
@@ -231,6 +238,31 @@ def test_layerwise_offload_uses_normal_tensors_under_inference_mode(monkeypatch)
assert model.blocks[0].bias._version >= 0
+def test_layerwise_offload_does_not_capture_nested_same_named_layers(monkeypatch):
+ monkeypatch.setattr(
+ layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+ )
+ monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+
+ model = _NestedSameNamedBlocksModel()
+ refiner_weight = model.token_refiner.blocks[0].weight.detach().clone()
+ manager = LayerwiseOffloadManager(
+ model=model,
+ layers_attr_str="blocks",
+ num_layers=1,
+ enabled=True,
+ pin_cpu_memory=False,
+ prefetch_size=1,
+ )
+
+ managed_names = {
+ name for metadata in manager._weight_metadata.values() for name in metadata
+ }
+ assert managed_names
+ assert all(name.startswith("blocks.") for name in managed_names)
+ assert torch.equal(model.token_refiner.blocks[0].weight, refiner_weight)
+
+
def test_layerwise_offload_keeps_shared_buffers_resident(monkeypatch):
monkeypatch.setattr(
layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
@@ -578,6 +610,10 @@ class _ResidentComponent(torch.nn.Module, LayerwiseOffloadableModuleMixin):
self.blocks = torch.nn.ModuleList([_DummyBlock() for _ in range(n)])
+class _AuxiliaryResidentComponent(_ResidentComponent):
+ layerwise_offload_dit_group_enabled = False
+
+
def _patch_fake_device(monkeypatch):
monkeypatch.setattr(
layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
@@ -691,6 +727,21 @@ def test_configure_resolves_resident_layers_ratio(monkeypatch):
assert comp.layerwise_offload_managers[0].resident_layers == 4
+def test_auxiliary_layerwise_components_ignore_dit_tuning(monkeypatch):
+ _patch_fake_device(monkeypatch)
+ comp = _AuxiliaryResidentComponent(8)
+ comp.configure_layerwise_offload(
+ _server_args(
+ dit_offload_prefetch_size=3,
+ dit_layerwise_resident_layers=0.5,
+ )
+ )
+
+ manager = comp.layerwise_offload_managers[0]
+ assert manager.prefetch_size == 1
+ assert manager.resident_layers == 0
+
+
class _MixinBlock(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
new file mode 100644
index 000000000..229f109be
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+"""High-value task, partition, and public request admission contracts."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
+ MiniMaxH3PipelineConfig,
+)
+from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+ VideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
+ MiniMaxH3PartitionAdmissionStage,
+ MiniMaxH3ReleaseMetadata,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.request_validation import (
+ minimax_h3_validate_canonical_request,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
+ minimax_h3_resolve_plan,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
+ partition_for_task,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args.server_args import Backend
+
+TARGET = {
+ "short_edge": 768,
+ "aspect_ratio": "16:9",
+ "duration_seconds": 5.0,
+}
+
+
+@pytest.mark.parametrize(
+ ("task", "conditions", "partition", "visual", "audio", "chains"),
+ [
+ ("t2va", [], "fl2va", [], [], []),
+ (
+ "fl2va",
+ [
+ {
+ "type": "image",
+ "uri": "file:///first.png",
+ "role": "keyframe",
+ "frame_index": 0,
+ },
+ {
+ "type": "image",
+ "uri": "file:///last.png",
+ "role": "keyframe",
+ "frame_index": -1,
+ },
+ ],
+ "fl2va",
+ [0, 1],
+ [],
+ ["image.target_canvas", "image.target_canvas"],
+ ),
+ (
+ "ref2va",
+ [
+ {
+ "type": "image",
+ "uri": "file:///image.png",
+ "role": "reference",
+ },
+ {
+ "type": "video",
+ "uri": "file:///video.mp4",
+ "role": "reference",
+ "start_time_seconds": 12.5,
+ },
+ {
+ "type": "audio",
+ "uri": "file:///audio.wav",
+ "role": "reference",
+ },
+ {
+ "type": "video_audio",
+ "uri": "file:///av.mp4",
+ "role": "reference",
+ },
+ ],
+ "ref2va",
+ [0, 1, 3],
+ [1, 2, 3],
+ [
+ "image.reference_preserve",
+ "video.reference_preserve",
+ "audio",
+ "video_audio.reference_preserve",
+ ],
+ ),
+ ],
+)
+def test_public_tasks_resolve_to_exact_partition_and_encoder_plan(
+ task, conditions, partition, visual, audio, chains
+):
+ canonical = minimax_h3_validate_canonical_request(
+ task=task,
+ prompt="contract",
+ conditions=conditions,
+ target=TARGET,
+ seed=0,
+ )
+ plan = minimax_h3_resolve_plan(canonical)
+
+ assert partition_for_task(task) == partition
+ assert plan.task == task
+ assert plan.encoders["visual"] == visual
+ assert plan.encoders["audio"] == audio
+ assert [material.material_chain for material in plan.materials] == chains
+ if task == "ref2va":
+ assert plan.materials[1].start_time_seconds == 12.5
+ assert plan.shape["frame_count"] == 124
+ assert plan.shape["video_latent_t"] == 37
+
+
+@pytest.mark.parametrize(
+ ("partition", "tasks"),
+ [("fl2va", ["t2va", "fl2va"]), ("ref2va", ["ref2va"])],
+)
+def test_loaded_weight_partition_admits_only_its_declared_tasks(partition, tasks):
+ metadata = MiniMaxH3ReleaseMetadata.from_model_index(
+ {
+ "_minimax_h3": {
+ "schema_version": 1,
+ "partition": partition,
+ "tasks": tasks,
+ "task_aliases": {},
+ "sigma_shift_scales": {"video": 12.0, "audio": 3.0},
+ }
+ }
+ )
+
+ assert [metadata.canonical_task(task) for task in tasks] == tasks
+ rejected = "ref2va" if partition == "fl2va" else "t2va"
+ with pytest.raises(ValueError):
+ metadata.canonical_task(rejected)
+
+
+def test_duration_admission_accepts_released_4_to_15_second_range():
+ for duration in (4.0, 15.0):
+ target = {**TARGET, "duration_seconds": duration}
+ canonical = minimax_h3_validate_canonical_request(
+ task="t2va",
+ prompt="duration contract",
+ conditions=[],
+ target=target,
+ seed=0,
+ )
+ assert canonical["target"]["duration_seconds"] == duration
+
+ for duration in (3.9, 15.1):
+ target = {**TARGET, "duration_seconds": duration}
+ with pytest.raises(ValueError, match=r"\[4, 15\]"):
+ minimax_h3_validate_canonical_request(
+ task="t2va",
+ prompt="duration contract",
+ conditions=[],
+ target=target,
+ seed=0,
+ )
+
+
+def test_video_adapter_lowers_only_native_fields_and_rejects_cfg():
+ request = VideoGenerationsRequest(
+ prompt="contract",
+ task="t2va",
+ conditions=[],
+ target=TARGET,
+ flow_shift=8.0,
+ audio_flow_shift=2.0,
+ quality="high",
+ imgvid_cond_noise_aug_for_inference=0.75,
+ audio_cond_noise_aug_for_inference=0.5,
+ )
+ generic = {
+ "prompt": request.prompt,
+ "seed": request.seed,
+ "flow_shift": request.flow_shift,
+ }
+
+ lowered = MiniMaxH3SamplingParams.lower_video_request_kwargs(request, generic)
+ assert lowered == {
+ "prompt": "contract",
+ "seed": request.seed,
+ "task": "t2va",
+ "conditions": [],
+ "target": TARGET,
+ "flow_shift": 8.0,
+ "audio_flow_shift": 2.0,
+ "quality": "high",
+ "imgvid_cond_noise_aug_for_inference": 0.75,
+ "audio_cond_noise_aug_for_inference": 0.5,
+ }
+
+ with pytest.raises(ValueError):
+ MiniMaxH3SamplingParams.lower_video_request_kwargs(
+ request, {**generic, "guidance_scale": 7.5}
+ )
+
+
+class _HopperCapability:
+ def to_int(self) -> int:
+ return 90
+
+
+def _quality_server_args():
+ return SimpleNamespace(
+ attention_backend=None,
+ model_variant="fl2va",
+ num_gpus=4,
+ backend=Backend.AUTO,
+ component_attention_backends={},
+ enable_breakable_cuda_graph=False,
+ enable_torch_compile=False,
+ is_dit_layerwise_offload_selected=False,
+ performance_mode="speed",
+ quantization=None,
+ regional_compile=False,
+ ring_degree=1,
+ sp_degree=4,
+ tp_size=1,
+ ulysses_degree=4,
+ use_fsdp_inference=False,
+ )
+
+
+def test_quality_admission_fails_closed_outside_validated_request():
+ metadata = MiniMaxH3ReleaseMetadata.from_model_index(
+ {
+ "_minimax_h3": {
+ "schema_version": 1,
+ "partition": "fl2va",
+ "tasks": ["t2va", "fl2va"],
+ "task_aliases": {},
+ "sigma_shift_scales": {"video": 12.0, "audio": 3.0},
+ }
+ }
+ )
+ canonical = minimax_h3_validate_canonical_request(
+ task="t2va",
+ prompt="quality",
+ conditions=[],
+ target=TARGET,
+ seed=0,
+ )
+ plan = minimax_h3_resolve_plan(canonical)
+ batch = SimpleNamespace(
+ sampling_params=SimpleNamespace(task="t2va", quality="high"),
+ num_inference_steps=50,
+ is_warmup=False,
+ )
+ stage = MiniMaxH3PartitionAdmissionStage(metadata)
+ config = MiniMaxH3PipelineConfig()
+ server_args = _quality_server_args()
+ server_args.pipeline_config = config
+
+ with (
+ patch.object(current_platform, "is_cuda", return_value=True),
+ patch.object(current_platform, "get_device_name", return_value="NVIDIA H200"),
+ patch.object(
+ current_platform,
+ "get_device_capability",
+ return_value=_HopperCapability(),
+ ),
+ patch(
+ "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages."
+ "minimax_h3.release_metadata.minimax_h3_plan_from_batch",
+ return_value=plan,
+ ),
+ ):
+ assert stage.forward(batch, server_args) is batch
+ batch.num_inference_steps = 40
+ with pytest.raises(ValueError, match="validated only"):
+ stage.forward(batch, server_args)
+
+ batch.sampling_params.quality = "lossless"
+ batch.num_inference_steps = 50
+ server_args.attention_backend = "sage_attn"
+ with pytest.raises(ValueError, match="does not support SageAttention"):
+ stage.forward(batch, server_args)
+
+ batch.sampling_params.quality = "unsupported"
+ server_args.attention_backend = None
+ with pytest.raises(ValueError, match="unsupported MiniMax-H3 quality profile"):
+ stage.forward(batch, server_args)
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py
new file mode 100644
index 000000000..8ca111133
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py
@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical contract for request-static H3 denoise metadata."""
+
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+ MINIMAX_H3_ADALN_MODALITY_NUM,
+)
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_minimax_h3_euler_ancestral import (
+ _minimax_h3_euler_eta0_step,
+ _minimax_h3_rf_v_to_x0,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import (
+ MiniMaxH3DenoiseBranch,
+ _build_local_embedding_layout,
+ _minimax_h3_update_target_rows_,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import (
+ minimax_h3_packed_sequence,
+ minimax_h3_packed_sequence_ref2va_blocks,
+)
+
+
+def _branch(
+ mode: str, token_tags: torch.Tensor | None = None
+) -> MiniMaxH3DenoiseBranch:
+ common = dict(text_len=3, latent_t=2, latent_h=4, latent_w=4, audio_t=3)
+ if mode == "t2va":
+ packed = minimax_h3_packed_sequence(
+ **common,
+ include_keyframe_cond=False,
+ )
+ elif mode == "fl2va":
+ packed = minimax_h3_packed_sequence(
+ **common,
+ include_keyframe_cond=True,
+ keyframe_frame_indices=[0, -1],
+ frame_count=5,
+ )
+ else:
+ packed = minimax_h3_packed_sequence_ref2va_blocks(
+ **common,
+ ref_blocks=[
+ {"kind": "image", "latent_h": 4, "latent_w": 4},
+ {"kind": "audio", "ref_audio_t": 2},
+ ],
+ )
+ return MiniMaxH3DenoiseBranch(
+ packed=packed,
+ text_embeddings=torch.zeros(3, 5120),
+ token_tags=packed["token_tags"] if token_tags is None else token_tags,
+ device=torch.device("cpu"),
+ )
+
+
+def test_precomputed_timestep_plan_matches_full_unique_reference():
+ """Preplanning must preserve fp32 collisions and every packed row class."""
+
+ for mode in ("t2va", "fl2va", "ref2va"):
+ branch = _branch(mode)
+ assert branch.static_kwargs["skip_mask_out_condition"]
+ assert "token_tags" not in branch.static_kwargs
+ assert not bool((branch.static_kwargs["block_token_tags"] < 0).any())
+ torch.testing.assert_close(
+ branch.static_kwargs["img_pos_for_infer_output_info"]["position_ids"],
+ branch.img_pos_dev[branch.update_mask_dev],
+ rtol=0,
+ atol=0,
+ )
+ video_steps = [0.75, 0.1]
+ audio_steps = [0.625, 0.2]
+ plan = branch.prepare_timestep_plan(
+ video_timesteps=video_steps,
+ audio_timesteps=audio_steps,
+ imgvid_cond_noise_aug=0.6,
+ audio_ref_cond_noise_aug=0.4,
+ )
+
+ assert branch.static_kwargs["packed_seq_params"]["cu_seqlens_q_host"] == tuple(
+ int(value)
+ for value in branch.static_kwargs["packed_seq_params"][
+ "cu_seqlens_q"
+ ].tolist()
+ )
+ assert branch.static_kwargs["refiner_packed_seq_params"][
+ "cu_seqlens_q_host"
+ ] == (0, 3, 3)
+
+ for step, (video_t, audio_t) in enumerate(
+ zip(video_steps, audio_steps, strict=True)
+ ):
+ reference = torch.full((branch.seq_len,), video_t, dtype=torch.float32)
+ reference[branch.img_cond_seq_idx] = max(video_t, 0.6)
+ reference[branch.audio_target_seq_idx] = audio_t
+ reference[branch.audio_ref_seq_idx] = max(audio_t, 0.4)
+ expected = torch.unique(reference, sorted=True, return_inverse=True)
+ torch.testing.assert_close(plan[step][0], expected[0], rtol=0, atol=0)
+ torch.testing.assert_close(plan[step][1], expected[1], rtol=0, atol=0)
+ torch.testing.assert_close(
+ plan[step][2],
+ branch.static_kwargs["block_token_tags"]
+ + expected[1] * MINIMAX_H3_ADALN_MODALITY_NUM,
+ rtol=0,
+ atol=0,
+ )
+
+ repeated_plan = branch.prepare_timestep_plan(
+ video_timesteps=[0.0, 0.1, 0.2],
+ audio_timesteps=[0.0, 0.2, 0.4],
+ imgvid_cond_noise_aug=0.999,
+ audio_ref_cond_noise_aug=1.0,
+ )
+ assert repeated_plan[1][1] is repeated_plan[2][1]
+ assert repeated_plan[1][2] is repeated_plan[2][2]
+
+
+def test_inplace_target_update_matches_scheduler_math():
+ generator = torch.Generator().manual_seed(7)
+ for sigma_curr, sigma_next in ((1.0, 0.7), (0.2, 0.0), (0.0, 0.0)):
+ state = torch.randn(11, 32, generator=generator)
+ velocity = torch.randn(11, 32, generator=generator)
+ timestep = torch.tensor(1.0 - sigma_curr)
+ ratio = torch.tensor(0.0 if sigma_curr == 0.0 else sigma_next / sigma_curr)
+ denoised = _minimax_h3_rf_v_to_x0(state, velocity, timestep)
+ expected = _minimax_h3_euler_eta0_step(
+ state,
+ denoised,
+ sigma_curr=sigma_curr,
+ sigma_next=sigma_next,
+ sigma_ratio=ratio,
+ )
+
+ actual = state.clone()
+ velocity_scratch = velocity.clone()
+ _minimax_h3_update_target_rows_(
+ actual,
+ velocity_scratch,
+ sigma_t=1.0 - timestep,
+ sigma_curr=sigma_curr,
+ sigma_ratio=ratio,
+ one_minus_sigma_ratio=1.0 - ratio,
+ denoised_scratch=torch.empty_like(actual),
+ )
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_local_text_layout_is_a_contiguous_prefix_per_ulysses_rank():
+ for mode in ("t2va", "fl2va", "ref2va"):
+ branch = _branch(mode)
+ text_len = int(branch.static_kwargs["prompt_embeds"].shape[0])
+ for world_size in (1, 2, 4, 8):
+ for rank in range(world_size):
+ layout = _build_local_embedding_layout(
+ seq_len=branch.seq_len,
+ text_pos=torch.arange(text_len),
+ img_pos=branch.img_pos,
+ audio_pos=branch.audio_pos,
+ world_size=world_size,
+ rank=rank,
+ device=torch.device("cpu"),
+ )
+ start = int(layout["text_source_start"])
+ stop = int(layout["text_source_stop"])
+ row_start = rank * (branch.seq_len // world_size)
+ expected = torch.nonzero(
+ (torch.arange(text_len) >= row_start)
+ & (
+ torch.arange(text_len)
+ < row_start + branch.seq_len // world_size
+ )
+ ).view(-1)
+ assert expected.tolist() == list(range(start, stop))
+
+
+def test_rank_local_token_tags_match_reference_slice():
+ for mode in ("t2va", "fl2va", "ref2va"):
+ seq_len = _branch(mode).seq_len
+ token_tags = torch.arange(seq_len, dtype=torch.long) - seq_len // 2
+ for world_size in (1, 2, 4, 8):
+ for rank in range(world_size):
+ with patch(
+ "sglang.multimodal_gen.runtime.pipelines_core.stages."
+ "model_specific_stages.minimax_h3.denoise_loop._ulysses_ctx",
+ return_value=(world_size, rank),
+ ):
+ branch = _branch(mode, token_tags=token_tags)
+ local_rows = branch.seq_len // world_size
+ expected = token_tags[
+ rank * local_rows : (rank + 1) * local_rows
+ ].clamp(min=0)
+ torch.testing.assert_close(
+ branch.static_kwargs["block_token_tags"], expected, rtol=0, atol=0
+ )
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
new file mode 100644
index 000000000..0a1021d7f
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Mixed-precision weight and TP/Ulysses numerical contracts for H3 DiT."""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+ MiniMaxH3DiTArchConfig,
+ MiniMaxH3DiTConfig,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+ maybe_init_distributed_environment_and_model_parallel,
+ model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPAImpl
+from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import (
+ Fp8Config,
+ Fp8LinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.usp import _usp_input_all_to_all_packed_qkv
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
+ MINIMAX_H3_FP32_BUFFER_NAMES,
+ MINIMAX_H3_FP32_PARAM_NAMES,
+ MiniMaxH3DiTModel,
+ _copy_grouped_qkv_tp_shard,
+ _reorder_grouped_qkv_to_qkv,
+)
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
+ ensure_distributed_env_defaults,
+)
+
+
+def _ensure_single_process_parallel_runtime() -> None:
+ if model_parallel_is_initialized():
+ return
+ ensure_distributed_env_defaults()
+ maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+
+
+def test_native_weight_names_and_grouped_qkv_reorder():
+ arch = MiniMaxH3DiTArchConfig()
+ assert arch.param_names_mapping == {}
+ assert arch.reverse_param_names_mapping == {}
+ mapping = get_param_names_mapping(arch.param_names_mapping)
+ for key in (
+ "video_patch_proj.weight",
+ "token_refiner.blocks.0.attn.qkv_proj.weight",
+ "blocks.49.attn.qkv_proj.weight",
+ "final_layer.audio_out.weight",
+ ):
+ assert mapping(key) == (key, None, None)
+
+ weight = torch.arange(12, dtype=torch.float32).reshape(12, 1)
+ actual = _reorder_grouped_qkv_to_qkv(
+ weight,
+ num_query_groups=2,
+ heads_per_group=1,
+ head_dim=2,
+ )
+ expected = torch.tensor(
+ [0, 1, 6, 7, 2, 3, 8, 9, 4, 5, 10, 11],
+ dtype=torch.float32,
+ ).reshape(12, 1)
+ torch.testing.assert_close(actual, expected)
+
+ grouped = torch.arange(48, dtype=torch.int16).reshape(24, 2).view(torch.bfloat16)
+ reordered = _reorder_grouped_qkv_to_qkv(
+ grouped,
+ num_query_groups=4,
+ heads_per_group=1,
+ head_dim=2,
+ )
+ for tp_size in (1, 2, 4):
+ local_rows = 8 // tp_size
+ for tp_rank in range(tp_size):
+ start = tp_rank * local_rows
+ param = torch.nn.Parameter(
+ torch.empty(3 * local_rows, 2, dtype=torch.bfloat16),
+ requires_grad=False,
+ )
+ param.output_dim = 0
+ assert _copy_grouped_qkv_tp_shard(
+ param,
+ grouped,
+ num_query_groups=4,
+ head_dim=2,
+ tp_rank=tp_rank,
+ tp_size=tp_size,
+ )
+ expected_shard = reordered.view(3, 8, 2)[
+ :, start : start + local_rows
+ ].reshape(-1, 2)
+ assert torch.equal(
+ param.view(torch.int16), expected_shard.view(torch.int16)
+ )
+
+
+def test_tp_and_ulysses_admission_uses_tp_local_shapes():
+ arch = MiniMaxH3DiTArchConfig()
+ model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel)
+ model._validate_tp_config(arch=arch, tp_size=2)
+ model._validate_tp_config(arch=arch, tp_size=8)
+ MiniMaxH3DiTModel._validate_sequence_parallel_config(
+ arch=arch,
+ tp_size=2,
+ ulysses_size=4,
+ ring_size=1,
+ )
+
+ with pytest.raises(ValueError):
+ model._validate_tp_config(arch=arch, tp_size=3)
+ with pytest.raises(ValueError):
+ MiniMaxH3DiTModel._validate_sequence_parallel_config(
+ arch=arch,
+ tp_size=4,
+ ulysses_size=4,
+ ring_size=1,
+ )
+ with pytest.raises(NotImplementedError):
+ MiniMaxH3DiTModel._validate_sequence_parallel_config(
+ arch=arch,
+ tp_size=1,
+ ulysses_size=1,
+ ring_size=2,
+ )
+
+
+def test_meta_model_enforces_mixed_precision_weight_contract():
+ expected_fp32 = set(MINIMAX_H3_FP32_PARAM_NAMES) | set(MINIMAX_H3_FP32_BUFFER_NAMES)
+ _ensure_single_process_parallel_runtime()
+ with torch.device("meta"):
+ model = MiniMaxH3DiTModel(
+ config=MiniMaxH3DiTConfig(),
+ hf_config={},
+ quant_config=None,
+ )
+
+ assert model._fsdp_mixed_dtype_params
+ for name, tensor in model.state_dict().items():
+ if name in expected_fp32:
+ assert tensor.dtype == torch.float32, name
+ elif tensor.is_floating_point():
+ assert tensor.dtype == torch.bfloat16, name
+
+
+def test_online_fp8_keeps_fp32_boundaries_and_ignored_layers_unquantized():
+ _ensure_single_process_parallel_runtime()
+ with torch.device("meta"):
+ model = MiniMaxH3DiTModel(
+ config=MiniMaxH3DiTConfig(),
+ hf_config={},
+ quant_config=Fp8Config(ignored_layers=["blocks.0.attn.out_proj"]),
+ )
+
+ assert isinstance(model.blocks[0].attn.qkv_proj.quant_method, Fp8LinearMethod)
+ assert isinstance(
+ model.blocks[0].attn.out_proj.quant_method, UnquantizedLinearMethod
+ )
+ for layer in (
+ model.video_patch_proj,
+ model.audio_patch_proj,
+ model.time_embedder.proj_in,
+ model.time_embedder.proj_out,
+ model.final_layer.video_out,
+ model.final_layer.audio_out,
+ ):
+ assert isinstance(layer.quant_method, UnquantizedLinearMethod)
+
+
+def test_sdpa_varlen_fallback_matches_naive_packed_reference():
+ torch.manual_seed(0)
+ heads, dim = 2, 8
+ bounds = (0, 5, 6, 13)
+ cu = torch.tensor(bounds, dtype=torch.int32)
+ q = torch.randn(bounds[-1], heads, dim)
+ k = torch.randn_like(q)
+ v = torch.randn_like(q)
+ scale = dim**-0.5
+
+ attention = SDPAImpl(
+ num_heads=heads,
+ head_size=dim,
+ causal=False,
+ softmax_scale=scale,
+ )
+ out = attention.forward_varlen(
+ q,
+ k,
+ v,
+ cu_seqlens=cu,
+ cu_seqlens_host=bounds,
+ max_seqlen=7,
+ )
+ for start, stop in zip(bounds[:-1], bounds[1:], strict=True):
+ seg_q = q[start:stop].transpose(0, 1)
+ seg_k = k[start:stop].transpose(0, 1)
+ seg_v = v[start:stop].transpose(0, 1)
+ expected = (
+ torch.softmax(seg_q @ seg_k.transpose(-1, -2) * scale, dim=-1) @ seg_v
+ ).transpose(0, 1)
+ torch.testing.assert_close(out[start:stop], expected, atol=1e-6, rtol=1e-6)
+
+
+@patch(
+ "sglang.multimodal_gen.runtime.layers.usp.get_ulysses_parallel_world_size",
+ return_value=2,
+)
+def test_packed_qkv_exchange_preserves_rank_and_head_order(_):
+ seq, heads, dim = 3, 4, 2
+ q_ranks = [
+ torch.arange(seq * heads * dim).reshape(seq, heads, dim) + rank * 1000
+ for rank in range(2)
+ ]
+ k_ranks = [q + 100 for q in q_ranks]
+ v_ranks = [q + 200 for q in q_ranks]
+ target_rank = 1
+
+ def packet(q, k, v, destination):
+ head_slice = slice(destination * 2, (destination + 1) * 2)
+ return torch.cat((q[:, head_slice], k[:, head_slice], v[:, head_slice]), dim=-1)
+
+ def fake_all_to_all(actual):
+ expected_input = torch.stack(
+ [
+ packet(q_ranks[0], k_ranks[0], v_ranks[0], destination)
+ for destination in range(2)
+ ]
+ )
+ torch.testing.assert_close(actual, expected_input)
+ return torch.stack(
+ [
+ packet(q, k, v, target_rank)
+ for q, k, v in zip(q_ranks, k_ranks, v_ranks, strict=True)
+ ]
+ )
+
+ with patch(
+ "sglang.multimodal_gen.runtime.layers.usp._usp_all_to_all_single",
+ side_effect=fake_all_to_all,
+ ):
+ actual = _usp_input_all_to_all_packed_qkv(q_ranks[0], k_ranks[0], v_ranks[0])
+
+ for index, tensors in enumerate((q_ranks, k_ranks, v_ranks)):
+ expected = torch.cat(
+ [tensor[:, target_rank * 2 : (target_rank + 1) * 2] for tensor in tensors]
+ )
+ torch.testing.assert_close(actual[index], expected)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_cuda_ulysses_qkv_pack_is_bit_exact():
+ from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
+ pack_qkv_destination_major,
+ )
+
+ torch.manual_seed(23)
+ rows, world_size, heads, head_size = 65, 8, 56, 128
+ qkv = torch.randn(
+ rows,
+ 3 * heads * head_size,
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ q, k, v = (
+ tensor.reshape(rows, heads, head_size)
+ for tensor in qkv.split(heads * head_size, dim=-1)
+ )
+ local_heads = heads // world_size
+ expected = torch.empty(
+ world_size,
+ rows,
+ local_heads,
+ 3 * head_size,
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ for index, tensor in enumerate((q, k, v)):
+ shards = tensor.view(rows, world_size, local_heads, head_size).permute(
+ 1, 0, 2, 3
+ )
+ expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
+
+ actual = pack_qkv_destination_major(q.contiguous(), k.contiguous(), v, world_size)
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py
new file mode 100644
index 000000000..c40ce302f
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_media.py
@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical boundaries for the one-pass Ref2VA media path."""
+
+import json
+import subprocess
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3 import (
+ material_io,
+ reference_encoding,
+)
+
+
+def test_ffprobe_falls_back_when_stream_side_data_is_unknown(monkeypatch):
+ material_io._ffprobe_entries = None
+ calls = []
+
+ def run(command, **_kwargs):
+ calls.append(command)
+ entries = command[command.index("-show_entries") + 1]
+ if "stream_side_data" in entries:
+ raise subprocess.CalledProcessError(
+ 1,
+ command,
+ stderr="ffprobe: No match for section 'stream_side_data'",
+ )
+ return SimpleNamespace(
+ stdout=json.dumps(
+ {
+ "streams": [{"codec_type": "audio", "sample_rate": "44100"}],
+ "format": {"format_name": "mp3", "duration": "1.0"},
+ }
+ )
+ )
+
+ monkeypatch.setattr(subprocess, "run", run)
+ payload = material_io._ffprobe_media("/input/ref.mp3")
+
+ assert len(calls) == 2
+ assert "stream_side_data" in calls[0][calls[0].index("-show_entries") + 1]
+ assert "stream_side_data" not in calls[1][calls[1].index("-show_entries") + 1]
+ assert payload["format"]["format_name"] == "mp3"
+ assert material_io._ffprobe_entries is not None
+ assert "stream_side_data" not in material_io._ffprobe_entries
+
+
+def test_video_transform_runs_once_and_qwen_samples_shared_rgb(monkeypatch):
+ expected = np.arange(25 * 4 * 6 * 3, dtype=np.uint8).reshape(25, 4, 6, 3)
+ commands = []
+
+ def run(command, **_kwargs):
+ commands.append(command)
+ return SimpleNamespace(stdout=expected.tobytes())
+
+ monkeypatch.setattr(subprocess, "run", run)
+ frames = reference_encoding.minimax_h3_decode_reference_video_frames(
+ "/input/ref.mp4",
+ target_width=6,
+ target_height=4,
+ target_frame_count=25,
+ fps=24.0,
+ start_time_seconds=2.25,
+ )
+ sampled = reference_encoding.minimax_h3_sample_reference_video_frames(frames)
+
+ assert len(commands) == 1
+ command = commands[0]
+ assert command[command.index("-vf") + 1] == (
+ "fps=24,scale=6:4:flags=lanczos,setsar=1"
+ )
+ assert command[command.index("-frames:v") + 1] == "25"
+ assert command[command.index("-ss") + 1] == "2.25"
+ assert command.index("-ss") < command.index("-i")
+ assert command[-5:] == ["-f", "rawvideo", "-pix_fmt", "rgb24", "pipe:1"]
+ assert "libx264" not in command
+ assert all(np.shares_memory(frame, frames) for frame in sampled["frames"])
+ assert [int(frame[0, 0, 0]) for frame in sampled["frames"]] == [
+ int(expected[index, 0, 0, 0]) for index in (0, 12, 24)
+ ]
+ assert sampled["block_timestamps"] == [0.25, 1.0]
+
+
+def test_audio_decode_is_bounded_float_pcm_without_temp_files(monkeypatch):
+ pcm = torch.arange(8, dtype=torch.float32).numpy().tobytes()
+ commands = []
+
+ def run(command, **_kwargs):
+ commands.append(command)
+ if command[0] == "ffprobe":
+ return SimpleNamespace(
+ stdout=json.dumps(
+ {"streams": [{"channels": 6, "sample_rate": "44100"}]}
+ )
+ )
+ return SimpleNamespace(stdout=pcm)
+
+ monkeypatch.setattr(subprocess, "run", run)
+ waveform, source_rate = reference_encoding._load_waveform(
+ "/input/ref.mp4",
+ material_chain="video.reference_preserve",
+ max_duration_seconds=3.5,
+ start_time_seconds=2.25,
+ )
+
+ ffmpeg = next(command for command in commands if command[0] == "ffmpeg")
+ assert source_rate == 44100
+ torch.testing.assert_close(
+ waveform,
+ torch.tensor([[0, 2, 4, 6], [1, 3, 5, 7]], dtype=torch.float32),
+ )
+ assert ffmpeg[ffmpeg.index("-t") + 1] == "3.5"
+ assert ffmpeg[ffmpeg.index("-ss") + 1] == "2.25"
+ assert ffmpeg.index("-ss") < ffmpeg.index("-i")
+ assert ffmpeg[-3:] == ["-f", "f32le", "pipe:1"]
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_packed_sequence.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_packed_sequence.py
new file mode 100644
index 000000000..f5868409b
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_packed_sequence.py
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Numerical contracts for MiniMax-H3 packed-sequence layouts."""
+
+import unittest
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import (
+ minimax_h3_packed_sequence,
+ minimax_h3_packed_sequence_ref2va_blocks,
+)
+
+
+class TestMiniMaxH3PackedSequence(unittest.TestCase):
+ def test_t2va_structure(self):
+ built = minimax_h3_packed_sequence(
+ text_len=97,
+ latent_t=62,
+ latent_h=48,
+ latent_w=76,
+ audio_t=348,
+ include_keyframe_cond=False,
+ )
+ self.assertTrue(built["update_mask"].all())
+ self.assertEqual(int(built["img_pos"].shape[0]), 62 * 24 * 38)
+ self.assertEqual(int(built["seq_len"]) % 64, 0)
+ self.assertEqual(built["token_tags"][built["audio_pos"]].unique().tolist(), [2])
+
+ def test_fl2va_first_last_cond_blocks_use_exact_rope_span(self):
+ text_len = 11
+ latent_t = 37
+ built = minimax_h3_packed_sequence(
+ text_len=text_len,
+ latent_t=latent_t,
+ latent_h=48,
+ latent_w=76,
+ audio_t=203,
+ include_keyframe_cond=True,
+ keyframe_frame_indices=[0, -1],
+ frame_count=124,
+ )
+ frame_rows = 24 * 38
+ cond_rows = 2 * frame_rows
+ self.assertEqual(int((~built["update_mask"]).sum()), cond_rows)
+ self.assertEqual(
+ int(built["img_pos"].shape[0]),
+ (2 + latent_t) * frame_rows,
+ )
+ cond_pos = built["img_pos"][:cond_rows].reshape(2, frame_rows)
+ cond_t = [
+ float(built["img_position_ids"][positions, 0].unique().item())
+ for positions in cond_pos
+ ]
+ frame_rescale = 5.0 / 3.0
+ temporal_span = sum(
+ frame_rescale * (1, 4, 4, 4, 4)[index % 5] for index in range(latent_t)
+ )
+ self.assertEqual(cond_t[0], float(text_len))
+ self.assertAlmostEqual(
+ cond_t[1],
+ float(text_len) + temporal_span - frame_rescale,
+ places=12,
+ )
+ self.assertFalse(built["update_mask"][:cond_rows].any())
+ self.assertTrue(built["update_mask"][cond_rows:].all())
+
+ def test_i2va_and_l2va_single_cond_blocks_use_endpoint_rope(self):
+ text_len = 11
+ latent_t = 37
+ frame_count = 124
+ frame_rescale = 5.0 / 3.0
+ temporal_span = sum(
+ frame_rescale * (1, 4, 4, 4, 4)[index % 5] for index in range(latent_t)
+ )
+ for semantic_index, expected_t in (
+ (0, float(text_len)),
+ (-1, float(text_len) + temporal_span - frame_rescale),
+ ):
+ with self.subTest(semantic_index=semantic_index):
+ built = minimax_h3_packed_sequence(
+ text_len=text_len,
+ latent_t=latent_t,
+ latent_h=48,
+ latent_w=76,
+ audio_t=203,
+ include_keyframe_cond=True,
+ keyframe_frame_indices=[semantic_index],
+ frame_count=frame_count,
+ )
+ frame_rows = 24 * 38
+ self.assertEqual(int((~built["update_mask"]).sum()), frame_rows)
+ positions = built["img_pos"][:frame_rows]
+ cond_t = float(built["img_position_ids"][positions, 0].unique().item())
+ self.assertAlmostEqual(cond_t, expected_t, places=12)
+
+ def test_fl2va_keyframe_index_validation_is_defensive(self):
+ common = dict(
+ text_len=11,
+ latent_t=37,
+ latent_h=48,
+ latent_w=76,
+ audio_t=203,
+ include_keyframe_cond=True,
+ frame_count=124,
+ )
+ for frame_indices in (None, [1], [0, 52, -1]):
+ with (
+ self.subTest(frame_indices=frame_indices),
+ self.assertRaises(ValueError),
+ ):
+ minimax_h3_packed_sequence(
+ **common,
+ keyframe_frame_indices=frame_indices,
+ )
+
+ def test_ref2va_structure(self):
+ built = minimax_h3_packed_sequence_ref2va_blocks(
+ text_len=97,
+ latent_t=112,
+ latent_h=48,
+ latent_w=84,
+ audio_t=631,
+ ref_blocks=[
+ {"kind": "image", "latent_h": 64, "latent_w": 48},
+ {"kind": "audio", "ref_audio_t": 582},
+ ],
+ )
+
+ self.assertEqual(int(built["seq_len"]) % 64, 0)
+ self.assertEqual(int((~built["update_mask"]).sum()), 32 * 24)
+ self.assertEqual(int((~built["audio_update_mask"]).sum()), 582 * 2)
+ self.assertEqual(built["token_tags"][built["audio_pos"]].unique().tolist(), [2])
+
+ def test_ref2va_mixed_media_preserves_temporal_origin(self):
+ built = minimax_h3_packed_sequence_ref2va_blocks(
+ text_len=5,
+ latent_t=2,
+ latent_h=4,
+ latent_w=4,
+ audio_t=5,
+ ref_blocks=[
+ {"kind": "image", "latent_h": 4, "latent_w": 4},
+ {
+ "kind": "video_audio",
+ "ref_audio_t": 3,
+ "latent_t": 2,
+ "latent_h": 4,
+ "latent_w": 4,
+ },
+ {"kind": "audio", "ref_audio_t": 1},
+ ],
+ )
+
+ self.assertEqual(int((~built["update_mask"]).sum()), 12)
+ self.assertEqual(int((~built["audio_update_mask"]).sum()), 8)
+ target_video_t0 = built["img_position_ids"][built["img_pos"][12], 0]
+ target_audio_t0 = built["img_position_ids"][built["audio_pos"][8], 0]
+ self.assertEqual(float(target_audio_t0), float(target_video_t0))
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py
new file mode 100644
index 000000000..d25fb8371
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_vae_parallel_modes.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MiniMax-H3 released VAE decode contract."""
+
+from unittest import mock
+
+import pytest
+
+from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
+ MiniMaxH3VideoVAEConfig,
+)
+from sglang.multimodal_gen.runtime.models.vaes.minimax_h3 import MiniMaxH3VideoVAE
+from sglang.multimodal_gen.runtime.models.vaes.minimax_h3_video_vae import (
+ AutoencoderKLLegacy,
+)
+
+
+def _init_kwargs(config: MiniMaxH3VideoVAEConfig):
+ with mock.patch.object(
+ AutoencoderKLLegacy, "__init__", autospec=True, return_value=None
+ ) as init:
+ model = MiniMaxH3VideoVAE(config)
+ return model, init.call_args.kwargs
+
+
+@pytest.mark.parametrize(
+ "mode",
+ [
+ None,
+ "auto",
+ "tiled",
+ ],
+)
+def test_decode_mode_uses_released_tiled_recipe(mode):
+ config = (
+ MiniMaxH3VideoVAEConfig()
+ if mode is None
+ else MiniMaxH3VideoVAEConfig(parallel_decode_mode=mode)
+ )
+ model, kwargs = _init_kwargs(config)
+
+ assert model.parallel_decode_mode == "tiled"
+ assert kwargs["decoder_tiling"] is True
+ assert kwargs["parallel_tiling"] is True
+ assert kwargs["decoder_parallel"] is False
+
+
+@pytest.mark.parametrize("mode", ["spatial", "spatial_shard", "patch"])
+def test_unvalidated_decode_modes_are_rejected(mode):
+ config = MiniMaxH3VideoVAEConfig(parallel_decode_mode=mode)
+ with pytest.raises(ValueError, match="use tiled"):
+ config.resolved_parallel_decode_mode()
diff --git a/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py b/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py
index 3f0edafb3..9963b663f 100644
--- a/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py
+++ b/python/sglang/multimodal_gen/test/unit/test_multi_output_grouping.py
@@ -19,6 +19,7 @@ class CountingDedupStage(PipelineStage):
deduplicated_output_fields = ("prompt_embeds",)
deduplicated_tensor_tree_output_fields = ("timesteps",)
deduplicated_deepcopy_output_fields = ("scheduler",)
+ deduplicated_extra_output_keys = ("shared",)
deduplicated_extra_tensor_tree_output_keys = ("mu",)
def __init__(self):
@@ -36,6 +37,7 @@ class CountingDedupStage(PipelineStage):
batch.prompt_embeds = [torch.tensor([value])]
batch.timesteps = torch.tensor([value])
batch.scheduler = {"state": [value]}
+ batch.extra["shared"] = {"tensor": batch.prompt_embeds[0]}
batch.extra["mu"] = torch.tensor([value])
return batch
@@ -159,12 +161,20 @@ class TestMultiOutputGrouping(unittest.TestCase):
self.assertTrue(torch.equal(req.prompt_embeds[0], torch.tensor([1.0])))
self.assertTrue(torch.equal(req.timesteps, torch.tensor([1.0])))
self.assertEqual(req.scheduler, {"state": [1.0]})
+ self.assertTrue(
+ torch.equal(req.extra["shared"]["tensor"], torch.tensor([1.0]))
+ )
self.assertTrue(torch.equal(req.extra["mu"], torch.tensor([1.0])))
self.assertIsNot(reqs[0].prompt_embeds, reqs[1].prompt_embeds)
self.assertIs(reqs[0].prompt_embeds[0], reqs[1].prompt_embeds[0])
self.assertIsNot(reqs[0].timesteps, reqs[1].timesteps)
self.assertIsNot(reqs[0].scheduler, reqs[1].scheduler)
+ self.assertIsNot(reqs[0].extra["shared"], reqs[1].extra["shared"])
+ self.assertIs(
+ reqs[0].extra["shared"]["tensor"],
+ reqs[1].extra["shared"]["tensor"],
+ )
self.assertIsNot(reqs[0].extra["mu"], reqs[1].extra["mu"])
def test_declarative_stage_dedup_runs_distinct_fingerprints_separately(self):
diff --git a/python/sglang/multimodal_gen/test/unit/test_platform_detection.py b/python/sglang/multimodal_gen/test/unit/test_platform_detection.py
new file mode 100644
index 000000000..8d198742e
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_platform_detection.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime import platforms
+
+
+class NVMLUnavailableError(Exception):
+ pass
+
+
+class TestCudaPlatformDetection(unittest.TestCase):
+ def test_torch_fallback_excludes_hip(self):
+ cases = (
+ ("6.0", None),
+ (
+ None,
+ "sglang.multimodal_gen.runtime.platforms.cuda.CudaPlatform",
+ ),
+ )
+
+ for hip_version, expected in cases:
+ with (
+ self.subTest(hip_version=hip_version),
+ patch(
+ "sglang.multimodal_gen.utils.import_pynvml",
+ side_effect=NVMLUnavailableError,
+ ),
+ patch.object(platforms.os.path, "isfile", return_value=False),
+ patch.object(platforms.os.path, "exists", return_value=False),
+ patch.object(torch.version, "hip", hip_version, create=True),
+ patch.object(torch.cuda, "is_available", return_value=True),
+ patch.object(torch.cuda, "device_count", return_value=1),
+ ):
+ self.assertEqual(platforms.cuda_platform_plugin(), expected)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py
index a6b2e6e49..7beca3bd2 100644
--- a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py
+++ b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py
@@ -41,6 +41,13 @@ class TestSamplingParamsValidate(unittest.TestCase):
with self.assertRaisesRegex(ValueError, r"num_outputs_per_prompt"):
SamplingParams(num_outputs_per_prompt=0)
+ def test_quality_must_be_a_non_empty_profile_name(self):
+ self.assertEqual(SamplingParams(quality="high").quality, "high")
+ with self.assertRaisesRegex(ValueError, r"quality must be a non-empty string"):
+ SamplingParams(quality="")
+ with self.assertRaisesRegex(ValueError, r"quality must be a non-empty string"):
+ SamplingParams(quality=True) # type: ignore[arg-type]
+
def test_seed_accepts_int_or_non_empty_int_list(self):
self.assertEqual(SamplingParams(seed=7).seed, 7)
self.assertEqual(SamplingParams(seed=[7, 8]).seed, [7, 8])
@@ -225,6 +232,12 @@ class TestSamplingParamsCliArgs(unittest.TestCase):
[7, 8],
)
+ def test_quality_is_request_scoped_cli_arg(self):
+ self.assertNotIn("quality", self._parse_cli_kwargs([]))
+ self.assertEqual(
+ self._parse_cli_kwargs(["--quality", "medium"])["quality"], "medium"
+ )
+
def test_qwen_image_cli_path_preserves_model_defaults(self):
params = self._make_qwen_image_params([])
diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py
index 5b2335312..7a94c276f 100644
--- a/python/sglang/multimodal_gen/test/unit/test_server_args.py
+++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py
@@ -20,6 +20,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
LTX2PipelineConfig,
LTX23PipelineConfig,
)
+from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
+ MiniMaxH3PipelineConfig,
+)
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
QwenImagePipelineConfig,
@@ -39,10 +42,17 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
WanT2V720PConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
-from sglang.multimodal_gen.registry import _get_config_info
+from sglang.multimodal_gen.registry import (
+ _get_config_info,
+ get_non_diffusers_pipeline_name,
+ is_known_non_diffusers_multimodal_model,
+)
from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
QwenImageTransformer2DModel,
)
+from sglang.multimodal_gen.runtime.pipelines.minimax_h3_pipeline import (
+ MiniMaxH3Pipeline,
+)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.utils import FlexibleArgumentParser
@@ -687,6 +697,29 @@ class TestWarmupImageIsModelValid(unittest.TestCase):
self.assertGreaterEqual(height, 64)
+class TestMiniMaxH3Routing(unittest.TestCase):
+ def test_semantic_variants_map_to_checkpoint_partitions(self):
+ self.assertEqual(
+ MiniMaxH3Pipeline.model_subfolder_for_variant("fl2va"), "FL2VA"
+ )
+ self.assertEqual(
+ MiniMaxH3Pipeline.model_subfolder_for_variant("ref2va"), "Ref2VA"
+ )
+ with self.assertRaisesRegex(ValueError, "unsupported MiniMax H3 model variant"):
+ MiniMaxH3Pipeline.model_subfolder_for_variant("v2v")
+
+ def test_modelscope_id_resolves_to_the_huggingface_config(self):
+ expected = _get_config_info("MiniMaxAI/MiniMax-H3")
+ actual = _get_config_info("MiniMax/MiniMax-H3")
+ self.assertIsNotNone(expected)
+ self.assertIs(actual, expected)
+ self.assertTrue(is_known_non_diffusers_multimodal_model("MiniMax/MiniMax-H3"))
+ self.assertEqual(
+ get_non_diffusers_pipeline_name("MiniMax/MiniMax-H3"),
+ "MiniMaxH3Pipeline",
+ )
+
+
class TestOffloadDefaults(unittest.TestCase):
def test_wan_decode_precision_defaults(self):
for pipeline_config in (
@@ -1506,6 +1539,73 @@ class TestOffloadDefaults(unittest.TestCase):
)
self.assertFalse(args.vae_cpu_offload)
+ def test_auto_minimax_h3_keeps_large_components_resident_with_headroom(self):
+ args = self._from_dict_with_pipeline_config(
+ MiniMaxH3PipelineConfig(),
+ memory_gb=141,
+ available_memory_gb=130,
+ kwargs={
+ "model_path": "MiniMaxAI/MiniMax-H3",
+ "num_gpus": 8,
+ "ulysses_degree": 8,
+ "performance_mode": "auto",
+ },
+ )
+
+ self.assertFalse(args.dit_cpu_offload)
+ self.assertFalse(args.text_encoder_cpu_offload)
+ self.assertFalse(args.vae_cpu_offload)
+ self.assertNotIn("text_encoder", args.layerwise_offload_components or [])
+ self.assertNotIn("vae", args.layerwise_offload_components or [])
+
+ def test_auto_minimax_h3_keeps_memory_policy_below_residency_threshold(self):
+ args = self._from_dict_with_pipeline_config(
+ MiniMaxH3PipelineConfig(),
+ memory_gb=96,
+ available_memory_gb=90,
+ kwargs={
+ "model_path": "MiniMaxAI/MiniMax-H3",
+ "num_gpus": 8,
+ "ulysses_degree": 8,
+ "performance_mode": "auto",
+ },
+ )
+
+ self.assertTrue(args.dit_cpu_offload)
+ self.assertIn("text_encoder", args.layerwise_offload_components or [])
+ self.assertIn("vae", args.layerwise_offload_components or [])
+
+ def test_memory_minimax_h3_combines_fsdp_with_aux_layerwise_offload(self):
+ args = self._from_dict_with_pipeline_config(
+ MiniMaxH3PipelineConfig(),
+ kwargs={
+ "model_path": "MiniMaxAI/MiniMax-H3",
+ "num_gpus": 8,
+ "ulysses_degree": 8,
+ "performance_mode": "memory",
+ "use_fsdp_inference": True,
+ },
+ )
+
+ self.assertTrue(args.use_fsdp_inference)
+ self.assertFalse(args.dit_cpu_offload)
+ self.assertFalse(args.dit_layerwise_offload)
+ self.assertIn("text_encoder", args.layerwise_offload_components or [])
+ self.assertIn("vae", args.layerwise_offload_components or [])
+
+ def test_minimax_h3_rejects_explicit_cfg_parallel(self):
+ with self.assertRaisesRegex(
+ ValueError, "MiniMaxH3PipelineConfig does not support CFG parallelism"
+ ):
+ self._from_dict_with_pipeline_config(
+ MiniMaxH3PipelineConfig(),
+ kwargs={
+ "model_path": "MiniMaxAI/MiniMax-H3",
+ "num_gpus": 4,
+ "cfg_parallel_degree": 2,
+ },
+ )
+
def test_speed_mode_single_gpu_disables_offload(self):
args = self._from_dict_with_pipeline_config(
QwenImagePipelineConfig(),
@@ -1560,6 +1660,20 @@ class TestOffloadDefaults(unittest.TestCase):
self.assertFalse(args.enable_torch_compile)
+ def test_speed_mode_uses_minimax_h3_compile_policy(self):
+ for explicit, expected in ((None, False), (True, True)):
+ kwargs = {
+ "model_path": "MiniMaxAI/MiniMax-H3",
+ "performance_mode": "speed",
+ }
+ if explicit is not None:
+ kwargs["enable_torch_compile"] = explicit
+ with self.subTest(explicit=explicit):
+ args = self._from_dict_with_pipeline_config(
+ MiniMaxH3PipelineConfig(), kwargs=kwargs
+ )
+ self.assertEqual(args.enable_torch_compile, expected)
+
def test_auto_mode_leaves_torch_compile_off(self):
args = self._from_dict_with_pipeline_config(
QwenImagePipelineConfig(),
diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
index 4c04bffa3..caae1c033 100644
--- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
+++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
@@ -7,6 +7,9 @@ import transformers
from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
TextEncoderLoader,
)
+from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
+ MiniMaxH3Qwen3VLEncoder,
+)
class TestTextEncoderClassResolution(unittest.TestCase):
@@ -79,5 +82,25 @@ class TestTextEncoderClassResolution(unittest.TestCase):
self.assertIs(cls, transformers.AutoModel)
+class TestMiniMaxH3CheckpointFilter(unittest.TestCase):
+ def test_only_known_unconsumed_weights_are_filtered(self):
+ should_load = MiniMaxH3Qwen3VLEncoder.should_materialize_checkpoint_weight
+ expected = {
+ "model.language_model.layers.49.self_attn.q_proj.weight": True,
+ "model.language_model.layers.50.self_attn.q_proj.weight": False,
+ "model.language_model.layers.63.mlp.down_proj.weight": False,
+ "model.language_model.norm.weight": False,
+ "lm_head.weight": False,
+ "model.language_model.rotary_emb.inv_freq": False,
+ "model.visual.blocks.0.attn.qkv.weight": True,
+ "language_model.layers.63.mlp.down_proj.weight": True,
+ "module.model.language_model.layers.63.mlp.down_proj.weight": True,
+ }
+ self.assertEqual(
+ {name: should_load(name) for name in expected},
+ expected,
+ )
+
+
if __name__ == "__main__":
unittest.main()
diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
index 14a18f9b9..8c98efeaf 100644
--- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
@@ -55,10 +55,15 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
ModelOptFp8Config,
_prepare_nvfp4_weight_bytes,
)
+from sglang.multimodal_gen.runtime.loader.component_loaders import transformer_loader
+from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
+ _warn_if_expected_param_dtype_missing,
+)
from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
_filter_duplicate_precision_variant_safetensors,
_Flux2Nvfp4FallbackAdapter,
_needs_device_weight_postprocess,
+ _resolve_quant_config,
resolve_transformer_quant_load_spec,
resolve_transformer_safetensors_to_load,
)
@@ -107,6 +112,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
),
nunchaku_config=None,
quantization=None,
+ quantization_ignored_layers=None,
tp_size=1,
dit_cpu_offload=False,
text_encoder_cpu_offload=False,
@@ -211,6 +217,24 @@ class TestTransformerQuantHelpers(unittest.TestCase):
self.assertEqual(plan.weight_postprocess_device, device)
self.assertTrue(plan.defer_component_cpu_offload)
+ def test_mixed_model_with_expected_dtype_does_not_warn(self):
+ model = torch.nn.Module()
+ model.fp32 = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32))
+ model.bf16 = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16))
+
+ with patch.object(transformer_loader.logger, "warning") as warning:
+ _warn_if_expected_param_dtype_missing(model, torch.bfloat16)
+
+ warning.assert_not_called()
+
+ def test_model_without_expected_dtype_warns(self):
+ model = torch.nn.Linear(1, 1, dtype=torch.float32)
+
+ with patch.object(transformer_loader.logger, "warning") as warning:
+ _warn_if_expected_param_dtype_missing(model, torch.bfloat16)
+
+ warning.assert_called_once()
+
def test_online_fp8_needs_device_weight_postprocess(self):
self.assertTrue(_needs_device_weight_postprocess(Fp8Config()))
self.assertFalse(
@@ -233,6 +257,23 @@ class TestTransformerQuantHelpers(unittest.TestCase):
)
)
+ def test_online_fp8_receives_cli_ignored_layer_patterns(self):
+ ignored_layers = ["blocks.0.attn.out_proj", "condition_proj"]
+ server_args = self._make_server_args(
+ quantization="fp8",
+ quantization_ignored_layers=ignored_layers,
+ )
+
+ quant_config = _resolve_quant_config(
+ hf_config={},
+ server_args=server_args,
+ safetensors_list=[],
+ component_model_path="/unused/component/path",
+ )
+
+ self.assertIsInstance(quant_config, Fp8Config)
+ self.assertEqual(quant_config.ignored_layers, ignored_layers)
+
@patch(
"sglang.multimodal_gen.runtime.loader.transformer_load_utils.build_nvfp4_config_from_safetensors_list",
return_value=None,
diff --git a/python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py b/python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py
new file mode 100644
index 000000000..cda5891e1
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_video_api_profiling.py
@@ -0,0 +1,50 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
+ VideoGenerationsRequest,
+)
+from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import (
+ _build_video_sampling_params,
+)
+
+
+def test_video_api_forwards_profiling_options():
+ request = VideoGenerationsRequest(
+ prompt="profile this request",
+ task="t2va",
+ conditions=[],
+ target={
+ "short_edge": 768,
+ "aspect_ratio": "16:9",
+ "duration_seconds": 5.0,
+ },
+ profile=True,
+ num_profiled_timesteps=3,
+ profile_all_stages=False,
+ )
+ server_args = SimpleNamespace(
+ backend="auto",
+ model_id=None,
+ model_path="MiniMaxAI/MiniMax-H3",
+ pipeline_class_name="MiniMaxH3Pipeline",
+ pipeline_config=object(),
+ )
+
+ with (
+ patch(
+ "sglang.multimodal_gen.runtime.entrypoints.openai.video_api."
+ "get_global_server_args",
+ return_value=server_args,
+ ),
+ patch(
+ "sglang.multimodal_gen.runtime.entrypoints.openai.video_api."
+ "build_sampling_params",
+ side_effect=lambda request_id, **kwargs: kwargs,
+ ),
+ ):
+ kwargs = _build_video_sampling_params("profile-request", request)
+
+ assert kwargs["profile"] is True
+ assert kwargs["num_profiled_timesteps"] == 3
+ assert kwargs["profile_all_stages"] is False
diff --git a/python/sglang/utils.py b/python/sglang/utils.py
index 1bf5eaa8c..9fb40b019 100644
--- a/python/sglang/utils.py
+++ b/python/sglang/utils.py
@@ -32,6 +32,10 @@ from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: dict[str, str] = {
+ "minimaxai/minimax-h3": "MiniMaxH3Pipeline",
+ "minimaxai--minimax-h3": "MiniMaxH3Pipeline",
+ "minimax/minimax-h3": "MiniMaxH3Pipeline",
+ "minimax--minimax-h3": "MiniMaxH3Pipeline",
"lerobot/pi05": "Pi05Pipeline",
"lerobot--pi05": "Pi05Pipeline",
"pi05": "Pi05Pipeline",
diff --git a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py
index 5033351e8..2030e78ce 100644
--- a/test/registered/kernels/ops/diffusion/test_qknorm_rope.py
+++ b/test/registered/kernels/ops/diffusion/test_qknorm_rope.py
@@ -122,7 +122,7 @@ def test_qknorm_rope(
if is_neox:
elems_per_thread = head_dim // 32
rotary_lanes = rope_dim // elems_per_thread
- if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
+ if rotary_lanes < 2 or rotary_lanes % 2:
continue
q = torch.randn(batch_size, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
@@ -150,5 +150,60 @@ def test_qknorm_rope(
triton.testing.assert_close(k_ref, k_fused, atol=ATOL, rtol=RTOL)
+def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
+ from sgl_kernel import rotary_embedding
+
+ from sglang.kernels.ops.diffusion.qknorm_rope import (
+ fused_inplace_qknorm_rope,
+ )
+ from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+
+ num_tokens, num_heads, head_dim, rope_dim = 257, 28, 128, 96
+ inner_dim = num_heads * head_dim
+ qkv = torch.randn(
+ num_tokens,
+ 3 * inner_dim,
+ device=DEVICE,
+ dtype=DTYPE,
+ )
+ q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+ k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+ positions = torch.arange(num_tokens, device=DEVICE, dtype=torch.int64)
+ cos_sin_cache = create_cos_sin_cache(rope_dim, num_tokens).to(DTYPE)
+
+ qkv_ref, qkv_fused = qkv.clone(), qkv.clone()
+ q_ref, k_ref, _ = qkv_ref.split(inner_dim, dim=-1)
+ q_fused, k_fused, _ = qkv_fused.split(inner_dim, dim=-1)
+ q_ref = q_ref.view(num_tokens, num_heads, head_dim)
+ k_ref = k_ref.view(num_tokens, num_heads, head_dim)
+ q_fused = q_fused.view(num_tokens, num_heads, head_dim)
+ k_fused = k_fused.view(num_tokens, num_heads, head_dim)
+
+ fused_inplace_qknorm(q_ref, k_ref, q_weight, k_weight, eps=1e-5)
+ rotary_embedding(
+ positions,
+ q_ref.view(num_tokens, -1),
+ k_ref.view(num_tokens, -1),
+ head_dim,
+ cos_sin_cache,
+ True,
+ )
+ fused_inplace_qknorm_rope(
+ q_fused,
+ k_fused,
+ q_weight,
+ k_weight,
+ cos_sin_cache,
+ positions,
+ is_neox=True,
+ eps=1e-5,
+ rope_dim=rope_dim,
+ round_norm_before_rope=True,
+ )
+
+ assert torch.equal(q_ref, q_fused)
+ assert torch.equal(k_ref, k_fused)
+
+
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
diff --git a/test/registered/kernels/ops/diffusion/test_usp_relayout.py b/test/registered/kernels/ops/diffusion/test_usp_relayout.py
new file mode 100644
index 000000000..588524865
--- /dev/null
+++ b/test/registered/kernels/ops/diffusion/test_usp_relayout.py
@@ -0,0 +1,69 @@
+"""Bitwise tests for the generic Ulysses output head-merge fast path."""
+
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.usp_relayout import (
+ can_use_usp_merge_heads,
+ usp_merge_heads,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+
+
+@pytest.mark.parametrize(
+ "world,seq,batch,h_local,head_dim",
+ [
+ (4, 7936, 1, 14, 128), # H3 768p production shape (Ulysses 4)
+ (2, 64, 3, 4, 64), # batched
+ (4, 33, 2, 4, 100), # scalar fallback inside the CUDA kernel
+ ],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+def test_usp_merge_heads_bitwise(dtype, world, seq, batch, h_local, head_dim):
+ generator = torch.Generator(device=DEVICE).manual_seed(4321)
+ x = torch.randn(
+ world,
+ seq,
+ batch,
+ h_local,
+ head_dim,
+ dtype=dtype,
+ device=DEVICE,
+ generator=generator,
+ )
+ assert can_use_usp_merge_heads(x)
+ out = usp_merge_heads(x)
+ ref = x.permute(2, 1, 0, 3, 4).contiguous()
+ assert out.shape == ref.shape
+ assert torch.equal(out, ref)
+
+
+def test_usp_merge_heads_unsupported_inputs_use_exact_fallback():
+ x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
+ unsupported = [x.transpose(0, 1), x[:0]]
+
+ for value in unsupported:
+ assert not can_use_usp_merge_heads(value)
+ assert torch.equal(
+ usp_merge_heads(value), value.permute(2, 1, 0, 3, 4).contiguous()
+ )
+
+ with patch.object(torch.version, "hip", "6.3"):
+ assert not can_use_usp_merge_heads(x)
+ assert torch.equal(usp_merge_heads(x), x.permute(2, 1, 0, 3, 4).contiguous())
+
+
+def test_usp_merge_heads_fast_path_rejects_wrong_rank():
+ x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
+ assert not can_use_usp_merge_heads(x[0])
+
+
+if __name__ == "__main__":
+ sys.exit(pytest.main([__file__, "-v"]))
diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py
index c3b2809a6..ceb843287 100644
--- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py
+++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py
@@ -29,7 +29,7 @@ def _make_processor(server_mode: str = "full") -> SchedulerBatchResultProcessor:
enable_return_hidden_states=True,
return_hidden_states_mode=server_mode,
),
- model_config=SimpleNamespace(think_end_id=None),
+ model_config=SimpleNamespace(think_end_ids=None),
token_to_kv_pool_allocator=Mock(),
tree_cache=None,
hisparse_coordinator=None,