diff --git a/.claude/skills/benchmark-model-kernels b/.claude/skills/benchmark-model-kernels
new file mode 120000
index 00000000000..1bfd1fefe85
--- /dev/null
+++ b/.claude/skills/benchmark-model-kernels
@@ -0,0 +1 @@
+../../.agents/skills/benchmark-model-kernels
\ No newline at end of file
diff --git a/examples/speculative_decoding/distributed_generate/launch.sh b/examples/speculative_decoding/distributed_generate/launch.sh
index 463f4c2a387..c93f59906e4 100644
--- a/examples/speculative_decoding/distributed_generate/launch.sh
+++ b/examples/speculative_decoding/distributed_generate/launch.sh
@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+set -euo pipefail
+
 if [ $# -lt 9 ]; then
     echo "Usage: $0 <job_id> <backend> <model_path> <data_path> <output_path> <scripts_path> <start_shard> <jobs_per_node> <comma_separated_node_names> <system_prompt>"
     echo "Example: $0 245387 vllm /model/ /input_data/ /output_data/ /scripts/ 0 20 cluster-01,cluster-02 "\"You are a helpful assistant.\"""
@@ -32,6 +34,9 @@ NODE_NAME=$9
 SYSTEM_PROMPT="${10:-}"
 IFS=',' read -r -a NODE_LIST <<< "$NODE_NAME"
 
+# Pyxis requires bind-mount sources to exist before srun creates the container.
+mkdir -p "$OUTPUT_PATH"
+
 # backend needs to be either vllm or sglang
 if [ "$BACKEND" != "vllm" ] && [ "$BACKEND" != "sglang" ]; then
     echo "Invalid backend: $BACKEND"
@@ -39,22 +44,37 @@ if [ "$BACKEND" != "vllm" ] && [ "$BACKEND" != "sglang" ]; then
 fi
 
 if [ "$BACKEND" == "vllm" ]; then
-    CONTAINER_IMAGE="vllm/vllm-openai:v0.8.5"
+    DEFAULT_CONTAINER_IMAGE="vllm/vllm-openai:v0.24.0"
 else
-    CONTAINER_IMAGE="lmsysorg/sglang:v0.4.6.post2-cu124"
+    DEFAULT_CONTAINER_IMAGE="lmsysorg/sglang:v0.5.3-cu129"
 fi
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-$DEFAULT_CONTAINER_IMAGE}
 counter=$START_SHARD
+worker_pids=()
 for node in "${NODE_LIST[@]}"; do
     echo "Processing node: $node"
-    srun --output=srun_worker_${node}.log --jobid=$JOB_ID -N 1 --ntasks=1 --ntasks-per-node=1 -w $node \
-        --mpi pmix --overlap --container-image=$CONTAINER_IMAGE \
-        --container-mounts=$MODEL_PATH:/model/,$DATA_PATH:/input_data/,$OUTPUT_PATH:/output_data/,$SCRIPTS_PATH:/scripts/ \
-        bash /scripts/distributed_generate/worker.sh $counter $BACKEND $JOBS_PER_NODE "$SYSTEM_PROMPT" &
+    srun --output="srun_worker_${node}.log" --jobid="$JOB_ID" -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
+        --mpi pmix --overlap --container-image="$CONTAINER_IMAGE" \
+        --container-mounts="$MODEL_PATH":/model/,"$DATA_PATH":/input_data/,"$OUTPUT_PATH":/output_data/,"$SCRIPTS_PATH":/scripts/ \
+        bash /scripts/distributed_generate/worker.sh "$counter" "$BACKEND" "$JOBS_PER_NODE" "$SYSTEM_PROMPT" &
 
     echo "srun command for node $node started with PID $!" >> srun_launch.log
+    worker_pids+=("$!")
 
     # increment counter by JOBS_PER_NODE
     counter=$((counter + JOBS_PER_NODE))
 done
 
 echo "Started workers, each processing $JOBS_PER_NODE shards of data. Will process shards $START_SHARD through $((counter - 1))."
+
+worker_status=0
+for worker_pid in "${worker_pids[@]}"; do
+    if ! wait "$worker_pid"; then
+        worker_status=1
+    fi
+done
+
+if [ "$worker_status" -ne 0 ]; then
+    echo "ERROR: one or more workers failed." >&2
+fi
+exit "$worker_status"
diff --git a/examples/speculative_decoding/distributed_generate/launch_multimodal.sh b/examples/speculative_decoding/distributed_generate/launch_multimodal.sh
new file mode 100755
index 00000000000..23558644c2a
--- /dev/null
+++ b/examples/speculative_decoding/distributed_generate/launch_multimodal.sh
@@ -0,0 +1,96 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+if [ $# -lt 10 ]; then
+    echo "Usage: $0 <job_id> <backend> <model_path> <shard_path> <output_path> <scripts_path> <start_shard> <jobs_per_node> <media_path> [num_frames] <comma_separated_node_names> [system_prompt]"
+    echo "Also accepted: $0 <job_id> <backend> <model_path> <shard_path> <output_path> <scripts_path> <start_shard> <jobs_per_node> <comma_separated_node_names> <media_path> [num_frames] [system_prompt]"
+    echo "Example: $0 245387 sglang /model/ /shards/ /output/ /scripts/ 0 10 /media/ 16 cluster-01"
+    echo "Optional env: SGLANG_TP_SIZE=8 NUM_TEMPERATURES=8 NUM_THREADS=8 SGLANG_EXTRA_ARGS='--mem-fraction-static 0.75'"
+    exit 1
+fi
+
+JOB_ID=$1
+BACKEND=$2
+MODEL_PATH=$3
+DATA_PATH=$4
+OUTPUT_PATH=$5
+SCRIPTS_PATH=$6
+START_SHARD=$7
+JOBS_PER_NODE=$8
+ARG9=$9
+ARG10=${10:-}
+ARG11=${11:-}
+ARG12=${12:-}
+
+if [[ "$ARG9" == */* || "$ARG9" == .* ]]; then
+    MEDIA_PATH=$ARG9
+    NUM_FRAMES="${ARG10:-32}"
+    NODE_NAME=$ARG11
+    SYSTEM_PROMPT="$ARG12"
+else
+    NODE_NAME=$ARG9
+    MEDIA_PATH=$ARG10
+    NUM_FRAMES="${ARG11:-32}"
+    SYSTEM_PROMPT="$ARG12"
+fi
+
+if [ -z "${NODE_NAME:-}" ] || [ -z "${MEDIA_PATH:-}" ]; then
+    echo "ERROR: both media_path and node_name are required." >&2
+    exit 1
+fi
+
+IFS=',' read -r -a NODE_LIST <<< "$NODE_NAME"
+
+if [ "$BACKEND" != "sglang" ]; then
+    echo "Multimodal generation currently supports backend=sglang."
+    exit 1
+fi
+
+mkdir -p "$OUTPUT_PATH"
+
+# Set CONTAINER_IMAGE to a local .sqsh image to avoid pulling from the registry.
+DEFAULT_CONTAINER_IMAGE="lmsysorg/sglang:v0.5.3-cu129"
+CONTAINER_IMAGE="${CONTAINER_IMAGE:-$DEFAULT_CONTAINER_IMAGE}"
+
+counter=$START_SHARD
+worker_pids=()
+for node in "${NODE_LIST[@]}"; do
+    echo "Processing node: $node"
+    srun --output=srun_vlm_worker_${node}.log --jobid=$JOB_ID -N 1 --ntasks=1 --ntasks-per-node=1 -w "$node" \
+        --mpi pmix --overlap --container-image="$CONTAINER_IMAGE" \
+        --container-mounts="$MODEL_PATH":/model/,"$DATA_PATH":/input_data/,"$OUTPUT_PATH":/output_data/,"$SCRIPTS_PATH":/scripts/,"$MEDIA_PATH":/media_data/ \
+        bash /scripts/distributed_generate/worker_multimodal.sh "$counter" "$BACKEND" "$JOBS_PER_NODE" "$NUM_FRAMES" "$SYSTEM_PROMPT" &
+
+    echo "srun multimodal command for node $node started with PID $!" >> srun_launch_multimodal.log
+    worker_pids+=("$!")
+    counter=$((counter + JOBS_PER_NODE))
+done
+
+echo "Started multimodal workers, each processing $JOBS_PER_NODE shards of data. Will process shards $START_SHARD through $((counter - 1))."
+
+worker_status=0
+for worker_pid in "${worker_pids[@]}"; do
+    if ! wait "$worker_pid"; then
+        worker_status=1
+    fi
+done
+
+if [ "$worker_status" -ne 0 ]; then
+    echo "ERROR: one or more multimodal workers failed." >&2
+fi
+exit "$worker_status"
diff --git a/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py b/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py
new file mode 100755
index 00000000000..6987c5b09d7
--- /dev/null
+++ b/examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py
@@ -0,0 +1,589 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generate multimodal SFT data from video prompts using SGLang native video input."""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures
+import json
+import os
+import sys
+import traceback
+from pathlib import Path
+from typing import Any
+from urllib.parse import quote
+
+import tqdm
+
+QWEN_IMAGE_TOKEN = "<|vision_start|><|image_pad|><|vision_end|>"
+_UNRESOLVED_MEDIA_PATHS: set[str] = set()
+
+
+def _load_json_or_jsonl(path: str) -> list[dict[str, Any]]:
+    if path.endswith("jsonl"):
+        with open(path, encoding="utf-8") as f:
+            return [json.loads(line) for line in f if line.strip()]
+    with open(path, encoding="utf-8") as f:
+        data = json.load(f)
+    if not isinstance(data, list):
+        raise ValueError(f"Expected list data in {path}")
+    return data
+
+
+def _first_user_message(sample: dict[str, Any]) -> dict[str, Any]:
+    messages = sample.get("messages") or sample.get("conversations") or sample.get("conversation")
+    if not isinstance(messages, list):
+        raise ValueError(f"Sample has no messages/conversations list: keys={sorted(sample)}")
+    for message in messages:
+        role = (message.get("role") or message.get("from") or "").lower()
+        if role in ("user", "human"):
+            return message
+    raise ValueError("Sample has no user message")
+
+
+def _extract_message_text_and_media(sample: dict[str, Any]) -> tuple[str, str | None, str | None]:
+    video_path = sample.get("video_path")
+    image_path = sample.get("image_path") or sample.get("image")
+
+    message = _first_user_message(sample)
+    content = message.get("content") or message.get("value")
+    text_parts: list[str] = []
+
+    if isinstance(content, str):
+        text_parts.append(content)
+    elif isinstance(content, list):
+        for item in content:
+            if not isinstance(item, dict):
+                continue
+            item_type = item.get("type")
+            if item_type == "text" and isinstance(item.get("text"), str):
+                text_parts.append(item["text"])
+            elif item_type == "video":
+                video_path = item.get("video") or video_path
+            elif item_type == "image":
+                image_path = item.get("image") or image_path
+    else:
+        raise ValueError(f"Unsupported user content: {content!r}")
+
+    prompt = "\n".join(part.strip() for part in text_parts if part.strip())
+    if not prompt:
+        raise ValueError("Could not extract text prompt from sample")
+    return prompt, video_path, image_path
+
+
+def _extract_text_and_media(sample: dict[str, Any]) -> tuple[str, str | None, str | None]:
+    prompt = sample.get("prompt")
+    text, video_path, image_path = _extract_message_text_and_media(sample)
+    if isinstance(prompt, str) and prompt.strip():
+        text = prompt.strip()
+    image_path = sample.get("image_path") or sample.get("image") or image_path
+    if image_path:
+        return text, None, image_path
+    return text, sample.get("video_path") or video_path, None
+
+
+def _resolve_media_path(
+    path: str | None, media_root: str | None, input_root: str | None
+) -> str | None:
+    if not path:
+        return None
+    if path.startswith(("http://", "https://", "data:")):
+        return path
+    candidate = Path(path)
+    if candidate.is_absolute() and candidate.exists():
+        return str(candidate)
+    if input_root:
+        rooted = Path(input_root) / path
+        if rooted.exists():
+            return str(rooted)
+    if media_root:
+        rooted = Path(media_root) / path
+        if rooted.exists():
+            return str(rooted)
+        # If the record stores an absolute host path, preserve its suffix under media_root.
+        parts = candidate.parts
+        if "videos" in parts:
+            suffix = Path(*parts[parts.index("videos") :])
+            rooted = Path(media_root) / suffix
+            if rooted.exists():
+                return str(rooted)
+    if candidate.exists():
+        return str(candidate)
+    if path not in _UNRESOLVED_MEDIA_PATHS:
+        print(f"WARNING: could not resolve media path: {path}")
+        _UNRESOLVED_MEDIA_PATHS.add(path)
+    return None
+
+
+def _as_openai_media_value(
+    path: str, media_url_base: str | None, media_root: str | None, input_root: str | None
+) -> str:
+    if path.startswith(("http://", "https://", "data:")):
+        return path
+    if media_url_base:
+        candidate = Path(path)
+        if candidate.is_absolute():
+            if media_root:
+                try:
+                    path = str(candidate.relative_to(media_root))
+                except ValueError:
+                    # The local HTTP server deliberately exposes only media_root.
+                    # Leave paths outside it local for SGLang to resolve directly.
+                    return path
+            else:
+                return f"{media_url_base.rstrip('/')}{quote(path, safe='/')}"
+        return f"{media_url_base.rstrip('/')}/{quote(path, safe='/')}"
+    # Do not convert local paths to file://. This SGLang build falls through to
+    # the base64 loader for file:// videos and raises "Incorrect padding".
+    return path
+
+
+def _openai_chat_url(url: str) -> str:
+    url = url.rstrip("/")
+    if url.endswith("/v1/chat/completions"):
+        return url
+    if url.endswith("/v1"):
+        return f"{url}/chat/completions"
+    return f"{url}/v1/chat/completions"
+
+
+def _messages_for_output(sample: dict[str, Any], answer: str) -> list[dict[str, Any]]:
+    messages = sample.get("messages")
+    if isinstance(messages, list):
+        output_messages = list(messages)
+    else:
+        prompt, video_path, image_path = _extract_text_and_media(sample)
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        if image_path:
+            content.append({"type": "image", "image": image_path})
+        elif video_path:
+            content.append({"type": "video", "video": video_path, "fps": 4})
+        output_messages = [{"role": "user", "content": content}]
+
+    output_messages.append({"role": "assistant", "content": answer})
+    return output_messages
+
+
+def _coerce_text(value: Any) -> str:
+    if value is None:
+        return ""
+    if type(value).__name__ == "ProgramState":
+        return ""
+    if isinstance(value, str):
+        return value.strip()
+    if isinstance(value, dict):
+        for key in ("answer", "text", "value", "content"):
+            text = _coerce_text(value.get(key))
+            if text:
+                return text
+        return ""
+    for attr_name in ("text", "value", "content"):
+        attr = getattr(value, attr_name, None)
+        if callable(attr):
+            try:
+                text = _coerce_text(attr())
+                if text:
+                    return text
+            except Exception:
+                pass
+        else:
+            text = _coerce_text(attr)
+            if text:
+                return text
+    return ""
+
+
+def _answer_from_messages(messages: Any) -> str:
+    if callable(messages):
+        try:
+            messages = messages()
+        except Exception:
+            return ""
+    if not isinstance(messages, list):
+        return ""
+    for message in reversed(messages):
+        if not isinstance(message, dict):
+            continue
+        role = (message.get("role") or "").lower()
+        if role != "assistant":
+            continue
+        return _coerce_text(message.get("content"))
+    return ""
+
+
+def _state_answer(state: Any) -> str:
+    # SGLang ProgramState normally supports state["answer"]. Some versions or
+    # failure paths expose variables through helper methods/attributes instead.
+    try:
+        text = _coerce_text(state["answer"])
+        if text:
+            return text
+    except Exception:
+        pass
+
+    if isinstance(state, dict):
+        return _coerce_text(state.get("answer") or state.get("value") or state.get("text"))
+
+    for method_name in ("get", "get_var", "get_variable", "var"):
+        method = getattr(state, method_name, None)
+        if not callable(method):
+            continue
+        try:
+            text = _coerce_text(method("answer"))
+            if text:
+                return text
+        except Exception:
+            pass
+
+    text = _answer_from_messages(getattr(state, "messages", None))
+    if text:
+        return text
+
+    for attr_name in ("answer", "variables", "vars"):
+        text = _coerce_text(getattr(state, attr_name, None))
+        if text:
+            return text
+
+    return ""
+
+
+def _prompt_with_vision_token(prompt: str, media_type: str, token_format: str) -> str:
+    if token_format == "none":
+        return prompt
+    if token_format != "qwen_vl":
+        raise ValueError(f"Unsupported vision token format: {token_format}")
+    if media_type == "video":
+        # SGLang's native sgl.video(...) transport binds video data using its
+        # own placeholder path. Adding a literal Qwen <|video_pad|> token here
+        # makes the server look for an extra unbound video iterator.
+        return prompt
+
+    known_tokens = (
+        "<|image_pad|>",
+        "<|video_pad|>",
+        "<image>",
+        "<video>",
+    )
+    if any(existing_token in prompt for existing_token in known_tokens):
+        return prompt
+    return f"{QWEN_IMAGE_TOKEN}\n{prompt}"
+
+
+def _openai_generate_one(
+    request: dict[str, Any],
+    url: str,
+    model_name: str,
+    temperature: float,
+    max_tokens: int,
+    timeout: int,
+    media_url_base: str | None,
+    media_root: str | None,
+    input_root: str | None,
+) -> tuple[str, str | None]:
+    try:
+        import requests
+    except ImportError as e:
+        raise ImportError("OpenAI API mode requires the requests package.") from e
+
+    content = [{"type": "text", "text": request["prompt"]}]
+    if request.get("image_path"):
+        content.append(
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": _as_openai_media_value(
+                        request["image_path"], media_url_base, media_root, input_root
+                    )
+                },
+            }
+        )
+    else:
+        content.append(
+            {
+                "type": "video_url",
+                "video_url": {
+                    "url": _as_openai_media_value(
+                        request["video_path"], media_url_base, media_root, input_root
+                    )
+                },
+            }
+        )
+
+    payload = {
+        "model": model_name,
+        "messages": [{"role": "user", "content": content}],
+        "temperature": temperature,
+        "max_tokens": max_tokens,
+    }
+
+    try:
+        response = requests.post(url, json=payload, timeout=timeout)
+        if response.status_code != 200:
+            return "", f"HTTP {response.status_code}: {response.text[:1000]}"
+        data = response.json()
+        choices = data.get("choices") or []
+        if not choices:
+            return "", f"No choices in OpenAI response: {data}"
+        message = choices[0].get("message") or {}
+        answer = _coerce_text(message.get("content"))
+        if not answer:
+            return "", f"Could not extract assistant content from OpenAI response: {data}"
+        return answer, None
+    except Exception:
+        return "", traceback.format_exc()
+
+
+def _write_generation(
+    outfile,
+    idx: int,
+    sample: dict[str, Any],
+    answer: str,
+    generation_error: str | None,
+) -> None:
+    out = {
+        "conversation_id": idx,
+        "messages": _messages_for_output(sample, answer),
+    }
+    if generation_error:
+        out["generation_error"] = generation_error
+    for key in (
+        "id",
+        "category",
+        "subcategory",
+        "dataset",
+        "answer",
+        "reference_answer",
+        "index2ans",
+        "video_path",
+        "image_path",
+        "source_image",
+        "source_video_path",
+        "source_mcap_path",
+        "source_subset",
+        "camera",
+        "camera_topic",
+        "task_path",
+        "skill",
+        "stage",
+        "major_task",
+        "task_id",
+        "episode_id",
+        "task_name",
+        "source_split",
+        "data_subtype",
+        "question_id",
+        "image_id",
+        "question_type",
+        "answer_type",
+        "multiple_choice_answer",
+        "answer_counts",
+        "answers",
+        "window_idx",
+        "windows_per_mcap",
+        "window_message_start",
+        "window_message_end",
+    ):
+        if key in sample:
+            out[key] = sample[key]
+    outfile.write(json.dumps(out, ensure_ascii=False) + "\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_path", required=True)
+    parser.add_argument("--output_path", required=True)
+    parser.add_argument("--input_root", default=None, help="Mounted shard root, e.g. /input_data")
+    parser.add_argument("--media_root", default=None, help="Mounted dataset root, e.g. /media_data")
+    parser.add_argument("--num_threads", type=int, default=64)
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--max_tokens", type=int, default=2048)
+    parser.add_argument("--num_frames", type=int, default=32)
+    parser.add_argument("--url", type=str, default="http://localhost:8000")
+    parser.add_argument(
+        "--media_url_base",
+        default=None,
+        help="HTTP base URL serving media_root/input_root for OpenAI multimodal requests.",
+    )
+    parser.add_argument(
+        "--api_mode",
+        choices=("openai", "native"),
+        default="native",
+        help="Use SGLang's OpenAI-compatible API or SGLang native DSL client.",
+    )
+    parser.add_argument("--model_name", default="model")
+    parser.add_argument("--request_timeout", type=int, default=1800)
+    parser.add_argument(
+        "--vision_token_format",
+        choices=("qwen_vl", "none"),
+        default="qwen_vl",
+        help="Insert explicit Qwen/Cosmos-VL image/video placeholder tokens into prompts.",
+    )
+    parser.add_argument("--log_empty_conversations", action="store_true")
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Remove output_path before generating. Useful after a bad/partial run.",
+    )
+    args = parser.parse_args()
+
+    url = args.url.removesuffix("/v1")
+
+    data = _load_json_or_jsonl(args.data_path)
+
+    if args.overwrite and os.path.exists(args.output_path):
+        os.remove(args.output_path)
+
+    finished_ids: set[int] = set()
+    done = False
+    if os.path.exists(args.output_path):
+        with open(args.output_path, encoding="utf-8") as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                outdata = json.loads(line)
+                finished_ids.add(outdata.get("conversation_id", -1))
+                if outdata.get("finished", False):
+                    done = True
+                    break
+    if done:
+        print("All conversations already generated")
+        sys.exit()
+
+    batch_args = []
+    batch_meta = []
+    for idx, sample in enumerate(data):
+        if idx in finished_ids:
+            continue
+        prompt, video_path, image_path = _extract_text_and_media(sample)
+        resolved_image = _resolve_media_path(image_path, args.media_root, args.input_root)
+        resolved_video = None
+        if not resolved_image:
+            resolved_video = _resolve_media_path(video_path, args.media_root, args.input_root)
+        if not resolved_video and not resolved_image:
+            print(f"Skipping sample {idx}: no video_path or image_path")
+            continue
+        request = {
+            "prompt": prompt,
+            "video_path": resolved_video,
+            "image_path": resolved_image,
+            "num_frames": args.num_frames,
+            "vision_token_format": args.vision_token_format,
+        }
+        batch_args.append([request])
+        batch_meta.append((idx, sample))
+
+    if not batch_args:
+        print("No pending multimodal conversations to generate")
+        if args.log_empty_conversations:
+            os.makedirs(os.path.dirname(args.output_path) or ".", exist_ok=True)
+            with open(args.output_path, "a", encoding="utf-8") as f:
+                f.write(json.dumps({"finished": True}) + "\n")
+        return
+
+    if args.api_mode == "openai":
+        chat_url = _openai_chat_url(args.url)
+        os.makedirs(os.path.dirname(args.output_path) or ".", exist_ok=True)
+        with open(args.output_path, "a", encoding="utf-8") as f:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_threads) as executor:
+                future_to_meta = {
+                    executor.submit(
+                        _openai_generate_one,
+                        request_args[0],
+                        chat_url,
+                        args.model_name,
+                        args.temperature,
+                        args.max_tokens,
+                        args.request_timeout,
+                        args.media_url_base,
+                        args.media_root,
+                        args.input_root,
+                    ): meta
+                    for request_args, meta in zip(batch_args, batch_meta)
+                }
+                for future in tqdm.tqdm(
+                    concurrent.futures.as_completed(future_to_meta),
+                    total=len(future_to_meta),
+                    desc="Generating and writing outputs",
+                ):
+                    idx, sample = future_to_meta[future]
+                    answer, error = future.result()
+                    if error:
+                        print(f"WARNING: generation failed for sample {idx}: {error[:500]}")
+                    if not answer and not args.log_empty_conversations:
+                        continue
+                    _write_generation(f, idx, sample, answer, error)
+
+            if args.log_empty_conversations:
+                f.write(json.dumps({"finished": True}) + "\n")
+        return
+
+    try:
+        import sglang as sgl
+    except ImportError as e:
+        raise ImportError("server_generate_vlm_sglang.py requires sglang in native mode.") from e
+
+    sgl.set_default_backend(sgl.RuntimeEndpoint(url))
+
+    @sgl.function
+    def generate_answer(s, request):
+        if request.get("image_path"):
+            prompt = _prompt_with_vision_token(
+                request["prompt"], "image", request["vision_token_format"]
+            )
+            s += sgl.user(sgl.image(request["image_path"]) + prompt)
+        else:
+            prompt = _prompt_with_vision_token(
+                request["prompt"], "video", request["vision_token_format"]
+            )
+            s += sgl.user(
+                sgl.video(request["video_path"], num_frames=request["num_frames"]) + prompt
+            )
+        s += sgl.assistant(sgl.gen("answer"))
+
+    states = generate_answer.run_batch(
+        batch_args,
+        temperature=args.temperature,
+        max_new_tokens=args.max_tokens,
+        num_threads=args.num_threads,
+        progress_bar=True,
+    )
+
+    os.makedirs(os.path.dirname(args.output_path) or ".", exist_ok=True)
+    with open(args.output_path, "a", encoding="utf-8") as f:
+        for state, (idx, sample) in tqdm.tqdm(
+            zip(states, batch_meta), total=len(batch_meta), desc="Writing outputs"
+        ):
+            answer = _state_answer(state)
+            if not answer:
+                print(
+                    "WARNING: Could not extract sgl.gen('answer') from state "
+                    f"type={type(state).__name__}, repr={repr(state)[:500]}"
+                )
+            if not answer and not args.log_empty_conversations:
+                continue
+            error = None
+            if not answer:
+                error = "Could not extract sgl.gen('answer') from SGLang ProgramState."
+            _write_generation(f, idx, sample, answer, error)
+
+        if args.log_empty_conversations:
+            f.write(json.dumps({"finished": True}) + "\n")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/speculative_decoding/distributed_generate/worker.sh b/examples/speculative_decoding/distributed_generate/worker.sh
index 97bf14c014a..7ae6aa495c9 100644
--- a/examples/speculative_decoding/distributed_generate/worker.sh
+++ b/examples/speculative_decoding/distributed_generate/worker.sh
@@ -14,52 +14,198 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-_CURRENT_COUNTER="$1"
+set -euo pipefail
 
+_CURRENT_COUNTER="$1"
 BACKEND="$2"
 JOBS_PER_NODE="$3"
 SYSTEM_PROMPT="$4"
 
-if [ "$BACKEND" == "vllm" ]; then
-    vllm serve /model/  --tensor-parallel-size 8 --served-model-name model --port 8000 --host 0.0.0.0 --trust-remote-code &
+BASE_PORT=${BASE_PORT:-8000}
+SGLANG_TP_SIZE=${SGLANG_TP_SIZE:-1}
+NUM_TEMPERATURES=${NUM_TEMPERATURES:-8}
+MAX_TOKENS=${MAX_TOKENS:-4096}
+MODEL_NAME=${MODEL_NAME:-model}
+STARTUP_TIMEOUT_SECONDS=${STARTUP_TIMEOUT_SECONDS:-600}
+GPU_COUNT=${GPU_COUNT:-$(nvidia-smi -L 2>/dev/null | wc -l)}
+
+if [ "$BACKEND" != "vllm" ] && [ "$BACKEND" != "sglang" ]; then
+    echo "ERROR: backend must be vllm or sglang, got: $BACKEND" >&2
+    exit 1
+fi
+if [ "$GPU_COUNT" -le 0 ]; then
+    echo "ERROR: no GPUs are visible inside the container." >&2
+    exit 1
+fi
+if [ "$SGLANG_TP_SIZE" -le 0 ] || [ "$SGLANG_TP_SIZE" -gt "$GPU_COUNT" ]; then
+    echo "ERROR: SGLANG_TP_SIZE=$SGLANG_TP_SIZE is invalid for GPU_COUNT=$GPU_COUNT." >&2
+    exit 1
+fi
+if [ "$NUM_TEMPERATURES" -le 0 ]; then
+    echo "ERROR: NUM_TEMPERATURES must be positive." >&2
+    exit 1
+fi
+if [ "$SGLANG_TP_SIZE" -eq 1 ] && [ "$NUM_TEMPERATURES" -gt "$GPU_COUNT" ]; then
+    echo "ERROR: NUM_TEMPERATURES=$NUM_TEMPERATURES exceeds GPU_COUNT=$GPU_COUNT for TP=1 serving." >&2
+    exit 1
+fi
+if ! [[ "$STARTUP_TIMEOUT_SECONDS" =~ ^[1-9][0-9]*$ ]]; then
+    echo "ERROR: STARTUP_TIMEOUT_SECONDS must be a positive integer." >&2
+    exit 1
+fi
+
+if [ "$SGLANG_TP_SIZE" -eq 1 ]; then
+    TEXT_NUM_THREADS=${TEXT_NUM_THREADS:-64}
 else
-    python3 -m sglang.launch_server --model-path /model --served-model-name model --tp 8 --port 8000 --host 0.0.0.0 --trust-remote-code &
+    TEXT_NUM_THREADS=${TEXT_NUM_THREADS:-320}
 fi
-# Wait for server to start up by polling the health endpoint
-echo "Waiting for server to start..."
-while true; do
-    response=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:8000/health" || true)
-    if [ "$response" -eq 200 ]; then
-        echo "Server is up!"
-        break
+
+SERVER_PIDS=()
+GENERATION_PIDS=()
+SERVER_PORTS=()
+
+cleanup() {
+    local attempt pid still_running
+    for pid in "${SERVER_PIDS[@]}"; do
+        kill "$pid" 2>/dev/null || true
+    done
+    for attempt in $(seq 1 10); do
+        still_running=0
+        for pid in "${SERVER_PIDS[@]}"; do
+            if kill -0 "$pid" 2>/dev/null; then
+                still_running=1
+            fi
+        done
+        [ "$still_running" -eq 0 ] && break
+        sleep 1
+    done
+    for pid in "${SERVER_PIDS[@]}"; do
+        if kill -0 "$pid" 2>/dev/null; then
+            kill -KILL "$pid" 2>/dev/null || true
+        fi
+    done
+    for pid in "${SERVER_PIDS[@]}"; do
+        wait "$pid" 2>/dev/null || true
+    done
+}
+
+trap cleanup EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+if [ "$SGLANG_TP_SIZE" -eq 1 ]; then
+    for gpu in $(seq 0 $((NUM_TEMPERATURES - 1))); do
+        port=$((BASE_PORT + gpu))
+        SERVER_PORTS+=("$port")
+        if [ "$BACKEND" = "vllm" ]; then
+            CUDA_VISIBLE_DEVICES=$gpu vllm serve /model/ \
+                --tensor-parallel-size 1 --served-model-name "$MODEL_NAME" \
+                --port "$port" --host 0.0.0.0 --trust-remote-code &
+        else
+            CUDA_VISIBLE_DEVICES=$gpu python3 -m sglang.launch_server \
+                --model-path /model --served-model-name "$MODEL_NAME" --tp 1 \
+                --port "$port" --host 0.0.0.0 ${SGLANG_EXTRA_ARGS:-} --trust-remote-code &
+        fi
+        SERVER_PIDS+=("$!")
+    done
+else
+    GPU_LIST=$(seq -s, 0 $((SGLANG_TP_SIZE - 1)))
+    SERVER_PORTS=("$BASE_PORT")
+    if [ "$BACKEND" = "vllm" ]; then
+        CUDA_VISIBLE_DEVICES=$GPU_LIST vllm serve /model/ \
+            --tensor-parallel-size "$SGLANG_TP_SIZE" --served-model-name "$MODEL_NAME" \
+            --port "$BASE_PORT" --host 0.0.0.0 --trust-remote-code &
+    else
+        CUDA_VISIBLE_DEVICES=$GPU_LIST python3 -m sglang.launch_server \
+            --model-path /model --served-model-name "$MODEL_NAME" --tp "$SGLANG_TP_SIZE" \
+            --port "$BASE_PORT" --host 0.0.0.0 ${SGLANG_EXTRA_ARGS:-} --trust-remote-code &
     fi
-    echo "Server not ready yet, retrying in 10 seconds..."
-    sleep 10
-done
+    SERVER_PIDS+=("$!")
+fi
+
+wait_for_servers() {
+    local deadline=$((SECONDS + STARTUP_TIMEOUT_SECONDS))
+    local index pid port ready response
 
+    echo "Waiting for ${#SERVER_PORTS[@]} server(s) to start..."
+    while (( SECONDS < deadline )); do
+        for index in "${!SERVER_PIDS[@]}"; do
+            pid=${SERVER_PIDS[$index]}
+            if ! kill -0 "$pid" 2>/dev/null; then
+                echo "ERROR: server process $pid exited before becoming ready." >&2
+                return 1
+            fi
+        done
 
-native_mpi_rank=$OMPI_COMM_WORLD_RANK
-# Works with Slurm launching with `--mpi=pmix`
+        ready=0
+        for port in "${SERVER_PORTS[@]}"; do
+            response=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:${port}/health" || true)
+            [ "$response" = "200" ] && ready=$((ready + 1))
+        done
+        if [ "$ready" -eq "${#SERVER_PORTS[@]}" ]; then
+            echo "All servers are up!"
+            return 0
+        fi
+        echo "$ready/${#SERVER_PORTS[@]} servers ready; retrying in 5 seconds..."
+        sleep 5
+    done
+
+    echo "ERROR: servers did not become ready within ${STARTUP_TIMEOUT_SECONDS} seconds." >&2
+    return 1
+}
+
+wait_for_servers
+
+native_mpi_rank=${OMPI_COMM_WORLD_RANK:-0}
 mpi_rank=${PMIX_RANK:-$native_mpi_rank}
 echo "Rank: $mpi_rank"
 echo "Counter: $_CURRENT_COUNTER"
 
-
 if [ "$mpi_rank" -eq 0 ]; then
-    start_shard=$((_CURRENT_COUNTER + 0))
+    start_shard=$_CURRENT_COUNTER
     end_shard=$((_CURRENT_COUNTER + JOBS_PER_NODE - 1))
 
-    for i in $(seq $start_shard $end_shard); do
-        echo "Processing shard: $i"
-        shard=$(printf "/input_data/train-%05d-%05d.jsonl" $i $i)
-        OUTPUT=$(printf "/output_data/output-%05d-%05d.jsonl" $i $i)
+    run_temperature() {
+        temp_id=$1
+        port=$2
+        temperature=$(printf "0.%d" "$temp_id")
+        for i in $(seq "$start_shard" "$end_shard"); do
+            echo "Temperature $temperature processing shard: $i on port $port"
+            shard=$(printf "/input_data/train-%05d-%05d.jsonl" "$i" "$i")
+            if [ ! -s "$shard" ]; then
+                echo "Skipping missing shard: $shard"
+                continue
+            fi
+            output=$(printf "/output_data/output-%05d-%05d-temp-%s.jsonl" "$i" "$i" "$temperature")
+            cmd=(
+                python3 /scripts/scripts/server_generate.py
+                --data_path "$shard"
+                --output_path "$output"
+                --num_threads "$TEXT_NUM_THREADS"
+                --max_tokens "$MAX_TOKENS"
+                --temperature "$temperature"
+                --url "http://localhost:$port/v1"
+                --log_empty_conversations
+            )
+            if [ -n "$SYSTEM_PROMPT" ]; then
+                cmd+=(--system_prompt "$SYSTEM_PROMPT")
+            fi
+            echo "Running: ${cmd[*]}"
+            "${cmd[@]}"
+        done
+    }
 
-        cmd="python3 /scripts/server_generate.py --data_path $shard --output_path $OUTPUT --num_threads 320 --max_tokens 4096 --temperature 0.0 --chat --log_empty_conversations"
-
-        if [ -n "$SYSTEM_PROMPT" ]; then
-            cmd+=" --system_prompt $SYSTEM_PROMPT"
-        fi
-        echo "Running command: $cmd"
-        eval $cmd
-    done
+    if [ "$SGLANG_TP_SIZE" -eq 1 ]; then
+        for temp_id in $(seq 0 $((NUM_TEMPERATURES - 1))); do
+            run_temperature "$temp_id" "$((BASE_PORT + temp_id))" &
+            GENERATION_PIDS+=("$!")
+        done
+        for generation_pid in "${GENERATION_PIDS[@]}"; do
+            wait "$generation_pid"
+        done
+    else
+        for temp_id in $(seq 0 $((NUM_TEMPERATURES - 1))); do
+            run_temperature "$temp_id" "$BASE_PORT"
+        done
+    fi
 fi
diff --git a/examples/speculative_decoding/distributed_generate/worker_multimodal.sh b/examples/speculative_decoding/distributed_generate/worker_multimodal.sh
new file mode 100755
index 00000000000..2db2b3c4aba
--- /dev/null
+++ b/examples/speculative_decoding/distributed_generate/worker_multimodal.sh
@@ -0,0 +1,269 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+_CURRENT_COUNTER="$1"
+BACKEND="$2"
+JOBS_PER_NODE="$3"
+NUM_FRAMES="${4:-32}"
+SYSTEM_PROMPT="${5:-}"
+
+if [ "$BACKEND" != "sglang" ]; then
+    echo "Multimodal generation currently uses the SGLang native video client; backend must be sglang."
+    exit 1
+fi
+
+if [ "${INSTALL_OPENCV_HEADLESS:-0}" = "1" ]; then
+    python3 -c "import cv2" || python3 -m pip install --user opencv-python-headless
+fi
+
+BASE_PORT=${BASE_PORT:-8000}
+SGLANG_TP_SIZE=${SGLANG_TP_SIZE:-8}
+NUM_TEMPERATURES=${NUM_TEMPERATURES:-8}
+API_MODE=${API_MODE:-openai}
+MEDIA_HTTP_PORT=${MEDIA_HTTP_PORT:-18080}
+MEDIA_URL_BASE=${MEDIA_URL_BASE:-http://127.0.0.1:${MEDIA_HTTP_PORT}}
+STARTUP_TIMEOUT_SECONDS=${STARTUP_TIMEOUT_SECONDS:-600}
+GPU_COUNT=${GPU_COUNT:-$(nvidia-smi -L 2>/dev/null | wc -l || true)}
+
+if [ "$GPU_COUNT" -le 0 ]; then
+    echo "ERROR: no GPUs are visible inside the container." >&2
+    exit 1
+fi
+if [ "$SGLANG_TP_SIZE" -gt "$GPU_COUNT" ]; then
+    echo "ERROR: SGLANG_TP_SIZE=$SGLANG_TP_SIZE exceeds GPU_COUNT=$GPU_COUNT." >&2
+    exit 1
+fi
+if [ "$SGLANG_TP_SIZE" -eq 1 ] && [ "$NUM_TEMPERATURES" -gt "$GPU_COUNT" ]; then
+    echo "ERROR: NUM_TEMPERATURES=$NUM_TEMPERATURES exceeds GPU_COUNT=$GPU_COUNT for single-GPU serving." >&2
+    exit 1
+fi
+if ! [[ "$STARTUP_TIMEOUT_SECONDS" =~ ^[1-9][0-9]*$ ]]; then
+    echo "ERROR: STARTUP_TIMEOUT_SECONDS must be a positive integer." >&2
+    exit 1
+fi
+
+SERVER_PIDS=()
+GENERATION_PIDS=()
+MEDIA_HTTP_PID=""
+
+cleanup() {
+    local attempt pid still_running
+    for pid in "${SERVER_PIDS[@]}" "${MEDIA_HTTP_PID:-}"; do
+        [ -n "${pid:-}" ] && kill "$pid" 2>/dev/null || true
+    done
+    for attempt in $(seq 1 10); do
+        still_running=0
+        for pid in "${SERVER_PIDS[@]}" "${MEDIA_HTTP_PID:-}"; do
+            if [ -n "${pid:-}" ] && kill -0 "$pid" 2>/dev/null; then
+                still_running=1
+            fi
+        done
+        [ "$still_running" -eq 0 ] && break
+        sleep 1
+    done
+    for pid in "${SERVER_PIDS[@]}" "${MEDIA_HTTP_PID:-}"; do
+        if [ -n "${pid:-}" ] && kill -0 "$pid" 2>/dev/null; then
+            kill -KILL "$pid" 2>/dev/null || true
+        fi
+    done
+    for pid in "${SERVER_PIDS[@]}" "${MEDIA_HTTP_PID:-}"; do
+        [ -n "${pid:-}" ] && wait "$pid" 2>/dev/null || true
+    done
+}
+
+trap cleanup EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+if [ "$API_MODE" = "openai" ]; then
+    python3 -m http.server "$MEDIA_HTTP_PORT" --bind 127.0.0.1 --directory /media_data \
+        >/tmp/multimodal_media_http_${MEDIA_HTTP_PORT}.log 2>&1 &
+    MEDIA_HTTP_PID=$!
+
+    echo "Waiting for media HTTP server at $MEDIA_URL_BASE..."
+    for _ in $(seq 1 30); do
+        if curl -fsS "$MEDIA_URL_BASE/" >/dev/null 2>&1; then
+            echo "Media HTTP server is up."
+            break
+        fi
+        sleep 1
+    done
+    if ! curl -fsS "$MEDIA_URL_BASE/" >/dev/null 2>&1; then
+        echo "ERROR: media HTTP server did not start. See /tmp/multimodal_media_http_${MEDIA_HTTP_PORT}.log" >&2
+        exit 1
+    fi
+fi
+
+SERVER_PORTS=()
+if [ "$SGLANG_TP_SIZE" -eq 1 ]; then
+    for gpu in $(seq 0 $((NUM_TEMPERATURES - 1))); do
+        port=$((BASE_PORT + gpu))
+        SERVER_PORTS+=("$port")
+        CUDA_VISIBLE_DEVICES=$gpu \
+            python3 -m sglang.launch_server \
+            --model-path /model \
+            --served-model-name model \
+            --tp 1 \
+            --port "$port" \
+            --host 0.0.0.0 \
+            ${SGLANG_EXTRA_ARGS:-} \
+            --trust-remote-code &
+        SERVER_PIDS+=("$!")
+    done
+else
+    gpu_list=$(seq -s, 0 $((SGLANG_TP_SIZE - 1)))
+    SERVER_PORTS=("$BASE_PORT")
+    CUDA_VISIBLE_DEVICES=$gpu_list \
+        python3 -m sglang.launch_server \
+        --model-path /model \
+        --served-model-name model \
+        --tp "$SGLANG_TP_SIZE" \
+        --port "$BASE_PORT" \
+        --host 0.0.0.0 \
+        ${SGLANG_EXTRA_ARGS:-} \
+        --trust-remote-code &
+    SERVER_PIDS+=("$!")
+fi
+
+wait_for_servers() {
+    local deadline=$((SECONDS + STARTUP_TIMEOUT_SECONDS))
+    local index pid port ready response
+
+    echo "Waiting for server to start..."
+    while (( SECONDS < deadline )); do
+        for index in "${!SERVER_PIDS[@]}"; do
+            pid=${SERVER_PIDS[$index]}
+            if ! kill -0 "$pid" 2>/dev/null; then
+                echo "ERROR: server process $pid exited before becoming ready." >&2
+                return 1
+            fi
+        done
+
+        ready=0
+        for port in "${SERVER_PORTS[@]}"; do
+            response=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:${port}/health" || true)
+            if [ "$response" -eq 200 ]; then
+                ready=$((ready + 1))
+            fi
+        done
+        if [ "$ready" -eq "${#SERVER_PORTS[@]}" ]; then
+            echo "All servers are up!"
+            return 0
+        fi
+        echo "$ready/${#SERVER_PORTS[@]} servers ready, retrying in 5 seconds..."
+        sleep 5
+    done
+
+    echo "ERROR: servers did not become ready within ${STARTUP_TIMEOUT_SECONDS} seconds." >&2
+    return 1
+}
+
+wait_for_servers
+
+native_mpi_rank=${OMPI_COMM_WORLD_RANK:-0}
+mpi_rank=${PMIX_RANK:-$native_mpi_rank}
+echo "Rank: $mpi_rank"
+echo "Counter: $_CURRENT_COUNTER"
+
+if [ "$API_MODE" = "openai" ]; then
+    NUM_THREADS=${NUM_THREADS:-4}
+elif [ "$SGLANG_TP_SIZE" -eq 1 ]; then
+    NUM_THREADS=${NUM_THREADS:-64}
+else
+    NUM_THREADS=${NUM_THREADS:-8}
+fi
+MAX_TOKENS=${MAX_TOKENS:-6144}
+VISION_TOKEN_FORMAT=${VISION_TOKEN_FORMAT:-qwen_vl}
+MODEL_NAME=${MODEL_NAME:-model}
+
+if [ "$mpi_rank" -eq 0 ]; then
+    start_shard=$((_CURRENT_COUNTER + 0))
+    end_shard=$((_CURRENT_COUNTER + JOBS_PER_NODE - 1))
+
+    run_temperature() {
+        server_id=$1
+        port=$2
+        temperature=$(printf "0.%d" "$server_id")
+
+        echo "Starting multimodal generation on port $port, temperature=$temperature"
+
+        for i in $(seq $start_shard $end_shard); do
+            echo "Temperature $temperature processing shard $i"
+
+            shard=$(printf "/input_data/train-%05d-%05d.jsonl" "$i" "$i")
+            if [ ! -s "$shard" ]; then
+                echo "Skipping missing shard: $shard"
+                continue
+            fi
+
+            output=$(printf \
+                "/output_data/output-%05d-%05d-temp-%s.jsonl" \
+                "$i" "$server_id" "$temperature")
+
+            cmd=(
+                python3 /scripts/distributed_generate/server_generate_vlm_sglang.py
+                --data_path "$shard"
+                --output_path "$output"
+                --input_root /input_data
+                --media_root /media_data
+                --num_threads "$NUM_THREADS"
+                --max_tokens "$MAX_TOKENS"
+                --num_frames "$NUM_FRAMES"
+                --temperature "$temperature"
+                --api_mode "$API_MODE"
+                --model_name "$MODEL_NAME"
+                --media_url_base "$MEDIA_URL_BASE"
+                --vision_token_format "$VISION_TOKEN_FORMAT"
+                --log_empty_conversations
+                --url "http://localhost:$port"
+            )
+
+            if [ "${OVERWRITE_OUTPUT:-0}" = "1" ]; then
+                cmd+=(--overwrite)
+            fi
+
+            if [ -n "$SYSTEM_PROMPT" ]; then
+                echo "WARNING: SYSTEM_PROMPT is ignored for multimodal generation; include it in the shard prompts instead."
+            fi
+
+            printf 'Running:'
+            printf ' %q' "${cmd[@]}"
+            printf '\n'
+            "${cmd[@]}"
+        done
+    }
+
+    if [ "$SGLANG_TP_SIZE" -eq 1 ]; then
+        for server_id in $(seq 0 $((NUM_TEMPERATURES - 1))); do
+            (
+                port=$((BASE_PORT + server_id))
+                run_temperature "$server_id" "$port"
+            ) &
+            GENERATION_PIDS+=("$!")
+        done
+        for generation_pid in "${GENERATION_PIDS[@]}"; do
+            wait "$generation_pid"
+        done
+    else
+        # A tensor-parallel VLM uses all GPUs for one server, so run the
+        # temperature sweep sequentially against that server.
+        for server_id in $(seq 0 $((NUM_TEMPERATURES - 1))); do
+            run_temperature "$server_id" "$BASE_PORT"
+        done
+    fi
+fi
diff --git a/examples/speculative_decoding/recipes/merge_dflash_datasets.py b/examples/speculative_decoding/recipes/merge_dflash_datasets.py
new file mode 100644
index 00000000000..ff734961e0a
--- /dev/null
+++ b/examples/speculative_decoding/recipes/merge_dflash_datasets.py
@@ -0,0 +1,577 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Merge DFlash training sources with conservative, multimodal-safe deduplication."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+from array import array
+from collections import Counter, OrderedDict, defaultdict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+WORD_PATTERN = re.compile(
+    r"[\u3400-\u4dbf\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]|[^\W_]+(?:['\u2019][^\W_]+)*",
+    flags=re.UNICODE,
+)
+OUTPUT_PREFIX_PATTERN = re.compile(r"^output-(\d+)-")
+
+
+@dataclass(frozen=True)
+class Source:
+    """An input JSONL file or directory and the root for its relative media paths."""
+
+    name: str
+    input_path: Path
+    media_root: Path | None
+    files: tuple[Path, ...]
+
+
+@dataclass
+class Partition:
+    """Source files assigned to one independent merge worker."""
+
+    files: dict[str, list[Path]] = field(default_factory=lambda: defaultdict(list))
+
+
+class MediaResolutionError(ValueError):
+    """Raised when a multimodal record cannot be converted to absolute media paths."""
+
+
+def parse_source(value: str) -> tuple[str, Path]:
+    """Parse a ``NAME=PATH`` command-line source argument."""
+
+    name, separator, path = value.partition("=")
+    if not separator or not name or not path:
+        raise argparse.ArgumentTypeError("sources must use NAME=PATH")
+    return name, Path(path)
+
+
+def parse_media_root(value: str) -> tuple[str, Path]:
+    """Parse a ``NAME=PATH`` command-line media-root argument."""
+
+    return parse_source(value)
+
+
+def source_files(input_path: Path) -> Iterator[Path]:
+    """Yield JSONL input files in deterministic order."""
+
+    if input_path.is_file():
+        if input_path.suffix != ".jsonl":
+            raise ValueError(f"Expected a .jsonl file: {input_path}")
+        yield input_path
+        return
+    if not input_path.is_dir():
+        raise FileNotFoundError(f"Missing input source: {input_path}")
+    files = sorted(input_path.rglob("*.jsonl"))
+    if not files:
+        raise FileNotFoundError(f"No .jsonl files found in: {input_path}")
+    yield from files
+
+
+def listed_source_files(file_list: Path) -> tuple[Path, ...]:
+    """Read an explicit, ordered list of JSONL input files."""
+
+    if not file_list.is_file():
+        raise FileNotFoundError(f"Missing source-file list: {file_list}")
+    files = tuple(
+        Path(line.strip())
+        for line in file_list.read_text(encoding="utf-8").splitlines()
+        if line.strip()
+    )
+    if not files:
+        raise ValueError(f"Source-file list is empty: {file_list}")
+    for path in files:
+        if path.suffix != ".jsonl" or not path.is_file():
+            raise FileNotFoundError(f"Invalid JSONL file listed in {file_list}: {path}")
+    return files
+
+
+def resolve_media_path(value: str, media_root: Path | None, resolved_paths: dict[str, str]) -> str:
+    """Return a checked absolute media path for an image or video content part."""
+
+    cache_key = f"{media_root}\0{value}"
+    if cache_key in resolved_paths:
+        return resolved_paths[cache_key]
+    path = Path(value)
+    if not path.is_absolute():
+        if media_root is None:
+            raise MediaResolutionError(f"Relative media path without a media root: {value}")
+        path = media_root / path
+    path = path.resolve()
+    if not path.is_file():
+        raise MediaResolutionError(f"Missing media file: {path}")
+    resolved_paths[cache_key] = str(path)
+    return resolved_paths[cache_key]
+
+
+def has_text(content: str | list[dict]) -> bool:
+    """Return whether a message content value contains non-empty text."""
+
+    if isinstance(content, str):
+        return bool(content.strip())
+    return any(part.get("type") == "text" and str(part.get("text", "")).strip() for part in content)
+
+
+def normalize_messages(
+    record: dict, media_root: Path | None, resolved_paths: dict[str, str]
+) -> list[dict] | None:
+    """Validate a record and resolve its image/video references."""
+
+    messages = record.get("messages") or record.get("conversations")
+    if not isinstance(messages, list) or record.get("generation_error"):
+        return None
+
+    normalized = []
+    for message in messages:
+        if not isinstance(message, dict):
+            return None
+        role = message.get("role")
+        content = message.get("content")
+        if role not in {"system", "user", "assistant"} or not isinstance(content, (str, list)):
+            return None
+        if isinstance(content, str):
+            normalized.append({"role": role, "content": content})
+            continue
+
+        parts = []
+        for part in content:
+            if not isinstance(part, dict) or part.get("type") not in {"text", "image", "video"}:
+                return None
+            normalized_part = dict(part)
+            part_type = normalized_part["type"]
+            if part_type == "text":
+                if not isinstance(normalized_part.get("text"), str):
+                    return None
+            else:
+                media_value = normalized_part.get(part_type)
+                if not isinstance(media_value, str):
+                    return None
+                normalized_part[part_type] = resolve_media_path(
+                    media_value, media_root, resolved_paths
+                )
+            parts.append(normalized_part)
+        normalized.append({"role": role, "content": parts})
+
+    if not any(
+        message["role"] == "user" and has_text(message["content"]) for message in normalized
+    ):
+        return None
+    if not any(
+        message["role"] == "assistant" and has_text(message["content"]) for message in normalized
+    ):
+        return None
+    return normalized
+
+
+def canonical_json(value: object) -> bytes:
+    """Serialize identity values stably and without whitespace."""
+
+    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode(
+        "utf-8"
+    )
+
+
+def candidate_key(messages: list[dict]) -> bytes:
+    """Hash the fixed context used only to find likely duplicate candidates."""
+
+    context = [message for message in messages if message["role"] != "assistant"]
+    return hashlib.sha256(canonical_json(context)).digest()
+
+
+def message_words(messages: list[dict]) -> array:
+    """Return sorted, multiplicity-preserving hashes of all textual message words."""
+
+    words: list[int] = []
+    for message in messages:
+        content = message["content"]
+        text_values = (
+            [content]
+            if isinstance(content, str)
+            else [part["text"] for part in content if part["type"] == "text"]
+        )
+        for text in text_values:
+            words.extend(
+                hash(word) & ((1 << 64) - 1) for word in WORD_PATTERN.findall(text.casefold())
+            )
+    return array("Q", sorted(words))
+
+
+def has_required_overlap(first: array, second: array, threshold: float) -> bool:
+    """Check mutual multiset-word overlap with a two-pointer scan."""
+
+    if not first or not second:
+        return False
+    minimum = int(threshold * len(first) + 0.999999999)
+    if len(second) < minimum or len(second) * threshold > len(first):
+        return False
+
+    first_index = 0
+    second_index = 0
+    shared = 0
+    while first_index < len(first) and second_index < len(second):
+        first_word = first[first_index]
+        second_word = second[second_index]
+        if first_word == second_word:
+            shared += 1
+            first_index += 1
+            second_index += 1
+        elif first_word < second_word:
+            first_index += 1
+        else:
+            second_index += 1
+    return shared >= minimum and shared >= int(threshold * len(second) + 0.999999999)
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse merge configuration."""
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--source",
+        action="append",
+        default=[],
+        metavar="NAME=PATH",
+        help="JSONL file or directory to include; may be specified multiple times.",
+    )
+    parser.add_argument(
+        "--media-root",
+        action="append",
+        default=[],
+        metavar="NAME=PATH",
+        help="Root used to resolve relative image/video paths for a source.",
+    )
+    parser.add_argument(
+        "--source-file-list",
+        action="append",
+        default=[],
+        metavar="NAME=PATH",
+        help="Ordered JSONL file list to use instead of recursively scanning a source directory.",
+    )
+    parser.add_argument("--output", required=True, type=Path, help="Merged JSONL destination.")
+    parser.add_argument(
+        "--word-overlap",
+        type=float,
+        default=0.90,
+        help="Mutual multiset-word overlap required to remove a candidate (default: 0.90).",
+    )
+    parser.add_argument(
+        "--cache-contexts",
+        type=int,
+        default=25_000,
+        help="Maximum recently seen prompt/media contexts retained for fuzzy comparisons.",
+    )
+    parser.add_argument(
+        "--jobs",
+        type=int,
+        default=1,
+        help=(
+            "Merge-worker count. Output files with the same synthetic shard prefix stay together; "
+            "deduplication is local to each worker when this is greater than one."
+        ),
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true", help="Replace an existing output after success."
+    )
+    parser.add_argument(
+        "--keep-work-dir", action="store_true", help="Keep parallel worker manifests and outputs."
+    )
+    parser.add_argument(
+        "--dry-run", action="store_true", help="Print parallel worker assignments without merging."
+    )
+    return parser.parse_args()
+
+
+def load_sources(args: argparse.Namespace) -> list[Source]:
+    """Validate source arguments and materialize their ordered JSONL files."""
+
+    media_roots = dict(parse_media_root(value) for value in args.media_root)
+    file_lists = dict(parse_source(value) for value in args.source_file_list)
+    source_arguments = list(map(parse_source, args.source))
+    source_names = {name for name, _ in source_arguments}
+    unknown_file_lists = set(file_lists) - source_names
+    if unknown_file_lists:
+        raise ValueError(f"--source-file-list names are not sources: {sorted(unknown_file_lists)}")
+    sources = [
+        Source(
+            name,
+            input_path,
+            media_roots.get(name),
+            listed_source_files(file_lists[name])
+            if name in file_lists
+            else tuple(source_files(input_path)),
+        )
+        for name, input_path in source_arguments
+    ]
+    if len({source.name for source in sources}) != len(sources):
+        raise ValueError("Each --source name must be unique")
+    for source in sources:
+        if source.media_root is not None and not source.media_root.is_dir():
+            raise FileNotFoundError(f"Missing media root for {source.name}: {source.media_root}")
+
+    return sources
+
+
+def merge_sources(
+    sources: list[Source], output: Path, word_overlap: float, cache_contexts: int
+) -> None:
+    """Merge source records into one atomic DFlash-training JSONL file."""
+
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = output.with_name(f".{output.name}.{os.getpid()}.tmp")
+    candidates: OrderedDict[bytes, list[array]] = OrderedDict()
+    written = Counter()
+    skipped = Counter()
+    processed = 0
+
+    try:
+        with temporary.open("x", encoding="utf-8") as destination:
+            for source in sources:
+                resolved_paths: dict[str, str] = {}
+                for path in source.files:
+                    with path.open(encoding="utf-8") as input_file:
+                        for line_number, line in enumerate(input_file, start=1):
+                            if not line.strip():
+                                continue
+                            processed += 1
+                            try:
+                                record = json.loads(line)
+                            except json.JSONDecodeError as error:
+                                raise ValueError(f"Invalid JSON at {path}:{line_number}") from error
+                            try:
+                                messages = normalize_messages(
+                                    record, source.media_root, resolved_paths
+                                )
+                            except MediaResolutionError:
+                                skipped["missing_media"] += 1
+                                continue
+                            if messages is None:
+                                skipped["invalid_or_failed"] += 1
+                                continue
+
+                            key = candidate_key(messages)
+                            words = message_words(messages)
+                            retained = candidates.get(key, [])
+                            if any(
+                                has_required_overlap(words, previous, word_overlap)
+                                for previous in retained
+                            ):
+                                skipped["near_duplicate"] += 1
+                                continue
+
+                            candidates.setdefault(key, []).append(words)
+                            candidates.move_to_end(key)
+                            while len(candidates) > cache_contexts:
+                                candidates.popitem(last=False)
+
+                            json.dump(
+                                {
+                                    "id": record.get("id", record.get("conversation_id")),
+                                    "dataset": record.get("dataset", source.name),
+                                    "messages": messages,
+                                },
+                                destination,
+                                ensure_ascii=False,
+                            )
+                            destination.write("\n")
+                            written[source.name] += 1
+                            if processed % 100_000 == 0:
+                                print(
+                                    f"Processed {processed:,}; wrote {sum(written.values()):,}; "
+                                    f"skipped {dict(skipped)}",
+                                    flush=True,
+                                )
+        os.replace(temporary, output)
+    except BaseException:
+        if temporary.exists():
+            temporary.unlink()
+        raise
+
+    print(f"Wrote {sum(written.values()):,} conversations to {output}")
+    print(f"By source: {dict(written)}; skipped: {dict(skipped)}")
+
+
+def assign_source_files(partitions: list[Partition], source: Source) -> None:
+    """Assign synthetic shard variants together and distribute other JSONL files evenly."""
+
+    if source.input_path.is_file():
+        partitions[0].files[source.name].extend(source.files)
+        return
+
+    grouped_by_prefix: dict[int, list[Path]] = defaultdict(list)
+    ungrouped = []
+    for path in source.files:
+        match = OUTPUT_PREFIX_PATTERN.match(path.name)
+        if match is None:
+            ungrouped.append(path)
+        else:
+            grouped_by_prefix[int(match.group(1))].append(path)
+
+    for prefix, files in sorted(grouped_by_prefix.items()):
+        partitions[prefix % len(partitions)].files[source.name].extend(files)
+    for index, path in enumerate(ungrouped):
+        partitions[index % len(partitions)].files[source.name].append(path)
+
+
+def write_manifest(path: Path, files: list[Path]) -> None:
+    """Write one worker's ordered source-file manifest."""
+
+    path.write_text("".join(f"{file}\n" for file in files), encoding="utf-8")
+
+
+def worker_command(
+    worker: int,
+    partition: Partition,
+    source_by_name: dict[str, Source],
+    args: argparse.Namespace,
+    manifests_dir: Path,
+    parts_dir: Path,
+) -> list[str]:
+    """Build a single-worker invocation of this script for one file partition."""
+
+    command = [
+        sys.executable,
+        str(Path(__file__).resolve()),
+        "--output",
+        str(parts_dir / f"part-{worker:02d}.jsonl"),
+        "--word-overlap",
+        str(args.word_overlap),
+        "--cache-contexts",
+        str(args.cache_contexts),
+    ]
+    for source_name, files in sorted(partition.files.items()):
+        source = source_by_name[source_name]
+        manifest = manifests_dir / f"{source_name}-{worker:02d}.txt"
+        write_manifest(manifest, files)
+        command.extend(("--source", f"{source.name}={source.input_path}"))
+        command.extend(("--source-file-list", f"{source.name}={manifest}"))
+        if source.media_root is not None:
+            command.extend(("--media-root", f"{source.name}={source.media_root}"))
+    return command
+
+
+def concatenate_parts(parts_dir: Path, output: Path, overwrite: bool) -> None:
+    """Join successful worker outputs into one atomic final JSONL file."""
+
+    if output.exists() and not overwrite:
+        raise FileExistsError(f"Output already exists (pass --overwrite to replace it): {output}")
+    temporary = output.with_name(f".{output.name}.{os.getpid()}.tmp")
+    try:
+        with temporary.open("xb") as destination:
+            for part in sorted(parts_dir.glob("part-*.jsonl")):
+                with part.open("rb") as source:
+                    shutil.copyfileobj(source, destination, length=16 * 1024 * 1024)
+        os.replace(temporary, output)
+    except BaseException:
+        if temporary.exists():
+            temporary.unlink()
+        raise
+
+
+def merge_in_parallel(sources: list[Source], args: argparse.Namespace) -> None:
+    """Merge source files concurrently while keeping synthetic shard variants together."""
+
+    partitions = [Partition() for _ in range(args.jobs)]
+    for source in sources:
+        assign_source_files(partitions, source)
+    for worker, partition in enumerate(partitions):
+        counts = {source: len(files) for source, files in sorted(partition.files.items())}
+        print(f"Worker {worker:02d}: {counts}")
+    if args.dry_run:
+        return
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    work_dir = args.output.with_name(f".{args.output.name}.parallel-{os.getpid()}")
+    manifests_dir = work_dir / "manifests"
+    parts_dir = work_dir / "parts"
+    manifests_dir.mkdir(parents=True)
+    parts_dir.mkdir()
+    source_by_name = {source.name: source for source in sources}
+    commands = [
+        (worker, worker_command(worker, partition, source_by_name, args, manifests_dir, parts_dir))
+        for worker, partition in enumerate(partitions)
+        if partition.files
+    ]
+
+    processes = []
+    merged = False
+    try:
+        for worker, command in commands:
+            print(f"Starting worker {worker:02d}", flush=True)
+            processes.append(subprocess.Popen(command))
+        failed_workers = [
+            worker for (worker, _), process in zip(commands, processes) if process.wait()
+        ]
+        if failed_workers:
+            raise RuntimeError(f"Merge workers failed: {failed_workers}")
+        concatenate_parts(parts_dir, args.output, args.overwrite)
+        merged = True
+    except BaseException:
+        for process in processes:
+            if process.poll() is None:
+                process.terminate()
+        for process in processes:
+            process.wait()
+        raise
+    finally:
+        if merged and not args.keep_work_dir:
+            shutil.rmtree(work_dir)
+
+    print(f"Wrote parallel merge to {args.output}")
+
+
+def main() -> None:
+    """Merge DFlash sources with one or more independent workers."""
+
+    args = parse_args()
+    if not args.source:
+        raise ValueError("At least one --source is required")
+    if not 0 < args.word_overlap <= 1:
+        raise ValueError("--word-overlap must be in (0, 1]")
+    if args.cache_contexts <= 0:
+        raise ValueError("--cache-contexts must be positive")
+    if args.jobs <= 0:
+        raise ValueError("--jobs must be positive")
+    if args.output.exists() and not args.overwrite:
+        raise FileExistsError(
+            f"Output already exists (pass --overwrite to replace it): {args.output}"
+        )
+
+    sources = load_sources(args)
+    if args.jobs == 1:
+        if args.dry_run:
+            for source in sources:
+                print(f"Source {source.name}: {len(source.files)} file(s)")
+            return
+        merge_sources(sources, args.output, args.word_overlap, args.cache_contexts)
+    else:
+        merge_in_parallel(sources, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py b/examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py
new file mode 100755
index 00000000000..680e8592cf2
--- /dev/null
+++ b/examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py
@@ -0,0 +1,520 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Create synthetic-generation shards from supported benchmark layouts.
+
+Supported adapters:
+  * pai_understanding: PAI-Bench-U parquet metadata plus a local video tree.
+  * vqa_v2: VQA v2 question/annotation JSON plus COCO image directories.
+  * specdec_multilingual_prompt: OpenAI-style text prompts from NVIDIA's dataset.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import random
+import shutil
+import subprocess
+import tempfile
+from collections import Counter
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+THINKING_SUFFIX = (
+    "\n\nAnswer using the following format:\n\n<think>\nYour reasoning grounded in the visual input.\n"
+    "</think>\n\nWrite the final answer immediately after the </think> tag."
+)
+
+VQA_SPLITS = {
+    "train": (
+        "train2014",
+        "COCO_train2014",
+        "v2_OpenEnded_mscoco_train2014_questions.json",
+        "v2_mscoco_train2014_annotations.json",
+    ),
+    "val": (
+        "val2014",
+        "COCO_val2014",
+        "v2_OpenEnded_mscoco_val2014_questions.json",
+        "v2_mscoco_val2014_annotations.json",
+    ),
+    "test": ("test2015", "COCO_test2015", "v2_OpenEnded_mscoco_test2015_questions.json", None),
+}
+
+
+def prompt(question: str, style: str, options: dict[str, str] | None = None) -> str:
+    text = question.strip()
+    if options:
+        choices = "\n".join(f"{key}. {value}" for key, value in sorted(options.items()))
+        text += f"\n\nOptions:\n{choices}"
+    if style == "plain":
+        return text
+    if options:
+        text += "\n\nSelect the single best option based only on the visual input."
+    return text + THINKING_SUFFIX
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    with path.open(encoding="utf-8") as file:
+        value = json.load(file)
+    if not isinstance(value, dict):
+        raise ValueError(f"Expected a JSON object: {path}")
+    return value
+
+
+def find_json(root: Path, filename: str) -> Path:
+    candidates = (root / filename, root / "Questions" / filename, root / "Annotations" / filename)
+    for candidate in candidates:
+        if candidate.is_file():
+            return candidate
+    raise FileNotFoundError(f"Could not find {filename} below {root}")
+
+
+def image_path(
+    image_root: Path, data_subtype: str, prefix: str, image_id: int
+) -> tuple[Path, str] | None:
+    filename = f"{prefix}_{image_id:012d}.jpg"
+    for candidate in (image_root / data_subtype / filename, image_root / filename):
+        if candidate.is_file():
+            return candidate, str(candidate.relative_to(image_root))
+    return None
+
+
+def media_path_within_root(media_root: Path, relative_path: str) -> Path | None:
+    """Resolve a media reference only when it remains below its declared root."""
+
+    media_root = media_root.resolve()
+    candidate = (media_root / relative_path).resolve()
+    try:
+        candidate.relative_to(media_root)
+    except ValueError:
+        return None
+    return candidate
+
+
+def normalize_options(value: Any) -> dict[str, str]:
+    if not isinstance(value, dict):
+        return {}
+    return {
+        str(key): str(answer)
+        for key, answer in value.items()
+        if answer is not None and str(answer).strip()
+    }
+
+
+def make_mosaic(video: Path, output_dir: Path, num_frames: int) -> str | None:
+    """Use ffmpeg only: preparation stays independent of host OpenCV builds."""
+    if not shutil.which("ffmpeg"):
+        raise RuntimeError("--media_mode image_mosaic requires ffmpeg on the preparation host.")
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ImportError("--media_mode image_mosaic requires Pillow.") from exc
+    digest = hashlib.sha1(str(video).encode("utf-8")).hexdigest()[:16]
+    frame_dir = output_dir / "media_cache" / digest
+    mosaic_path = frame_dir / "mosaic.png"
+    if mosaic_path.is_file():
+        return str(mosaic_path.relative_to(output_dir))
+    frame_dir.mkdir(parents=True, exist_ok=True)
+    for frame in frame_dir.glob("frame_*.png"):
+        frame.unlink()
+    subprocess.run(
+        [
+            "ffmpeg",
+            "-y",
+            "-i",
+            str(video),
+            "-vf",
+            f"fps={max(1, min(num_frames, 32))}",
+            "-frames:v",
+            str(num_frames),
+            str(frame_dir / "frame_%05d.png"),
+        ],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        check=False,
+    )
+    frames = sorted(frame_dir.glob("frame_*.png"))[:num_frames]
+    if not frames:
+        return None
+    columns = math.ceil(math.sqrt(len(frames)))
+    rows = math.ceil(len(frames) / columns)
+    canvas = Image.new("RGB", (columns * 224, rows * 224), (0, 0, 0))
+    for index, frame in enumerate(frames):
+        image = Image.open(frame).convert("RGB").resize((224, 224))
+        row, column = divmod(index, columns)
+        canvas.paste(image, (column * 224, row * 224))
+    canvas.save(mosaic_path, "PNG")
+    return str(mosaic_path.relative_to(output_dir))
+
+
+def download_pai_bench(dataset_dir: Path, repo_id: str, force_download: bool) -> None:
+    """Materialize the complete PAI-Bench-U Hugging Face dataset, including videos."""
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as exc:
+        raise ImportError(
+            "--download requires huggingface_hub. Install it in the preparation environment."
+        ) from exc
+    dataset_dir.mkdir(parents=True, exist_ok=True)
+    snapshot_download(
+        repo_id=repo_id,
+        repo_type="dataset",
+        local_dir=str(dataset_dir),
+        force_download=force_download,
+    )
+
+
+def pai_records(args: argparse.Namespace, output_dir: Path) -> Iterator[dict[str, Any]]:
+    if not args.dataset_dir:
+        raise ValueError("--dataset_dir is required with --dataset pai_understanding")
+    dataset_dir = Path(args.dataset_dir).expanduser().resolve()
+    if args.download:
+        download_pai_bench(dataset_dir, args.repo_id, args.force_download)
+    media_root = Path(args.media_root or dataset_dir).expanduser().resolve()
+    try:
+        from datasets import load_dataset
+    except ImportError as exc:
+        raise ImportError("PAI-Bench-U preparation requires datasets and pyarrow.") from exc
+    parquet_files = sorted((dataset_dir / "data").glob("test-*.parquet"))
+    if not parquet_files:
+        raise FileNotFoundError(f"No data/test-*.parquet found below {dataset_dir}")
+    rows = load_dataset(
+        "parquet", data_files={"test": [str(path) for path in parquet_files]}, split="test"
+    )
+    for row_index, row in enumerate(rows):
+        category = row.get("category") or ""
+        if args.category and category != args.category:
+            continue
+        question, relative_video = row.get("question"), row.get("video_path")
+        if (
+            not isinstance(question, str)
+            or not isinstance(relative_video, str)
+            or not relative_video.strip()
+        ):
+            continue
+        relative_video = relative_video.strip()
+        video = media_path_within_root(media_root, relative_video)
+        if video is None:
+            yield {"_missing_media": f"outside media root: {relative_video}"}
+            continue
+        if not video.is_file():
+            yield {"_missing_media": str(video)}
+            continue
+        options = normalize_options(row.get("index2ans"))
+        user_content: list[dict[str, Any]] = [
+            {"type": "text", "text": prompt(question, args.prompt_style, options)}
+        ]
+        media: dict[str, str] = {"video_path": relative_video}
+        if args.media_mode == "image_mosaic":
+            image = make_mosaic(video, output_dir, args.num_frames)
+            if not image:
+                yield {"_missing_media": str(video)}
+                continue
+            user_content.append({"type": "image", "image": image})
+            media.update({"image_path": image, "source_video_path": relative_video})
+        else:
+            user_content.append({"type": "video", "video": relative_video, "fps": 4})
+        yield {
+            "id": f"pai_understanding_{row_index:05d}",
+            "dataset": "pai_understanding",
+            "messages": [{"role": "user", "content": user_content}],
+            "prompt": prompt(question, args.prompt_style, options),
+            "question": question,
+            "reference_answer": row.get("answer"),
+            "index2ans": options,
+            "category": category,
+            "subcategory": row.get("subcategory") or "",
+            **media,
+        }
+
+
+def vqa_records(args: argparse.Namespace) -> Iterator[dict[str, Any]]:
+    if not args.vqa_root or not args.image_root:
+        raise ValueError("--vqa_root and --image_root are required with --dataset vqa_v2")
+    vqa_root = Path(args.vqa_root).expanduser().resolve()
+    image_root = Path(args.image_root).expanduser().resolve()
+    for split in [value.strip() for value in args.vqa_splits.split(",") if value.strip()]:
+        if split not in VQA_SPLITS:
+            raise ValueError(f"Unsupported VQA split {split!r}; choose from {sorted(VQA_SPLITS)}")
+        subtype, prefix, question_name, annotation_name = VQA_SPLITS[split]
+        questions = load_json(find_json(vqa_root, question_name)).get("questions")
+        if not isinstance(questions, list):
+            raise ValueError(f"questions is not a list for split={split}")
+        annotations: dict[int, dict[str, Any]] = {}
+        if annotation_name:
+            annotation_rows = load_json(find_json(vqa_root, annotation_name)).get("annotations")
+            if not isinstance(annotation_rows, list):
+                raise ValueError(f"annotations is not a list for split={split}")
+            annotations = {
+                item["question_id"]: item
+                for item in annotation_rows
+                if isinstance(item, dict) and isinstance(item.get("question_id"), int)
+            }
+        elif not args.include_unannotated_vqa:
+            raise ValueError(
+                f"VQA split={split} has no public annotations; pass --include_unannotated_vqa to use it."
+            )
+        for question_row in questions:
+            if not isinstance(question_row, dict):
+                continue
+            question_id, image_id, question = (
+                question_row.get("question_id"),
+                question_row.get("image_id"),
+                question_row.get("question"),
+            )
+            if (
+                not isinstance(question_id, int)
+                or not isinstance(image_id, int)
+                or not isinstance(question, str)
+            ):
+                continue
+            resolved = image_path(image_root, subtype, prefix, image_id)
+            if not resolved:
+                yield {"_missing_media": f"{subtype}/{prefix}_{image_id:012d}.jpg"}
+                continue
+            _, relative_image = resolved
+            annotation = annotations.get(question_id, {})
+            counts = Counter(
+                item.get("answer", "").strip()
+                for item in (annotation.get("answers") or [])
+                if isinstance(item, dict)
+                and isinstance(item.get("answer"), str)
+                and item["answer"].strip()
+            )
+            record: dict[str, Any] = {
+                "id": f"vqa_v2_{question_id}",
+                "dataset": "vqa_v2",
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt(question, args.prompt_style)},
+                            {"type": "image", "image": relative_image},
+                        ],
+                    }
+                ],
+                "prompt": prompt(question, args.prompt_style),
+                "question": question,
+                "image_path": relative_image,
+                "question_id": question_id,
+                "image_id": image_id,
+                "source_split": split,
+            }
+            if annotation:
+                record.update(
+                    {
+                        "reference_answer": annotation.get("multiple_choice_answer"),
+                        "question_type": annotation.get("question_type"),
+                        "answer_type": annotation.get("answer_type"),
+                        "answer_counts": dict(counts),
+                    }
+                )
+                if args.include_all_answers:
+                    record["answers"] = annotation.get("answers")
+            yield record
+
+
+def text_records(args: argparse.Namespace) -> Iterator[dict[str, Any]]:
+    if not args.text_data:
+        raise ValueError("--text_data is required with --dataset specdec_multilingual_prompt")
+    path = Path(args.text_data).expanduser().resolve()
+    if path.is_file():
+        files = [path]
+    else:
+        # The NVIDIA repository includes sample-*.jsonl convenience subsets.
+        # Prefer default.jsonl for a directory input so a full run neither
+        # misses most of the data nor duplicates the sampled records.
+        default_file = path / "default.jsonl"
+        files = [default_file] if default_file.is_file() else sorted(path.glob("*.jsonl"))
+    if not files:
+        raise FileNotFoundError(f"No JSONL input file found at {path}")
+
+    record_index = 0
+    for file_path in files:
+        with file_path.open(encoding="utf-8") as file:
+            for line_number, line in enumerate(file, start=1):
+                if not line.strip():
+                    continue
+                try:
+                    source = json.loads(line)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(f"Invalid JSON in {file_path}:{line_number}") from exc
+                messages = source.get("messages") or source.get("conversations")
+                if not isinstance(messages, list):
+                    continue
+                conversations: list[dict[str, str]] = []
+                for message in messages:
+                    if not isinstance(message, dict):
+                        continue
+                    role = message.get("role") or message.get("from")
+                    content = message.get("content") or message.get("value")
+                    if role in ("human", "user"):
+                        role = "user"
+                    elif role in ("gpt", "assistant"):
+                        role = "assistant"
+                    elif role in ("system", "developer"):
+                        role = "system"
+                    else:
+                        continue
+                    if not isinstance(content, str) or not content.strip():
+                        continue
+                    conversations.append({"role": role, "content": content})
+                if not any(message["role"] == "user" for message in conversations):
+                    continue
+                yield {
+                    "id": source.get("id", f"specdec_multilingual_prompt_{record_index:07d}"),
+                    "dataset": "nvidia/Speculative-Decoding-Multilingual-Prompt-v2",
+                    "source_file": file_path.name,
+                    "conversations": conversations,
+                }
+                record_index += 1
+
+
+def commit_staged_output(staging_dir: Path, output_dir: Path) -> None:
+    """Replace generated shards only after the complete staged build succeeds."""
+
+    for shard in output_dir.glob("train-*.jsonl"):
+        shard.unlink()
+    shutil.rmtree(output_dir / "media_cache", ignore_errors=True)
+    for path in staging_dir.iterdir():
+        path.replace(output_dir / path.name)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        choices=("pai_understanding", "vqa_v2", "specdec_multilingual_prompt"),
+        required=True,
+    )
+    parser.add_argument("--output_dir", required=True)
+    parser.add_argument("--dataset_dir", help="PAI-Bench-U metadata root")
+    parser.add_argument("--media_root", help="PAI-Bench-U video root")
+    parser.add_argument(
+        "--download",
+        action="store_true",
+        help="For pai_understanding, snapshot the complete Hugging Face dataset into --dataset_dir before sharding.",
+    )
+    parser.add_argument("--repo_id", default="shi-labs/physical-ai-bench-understanding")
+    parser.add_argument("--force_download", action="store_true")
+    parser.add_argument("--category", help="Optional exact PAI-Bench-U category")
+    parser.add_argument("--vqa_root", help="VQA v2 questions and annotations root")
+    parser.add_argument("--image_root", help="VQA v2 COCO image root")
+    parser.add_argument("--vqa_splits", default="train", help="Comma-separated VQA v2 splits")
+    parser.add_argument("--include_unannotated_vqa", action="store_true")
+    parser.add_argument("--include_all_answers", action="store_true")
+    parser.add_argument(
+        "--text_data",
+        help="JSONL file or directory for specdec_multilingual_prompt; rows must contain messages or conversations.",
+    )
+    parser.add_argument("--media_mode", choices=("native", "image_mosaic"), default="native")
+    parser.add_argument("--prompt_style", choices=("plain", "thinking"), default="thinking")
+    parser.add_argument("--num_frames", type=int, default=16)
+    parser.add_argument("--max_lines_per_shard", type=int, default=128)
+    parser.add_argument("--num_samples", type=int)
+    parser.add_argument("--start_index", type=int, default=0)
+    parser.add_argument(
+        "--shuffle_seed",
+        type=int,
+        help="Deterministically shuffle PAI or VQA records before applying --start_index/--num_samples.",
+    )
+    parser.add_argument("--missing_media", choices=("error", "skip"), default="error")
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+    if args.max_lines_per_shard <= 0 or args.num_frames <= 0 or args.start_index < 0:
+        raise ValueError(
+            "max_lines_per_shard/num_frames must be positive and start_index must be non-negative"
+        )
+
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    existing_shards = list(output_dir.glob("train-*.jsonl"))
+    if (existing_shards or (output_dir / "media_cache").exists()) and not args.overwrite:
+        raise FileExistsError(
+            f"Output already contains generated shards: {output_dir} (pass --overwrite to replace them)."
+        )
+
+    staging_dir = Path(
+        tempfile.mkdtemp(prefix=f".{output_dir.name}.staging-", dir=output_dir.parent)
+    )
+    output = None
+    written = missing = source_index = shard_index = lines = 0
+    try:
+        if args.dataset == "pai_understanding":
+            iterator = pai_records(args, staging_dir)
+        elif args.dataset == "vqa_v2":
+            iterator = vqa_records(args)
+        else:
+            iterator = text_records(args)
+        if args.shuffle_seed is not None:
+            if args.dataset == "specdec_multilingual_prompt":
+                raise ValueError(
+                    "--shuffle_seed is not supported for the full text dataset; it would materialize every prompt."
+                )
+            shuffled_records = list(iterator)
+            random.Random(args.shuffle_seed).shuffle(shuffled_records)
+            iterator = iter(shuffled_records)
+
+        for record in iterator:
+            if "_missing_media" in record:
+                missing += 1
+                print(f"Missing media: {record['_missing_media']}")
+                continue
+            if source_index < args.start_index:
+                source_index += 1
+                continue
+            source_index += 1
+            if args.num_samples is not None and written >= args.num_samples:
+                break
+            if output is None or lines >= args.max_lines_per_shard:
+                if output:
+                    output.close()
+                    shard_index += 1
+                output = (staging_dir / f"train-{shard_index:05d}-{shard_index:05d}.jsonl").open(
+                    "w", encoding="utf-8"
+                )
+                lines = 0
+            output.write(json.dumps(record, ensure_ascii=False) + "\n")
+            written += 1
+            lines += 1
+        if output:
+            output.close()
+            output = None
+        if missing and args.missing_media == "error":
+            raise FileNotFoundError(
+                f"{missing} referenced media files were missing; rerun with --missing_media skip "
+                "for an explicit partial dataset."
+            )
+        if not written:
+            raise RuntimeError("No usable records were written")
+        commit_staged_output(staging_dir, output_dir)
+    finally:
+        if output:
+            output.close()
+        shutil.rmtree(staging_dir, ignore_errors=True)
+
+    print(f"Wrote {written} {args.dataset} records to {shard_index + 1} shard(s) at {output_dir}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/speculative_decoding/recipes/run_multimodal_synthetic_generation.sh b/examples/speculative_decoding/recipes/run_multimodal_synthetic_generation.sh
new file mode 100755
index 00000000000..242b0bf5566
--- /dev/null
+++ b/examples/speculative_decoding/recipes/run_multimodal_synthetic_generation.sh
@@ -0,0 +1,92 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+if [ $# -lt 4 ] || [ $# -gt 5 ]; then
+    echo "Usage: MODEL_PATH=/path/to/model $0 <pai_understanding|vqa_v2|specdec_multilingual_prompt> [job_id] <start_shard> <jobs_per_node> <comma_separated_node_names>" >&2
+    echo "PAI env: DATASET_DIR=/dataset/root [MEDIA_ROOT=/video/root] [PREPARE_DOWNLOAD=1] [PAI_CATEGORY=...]" >&2
+    echo "VQA env: VQA_ROOT=/questions-and-annotations IMAGE_ROOT=/coco-images [VQA_SPLITS=train,val]" >&2
+    echo "Text env: TEXT_DATA=/path/to/sample-1K.jsonl [BACKEND=vllm|sglang]" >&2
+    exit 1
+fi
+
+DATASET=$1
+shift
+case "$DATASET" in pai_understanding|vqa_v2|specdec_multilingual_prompt) ;; *) echo "ERROR: unsupported dataset: $DATASET" >&2; exit 1;; esac
+if [ $# -eq 3 ]; then
+    JOB_ID=${SLURM_JOB_ID:-}
+    START_SHARD=$1
+    JOBS_PER_NODE=$2
+    NODE_NAMES=$3
+else
+    JOB_ID=$1
+    START_SHARD=$2
+    JOBS_PER_NODE=$3
+    NODE_NAMES=$4
+fi
+[ -n "${JOB_ID:-}" ] || { echo "ERROR: SLURM_JOB_ID is not set; pass job_id explicitly." >&2; exit 1; }
+[ -n "${MODEL_PATH:-}" ] || { echo "ERROR: MODEL_PATH must name the local VLM checkpoint." >&2; exit 1; }
+
+SPEC_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+SCRIPTS_PATH=${SCRIPTS_PATH:-$SPEC_ROOT}
+MEDIA_MODE=${MEDIA_MODE:-native}
+SHARD_PATH=${SHARD_PATH:-$SPEC_ROOT/CR3_data/${DATASET}_${MEDIA_MODE}_shards}
+OUTPUT_PATH=${OUTPUT_PATH:-$SPEC_ROOT/${DATASET}_synthetic_outputs}
+MAX_LINES_PER_SHARD=${MAX_LINES_PER_SHARD:-128}
+PROMPT_STYLE=${PROMPT_STYLE:-thinking}
+NUM_FRAMES=${NUM_FRAMES:-16}
+
+export API_MODE=${API_MODE:-native}
+export OVERWRITE_OUTPUT=${OVERWRITE_OUTPUT:-1}
+export SGLANG_TP_SIZE=${SGLANG_TP_SIZE:-1}
+export NUM_TEMPERATURES=${NUM_TEMPERATURES:-8}
+export NUM_THREADS=${NUM_THREADS:-8}
+
+prepare_args=(--dataset "$DATASET" --output_dir "$SHARD_PATH" --media_mode "$MEDIA_MODE" --prompt_style "$PROMPT_STYLE" --num_frames "$NUM_FRAMES" --max_lines_per_shard "$MAX_LINES_PER_SHARD" --missing_media "${MISSING_MEDIA:-error}")
+if [ "$DATASET" = "pai_understanding" ]; then
+    [ -n "${DATASET_DIR:-}" ] || { echo "ERROR: DATASET_DIR is required for pai_understanding." >&2; exit 1; }
+    MEDIA_ROOT=${MEDIA_ROOT:-$DATASET_DIR}
+    prepare_args+=(--dataset_dir "$DATASET_DIR" --media_root "$MEDIA_ROOT")
+    [ "${PREPARE_DOWNLOAD:-0}" = "1" ] && prepare_args+=(--download)
+    [ "${FORCE_DOWNLOAD:-0}" = "1" ] && prepare_args+=(--force_download)
+    [ -n "${PAI_REPO_ID:-}" ] && prepare_args+=(--repo_id "$PAI_REPO_ID")
+    [ -n "${PAI_CATEGORY:-}" ] && prepare_args+=(--category "$PAI_CATEGORY")
+    if [ "$MEDIA_MODE" = "image_mosaic" ]; then GENERATION_MEDIA_ROOT=$SHARD_PATH; else GENERATION_MEDIA_ROOT=$MEDIA_ROOT; fi
+elif [ "$DATASET" = "vqa_v2" ]; then
+    [ -n "${VQA_ROOT:-}" ] || { echo "ERROR: VQA_ROOT is required for vqa_v2." >&2; exit 1; }
+    [ -n "${IMAGE_ROOT:-}" ] || { echo "ERROR: IMAGE_ROOT is required for vqa_v2." >&2; exit 1; }
+    prepare_args+=(--vqa_root "$VQA_ROOT" --image_root "$IMAGE_ROOT" --vqa_splits "${VQA_SPLITS:-train}")
+    [ "${VQA_INCLUDE_UNANNOTATED:-0}" = "1" ] && prepare_args+=(--include_unannotated_vqa)
+    [ "${VQA_INCLUDE_ALL_ANSWERS:-0}" = "1" ] && prepare_args+=(--include_all_answers)
+    GENERATION_MEDIA_ROOT=$IMAGE_ROOT
+else
+    [ -n "${TEXT_DATA:-}" ] || { echo "ERROR: TEXT_DATA is required for specdec_multilingual_prompt." >&2; exit 1; }
+    prepare_args+=(--text_data "$TEXT_DATA")
+fi
+[ -n "${NUM_SAMPLES:-}" ] && prepare_args+=(--num_samples "$NUM_SAMPLES")
+[ -n "${START_INDEX:-}" ] && prepare_args+=(--start_index "$START_INDEX")
+[ "${PREPARE_OVERWRITE:-0}" = "1" ] && prepare_args+=(--overwrite)
+
+if [ "${PREPARE_SHARDS:-auto}" = "1" ] || ! compgen -G "$SHARD_PATH/train-*.jsonl" >/dev/null; then
+    python3 "$SCRIPTS_PATH/recipes/prepare_multimodal_synthetic_shards.py" "${prepare_args[@]}"
+fi
+
+if [ "$DATASET" = "specdec_multilingual_prompt" ]; then
+    bash "$SCRIPTS_PATH/distributed_generate/launch.sh" "$JOB_ID" "${BACKEND:-vllm}" "$MODEL_PATH" "$SHARD_PATH" "$OUTPUT_PATH" "$SCRIPTS_PATH" "$START_SHARD" "$JOBS_PER_NODE" "$NODE_NAMES" "${SYSTEM_PROMPT:-}"
+else
+    bash "$SCRIPTS_PATH/distributed_generate/launch_multimodal.sh" "$JOB_ID" sglang "$MODEL_PATH" "$SHARD_PATH" "$OUTPUT_PATH" "$SCRIPTS_PATH" "$START_SHARD" "$JOBS_PER_NODE" "$GENERATION_MEDIA_ROOT" "$NUM_FRAMES" "$NODE_NAMES"
+fi
diff --git a/examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb b/examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb
new file mode 100644
index 00000000000..6c020bcdf01
--- /dev/null
+++ b/examples/speculative_decoding/recipes/train_dflash_cosmos3_nano.ipynb
@@ -0,0 +1,732 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7fb27b941602401d91542211134fc71a",
+   "metadata": {},
+   "source": [
+    "# Training a DFlash Draft Model for Cosmos3 Nano\n",
+    "\n",
+    "This production runbook trains an online DFlash draft model with Cosmos3 Nano as the frozen vision-language target. DFlash learns a compact block-diffusion draft model.\n",
+    "\n",
+    "| Step | Description |\n",
+    "| :---: | :--- |\n",
+    "| 1 | Configure the training-data build |\n",
+    "| 2 | Build and merge the training data |\n",
+    "| 3 | Submit the eight-GPU training job |\n",
+    "| 4 | Export a checkpoint for deployment |\n",
+    "\n",
+    "> **Hardware** – the example requests eight GPUs by default.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "acae54e37e7d407bbb7b55eff062a284",
+   "metadata": {},
+   "source": [
+    "## Step 1 – Configure the Training-Data Build\n",
+    "\n",
+    "Build the merged Cosmos3 Nano DFlash JSONL from the sources below, then use this cell to set the model, data, and output locations shared by the remaining steps. Treat the provided source mix as a starting point and extend it with privacy-reviewed, task-relevant user data when available.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9a63283cbaf04dbcab1f6479b197f3a8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "REPO_ROOT = Path.cwd().resolve().parents[2]\n",
+    "SPEC_ROOT = REPO_ROOT / \"examples\" / \"speculative_decoding\"\n",
+    "if not (REPO_ROOT / \"pyproject.toml\").is_file():\n",
+    "    raise RuntimeError(\n",
+    "        f\"Run this notebook from examples/speculative_decoding/recipes; got {Path.cwd()}\"\n",
+    "    )\n",
+    "\n",
+    "TRAINING_DATA = Path(\n",
+    "    os.environ.get(\"TRAINING_DATA\", str(SPEC_ROOT / \"CR3_data\" / \"cosmos3_nano_dflash_train.jsonl\"))\n",
+    ").resolve()\n",
+    "MODEL_PATH = os.environ.get(\n",
+    "    \"MODEL_PATH\",\n",
+    "    \"/lustre/fsw/portfolios/coreai/users/skierat/ngc/ngc-cli/cosmos3-nano-reasoner_vbf16-final\",\n",
+    ")\n",
+    "OUTPUT_DIR = Path(\n",
+    "    os.environ.get(\"OUTPUT_DIR\", str(REPO_ROOT / \"ckpts\" / \"cosmos3-nano-dflash-first-run\"))\n",
+    ").resolve()\n",
+    "os.environ.update(\n",
+    "    REPO_ROOT=str(REPO_ROOT),\n",
+    "    TRAINING_DATA=str(TRAINING_DATA),\n",
+    "    MODEL_PATH=MODEL_PATH,\n",
+    "    OUTPUT_DIR=str(OUTPUT_DIR),\n",
+    ")\n",
+    "print(f\"Training data will be written to: {TRAINING_DATA}\")\n",
+    "if TRAINING_DATA.is_file() and TRAINING_DATA.stat().st_size:\n",
+    "    print(f\"Existing file size: {TRAINING_DATA.stat().st_size / 2**30:.1f} GiB\")\n",
+    "print(f\"Training output: {OUTPUT_DIR}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8edb47106e1a46a883d545849b8ab81b",
+   "metadata": {},
+   "source": [
+    "## Step 2 – Build the Training Data\n",
+    "\n",
+    "Prepare prompt shards, generate Cosmos3 Nano completions in a Slurm allocation, then merge the results into the `TRAINING_DATA` path configured above. The paths below default to this checkout rather than a prior run.\n",
+    "\n",
+    "These cells require a Jupyter environment with the example requirements installed. Run the configuration cell below before any data-preparation cell, and set `PLAIN_TEXT_INPUT` to curated Nemotron Chat or compatible user-data JSONL before merging.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10185d26023b46108eb7d9f57d49d2b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "\n",
+    "SPEC_ROOT = Path.cwd().resolve().parent\n",
+    "REPO_ROOT = SPEC_ROOT.parents[1]\n",
+    "DATA_ROOT = Path(os.environ.get(\"DATA_ROOT\", str(SPEC_ROOT / \"CR3_data\"))).resolve()\n",
+    "MODEL_PATH = os.environ.get(\n",
+    "    \"MODEL_PATH\",\n",
+    "    \"/lustre/fsw/portfolios/coreai/users/skierat/ngc/ngc-cli/cosmos3-nano-reasoner_vbf16-final\",\n",
+    ")\n",
+    "TRAINING_DATA = Path(\n",
+    "    os.environ.get(\"TRAINING_DATA\", str(DATA_ROOT / \"cosmos3_nano_dflash_train.jsonl\"))\n",
+    ").resolve()\n",
+    "OUTPUT_DIR = Path(\n",
+    "    os.environ.get(\"OUTPUT_DIR\", str(REPO_ROOT / \"ckpts\" / \"cosmos3-nano-dflash-first-run\"))\n",
+    ").resolve()\n",
+    "\n",
+    "PAI_SHARDS = Path(\n",
+    "    os.environ.get(\"PAI_SHARDS\", str(DATA_ROOT / \"pai_understanding_native_shards\"))\n",
+    ").resolve()\n",
+    "VQA_SHARD_PATH = Path(\n",
+    "    os.environ.get(\"VQA_SHARD_PATH\", str(DATA_ROOT / \"vqa_v2_train_shards\"))\n",
+    ").resolve()\n",
+    "TEXT_SHARDS = Path(\n",
+    "    os.environ.get(\"TEXT_SHARDS\", str(DATA_ROOT / \"specdec_multilingual_prompt_full_shards\"))\n",
+    ").resolve()\n",
+    "PAI_OUTPUT = Path(\n",
+    "    os.environ.get(\"PAI_OUTPUT\", str(SPEC_ROOT / \"pai_understanding_synthetic_outputs\"))\n",
+    ").resolve()\n",
+    "VQA_OUTPUT = Path(\n",
+    "    os.environ.get(\"VQA_OUTPUT\", str(SPEC_ROOT / \"vqa_v2_synthetic_outputs\"))\n",
+    ").resolve()\n",
+    "MULTILINGUAL_OUTPUT = Path(\n",
+    "    os.environ.get(\n",
+    "        \"MULTILINGUAL_OUTPUT\", str(SPEC_ROOT / \"specdec_multilingual_prompt_synthetic_outputs\")\n",
+    "    )\n",
+    ").resolve()\n",
+    "PLAIN_TEXT_INPUT = os.environ.get(\"PLAIN_TEXT_INPUT\", \"\")\n",
+    "\n",
+    "PAI_REVISION = os.environ.get(\"PAI_REVISION\", \"\")\n",
+    "TEXT_PROMPT_REVISION = os.environ.get(\"TEXT_PROMPT_REVISION\", \"\")\n",
+    "PAI_NUM_GENERATION_SHARDS = int(os.environ.get(\"PAI_NUM_GENERATION_SHARDS\", \"5\"))\n",
+    "PAI_LINES_PER_SHARD = int(os.environ.get(\"PAI_LINES_PER_SHARD\", \"128\"))\n",
+    "PAI_SHUFFLE_SEED = int(os.environ.get(\"PAI_SHUFFLE_SEED\", \"42\"))\n",
+    "VQA_NUM_SAMPLES = int(os.environ.get(\"VQA_NUM_SAMPLES\", \"20000\"))\n",
+    "VQA_SHUFFLE_SEED = int(os.environ.get(\"VQA_SHUFFLE_SEED\", \"42\"))\n",
+    "DEDUP_WORD_OVERLAP = float(os.environ.get(\"DEDUP_WORD_OVERLAP\", \"0.90\"))\n",
+    "DEDUP_CACHE_CONTEXTS = int(os.environ.get(\"DEDUP_CACHE_CONTEXTS\", \"25000\"))\n",
+    "MERGE_WORKERS = int(os.environ.get(\"MERGE_WORKERS\", \"8\"))\n",
+    "\n",
+    "if PAI_NUM_GENERATION_SHARDS <= 0 or PAI_LINES_PER_SHARD <= 0:\n",
+    "    raise ValueError(\"PAI_NUM_GENERATION_SHARDS and PAI_LINES_PER_SHARD must be positive.\")\n",
+    "if not 0 < DEDUP_WORD_OVERLAP <= 1:\n",
+    "    raise ValueError(\"DEDUP_WORD_OVERLAP must be in (0, 1].\")\n",
+    "if DEDUP_CACHE_CONTEXTS <= 0 or MERGE_WORKERS <= 0:\n",
+    "    raise ValueError(\"DEDUP_CACHE_CONTEXTS and MERGE_WORKERS must be positive.\")\n",
+    "\n",
+    "for name, value in {\n",
+    "    \"REPO_ROOT\": REPO_ROOT,\n",
+    "    \"DATA_ROOT\": DATA_ROOT,\n",
+    "    \"MODEL_PATH\": MODEL_PATH,\n",
+    "    \"TRAINING_DATA\": TRAINING_DATA,\n",
+    "    \"OUTPUT_DIR\": OUTPUT_DIR,\n",
+    "    \"PAI_SHARDS\": PAI_SHARDS,\n",
+    "    \"VQA_SHARD_PATH\": VQA_SHARD_PATH,\n",
+    "    \"TEXT_SHARDS\": TEXT_SHARDS,\n",
+    "    \"PAI_OUTPUT\": PAI_OUTPUT,\n",
+    "    \"VQA_OUTPUT\": VQA_OUTPUT,\n",
+    "    \"MULTILINGUAL_OUTPUT\": MULTILINGUAL_OUTPUT,\n",
+    "    \"PLAIN_TEXT_INPUT\": PLAIN_TEXT_INPUT,\n",
+    "    \"PAI_REVISION\": PAI_REVISION,\n",
+    "    \"TEXT_PROMPT_REVISION\": TEXT_PROMPT_REVISION,\n",
+    "    \"PYTHON_BIN\": sys.executable,\n",
+    "    \"PAI_NUM_GENERATION_SHARDS\": PAI_NUM_GENERATION_SHARDS,\n",
+    "    \"PAI_LINES_PER_SHARD\": PAI_LINES_PER_SHARD,\n",
+    "    \"PAI_SHUFFLE_SEED\": PAI_SHUFFLE_SEED,\n",
+    "    \"VQA_NUM_SAMPLES\": VQA_NUM_SAMPLES,\n",
+    "    \"VQA_SHUFFLE_SEED\": VQA_SHUFFLE_SEED,\n",
+    "    \"DEDUP_WORD_OVERLAP\": DEDUP_WORD_OVERLAP,\n",
+    "    \"DEDUP_CACHE_CONTEXTS\": DEDUP_CACHE_CONTEXTS,\n",
+    "    \"MERGE_WORKERS\": MERGE_WORKERS,\n",
+    "}.items():\n",
+    "    os.environ[name] = str(value)\n",
+    "\n",
+    "print(f\"Training data: {TRAINING_DATA}\")\n",
+    "print(f\"Training output: {OUTPUT_DIR}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8763a12b2bbd4a93a75aff182afb95dc",
+   "metadata": {},
+   "source": [
+    "### Sources and Workflow\n",
+    "\n",
+    "DFlash benefits most from target-model completions rather than a large collection of human-written answers. This build combines four complementary sources:\n",
+    "\n",
+    "1. **Representative usage samples** – PAI-Understanding video prompts.\n",
+    "2. **Visual reasoning data** – VQA v2 image-question prompts.\n",
+    "3. **Curated text** – either selected Nemotron Chat conversations or approved user data.\n",
+    "4. **High-quality text prompts** – multilingual prompt completions.\n",
+    "\n",
+    "The PAI and VQA sources retain their media; curated and multilingual text are text-only. The merge resolves media to absolute paths and conservatively deduplicates within matching prompt/media groups. This is why training uses `data.vlm_img_dir=/`.\n",
+    "\n",
+    "To create the Curated Nemotron Chat component, follow `recipes/train_eagle_head_cosmos_reason2.ipynb`: accept the Hugging Face licences for `nvidia/Nemotron-Post-Training-Dataset-v2`, then run `python ../prepare_input_conversations/add_nemotron_chat.py --mapping-file nemotron_mapping.bin` from that notebook. It writes `input_conversations/nemotron-chat.jsonl` with exactly 89,511 conversations. Set `PLAIN_TEXT_INPUT` to that file or to a privacy-reviewed JSONL of assistant-completed user conversations in the same `messages`/`conversations` format. To include both, combine them into one valid JSONL first.\n",
+    "\n",
+    "For DFlash, `training_seq_len` must be divisible by `dflash_block_size`; the production job uses `16384` and `8`.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7623eae2785240b9bd12b16a66d81610",
+   "metadata": {},
+   "source": [
+    "### Representative Usage: PAI-Understanding\n",
+    "\n",
+    "Use a modest representative sample. In production, replace this benchmark with a privacy-reviewed sample of real requests when available. The target model produces the assistant response used for training. The CPU-only setup cell creates a deterministic shuffled sample. Its default produces five shards; set `PAI_NUM_GENERATION_SHARDS` to a multiple of the number of allocated nodes.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7cdc8c89c7104fffa095e18ddfef8986",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "# CPU-only node: download PAI-Understanding and make native-video prompt shards.\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${PAI_NUM_GENERATION_SHARDS:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding}\n",
+    "PAI_SHARDS=${PAI_SHARDS:-$SPEC/CR3_data/pai_understanding_native_shards}\n",
+    "PAI_LINES_PER_SHARD=${PAI_LINES_PER_SHARD:-128}\n",
+    "PAI_NUM_SAMPLES=${PAI_NUM_SAMPLES:-$((PAI_NUM_GENERATION_SHARDS * PAI_LINES_PER_SHARD))}\n",
+    "PAI_REVISION_ARGS=()\n",
+    "[ -n \"${PAI_REVISION:-}\" ] && PAI_REVISION_ARGS+=(--revision \"$PAI_REVISION\")\n",
+    "HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache}\n",
+    "\n",
+    "# Invoke the CLI through the active kernel interpreter to avoid a stale `hf` executable on PATH.\n",
+    "\"$PYTHON_BIN\" -m huggingface_hub.cli.hf --version\n",
+    "# If networking on the CPU-only node stalls, retry this command with HF_HUB_DISABLE_XET=1.\n",
+    "\"$PYTHON_BIN\" -m huggingface_hub.cli.hf download shi-labs/physical-ai-bench-understanding \\\n",
+    "  --repo-type dataset \\\n",
+    "  --local-dir \"$PAI_ROOT\" \\\n",
+    "  \"${PAI_REVISION_ARGS[@]}\" \\\n",
+    "  --max-workers 8\n",
+    "\n",
+    "HF_DATASETS_CACHE=\"$HF_DATASETS_CACHE\" \"$PYTHON_BIN\" \"$SPEC/recipes/prepare_multimodal_synthetic_shards.py\" \\\n",
+    "  --dataset pai_understanding \\\n",
+    "  --dataset_dir \"$PAI_ROOT\" \\\n",
+    "  --media_root \"$PAI_ROOT\" \\\n",
+    "  --output_dir \"$PAI_SHARDS\" \\\n",
+    "  --max_lines_per_shard \"$PAI_LINES_PER_SHARD\" \\\n",
+    "  --num_samples \"$PAI_NUM_SAMPLES\" \\\n",
+    "  --shuffle_seed \"$PAI_SHUFFLE_SEED\" \\\n",
+    "  --overwrite\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b118ea5561624da68c537baed56e602f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "# Compute node: run a representative PAI shard slice in an existing Slurm GPU allocation.\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${MODEL_PATH:?Set MODEL_PATH and run Step 2 before this compute-node cell.}\"\n",
+    "PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding}\n",
+    "export MODEL_PATH\n",
+    "export DATASET_DIR=\"$PAI_ROOT\"\n",
+    "export MEDIA_ROOT=\"$PAI_ROOT\"\n",
+    "export SHARD_PATH=${PAI_SHARDS:-$SPEC/CR3_data/pai_understanding_native_shards}\n",
+    "export PREPARE_SHARDS=0\nexport OUTPUT_PATH=\"${PAI_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "NODE_NAMES=\"$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | paste -sd, -)\"\n",
+    ": \"${PAI_NUM_GENERATION_SHARDS:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "PAI_START_SHARD=${PAI_START_SHARD:-0}\n",
+    "PAI_NUM_AVAILABLE_SHARDS=$(find \"$SHARD_PATH\" -maxdepth 1 -type f -name 'train-*.jsonl' | wc -l)\n",
+    "PAI_NUM_NODES=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | wc -l)\n",
+    "[ \"$PAI_NUM_AVAILABLE_SHARDS\" -gt 0 ] || { echo \"No PAI shards found in $SHARD_PATH\" >&2; exit 1; }\n",
+    "[ \"$PAI_NUM_NODES\" -gt 0 ] || { echo \"No allocated nodes found\" >&2; exit 1; }\n",
+    "(( PAI_START_SHARD + PAI_NUM_GENERATION_SHARDS <= PAI_NUM_AVAILABLE_SHARDS )) || { echo \"Requested PAI shard range exceeds $PAI_NUM_AVAILABLE_SHARDS available shards\" >&2; exit 1; }\n",
+    "(( PAI_NUM_GENERATION_SHARDS % PAI_NUM_NODES == 0 )) || { echo \"$PAI_NUM_GENERATION_SHARDS selected PAI shards cannot be distributed evenly over $PAI_NUM_NODES nodes\" >&2; exit 1; }\n",
+    "PAI_SHARDS_PER_NODE=$(( PAI_NUM_GENERATION_SHARDS / PAI_NUM_NODES ))\n",
+    "\n",
+    "\"$SPEC/recipes/run_multimodal_synthetic_generation.sh\" \\\n",
+    "  pai_understanding \"$SLURM_JOB_ID\" \"$PAI_START_SHARD\" \"$PAI_SHARDS_PER_NODE\" \"$NODE_NAMES\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "938c804e27f84196a10c8828c723f798",
+   "metadata": {},
+   "source": [
+    "### Visual Reasoning: VQA v2\n",
+    "\n",
+    "This component is larger than the representative sample. The default creates a deterministic sample of 20,000 train-split prompts; increase `VQA_NUM_SAMPLES` for more visual coverage.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "504fb2a444614c0babb325280ed9130a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "# CPU-only node: fetch the VQA v2 train questions, annotations, and COCO images.\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2}\n",
+    "IMAGE_ROOT=\"$VQA_ROOT/images\"\n",
+    ": \"${VQA_NUM_SAMPLES:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "VQA_SHARD_PATH=${VQA_SHARD_PATH:-$SPEC/CR3_data/vqa_v2_train_shards}\n",
+    "HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache}\n",
+    "VQA_QUESTIONS_URL=${VQA_QUESTIONS_URL:-https://cvmlp.s3.amazonaws.com/vqa/mscoco/vqa/v2_Questions_Train_mscoco.zip}\n",
+    "VQA_ANNOTATIONS_URL=${VQA_ANNOTATIONS_URL:-https://cvmlp.s3.amazonaws.com/vqa/mscoco/vqa/v2_Annotations_Train_mscoco.zip}\n",
+    "COCO_TRAIN_URL=${COCO_TRAIN_URL:-https://images.cocodataset.org/zips/train2014.zip}\n",
+    "\n",
+    "extract_zip() {\n",
+    "  local archive=$1 destination=$2 marker=$3\n",
+    "  if [ -f \"$marker\" ]; then\n",
+    "    echo \"Already extracted: $archive\"\n",
+    "    return\n",
+    "  fi\n",
+    "  if command -v unzip >/dev/null 2>&1; then\n",
+    "    unzip -n \"$archive\" -d \"$destination\"\n",
+    "  else\n",
+    "    # Minimal CPU-only node images may omit `unzip`; Python provides a compatible fallback.\n",
+    "    \"$PYTHON_BIN\" -m zipfile -e \"$archive\" \"$destination\"\n",
+    "  fi\n",
+    "  touch \"$marker\"\n",
+    "}\n",
+    "\n",
+    "mkdir -p \"$VQA_ROOT\" \"$IMAGE_ROOT\"\n",
+    "curl --proto '=https' -L --fail --retry 5 -C - -o \"$VQA_ROOT/v2_Questions_Train_mscoco.zip\" \\\n",
+    "  \"$VQA_QUESTIONS_URL\"\n",
+    "curl --proto '=https' -L --fail --retry 5 -C - -o \"$VQA_ROOT/v2_Annotations_Train_mscoco.zip\" \\\n",
+    "  \"$VQA_ANNOTATIONS_URL\"\n",
+    "curl --proto '=https' -L --fail --retry 5 -C - -o \"$IMAGE_ROOT/train2014.zip\" \\\n",
+    "  \"$COCO_TRAIN_URL\"\n",
+    "extract_zip \"$VQA_ROOT/v2_Questions_Train_mscoco.zip\" \"$VQA_ROOT\" \"$VQA_ROOT/.questions.extracted\"\n",
+    "extract_zip \"$VQA_ROOT/v2_Annotations_Train_mscoco.zip\" \"$VQA_ROOT\" \"$VQA_ROOT/.annotations.extracted\"\n",
+    "extract_zip \"$IMAGE_ROOT/train2014.zip\" \"$IMAGE_ROOT\" \"$IMAGE_ROOT/.train2014.extracted\"\n",
+    "\n",
+    "HF_DATASETS_CACHE=\"$HF_DATASETS_CACHE\" \"$PYTHON_BIN\" \"$SPEC/recipes/prepare_multimodal_synthetic_shards.py\" \\\n",
+    "  --dataset vqa_v2 \\\n",
+    "  --vqa_root \"$VQA_ROOT\" \\\n",
+    "  --image_root \"$IMAGE_ROOT\" \\\n",
+    "  --vqa_splits train \\\n",
+    "  --num_samples \"$VQA_NUM_SAMPLES\" \\\n",
+    "  --shuffle_seed \"$VQA_SHUFFLE_SEED\" \\\n",
+    "  --output_dir \"$VQA_SHARD_PATH\" \\\n",
+    "  --overwrite\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59bbdb311c014d738909a11f9e486628",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "# Compute node: distribute all remaining prepared VQA shards across the allocated nodes.\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${MODEL_PATH:?Set MODEL_PATH and run Step 2 before this compute-node cell.}\"\n",
+    "export MODEL_PATH\n",
+    "VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2}\n",
+    "IMAGE_ROOT=\"$VQA_ROOT/images\"\n",
+    "export VQA_ROOT IMAGE_ROOT\n",
+    "export SHARD_PATH=${VQA_SHARD_PATH:-$SPEC/CR3_data/vqa_v2_train_shards}\n",
+    "export PREPARE_SHARDS=0\nexport OUTPUT_PATH=\"${VQA_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "NODE_NAMES=\"$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | paste -sd, -)\"\n",
+    "VQA_NUM_SHARDS=$(find \"$SHARD_PATH\" -maxdepth 1 -type f -name '*.jsonl' | wc -l)\n",
+    "VQA_NUM_NODES=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | wc -l)\n",
+    "VQA_START_SHARD=${VQA_START_SHARD:-0}\n",
+    "[ \"$VQA_NUM_SHARDS\" -gt 0 ] || { echo \"No VQA shards found in $SHARD_PATH\" >&2; exit 1; }\n",
+    "[ \"$VQA_NUM_NODES\" -gt 0 ] || { echo \"No allocated nodes found\" >&2; exit 1; }\n",
+    "[[ \"$VQA_START_SHARD\" =~ ^[0-9]+$ ]] || { echo \"VQA_START_SHARD must be a non-negative integer: $VQA_START_SHARD\" >&2; exit 1; }\n",
+    "(( VQA_START_SHARD >= 0 && VQA_START_SHARD < VQA_NUM_SHARDS )) || { echo \"VQA_START_SHARD=$VQA_START_SHARD is outside 0 through $((VQA_NUM_SHARDS - 1))\" >&2; exit 1; }\n",
+    "VQA_REMAINING_SHARDS=$(( VQA_NUM_SHARDS - VQA_START_SHARD ))\n",
+    "# The multimodal worker skips its small rounded-up tail, so all remaining shards are covered.\n",
+    "VQA_SHARDS_PER_NODE=$(( (VQA_REMAINING_SHARDS + VQA_NUM_NODES - 1) / VQA_NUM_NODES ))\n",
+    "\n",
+    "\"$SPEC/recipes/run_multimodal_synthetic_generation.sh\" \\\n",
+    "  vqa_v2 \"$SLURM_JOB_ID\" \"$VQA_START_SHARD\" \"$VQA_SHARDS_PER_NODE\" \"$NODE_NAMES\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b43b363d81ae4b689946ece5c682cd59",
+   "metadata": {},
+   "source": [
+    "### High-Quality Text: Synthetic Multilingual Prompts\n",
+    "\n",
+    "Use a large text-only component unless the deployment is deliberately narrow. `nvidia/Speculative-Decoding-Multilingual-Prompt-v2` provides prompts and Cosmos3 Nano generates the assistant answers. Eight temperatures increase response diversity.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c3933fab20d04ec698c2621248eb3be0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "# CPU-only node: download the text prompts and create full prompt shards.\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "TEXT_ROOT=${TEXT_ROOT:-$DATA_ROOT/specdec_multilingual_prompt}\n",
+    "TEXT_SHARDS=${TEXT_SHARDS:-$SPEC/CR3_data/specdec_multilingual_prompt_full_shards}\n",
+    "TEXT_REVISION_ARGS=()\n",
+    "[ -n \"${TEXT_PROMPT_REVISION:-}\" ] && TEXT_REVISION_ARGS+=(--revision \"$TEXT_PROMPT_REVISION\")\n",
+    "HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-$DATA_ROOT/.hf_datasets_cache}\n",
+    "\n",
+    "# Only default.jsonl is used below; the repository's convenience samples are intentionally not downloaded.\n",
+    "# If networking on the CPU-only node stalls, retry this command with HF_HUB_DISABLE_XET=1.\n",
+    "\"$PYTHON_BIN\" -m huggingface_hub.cli.hf download nvidia/Speculative-Decoding-Multilingual-Prompt-v2 default.jsonl \\\n",
+    "  --repo-type dataset \\\n",
+    "  --local-dir \"$TEXT_ROOT\" \\\n",
+    "  \"${TEXT_REVISION_ARGS[@]}\" \\\n",
+    "  --max-workers 8\n",
+    "\n",
+    "HF_DATASETS_CACHE=\"$HF_DATASETS_CACHE\" \"$PYTHON_BIN\" \"$SPEC/recipes/prepare_multimodal_synthetic_shards.py\" \\\n",
+    "  --dataset specdec_multilingual_prompt \\\n",
+    "  --text_data \"$TEXT_ROOT\" \\\n",
+    "  --max_lines_per_shard 1024 \\\n",
+    "  --overwrite \\\n",
+    "  --output_dir \"$TEXT_SHARDS\"\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4dd4641cc4064e0191573fe9c69df29b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "# Compute node: use the allocation dynamically; no fixed node list is needed.\n",
+    "# START_SHARD=1059 is a resume example. Set it to 0 for a new full run.\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${MODEL_PATH:?Set MODEL_PATH and run Step 2 before this compute-node cell.}\"\n",
+    "export MODEL_PATH\n",
+    "TEXT_ROOT=${TEXT_ROOT:-$DATA_ROOT/specdec_multilingual_prompt}\n",
+    "export TEXT_DATA=\"$TEXT_ROOT\"\n",
+    "export SHARD_PATH=${TEXT_SHARDS:-$SPEC/CR3_data/specdec_multilingual_prompt_full_shards}\n",
+    "export OUTPUT_PATH=\"${MULTILINGUAL_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "export PREPARE_SHARDS=0\n",
+    "export BACKEND=vllm\n",
+    "export CONTAINER_IMAGE=vllm/vllm-openai:v0.24.0\n",
+    "export SGLANG_TP_SIZE=1\n",
+    "export NUM_TEMPERATURES=8\n",
+    "NODE_NAMES=\"$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | paste -sd, -)\"\n",
+    "START_SHARD=${START_SHARD:-0}\n",
+    "TEXT_NUM_SHARDS=$(find \"$SHARD_PATH\" -maxdepth 1 -type f -name 'train-*.jsonl' | wc -l)\n",
+    "TEXT_NUM_NODES=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | wc -l)\n",
+    "[ \"$TEXT_NUM_SHARDS\" -gt 0 ] || { echo \"No text shards found in $SHARD_PATH\" >&2; exit 1; }\n",
+    "[ \"$TEXT_NUM_NODES\" -gt 0 ] || { echo \"No allocated nodes found\" >&2; exit 1; }\n",
+    "[ \"$START_SHARD\" -lt \"$TEXT_NUM_SHARDS\" ] || { echo \"START_SHARD=$START_SHARD is beyond $TEXT_NUM_SHARDS text shards\" >&2; exit 1; }\n",
+    "TEXT_REMAINING_SHARDS=$(( TEXT_NUM_SHARDS - START_SHARD ))\n",
+    "# The text worker skips its small rounded-up tail, so all remaining shards are covered.\n",
+    "JOBS_PER_NODE=${JOBS_PER_NODE:-$(( (TEXT_REMAINING_SHARDS + TEXT_NUM_NODES - 1) / TEXT_NUM_NODES ))}\n",
+    "\n",
+    "\"$SPEC/recipes/run_multimodal_synthetic_generation.sh\" \\\n",
+    "  specdec_multilingual_prompt \"$SLURM_JOB_ID\" \"$START_SHARD\" \"$JOBS_PER_NODE\" \"$NODE_NAMES\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8309879909854d7188b41380fd92a7c3",
+   "metadata": {},
+   "source": [
+    "### Merge, Resolve Media, and Deduplicate\n",
+    "\n",
+    "Run this cell after the PAI, VQA v2, and multilingual synthesis jobs finish and `PLAIN_TEXT_INPUT` is set. It merges all four required sources into one atomic JSONL with absolute media paths. `MERGE_WORKERS` controls parallel merge workers; temperature variants for one generated shard always stay together, and deduplication is local to each worker partition. The merger never overwrites an existing final JSONL; choose a new `TRAINING_DATA` path when intentionally rebuilding.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5d1c9edacf9a4d2fa6b2e44b73dd4c20",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "set -euo pipefail\n",
+    "SPEC=\"$(cd .. && pwd)\"\n",
+    ": \"${DATA_ROOT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${PYTHON_BIN:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "PAI_ROOT=${PAI_ROOT:-$DATA_ROOT/pai_understanding}\n",
+    "VQA_ROOT=${VQA_ROOT:-$DATA_ROOT/vqa_v2}\n",
+    ": \"${PAI_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${VQA_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${PLAIN_TEXT_INPUT:?Set PLAIN_TEXT_INPUT in the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${MULTILINGUAL_OUTPUT:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${TRAINING_DATA:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${DEDUP_WORD_OVERLAP:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${DEDUP_CACHE_CONTEXTS:?Run the Step 2 configuration cell before this cell.}\"\n",
+    ": \"${MERGE_WORKERS:?Run the Step 2 configuration cell before this cell.}\"\n",
+    "\n",
+    "\"$PYTHON_BIN\" \"$SPEC/recipes/merge_dflash_datasets.py\" \\\n",
+    "  --source \"pai_understanding=$PAI_OUTPUT\" \\\n",
+    "  --source \"vqa_v2=$VQA_OUTPUT\" \\\n",
+    "  --source \"curated_text=$PLAIN_TEXT_INPUT\" \\\n",
+    "  --source \"specdec_multilingual_prompt=$MULTILINGUAL_OUTPUT\" \\\n",
+    "  --media-root \"pai_understanding=$PAI_ROOT\" \\\n",
+    "  --media-root \"vqa_v2=$VQA_ROOT/images\" \\\n",
+    "  --output \"$TRAINING_DATA\" \\\n",
+    "  --jobs \"$MERGE_WORKERS\" \\\n",
+    "  --word-overlap \"$DEDUP_WORD_OVERLAP\" \\\n",
+    "  --cache-contexts \"$DEDUP_CACHE_CONTEXTS\"\n",
+    "\n",
+    "wc -l \"$TRAINING_DATA\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb1e1581032b452c9409d6c6813c49d1",
+   "metadata": {},
+   "source": [
+    "## Step 3 – Submit the DFlash Training Job\n",
+    "\n",
+    "This self-contained cell submits the training job directly to Slurm. Its resource request, container setup, and training command live together here; no separate `job_*.sh` file is required. The example requests eight GPUs for four hours and saves every 1,000 updates. For a single-GPU run, request one GPU and raise `training.gradient_accumulation_steps` from `2` to `16` to keep the same global batch size.\n",
+    "\n",
+    "The defaults use the JSONL built in Step 2 and write to `ckpts/cosmos3-nano-dflash-first-run`. Set `MODEL_PATH`, `TRAINING_DATA`, or `OUTPUT_DIR` before submitting to override them. Remote code is disabled by default; set `TRUST_REMOTE_CODE=true` only for a model checkpoint you trust and that requires custom code.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "379cbbc1e968416e875cc15c1202d7eb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "set -euo pipefail\n",
+    "REPO_ROOT=\"$(cd ../../.. && pwd)\"\n",
+    "\n",
+    ": \"${MODEL_PATH:?Run Step 1 before this cell.}\"\n",
+    ": \"${TRAINING_DATA:?Run Steps 1 and 2 before this cell.}\"\n",
+    ": \"${OUTPUT_DIR:?Run Step 1 before this cell.}\"\n",
+    "[ -d \"$MODEL_PATH\" ] || { echo \"Missing MODEL_PATH: $MODEL_PATH\" >&2; exit 1; }\n",
+    "[ -s \"$TRAINING_DATA\" ] || { echo \"Missing training JSONL: $TRAINING_DATA\" >&2; exit 1; }\n",
+    "\n",
+    "export TRUST_REMOTE_CODE=\"${TRUST_REMOTE_CODE:-false}\"\n",
+    "export REPO_ROOT MODEL_PATH TRAINING_DATA OUTPUT_DIR TRUST_REMOTE_CODE\n",
+    "export HF_HOME=\"${HF_HOME:-/lustre/fsw/portfolios/coreai/users/skierat/cache_hf}\"\n",
+    "\n",
+    "# Submit this heredoc directly; no batch-script file is created.\n",
+    "sbatch --export=ALL <<'SBATCH'\n",
+    "#!/usr/bin/env bash\n",
+    "#SBATCH -p interactive\n",
+    "#SBATCH --nodes=1\n",
+    "#SBATCH --ntasks-per-node=1\n",
+    "#SBATCH --gpus-per-node=8\n",
+    "#SBATCH --time=04:00:00\n",
+    "#SBATCH --account=coreai_tritoninference_triton3\n",
+    "#SBATCH --job-name=cosmos3-dflash\n",
+    "#SBATCH --signal=TERM@120\n",
+    "\n",
+    "set -euo pipefail\n",
+    "trap \"kill -- -$$\" TERM INT HUP\n",
+    "\n",
+    "export NUM_GPUS=8\n",
+    "export PIP_CACHE_DIR=\"${PIP_CACHE_DIR:-/tmp/pip-cache-${SLURM_JOB_ID}}\"\n",
+    "export TRITON_CACHE_DIR=\"${TRITON_CACHE_DIR:-/tmp/triton-cache-${SLURM_JOB_ID}}\"\n",
+    "export VLM_MIN_PIXELS=\"${VLM_MIN_PIXELS:-50176}\"\n",
+    "export VLM_MAX_PIXELS=\"${VLM_MAX_PIXELS:-802816}\"\n",
+    "export VLM_MAX_ASSISTANT_TOKENS=\"${VLM_MAX_ASSISTANT_TOKENS:-2048}\"\n",
+    "export VLM_MAX_PROMPT_TOKENS=\"${VLM_MAX_PROMPT_TOKENS:-8192}\"\n",
+    "export VLM_VIDEO_MIN_PIXELS=\"${VLM_VIDEO_MIN_PIXELS:-100352}\"\n",
+    "export VLM_VIDEO_MAX_PIXELS=\"${VLM_VIDEO_MAX_PIXELS:-2097152}\"\n",
+    "\n",
+    "echo \"Model: ${MODEL_PATH}\"\n",
+    "echo \"Training data: ${TRAINING_DATA}\"\n",
+    "echo \"Output: ${OUTPUT_DIR}\"\n",
+    "\n",
+    "srun \\\n",
+    "  --ntasks=1 \\\n",
+    "  --gpus-per-task=\"${NUM_GPUS}\" \\\n",
+    "  --kill-on-bad-exit=1 \\\n",
+    "  --container-image=/lustre/fsw/portfolios/coreai/users/skierat/containers/recent_vllm.sqfs \\\n",
+    "  --container-workdir=/ \\\n",
+    "  --container-mounts=/lustre:/lustre \\\n",
+    "  bash -lc '\n",
+    "    set -euo pipefail\n",
+    "    export WANDB_MODE=disabled\n",
+    "    export TOKENIZERS_PARALLELISM=false\n",
+    "    unset NCCL_ASYNC_ERROR_HANDLING\n",
+    "    export TORCH_NCCL_ASYNC_ERROR_HANDLING=1\n",
+    "    export PYTHONPATH=\"${REPO_ROOT}:${PYTHONPATH:-}\"\n",
+    "    mkdir -p \"${PIP_CACHE_DIR}\" \"${TRITON_CACHE_DIR}\"\n",
+    "\n",
+    "    cd \"${REPO_ROOT}\"\n",
+    "    python3 -m pip install -r examples/speculative_decoding/requirements.txt\n",
+    "    python3 -m pip install torchcodec\n",
+    "    python3 -m pip install -e \".[hf]\"\n",
+    "\n",
+    "    python3 -m torch.distributed.run \\\n",
+    "      --standalone \\\n",
+    "      --nproc_per_node=\"${NUM_GPUS}\" \\\n",
+    "      examples/speculative_decoding/main.py \\\n",
+    "      --config modelopt_recipes/general/speculative_decoding/dflash.yaml \\\n",
+    "      model.model_name_or_path=\"${MODEL_PATH}\" \\\n",
+    "      model.trust_remote_code=\"${TRUST_REMOTE_CODE:-false}\" \\\n",
+    "      data.data_path=\"${TRAINING_DATA}\" \\\n",
+    "      data.vlm_processor=\"${MODEL_PATH}\" \\\n",
+    "      data.vlm_img_dir=/ \\\n",
+    "      training.output_dir=\"${OUTPUT_DIR}\" \\\n",
+    "      training.num_train_epochs=25 \\\n",
+    "      training.per_device_train_batch_size=1 \\\n",
+    "      training.gradient_accumulation_steps=2 \\\n",
+    "      training.training_seq_len=16384 \\\n",
+    "      training.answer_only_loss=true \\\n",
+    "      training.save_steps=1000 \\\n",
+    "      training.save_total_limit=10 \\\n",
+    "      training.logging_steps=10 \\\n",
+    "      training.dataloader_num_workers=2 \\\n",
+    "      training.dataloader_prefetch_factor=2 \\\n",
+    "      training.ddp_find_unused_parameters=false \\\n",
+    "      training.report_to=none \\\n",
+    "      dflash.dflash_block_size=8 \\\n",
+    "      dflash.dflash_num_anchors=128 \\\n",
+    "      dflash.dflash_loss_objective=decay \\\n",
+    "      dflash.dflash_loss_decay_factor=4 \\\n",
+    "      dflash.dflash_architecture_config.num_hidden_layers=5 \\\n",
+    "      dflash.dflash_architecture_config.num_attention_heads=32 \\\n",
+    "      dflash.dflash_architecture_config.num_key_value_heads=8 \\\n",
+    "      dflash.dflash_architecture_config.head_dim=128 \\\n",
+    "      dflash.dflash_architecture_config.intermediate_size=12288 \\\n",
+    "      dflash.dflash_architecture_config.max_position_embeddings=262144 \\\n",
+    "      dflash.dflash_architecture_config.rms_norm_eps=1e-06 \\\n",
+    "      dflash.dflash_architecture_config.rope_theta=5000000 \\\n",
+    "      dflash.dflash_mask_token_id=151669\n",
+    "  '\n",
+    "SBATCH\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "277c27b1587741f2af2001be3712ef0d",
+   "metadata": {},
+   "source": [
+    "## Step 4 – Export the DFlash Checkpoint\n",
+    "\n",
+    "Export one completed checkpoint to the compact DFlash Hugging Face format consumed by vLLM. Run this cell in the same container or Python environment used for training (or another environment with ModelOpt and the example dependencies installed).\n",
+    "\n",
+    "Set `CKPT_DIR` to a saved checkpoint directory such as `checkpoint-180000`, not the top-level run directory. The export reads `modelopt_state.pth` from that directory and writes a new deployment directory containing `config.json` and `model.safetensors`. The cell refuses to overwrite an existing export.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "db7b79bc585a40fcaf58bf750017e135",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "set -euo pipefail\n",
+    "REPO_ROOT=\"$(cd ../../.. && pwd)\"\n",
+    "RUN_DIR=${OUTPUT_DIR:-$REPO_ROOT/ckpts/cosmos3-nano-dflash-first-run}\n",
+    "CKPT_DIR=${CKPT_DIR:?Set CKPT_DIR to a saved checkpoint directory, for example $RUN_DIR/checkpoint-180000}\n",
+    "EXPORT_PATH=${EXPORT_PATH:-$RUN_DIR/export-$(basename \"$CKPT_DIR\" | sed 's/^checkpoint-//')}\n",
+    "\n",
+    "[ -f \"$CKPT_DIR/modelopt_state.pth\" ] || { echo \"Missing DFlash checkpoint: $CKPT_DIR\" >&2; exit 1; }\n",
+    "[ ! -e \"$EXPORT_PATH\" ] || { echo \"Export path already exists: $EXPORT_PATH\" >&2; exit 1; }\n",
+    "TRUST_REMOTE_CODE=\"${TRUST_REMOTE_CODE:-false}\"\n",
+    "EXPORT_ARGS=()\n",
+    "if [[ \"${TRUST_REMOTE_CODE,,}\" == \"true\" || \"$TRUST_REMOTE_CODE\" == \"1\" ]]; then\n",
+    "  EXPORT_ARGS+=(--trust_remote_code)\n",
+    "fi\n",
+    "python3 \"$REPO_ROOT/examples/speculative_decoding/scripts/export_hf_checkpoint.py\" \\\n",
+    "  --model_path \"$CKPT_DIR\" \\\n",
+    "  --export_path \"$EXPORT_PATH\" \\\n",
+    "  \"${EXPORT_ARGS[@]}\"\n",
+    "\n",
+    "[ -s \"$EXPORT_PATH/config.json\" ] && [ -s \"$EXPORT_PATH/model.safetensors\" ] \\\n",
+    "  || { echo \"Export is incomplete: $EXPORT_PATH\" >&2; exit 1; }\n",
+    "echo \"DFlash deployment checkpoint: $EXPORT_PATH\"\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "916684f9a58a4a2aa5f864670399430d",
+   "metadata": {},
+   "source": [
+    "## Deployment\n",
+    "\n",
+    "Serve the original Cosmos3 Nano target with the exported DFlash draft. A DFlash block size of eight yields seven speculative tokens: the remaining position is the context/bonus token.\n",
+    "\n",
+    "For a first deployment, use a vLLM environment compatible with this recipe. The command below is a single-GPU smoke test: it binds only to localhost and caps the context at 4096 tokens to keep startup and memory use bounded. Raise or remove `--max-model-len` only after sizing the desired context and concurrency. Add `--trust-remote-code` only when serving a model checkpoint you trust and that requires custom code.\n",
+    "\n",
+    "```bash\n",
+    "SERVED_MODEL_NAME=cosmos3-nano-dflash\n",
+    "vllm serve \"$MODEL_PATH\" \\\n",
+    "  --host 127.0.0.1 --port 8000 \\\n",
+    "  --max-model-len 4096 \\\n",
+    "  --served-model-name \"$SERVED_MODEL_NAME\" \\\n",
+    "  --speculative-config \"{\\\"method\\\":\\\"dflash\\\",\\\"model\\\":\\\"$EXPORT_PATH\\\",\\\"num_speculative_tokens\\\":7}\"\n",
+    "\n",
+    "# In another terminal, verify the server and make one request.\n",
+    "curl -fsS http://127.0.0.1:8000/health\n",
+    "curl -fsS http://127.0.0.1:8000/v1/models\n",
+    "curl -fsS http://127.0.0.1:8000/v1/chat/completions \\\n",
+    "  -H 'Content-Type: application/json' \\\n",
+    "  -d \"{\\\"model\\\":\\\"$SERVED_MODEL_NAME\\\",\\\"messages\\\":[{\\\"role\\\":\\\"user\\\",\\\"content\\\":\\\"Reply with exactly: DFlash deployment smoke test passed.\\\"}],\\\"temperature\\\":0,\\\"max_tokens\\\":16}\"\n",
+    "\n",
+    "# Stop the foreground server with Ctrl-C when the smoke test completes.\n",
+    "```\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/speculative_decoding/scripts/server_generate.py b/examples/speculative_decoding/scripts/server_generate.py
index 0fb71a0a0a1..a0516bc3922 100644
--- a/examples/speculative_decoding/scripts/server_generate.py
+++ b/examples/speculative_decoding/scripts/server_generate.py
@@ -75,6 +75,7 @@ def generate_data(messages, idx, system_prompt):
 
         if args.chat:
             output_messages = []
+            truncated = False
 
             if system_prompt and len(messages) > 0:
                 system_message = {"role": "system", "content": system_prompt}
@@ -106,15 +107,17 @@ def generate_data(messages, idx, system_prompt):
                         max_tokens=args.max_tokens,
                         temperature=args.temperature,
                     )
-                    if response.choices[0].finish_reason == "length":
-                        break
-                    response = response.choices[0].message.content.strip()
+                    choice = response.choices[0]
+                    response = (choice.message.content or "").strip()
                     output_messages.append(
                         {
                             "role": "assistant",
                             "content": response,
                         }
                     )
+                    if choice.finish_reason == "length":
+                        truncated = True
+                        break
                 except Exception as e:
                     print(e)
                     break
@@ -124,6 +127,8 @@ def generate_data(messages, idx, system_prompt):
                 to_write = {"conversation_id": idx}
             else:
                 to_write = {"conversation_id": idx, "conversations": output_messages}
+            if truncated:
+                to_write["truncated"] = True
             with open(args.output_path, "a") as f:
                 # write in share gpt format
                 f.write(json.dumps(to_write) + "\n")
diff --git a/modelopt/torch/utils/plugins/transformers_dataset.py b/modelopt/torch/utils/plugins/transformers_dataset.py
index 97ae4ea2d14..6a6a7a4e249 100644
--- a/modelopt/torch/utils/plugins/transformers_dataset.py
+++ b/modelopt/torch/utils/plugins/transformers_dataset.py
@@ -17,6 +17,7 @@
 
 import copy
 import itertools
+import json
 import os
 
 import torch
@@ -329,8 +330,48 @@ def __init__(
         local_image_path: str = "",
         return_labels: bool = False,
     ):
-        """Initialize the VisionLanguageDataset."""
-        self.processor = transformers.AutoProcessor.from_pretrained(processor)
+        """Initialize the collator and its Hugging Face processor."""
+        processor_kwargs = {}
+        for env_name, kwarg_name in (
+            ("VLM_MIN_PIXELS", "min_pixels"),
+            ("VLM_MAX_PIXELS", "max_pixels"),
+        ):
+            value = os.environ.get(env_name)
+            if value is not None:
+                try:
+                    processor_kwargs[kwarg_name] = int(value)
+                except ValueError as exc:
+                    raise ValueError(f"{env_name} must be an integer, got {value!r}") from exc
+
+        self.processor = transformers.AutoProcessor.from_pretrained(processor, **processor_kwargs)
+        if processor_kwargs:
+            print_rank_0(f"Loaded VLM processor with {processor_kwargs}")
+        self._configure_video_processor(self.processor, processor_kwargs)
+        max_prompt_tokens = os.environ.get("VLM_MAX_PROMPT_TOKENS")
+        self.max_prompt_tokens = None
+        if max_prompt_tokens is not None:
+            try:
+                self.max_prompt_tokens = int(max_prompt_tokens)
+            except ValueError as exc:
+                raise ValueError(
+                    f"VLM_MAX_PROMPT_TOKENS must be a positive integer, got {max_prompt_tokens!r}"
+                ) from exc
+            if self.max_prompt_tokens <= 0:
+                raise ValueError("VLM_MAX_PROMPT_TOKENS must be a positive integer.")
+        self._prompt_content_truncation_warned = False
+        max_assistant_tokens = os.environ.get("VLM_MAX_ASSISTANT_TOKENS")
+        self.max_assistant_tokens = None
+        if max_assistant_tokens is not None:
+            try:
+                self.max_assistant_tokens = int(max_assistant_tokens)
+            except ValueError as exc:
+                raise ValueError(
+                    "VLM_MAX_ASSISTANT_TOKENS must be a positive integer, got "
+                    f"{max_assistant_tokens!r}"
+                ) from exc
+            if self.max_assistant_tokens <= 0:
+                raise ValueError("VLM_MAX_ASSISTANT_TOKENS must be a positive integer.")
+        self._assistant_content_truncation_warned = False
         self.chat_template = chat_template
         self.local_image_path = local_image_path
         self._conversations_warned = False
@@ -344,6 +385,42 @@ def __init__(
             shift_labels=shift_labels,
             return_labels=return_labels,
         )
+        self.processor.chat_template = self.tokenizer.chat_template
+
+    @staticmethod
+    def _configure_video_processor(processor, image_pixel_bounds):
+        """Apply visual pixel bounds to a separate Qwen video processor."""
+        video_processor = getattr(processor, "video_processor", None)
+        video_pixel_bounds = {}
+        explicit_video_bounds = False
+        for env_name, size_key, image_key in (
+            ("VLM_VIDEO_MIN_PIXELS", "shortest_edge", "min_pixels"),
+            ("VLM_VIDEO_MAX_PIXELS", "longest_edge", "max_pixels"),
+        ):
+            value = os.environ.get(env_name)
+            if value is None:
+                value = image_pixel_bounds.get(image_key)
+            else:
+                explicit_video_bounds = True
+                try:
+                    value = int(value)
+                except ValueError as exc:
+                    raise ValueError(f"{env_name} must be an integer, got {value!r}") from exc
+            if value is not None:
+                video_pixel_bounds[size_key] = value
+
+        if explicit_video_bounds and video_processor is None:
+            raise ValueError(
+                "VLM_VIDEO_MIN_PIXELS/VLM_VIDEO_MAX_PIXELS were set, but the "
+                "configured VLM processor has no video_processor."
+            )
+        if video_processor is not None and video_pixel_bounds:
+            video_size = dict(video_processor.size)
+            video_size.update(video_pixel_bounds)
+            if video_size["shortest_edge"] > video_size["longest_edge"]:
+                raise ValueError("VLM_VIDEO_MIN_PIXELS must not exceed VLM_VIDEO_MAX_PIXELS.")
+            video_processor.size = video_size
+            print_rank_0(f"Configured VLM video processor with {video_size}")
 
     def _verify_generation_tags(self):
         """Accept VLM templates whose assistant spans have stable chat markers.
@@ -430,23 +507,115 @@ def _build_assistant_masks(self, tokenized_messages):
 
         return assistant_masks
 
-    def _process_multimodal_sample(self, examples):
-        derive_masks_from_markers = self.answer_only_loss and bool(self._assistant_marker_specs())
+    def _pad_sequence_tensors(self, tokenized_messages):
+        """Pad processor outputs to the fixed DFlash training sequence length."""
+        if self.train_len is None:
+            return tokenized_messages
+
+        input_ids = tokenized_messages.get("input_ids")
+        if not isinstance(input_ids, torch.Tensor) or input_ids.ndim != 2:
+            raise ValueError("VLM processor did not return rank-2 input_ids.")
+        if input_ids.shape[1] > self.train_len:
+            raise ValueError(
+                f"VLM processor returned seq_len {input_ids.shape[1]} above "
+                f"training_seq_len {self.train_len}; reduce visual or assistant-text caps."
+            )
+
+        pad_width = self.train_len - input_ids.shape[1]
+        if pad_width == 0:
+            return tokenized_messages
+
+        sequence_length = input_ids.shape[1]
+        batch_size = input_ids.shape[0]
+        for key, value in tokenized_messages.items():
+            if (
+                not isinstance(value, torch.Tensor)
+                or value.ndim != 2
+                or value.shape != (batch_size, sequence_length)
+            ):
+                continue
+            pad_value = self.tokenizer.pad_token_id if key == "input_ids" else 0
+            pad = value.new_full((batch_size, pad_width), pad_value)
+            tokenized_messages[key] = torch.cat((value, pad), dim=1)
+
+        return tokenized_messages
+
+    def _truncate_assistant_content(self, messages):
+        """Limit assistant text before image/video placeholders expand into tokens."""
+        if self.max_assistant_tokens is None:
+            return
+
+        for message in messages:
+            if message.get("role") != "assistant":
+                continue
+            for content in message.get("content", []):
+                if content.get("type") != "text" or not isinstance(content.get("text"), str):
+                    continue
+                token_ids = self.tokenizer(content["text"], add_special_tokens=False).input_ids
+                if len(token_ids) <= self.max_assistant_tokens:
+                    continue
+
+                suffix_tokens = min(128, max(1, self.max_assistant_tokens // 4))
+                prefix_tokens = self.max_assistant_tokens - suffix_tokens
+                content["text"] = self.tokenizer.decode(
+                    token_ids[:prefix_tokens] + token_ids[-suffix_tokens:],
+                    skip_special_tokens=True,
+                )
+                if not self._assistant_content_truncation_warned:
+                    print_rank_0(
+                        "Truncating assistant content to "
+                        f"{self.max_assistant_tokens} tokens before VLM tokenization."
+                    )
+                    self._assistant_content_truncation_warned = True
+
+    def _truncate_prompt_content(self, messages):
+        """Limit user/system text before image/video placeholders expand into tokens."""
+        if self.max_prompt_tokens is None:
+            return
+
+        for message in messages:
+            if message.get("role") not in {"system", "user"}:
+                continue
+            for content in message.get("content", []):
+                if content.get("type") != "text" or not isinstance(content.get("text"), str):
+                    continue
+                token_ids = self.tokenizer(content["text"], add_special_tokens=False).input_ids
+                if len(token_ids) <= self.max_prompt_tokens:
+                    continue
+
+                suffix_tokens = min(128, max(1, self.max_prompt_tokens // 4))
+                prefix_tokens = self.max_prompt_tokens - suffix_tokens
+                content["text"] = self.tokenizer.decode(
+                    token_ids[:prefix_tokens] + token_ids[-suffix_tokens:],
+                    skip_special_tokens=True,
+                )
+                if not self._prompt_content_truncation_warned:
+                    print_rank_0(
+                        "Truncating user/system content to "
+                        f"{self.max_prompt_tokens} tokens before VLM tokenization."
+                    )
+                    self._prompt_content_truncation_warned = True
+
+    def _apply_chat_template(self, examples):
+        """Tokenize VLM messages without passing text-only kwargs to the processor."""
+        derive_masks_from_markers = bool(self._assistant_marker_specs())
         tokenized_messages = self.processor.apply_chat_template(
             examples,
             tokenize=True,
             return_tensors="pt",
             return_dict=True,
-            padding="max_length",
-            truncation=True,
-            max_length=self.train_len,
             add_generation_prompt=self.add_generation_prompt,
             return_assistant_tokens_mask=self.answer_only_loss and not derive_masks_from_markers,
         )
-
-        if derive_masks_from_markers:
+        tokenized_messages = self._pad_sequence_tensors(tokenized_messages)
+        if self.answer_only_loss and derive_masks_from_markers:
             tokenized_messages["assistant_masks"] = self._build_assistant_masks(tokenized_messages)
 
+        return tokenized_messages
+
+    def _process_multimodal_sample(self, examples):
+        tokenized_messages = self._apply_chat_template(examples)
+
         if self.return_labels:
             input_ids = tokenized_messages["input_ids"]
             labels = input_ids.new_full(input_ids.shape, IGNORE_TOKEN_ID)
@@ -498,14 +667,63 @@ def __call__(self, examples):
             for msg in copy_messages:
                 if isinstance(msg["content"], str):
                     msg["content"] = [{"type": "text", "text": msg["content"]}]
+                elif isinstance(msg["content"], dict):
+                    # A top-level structured answer can legitimately contain
+                    # keys such as "image" or "video" (for example, a URL in
+                    # a JSON response).  It is media only when it explicitly
+                    # declares a supported multimodal type.
+                    if msg["content"].get("type") in {"text", "image", "video"}:
+                        msg["content"] = [msg["content"]]
+                    else:
+                        msg["content"] = [
+                            {
+                                "type": "text",
+                                "text": json.dumps(msg["content"], ensure_ascii=False),
+                            }
+                        ]
+                elif not isinstance(msg["content"], list):
+                    # Some synthetic rows use a scalar answer or prompt (for
+                    # example, an integer multiple-choice answer). Preserve it
+                    # as text rather than discarding an otherwise valid row.
+                    msg["content"] = [
+                        {
+                            "type": "text",
+                            "text": json.dumps(msg["content"], ensure_ascii=False),
+                        }
+                    ]
 
-                for ctn in msg["content"]:
+                for index, ctn in enumerate(msg["content"]):
+                    if isinstance(ctn, str):
+                        ctn = {"type": "text", "text": ctn}
+                        msg["content"][index] = ctn
+                    if not isinstance(ctn, dict):
+                        ctn = {"type": "text", "text": json.dumps(ctn, ensure_ascii=False)}
+                        msg["content"][index] = ctn
+                    if not {"type", "text", "image", "video", "fps"}.intersection(ctn):
+                        ctn = {"type": "text", "text": json.dumps(ctn, ensure_ascii=False)}
+                        msg["content"][index] = ctn
                     # Some JSONL producers use a fixed multimodal-part schema
                     # (text/image/video/fps on every part) so Arrow can load
                     # heterogeneous image and video datasets together.  Drop
                     # the inactive placeholders before handing a part to the
                     # processor, which expects only fields relevant to its type.
                     content_type = ctn.get("type")
+                    if content_type not in {"text", "image", "video"}:
+                        inferred_types = [
+                            part_type
+                            for key, part_type in (
+                                ("text", "text"),
+                                ("image", "image"),
+                                ("video", "video"),
+                            )
+                            if ctn.get(key) not in (None, "")
+                        ]
+                        if len(inferred_types) != 1:
+                            raise ValueError(
+                                f"Unable to infer a multimodal content type from {ctn!r}."
+                            )
+                        content_type = inferred_types[0]
+                        ctn["type"] = content_type
                     if content_type != "text" and ctn.get("text") == "":
                         del ctn["text"]
                     if content_type != "image" and ctn.get("image") == "":
@@ -525,6 +743,8 @@ def __call__(self, examples):
                     for k in keys_to_delete:
                         del ctn[k]
 
+            self._truncate_prompt_content(copy_messages)
+            self._truncate_assistant_content(copy_messages)
             batch.append(copy_messages)
 
         return self._process_multimodal_sample(batch)
diff --git a/tests/examples/speculative_decoding/test_merge_dflash_datasets.py b/tests/examples/speculative_decoding/test_merge_dflash_datasets.py
new file mode 100644
index 00000000000..525f0a5a0fc
--- /dev/null
+++ b/tests/examples/speculative_decoding/test_merge_dflash_datasets.py
@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the DFlash JSONL merger recipe."""
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).parents[3]
+MERGER = REPO_ROOT / "examples/speculative_decoding/recipes/merge_dflash_datasets.py"
+
+
+def write_record(path: Path, record_id: str, prompt: str, answer: str) -> None:
+    """Write a one-turn DFlash-training conversation."""
+
+    path.write_text(
+        json.dumps(
+            {
+                "id": record_id,
+                "messages": [
+                    {"role": "user", "content": prompt},
+                    {"role": "assistant", "content": answer},
+                ],
+            }
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+
+
+def test_merge_keeps_synthetic_variants_together(tmp_path):
+    """Single and parallel merges deduplicate matching variants from one synthetic shard."""
+
+    synthetic_output = tmp_path / "synthetic-output"
+    synthetic_output.mkdir()
+    write_record(
+        synthetic_output / "output-00000-00000-temp-0.0.jsonl",
+        "duplicate-first",
+        "Describe the image.",
+        "The image shows a blue square.",
+    )
+    write_record(
+        synthetic_output / "output-00000-00000-temp-0.1.jsonl",
+        "duplicate-second",
+        "Describe the image.",
+        "The image shows a blue square.",
+    )
+    write_record(
+        synthetic_output / "output-00001-00001-temp-0.0.jsonl",
+        "synthetic-unique",
+        "Describe the video.",
+        "The video shows a red circle.",
+    )
+    curated_text = tmp_path / "curated-text.jsonl"
+    write_record(curated_text, "curated-unique", "What is two plus two?", "Four.")
+    expected_ids = {"duplicate-first", "synthetic-unique", "curated-unique"}
+    for jobs in (1, 2):
+        output = tmp_path / f"merged-{jobs}.jsonl"
+        subprocess.run(
+            [
+                sys.executable,
+                str(MERGER),
+                "--source",
+                f"synthetic={synthetic_output}",
+                "--source",
+                f"curated={curated_text}",
+                "--output",
+                str(output),
+                "--jobs",
+                str(jobs),
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        records = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()]
+        assert {record["id"] for record in records} == expected_ids
+        assert not list(tmp_path.glob(f".{output.name}.parallel-*"))
+
+
+def test_parallel_merge_failure_preserves_work_dir_with_stale_output(tmp_path):
+    """A failed overwrite retains worker artifacts even when prior output exists."""
+
+    invalid_source = tmp_path / "invalid.jsonl"
+    invalid_source.write_text("not valid JSON\n", encoding="utf-8")
+    output = tmp_path / "merged.jsonl"
+    output.write_text("stale output\n", encoding="utf-8")
+
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(MERGER),
+            "--source",
+            f"invalid={invalid_source}",
+            "--output",
+            str(output),
+            "--jobs",
+            "2",
+            "--overwrite",
+        ],
+        capture_output=True,
+        text=True,
+    )
+
+    assert result.returncode != 0
+    assert output.read_text(encoding="utf-8") == "stale output\n"
+    work_dirs = list(tmp_path.glob(f".{output.name}.parallel-*"))
+    assert len(work_dirs) == 1
+    assert list((work_dirs[0] / "manifests").glob("*.txt"))
diff --git a/tests/examples/speculative_decoding/test_prepare_multimodal_synthetic_shards.py b/tests/examples/speculative_decoding/test_prepare_multimodal_synthetic_shards.py
new file mode 100644
index 00000000000..33a06a24334
--- /dev/null
+++ b/tests/examples/speculative_decoding/test_prepare_multimodal_synthetic_shards.py
@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for multimodal synthetic-shard preparation helpers."""
+
+import importlib.util
+from pathlib import Path
+
+_SCRIPT_PATH = (
+    Path(__file__).parents[3]
+    / "examples/speculative_decoding/recipes/prepare_multimodal_synthetic_shards.py"
+)
+_spec = importlib.util.spec_from_file_location("prepare_multimodal_synthetic_shards", _SCRIPT_PATH)
+assert _spec is not None and _spec.loader is not None
+prepare_multimodal_synthetic_shards = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(prepare_multimodal_synthetic_shards)
+
+
+def test_media_path_within_root_rejects_path_traversal_and_absolute_paths(tmp_path):
+    """PAI metadata cannot point preparation outside its media root."""
+
+    media_root = tmp_path / "media"
+    media_path = media_root / "videos" / "clip.mp4"
+    media_path.parent.mkdir(parents=True)
+    media_path.touch()
+    outside_path = tmp_path / "outside.mp4"
+    outside_path.touch()
+
+    assert (
+        prepare_multimodal_synthetic_shards.media_path_within_root(media_root, "videos/clip.mp4")
+        == media_path
+    )
+    assert (
+        prepare_multimodal_synthetic_shards.media_path_within_root(media_root, "../outside.mp4")
+        is None
+    )
+    assert (
+        prepare_multimodal_synthetic_shards.media_path_within_root(media_root, str(outside_path))
+        is None
+    )
diff --git a/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py b/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py
new file mode 100644
index 00000000000..6a94d6d58c4
--- /dev/null
+++ b/tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py
@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the multimodal SGLang generation client."""
+
+import importlib.util
+from pathlib import Path
+
+_SCRIPT_PATH = (
+    Path(__file__).parents[3]
+    / "examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py"
+)
+_spec = importlib.util.spec_from_file_location("server_generate_vlm_sglang", _SCRIPT_PATH)
+assert _spec is not None and _spec.loader is not None
+server_generate_vlm_sglang = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(server_generate_vlm_sglang)
+
+
+def test_resolve_media_path_supports_local_and_remote_media(tmp_path):
+    """Existing local media and supported remote media remain usable."""
+
+    local_media = tmp_path / "image.jpg"
+    local_media.touch()
+    remote_media = "https://example.com/image.jpg"
+
+    assert server_generate_vlm_sglang._resolve_media_path("image.jpg", str(tmp_path), None) == str(
+        local_media
+    )
+    assert (
+        server_generate_vlm_sglang._resolve_media_path(remote_media, str(tmp_path), None)
+        == remote_media
+    )
+
+
+def test_resolve_media_path_returns_none_and_warns_once_for_missing_media(monkeypatch, capsys):
+    """Missing local media can fall back to video or be skipped by the caller."""
+
+    monkeypatch.setattr(server_generate_vlm_sglang, "_UNRESOLVED_MEDIA_PATHS", set())
+
+    assert server_generate_vlm_sglang._resolve_media_path("missing.mp4", None, None) is None
+    assert server_generate_vlm_sglang._resolve_media_path("missing.mp4", None, None) is None
+
+    assert capsys.readouterr().out == "WARNING: could not resolve media path: missing.mp4\n"
+
+
+def test_openai_media_value_is_relative_to_the_media_root(tmp_path):
+    """The local HTTP server exposes media files, not the container root."""
+
+    media_root = tmp_path / "media"
+    media_path = media_root / "videos" / "clip.mp4"
+    media_path.parent.mkdir(parents=True)
+    media_path.touch()
+    input_path = tmp_path / "input" / "private.json"
+    input_path.parent.mkdir()
+    input_path.touch()
+
+    assert (
+        server_generate_vlm_sglang._as_openai_media_value(
+            str(media_path), "http://127.0.0.1:18080", str(media_root), None
+        )
+        == "http://127.0.0.1:18080/videos/clip.mp4"
+    )
+    assert server_generate_vlm_sglang._as_openai_media_value(
+        str(input_path), "http://127.0.0.1:18080", str(media_root), None
+    ) == str(input_path)
diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
index 5abaa124d68..16408e2c3e8 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py
@@ -110,6 +110,421 @@ def test_vlm_data_collator_accepts_unshifted_labels(monkeypatch):
     assert collator.shift_labels is False
 
 
+class _WhitespaceTokenizer:
+    """Minimal tokenizer for testing VLM content truncation."""
+
+    pad_token_id = 0
+
+    def __call__(self, text, add_special_tokens=False):
+        del add_special_tokens
+        return types.SimpleNamespace(input_ids=list(range(len(text.split()))))
+
+    def decode(self, token_ids, skip_special_tokens=True):
+        del skip_special_tokens
+        return " ".join(f"token-{token_id}" for token_id in token_ids)
+
+
+def _bare_vlm_collator():
+    """Create a collator for helper-method tests without loading a processor."""
+
+    collator = object.__new__(transformers_dataset.VisionLanguageDataCollator)
+    collator.tokenizer = _WhitespaceTokenizer()
+    collator.max_prompt_tokens = None
+    collator.max_assistant_tokens = None
+    collator._prompt_content_truncation_warned = False
+    collator._assistant_content_truncation_warned = False
+    collator._conversations_warned = False
+    return collator
+
+
+def test_vlm_collator_configures_visual_limits(monkeypatch):
+    """VLM environment limits configure image, video, and text processing."""
+
+    tokenizer = get_tiny_tokenizer()
+    video_processor = types.SimpleNamespace(size={"shortest_edge": 32, "longest_edge": 256})
+    processor = types.SimpleNamespace(tokenizer=tokenizer, video_processor=video_processor)
+    from_pretrained = MagicMock(return_value=processor)
+    monkeypatch.setattr(
+        transformers_dataset.transformers.AutoProcessor, "from_pretrained", from_pretrained
+    )
+    monkeypatch.setenv("VLM_MIN_PIXELS", "64")
+    monkeypatch.setenv("VLM_MAX_PIXELS", "512")
+    monkeypatch.setenv("VLM_VIDEO_MIN_PIXELS", "128")
+    monkeypatch.setenv("VLM_VIDEO_MAX_PIXELS", "1024")
+    monkeypatch.setenv("VLM_MAX_PROMPT_TOKENS", "32")
+    monkeypatch.setenv("VLM_MAX_ASSISTANT_TOKENS", "16")
+
+    collator = transformers_dataset.VisionLanguageDataCollator(
+        processor="dummy-vlm-processor", chat_template="{{ messages }}"
+    )
+
+    from_pretrained.assert_called_once_with("dummy-vlm-processor", min_pixels=64, max_pixels=512)
+    assert video_processor.size == {"shortest_edge": 128, "longest_edge": 1024}
+    assert collator.max_prompt_tokens == 32
+    assert collator.max_assistant_tokens == 16
+    assert processor.chat_template == collator.tokenizer.chat_template
+
+
+def test_vlm_collator_rejects_video_limits_without_video_processor(monkeypatch):
+    """Explicit video limits require a processor that supports video input."""
+
+    monkeypatch.setenv("VLM_VIDEO_MIN_PIXELS", "128")
+    processor = types.SimpleNamespace(video_processor=None)
+
+    with pytest.raises(ValueError, match="has no video_processor"):
+        transformers_dataset.VisionLanguageDataCollator._configure_video_processor(processor, {})
+
+
+def test_vlm_collator_applies_image_limits_to_video_processor():
+    """Image pixel limits are used for video when no video-specific limits are set."""
+
+    video_processor = types.SimpleNamespace(size={"shortest_edge": 32, "longest_edge": 256})
+    processor = types.SimpleNamespace(video_processor=video_processor)
+
+    transformers_dataset.VisionLanguageDataCollator._configure_video_processor(
+        processor, {"min_pixels": 64, "max_pixels": 512}
+    )
+
+    assert video_processor.size == {"shortest_edge": 64, "longest_edge": 512}
+
+
+@pytest.mark.parametrize(
+    ("env_name", "value", "message"),
+    [
+        ("VLM_MIN_PIXELS", "not-an-integer", "VLM_MIN_PIXELS must be an integer"),
+        ("VLM_MAX_PROMPT_TOKENS", "0", "VLM_MAX_PROMPT_TOKENS must be a positive integer"),
+        (
+            "VLM_MAX_PROMPT_TOKENS",
+            "not-an-integer",
+            "VLM_MAX_PROMPT_TOKENS must be a positive integer",
+        ),
+        (
+            "VLM_MAX_ASSISTANT_TOKENS",
+            "not-an-integer",
+            "VLM_MAX_ASSISTANT_TOKENS must be a positive integer",
+        ),
+    ],
+)
+def test_vlm_collator_rejects_invalid_environment_limits(monkeypatch, env_name, value, message):
+    """Configuration errors identify the invalid VLM environment variable."""
+
+    processor = types.SimpleNamespace(tokenizer=get_tiny_tokenizer())
+    monkeypatch.setattr(
+        transformers_dataset.transformers.AutoProcessor,
+        "from_pretrained",
+        lambda *_args, **_kwargs: processor,
+    )
+    monkeypatch.setenv(env_name, value)
+
+    with pytest.raises(ValueError, match=message):
+        transformers_dataset.VisionLanguageDataCollator(
+            processor="dummy-vlm-processor", chat_template="{{ messages }}"
+        )
+
+
+def test_vlm_collator_rejects_non_integer_video_limit(monkeypatch):
+    """Video bounds must be integer pixel counts."""
+
+    monkeypatch.setenv("VLM_VIDEO_MIN_PIXELS", "not-an-integer")
+    processor = types.SimpleNamespace(video_processor=types.SimpleNamespace(size={}))
+
+    with pytest.raises(ValueError, match="VLM_VIDEO_MIN_PIXELS must be an integer"):
+        transformers_dataset.VisionLanguageDataCollator._configure_video_processor(processor, {})
+
+
+def test_vlm_collator_pads_template_output_and_builds_unshifted_labels():
+    """DFlash batches preserve current-token labels after fixed-length padding."""
+
+    collator = _bare_vlm_collator()
+    collator.train_len = 5
+    collator.answer_only_loss = True
+    collator.return_labels = True
+    collator.shift_labels = False
+    collator.add_generation_prompt = False
+    collator._assistant_marker_specs = list
+    collator.processor = types.SimpleNamespace(
+        apply_chat_template=MagicMock(
+            return_value={
+                "input_ids": torch.tensor([[1, 2, 3]]),
+                "attention_mask": torch.tensor([[1, 1, 1]]),
+                "assistant_masks": torch.tensor([[0, 1, 1]]),
+                "pixel_values": torch.tensor([1.0]),
+            }
+        )
+    )
+
+    output = collator._process_multimodal_sample([[{"role": "user", "content": "prompt"}]])
+
+    assert output["input_ids"].tolist() == [[1, 2, 3, 0, 0]]
+    assert output["attention_mask"].tolist() == [[1, 1, 1, 0, 0]]
+    assert output["assistant_masks"].tolist() == [[0, 1, 1, 0, 0]]
+    assert output["labels"].tolist() == [
+        [
+            transformers_dataset.IGNORE_TOKEN_ID,
+            2,
+            3,
+            transformers_dataset.IGNORE_TOKEN_ID,
+            transformers_dataset.IGNORE_TOKEN_ID,
+        ]
+    ]
+    assert output["pixel_values"].tolist() == [1.0]
+    collator.processor.apply_chat_template.assert_called_once_with(
+        [[{"role": "user", "content": "prompt"}]],
+        tokenize=True,
+        return_tensors="pt",
+        return_dict=True,
+        add_generation_prompt=False,
+        return_assistant_tokens_mask=True,
+    )
+
+
+def test_vlm_collator_preserves_processor_length_when_it_matches_training_length():
+    """No padding is added when the processor already returns train_len tokens."""
+
+    collator = _bare_vlm_collator()
+    collator.train_len = 3
+    tokenized = {"input_ids": torch.tensor([[1, 2, 3]])}
+
+    assert collator._pad_sequence_tensors(tokenized) is tokenized
+
+
+@pytest.mark.parametrize(
+    ("shift_labels", "assistant_mask", "expected_labels"),
+    [
+        (
+            True,
+            torch.tensor([[0, 1, 0]]),
+            [[2, transformers_dataset.IGNORE_TOKEN_ID, transformers_dataset.IGNORE_TOKEN_ID]],
+        ),
+        (
+            False,
+            torch.zeros((1, 3), dtype=torch.long),
+            [
+                [
+                    transformers_dataset.IGNORE_TOKEN_ID,
+                    transformers_dataset.IGNORE_TOKEN_ID,
+                    transformers_dataset.IGNORE_TOKEN_ID,
+                ]
+            ],
+        ),
+    ],
+)
+def test_vlm_collator_masks_shifted_and_empty_assistant_labels(
+    shift_labels, assistant_mask, expected_labels
+):
+    """Assistant masks align shifted labels and mask empty assistant spans."""
+
+    collator = _bare_vlm_collator()
+    collator.train_len = None
+    collator.answer_only_loss = True
+    collator.return_labels = True
+    collator.shift_labels = shift_labels
+    collator.add_generation_prompt = False
+    collator._assistant_marker_specs = list
+    collator.processor = types.SimpleNamespace(
+        apply_chat_template=MagicMock(
+            return_value={
+                "input_ids": torch.tensor([[1, 2, 3]]),
+                "assistant_masks": assistant_mask,
+            }
+        )
+    )
+
+    output = collator._process_multimodal_sample([[{"role": "assistant", "content": "answer"}]])
+
+    assert output["labels"].tolist() == expected_labels
+
+
+def test_vlm_collator_requires_assistant_masks_for_answer_only_loss():
+    """A processor without assistant masks cannot support answer-only loss."""
+
+    collator = _bare_vlm_collator()
+    collator.train_len = None
+    collator.answer_only_loss = True
+    collator.return_labels = True
+    collator.shift_labels = True
+    collator.add_generation_prompt = False
+    collator._assistant_marker_specs = list
+    collator.processor = types.SimpleNamespace(
+        apply_chat_template=MagicMock(return_value={"input_ids": torch.tensor([[1, 2, 3]])})
+    )
+
+    with pytest.raises(ValueError, match="requires assistant_masks"):
+        collator._process_multimodal_sample([[{"role": "assistant", "content": "answer"}]])
+
+
+def test_vlm_collator_derives_masks_from_chatml_markers():
+    """ChatML fallback derives masks after the processor tokenizes the batch."""
+
+    collator = _bare_vlm_collator()
+    collator.train_len = None
+    collator.answer_only_loss = True
+    collator.add_generation_prompt = False
+    collator._assistant_marker_specs = lambda: [([10, 11], [[99]])]
+    collator.processor = types.SimpleNamespace(
+        apply_chat_template=MagicMock(
+            return_value={
+                "input_ids": torch.tensor([[10, 11, 7, 99]]),
+                "attention_mask": torch.tensor([[1, 1, 1, 1]]),
+            }
+        )
+    )
+
+    tokenized = collator._apply_chat_template([[{"role": "assistant", "content": "answer"}]])
+
+    assert tokenized["assistant_masks"].tolist() == [[0, 0, 1, 0]]
+    assert (
+        collator.processor.apply_chat_template.call_args.kwargs["return_assistant_tokens_mask"]
+        is False
+    )
+
+
+def test_vlm_collator_normalizes_and_truncates_structured_messages(tmp_path):
+    """Structured contents become processor-compatible text, image, and video parts."""
+
+    collator = _bare_vlm_collator()
+    collator.local_image_path = str(tmp_path)
+    collator.max_prompt_tokens = 3
+    collator.max_assistant_tokens = 3
+    captured = {}
+    collator._process_multimodal_sample = lambda batch: captured.setdefault("batch", batch)
+
+    result = collator(
+        [
+            {
+                "messages": [
+                    {"role": "system", "content": {"format": "json"}},
+                    {
+                        "role": "user",
+                        "content": [
+                            "one two three four",
+                            {"text": "five six seven eight"},
+                            {"image": "image.png", "text": "", "video": "", "fps": 0},
+                        ],
+                    },
+                    {"role": "assistant", "content": 42},
+                    {"role": "assistant", "content": "nine ten eleven twelve"},
+                ]
+            }
+        ]
+    )
+
+    assert result == captured["batch"]
+    messages = captured["batch"][0]
+    assert messages[0]["content"] == [{"type": "text", "text": '{"format": "json"}'}]
+    assert messages[1]["content"][0]["text"] == "token-0 token-1 token-3"
+    assert messages[1]["content"][1]["text"] == "token-0 token-1 token-3"
+    assert messages[1]["content"][2] == {"type": "image", "image": str(tmp_path / "image.png")}
+    assert messages[2]["content"] == [{"type": "text", "text": "42"}]
+    assert messages[3]["content"] == [{"type": "text", "text": "token-0 token-1 token-3"}]
+
+
+def test_vlm_collator_normalizes_conversations_and_malformed_content():
+    """Deprecated conversations and malformed content remain processor-compatible."""
+
+    collator = _bare_vlm_collator()
+    captured = {}
+    collator._process_multimodal_sample = lambda batch: captured.setdefault("batch", batch)
+
+    result = collator(
+        [
+            {
+                "conversations": [
+                    {"role": "system", "content": {"type": "text", "text": "declared"}},
+                    {
+                        "role": "user",
+                        "content": [
+                            7,
+                            {"unrecognized": "data"},
+                            {"type": "text", "text": "kept", "image": "", "video": None},
+                        ],
+                    },
+                ]
+            }
+        ]
+    )
+
+    assert result == captured["batch"]
+    messages = captured["batch"][0]
+    assert collator._conversations_warned is True
+    assert messages[0]["content"] == [{"type": "text", "text": "declared"}]
+    assert messages[1]["content"] == [
+        {"type": "text", "text": "7"},
+        {"type": "text", "text": '{"unrecognized": "data"}'},
+        {"type": "text", "text": "kept"},
+    ]
+
+
+def test_vlm_collator_builds_assistant_masks_from_markers():
+    """ChatML boundaries provide assistant-only masks without generation tags."""
+
+    collator = _bare_vlm_collator()
+    collator._assistant_marker_specs = lambda: [([10, 11], [[99]])]
+    tokenized_messages = {
+        "input_ids": torch.tensor([[0, 10, 11, 7, 8, 99, 0]]),
+        "attention_mask": torch.tensor([[0, 1, 1, 1, 1, 1, 0]]),
+    }
+
+    assistant_masks = collator._build_assistant_masks(tokenized_messages)
+
+    assert assistant_masks.tolist() == [[0, 0, 0, 1, 1, 0, 0]]
+
+
+def test_vlm_collator_recognizes_chatml_markers_without_generation_tags():
+    """ChatML templates do not need Hugging Face generation tags."""
+
+    collator = _bare_vlm_collator()
+    collator.tokenizer.chat_template = "<|im_start|>{{ message }}<|im_end|>"
+    markers = {
+        "<|im_start|>assistant\n": [10, 11],
+        "<|im_end|>\n": [99],
+        "<|im_end|>": [99],
+    }
+    collator.tokenizer = MagicMock(
+        chat_template=collator.tokenizer.chat_template,
+        side_effect=lambda text, **_kwargs: {"input_ids": markers[text]},
+    )
+
+    collator._verify_generation_tags()
+
+    assert collator._assistant_marker_specs() == [([10, 11], [[99], [99]])]
+
+
+def test_vlm_collator_rejects_templates_without_generation_or_chatml_markers():
+    """Answer-only loss still rejects templates with no supported boundaries."""
+
+    collator = _bare_vlm_collator()
+    collator.tokenizer.chat_template = "{{ messages }}"
+
+    with pytest.raises(ValueError, match=r"requires \{\% generation \%\}"):
+        collator._verify_generation_tags()
+
+
+def test_vlm_collator_skips_unlimited_and_non_text_content_truncation():
+    """Disabled limits and media parts are left unchanged."""
+
+    collator = _bare_vlm_collator()
+    messages = [{"role": "assistant", "content": [{"type": "image", "image": "image.png"}]}]
+
+    collator._truncate_assistant_content(messages)
+    collator._truncate_prompt_content(messages)
+
+    collator.max_assistant_tokens = 1
+    collator._truncate_assistant_content(messages)
+
+    assert messages == [{"role": "assistant", "content": [{"type": "image", "image": "image.png"}]}]
+
+
+def test_vlm_collator_rejects_processor_output_above_training_length():
+    """Visual expansion must not silently truncate DFlash sequences."""
+
+    collator = _bare_vlm_collator()
+    collator.train_len = 2
+
+    with pytest.raises(ValueError, match="above training_seq_len"):
+        collator._pad_sequence_tensors({"input_ids": torch.tensor([[1, 2, 3]])})
+
+
 # ---------------------------------------------------------------------------
 # sample_size truncation tests
 # ---------------------------------------------------------------------------