[diffusion] model: support minimax-h3 (#33275)
Co-authored-by: zhenaozhenfu <zhenaozhenfu@minimaxi.com> Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: andyluo7 <andy.luo@amd.com> Co-authored-by: Zijie Xia <zijie_xia@icloud.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: chao-xue <877184285@qq.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Mick
co-authored by
zhenaozhenfu
BBuf
andyluo7
Zijie Xia
Claude Fable 5
chao-xue
Cursor
zhenaozhenfu
BBuf
andyluo7
Zijie Xia
Claude Fable 5
chao-xue
Cursor
parent
0877a0e2f1
commit
70fe2e0dd5
@@ -435,6 +435,55 @@ jobs:
|
||||
- uses: ./.github/actions/upload-cuda-coredumps
|
||||
if: failure()
|
||||
|
||||
multimodal-gen-test-4-h100:
|
||||
# Temporarily disabled while the 4-gpu-h100 runner is unstable
|
||||
if: ${{ false }}
|
||||
runs-on: 4-gpu-h100
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.git_ref || github.sha }}
|
||||
|
||||
- uses: ./.github/actions/check-pr-test-health
|
||||
|
||||
- uses: ./.github/actions/check-maintenance
|
||||
|
||||
- name: Download artifacts
|
||||
if: inputs.sgl_kernel == 'true'
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
path: python/sglang/kernels/aot/dist/
|
||||
merge-multiple: true
|
||||
pattern: wheel-python3.10-cuda*
|
||||
|
||||
- name: Install dependencies
|
||||
timeout-minutes: 20
|
||||
run: |
|
||||
CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
|
||||
|
||||
- name: Run MiniMax-H3 PR smoke test
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
RUNAI_STREAMER_MEMORY_LIMIT: 0
|
||||
SGLANG_TEST_WAIT_SECS: 1800
|
||||
SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
|
||||
run: |
|
||||
cd python
|
||||
python3 sglang/multimodal_gen/test/run_suite.py --suite 4-gpu-h100
|
||||
|
||||
- name: Upload diffusion failure artifacts
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: diffusion-failures-${{ github.job }}-${{ github.run_attempt }}
|
||||
path: diffusion-failures/
|
||||
if-no-files-found: ignore
|
||||
|
||||
- uses: ./.github/actions/upload-cuda-coredumps
|
||||
if: failure()
|
||||
|
||||
multimodal-gen-unit-test:
|
||||
if: |
|
||||
((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
|
||||
|
||||
@@ -15,6 +15,11 @@ on:
|
||||
description: "Docker Hub repo to push to. Use lmsysorg/sglang-staging for testing."
|
||||
required: false
|
||||
default: "lmsysorg/sglang"
|
||||
build_only:
|
||||
description: "Build and validate one Linux AMD64 CUDA 13 image locally on the runner without logging in or pushing."
|
||||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
overlay_dockerfile:
|
||||
description: "Optional extra Dockerfile lines appended after FROM <base> to build a layered image (e.g. 'RUN pip install ...'). Note: this job has no repo checkout, so only FROM + RUN work — you cannot COPY files from this repo. Leave empty to skip overlay."
|
||||
required: false
|
||||
@@ -135,6 +140,7 @@ jobs:
|
||||
|
||||
build-and-publish:
|
||||
needs: prepare
|
||||
if: ${{ !inputs.build_only }}
|
||||
uses: ./.github/workflows/_docker-build-and-publish.yml
|
||||
with:
|
||||
docker_target: framework_final
|
||||
@@ -144,9 +150,68 @@ jobs:
|
||||
image_repo: ${{ inputs.image_repo || 'lmsysorg/sglang' }}
|
||||
secrets: inherit
|
||||
|
||||
build-only:
|
||||
needs: prepare
|
||||
if: ${{ inputs.build_only && github.repository == 'sgl-project/sglang' }}
|
||||
runs-on: x64-docker-build-node
|
||||
env:
|
||||
IMAGE: sglang-diffusion-build-only:${{ github.run_id }}-${{ github.run_attempt }}
|
||||
SOURCE_DIR: source-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ needs.prepare.outputs.checkout_ref || github.ref }}
|
||||
path: ${{ env.SOURCE_DIR }}
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Build CUDA 13 diffusion image locally (no push)
|
||||
run: |
|
||||
cd "$SOURCE_DIR"
|
||||
docker buildx build \
|
||||
--target framework_final \
|
||||
--platform linux/amd64 \
|
||||
--load \
|
||||
-t "$IMAGE" \
|
||||
-f docker/Dockerfile \
|
||||
--build-arg CUDA_VERSION=13.0.1 \
|
||||
--build-arg BUILD_TYPE=all \
|
||||
--build-arg GRACE_BLACKWELL=0 \
|
||||
--build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
|
||||
${{ needs.prepare.outputs.extra_build_args }} \
|
||||
--no-cache \
|
||||
.
|
||||
|
||||
- name: Verify diffusion dependencies
|
||||
run: |
|
||||
docker run --rm "$IMAGE" python3 -c '
|
||||
import importlib
|
||||
modules = (
|
||||
"av",
|
||||
"cache_dit",
|
||||
"cv2",
|
||||
"diffusers",
|
||||
"imageio_ffmpeg",
|
||||
"moviepy",
|
||||
"msgpack",
|
||||
"sglang.multimodal_gen",
|
||||
"soundfile",
|
||||
"trimesh",
|
||||
)
|
||||
for module in modules:
|
||||
importlib.import_module(module)
|
||||
print("diffusion dependency smoke check passed")
|
||||
'
|
||||
|
||||
- name: Remove local test image
|
||||
if: always()
|
||||
run: docker image rm "$IMAGE" || true
|
||||
|
||||
cleanup-nightly:
|
||||
needs: build-and-publish
|
||||
if: ${{ !inputs.tag && !inputs.pr_number }}
|
||||
if: ${{ !inputs.build_only && !inputs.tag && !inputs.pr_number }}
|
||||
uses: ./.github/workflows/_docker-cleanup-nightly.yml
|
||||
with:
|
||||
tag_prefixes: '["nightly-dev", "nightly-dev-cu12", "nightly-dev-cu13"]'
|
||||
@@ -155,7 +220,7 @@ jobs:
|
||||
|
||||
build-overlay:
|
||||
needs: [prepare, build-and-publish]
|
||||
if: ${{ inputs.overlay_dockerfile != '' && github.repository == 'sgl-project/sglang' }}
|
||||
if: ${{ !inputs.build_only && inputs.overlay_dockerfile != '' && github.repository == 'sgl-project/sglang' }}
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
fail-fast: false
|
||||
|
||||
@@ -0,0 +1,807 @@
|
||||
---
|
||||
title: MiniMax-H3
|
||||
description: Run native MiniMax-H3 video-and-audio generation with SGLang Diffusion.
|
||||
metatags:
|
||||
description: "Serve MiniMax-H3 with SGLang Diffusion for text-to-video-and-audio, first/last-frame conditioning, video-to-video, and multimodal reference conditioning."
|
||||
---
|
||||
|
||||
## 1. Model introduction
|
||||
|
||||
[MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) generates a video and a synchronized stereo audio track in one request. SGLang Diffusion provides a native pipeline for the three public task profiles, split across the released FL2VA (First-and-Last-Frame-to-Video-and-Audio) and Ref2VA (Reference-to-Video-and-Audio) checkpoint partitions:
|
||||
|
||||
| Task | `task` value | Conditioning |
|
||||
| --- | --- | --- |
|
||||
| Text to video and audio | `t2va` | Text prompt only |
|
||||
| First/last frame to video and audio | `fl2va` | First frame, last frame, or both |
|
||||
| Reference to video and audio | `ref2va` | Image, video, and audio references |
|
||||
|
||||
Video-to-video (V2V) is a supported `ref2va` use case, not a fourth task
|
||||
value. Run the `Ref2VA` partition and provide a video reference in
|
||||
`conditions`.
|
||||
|
||||
Use the selected Hub's root model ID: `MiniMaxAI/MiniMax-H3` on Hugging Face
|
||||
or `MiniMax/MiniMax-H3` on ModelScope. Select the checkpoint variant with
|
||||
`--model-variant`: `fl2va` serves both `t2va` and `fl2va`, while `ref2va`
|
||||
serves reference-conditioned requests. SGLang owns the checkpoint-directory
|
||||
mapping; do not point `--model-path` at a manually downloaded subdirectory.
|
||||
|
||||
<Warning>
|
||||
Review the license and usage terms in the MiniMax-H3 model card before production or commercial use. SGLang support does not grant additional model usage rights.
|
||||
</Warning>
|
||||
|
||||
## 2. Installation
|
||||
|
||||
Install SGLang with the diffusion dependencies:
|
||||
|
||||
```bash Command
|
||||
uv pip install "sglang[diffusion]" --prerelease=allow
|
||||
```
|
||||
|
||||
For platform-specific setup, see the [SGLang Diffusion installation guide](/docs/sglang-diffusion/installation).
|
||||
|
||||
## 3. Serve MiniMax-H3
|
||||
|
||||
Use the interactive selector to choose a hardware platform, deployment profile,
|
||||
one of the two checkpoint partitions, a request mode, and deployment features.
|
||||
It generates Python and, where available, Docker launch forms. AMD selections
|
||||
use the Python form until an H3-capable ROCm image is validated. The **$ cURL**
|
||||
button follows the selected request mode and switches the payload across
|
||||
text-only, all three first/last-frame signatures, and the image/audio/video
|
||||
reference combinations listed below.
|
||||
Set **Outputs per prompt** in the picker’s **Env** panel to generate more than
|
||||
one output without mixing request sampling controls into the deployment
|
||||
matrix.
|
||||
|
||||
The Docker form does not assume the base SGLang image contains optional
|
||||
diffusion dependencies. It installs the platform-specific diffusion extra from
|
||||
the source bundled in the image before starting the server. Set **Host media
|
||||
directory** in the **Env** panel for FL2VA, V2V, or Ref2VA; the picker mounts
|
||||
that directory read-only at `/data/minimax-h3` inside the container.
|
||||
|
||||
Every hardware/topology cell in this picker has completed a real request on
|
||||
that exact GPU model. Approximate load-time features such as online
|
||||
quantization are called out separately in the generated command. Sampling
|
||||
behavior such as Cache-DiT is documented separately below.
|
||||
|
||||
**Deployment Profile** exposes resident and FSDP placement on B200, B300,
|
||||
H200, and H100. Resident is the latency-oriented default; FSDP reduces DiT
|
||||
weight residency at the cost of per-block parameter collectives. **Online
|
||||
Quantization** appears only on B200 and B300. AMD keeps its resident AITER
|
||||
recipe, while RTX 5090 uses its dedicated layerwise-offload profile.
|
||||
|
||||
import { Deployment } from "/src/snippets/_deployment.jsx";
|
||||
import { config } from "/src/snippets/configs/MiniMaxAI/minimax-h3.jsx";
|
||||
|
||||
<Deployment config={config} />
|
||||
|
||||
<Note>
|
||||
The ready-to-run request template lives behind the **$ cURL** button in the
|
||||
picker above. It regenerates as you change the selection, so the payload it
|
||||
shows always matches the serve command next to it.
|
||||
</Note>
|
||||
|
||||
The selector uses the verified Hugging Face ID. To use ModelScope through the
|
||||
same normal `sglang serve` path, prefix the copied command with
|
||||
`SGLANG_USE_MODELSCOPE=true` and replace the model path with
|
||||
`MiniMax/MiniMax-H3`; keep its selected variant and topology flags unchanged.
|
||||
|
||||
For a four-card H200 host, keep the full BF16/FP32 model resident by default.
|
||||
The model fits without FSDP, so this path avoids the per-block parameter
|
||||
all-gathers of the memory-oriented FSDP profile:
|
||||
|
||||
```bash 4×H200 resident
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant fl2va \
|
||||
--num-gpus 4 \
|
||||
--ulysses-degree 4 \
|
||||
--performance-mode speed \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
For 4×H100 80 GB, balance the large packed activation with resident weight
|
||||
sharding. TP2 + Ulysses2 was the fastest measured lossless topology while the
|
||||
Qwen encoder still folds across all four GPUs:
|
||||
|
||||
```bash 4×H100 fastest
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant fl2va \
|
||||
--num-gpus 4 \
|
||||
--tp-size 2 \
|
||||
--ulysses-degree 2 \
|
||||
--performance-mode speed \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
Pure Ulysses4 could not keep the full pipeline resident on 80 GB H100s. Use
|
||||
`--tp-size 4 --ulysses-degree 1` when lower resident memory matters more than
|
||||
the last few percent of latency. FSDP remains a verified capacity option, but
|
||||
its per-block weight all-gathers do not make it the H100 speed default:
|
||||
|
||||
```bash 4×H100 FSDP capacity
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant fl2va \
|
||||
--num-gpus 4 \
|
||||
--ulysses-degree 4 \
|
||||
--performance-mode speed \
|
||||
--use-fsdp-inference true \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
For a two-card RTX 5090 host, use TP2 and keep 20 DiT blocks
|
||||
resident. Layerwise placement is lossless: it changes parameter placement and
|
||||
transfer scheduling, not the BF16/FP32 denoising or VAE math. This is the
|
||||
fastest measured 32 GB operating point:
|
||||
|
||||
```bash 2×RTX 5090 fastest lossless
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant fl2va \
|
||||
--num-gpus 2 \
|
||||
--tp-size 2 \
|
||||
--ulysses-degree 1 \
|
||||
--performance-mode memory \
|
||||
--layerwise-offload-components dit,text_encoder,vae \
|
||||
--dit-offload-prefetch-size 1 \
|
||||
--dit-layerwise-resident-layers 20 \
|
||||
--enable-torch-compile false \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
The DiT residency and prefetch knobs apply only to the repeatedly executed DiT
|
||||
blocks. The text encoder and the video VAE decoder blocks use one-layer
|
||||
prefetch with zero resident layers. The video VAE encoder stays resident
|
||||
because its indexed down blocks cannot host executable layerwise hooks; the
|
||||
roughly 577 MiB audio VAE also stays resident because offloading it only adds
|
||||
transfer overhead. This exact recipe was validated on
|
||||
2× RTX 5090 (32 GB each) and a 377 GiB host; use a 384 GiB-class machine. The
|
||||
latency and memory comparison is collected in the benchmark section below.
|
||||
|
||||
The first launch downloads the model through the selected Hub. If the Hugging
|
||||
Face repository requires authentication, export a Hugging Face token in the
|
||||
server environment.
|
||||
|
||||
For MiniMax-H3, `--performance-mode speed` deliberately keeps the DiT eager. The current `torch.compile` path changes the model's numerical output, so it is not enabled implicitly by any recommended lossless preset. An explicit `--enable-torch-compile true` remains available for controlled experiments, but it should not be used to generate consistency ground truth.
|
||||
|
||||
## 4. Generate video and audio
|
||||
|
||||
MiniMax-H3 uses the asynchronous OpenAI-compatible video endpoint. Choose a
|
||||
generation mode below, submit a job, poll its status, and then download the
|
||||
completed MP4.
|
||||
|
||||
<Tabs>
|
||||
|
||||
<Tab title="T2VA">
|
||||
|
||||
MiniMax-H3 supports output durations from 4 through 15 seconds, inclusive. The
|
||||
following request keeps the verified 5-second profile at a 768-pixel short
|
||||
edge. MiniMax-H3 resolves the aligned output canvas and frame count from
|
||||
`target`.
|
||||
|
||||
```bash Command
|
||||
video_id=$(
|
||||
curl -sS -X POST http://127.0.0.1:30010/v1/videos \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "MiniMaxAI/MiniMax-H3",
|
||||
"prompt": "At night, while their owner sleeps in a bedroom, three cats march in loudly playing tiny brass instruments, then abruptly file out.",
|
||||
"seconds": 5,
|
||||
"task": "t2va",
|
||||
"conditions": [],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 5.0
|
||||
},
|
||||
"num_outputs_per_prompt": 1,
|
||||
"num_inference_steps": 50,
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
"seed": 1101
|
||||
}' |
|
||||
jq -r '.id'
|
||||
)
|
||||
|
||||
while true; do
|
||||
status=$(curl -sS "http://127.0.0.1:30010/v1/videos/${video_id}" | jq -r '.status')
|
||||
[ "$status" = "completed" ] && break
|
||||
[ "$status" = "failed" ] && exit 1
|
||||
sleep 1
|
||||
done
|
||||
|
||||
curl -sS -L "http://127.0.0.1:30010/v1/videos/${video_id}/content" \
|
||||
-o minimax-h3-t2va.mp4
|
||||
```
|
||||
|
||||
The output contract is an MP4 containing H.264 video at 24 fps and one AAC stereo audio stream at 32 kHz.
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="FL2VA">
|
||||
|
||||
For `fl2va`, provide one or two image conditions with role `keyframe`. The supported frame-index sets are `[0]`, `[-1]`, and `[0, -1]`.
|
||||
|
||||
The following request uses one server-local first frame. Use
|
||||
`frame_index: -1` for a last frame, or include both entries for first-and-last
|
||||
conditioning.
|
||||
|
||||
Choose FL2VA when the supplied image should be the actual first or last frame
|
||||
of the generated clip. Use image-based Ref2VA instead when the image should
|
||||
guide identity, style, or composition without being preserved as an endpoint;
|
||||
Ref2VA may recompose or crop the reference.
|
||||
|
||||
```bash Command
|
||||
curl -sS -X POST http://127.0.0.1:30010/v1/videos \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "MiniMaxAI/MiniMax-H3",
|
||||
"prompt": "The supplied frame continues with calm, natural motion and synchronized ambient sound.",
|
||||
"seconds": 5,
|
||||
"task": "fl2va",
|
||||
"conditions": [
|
||||
{
|
||||
"type": "image",
|
||||
"uri": "file:///data/minimax-h3/first-frame.png",
|
||||
"role": "keyframe",
|
||||
"frame_index": 0
|
||||
}
|
||||
],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "auto",
|
||||
"duration_seconds": 5.0
|
||||
},
|
||||
"num_outputs_per_prompt": 1,
|
||||
"num_inference_steps": 50,
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
"seed": 2101
|
||||
}'
|
||||
```
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="V2V">
|
||||
|
||||
V2V uses the reference-conditioning weights. Launch the server with
|
||||
`--model-variant ref2va`, keep the request `task` set to `ref2va`, and provide a video
|
||||
reference in `conditions`. There is no separate `v2v` task value.
|
||||
|
||||
Use `type: "video"` when the input may be silent. If the file has a soundtrack,
|
||||
H3 also uses it as an audio reference. Use `type: "video_audio"` only when both
|
||||
streams are required; that form rejects an input without audio. The prompt tag
|
||||
for the visual stream is `<Video 1>`; an available soundtrack is exposed as
|
||||
`<Audio 1>`.
|
||||
|
||||
<Note>
|
||||
Ref2VA treats the input video as reference material, not as a pixel-aligned
|
||||
edit source. It can resynthesize or reorder motion and cuts, and it does not
|
||||
expose a denoising-strength control. Do not rely on it to preserve every source
|
||||
frame or exact timing.
|
||||
</Note>
|
||||
|
||||
Set `conditions[].start_time_seconds` to select a segment from a longer source.
|
||||
The default is `0`. SGLang seeks the visual stream and soundtrack to the same
|
||||
offset, then decodes at most the requested target duration in one pass; the
|
||||
source is not re-encoded into an intermediate clip.
|
||||
|
||||
```bash Command
|
||||
curl -sS -X POST http://127.0.0.1:30010/v1/videos \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "MiniMaxAI/MiniMax-H3",
|
||||
"prompt": "Follow the motion and appearance of <Video 1>, changing the setting to a moonlit bedroom while preserving coherent timing.",
|
||||
"seconds": 5,
|
||||
"task": "ref2va",
|
||||
"conditions": [
|
||||
{
|
||||
"type": "video",
|
||||
"uri": "file:///data/minimax-h3/input.mp4",
|
||||
"role": "reference",
|
||||
"start_time_seconds": 35.0
|
||||
}
|
||||
],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 5.0
|
||||
},
|
||||
"num_outputs_per_prompt": 1,
|
||||
"num_inference_steps": 50,
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
"seed": 4101
|
||||
}'
|
||||
```
|
||||
|
||||
Use `conditions[].uri` for H3 V2V. The generic top-level `video_path`,
|
||||
`video_url`, and `video_reference` upload fields are not lowered into H3
|
||||
reference conditions.
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="Multimodal Ref2VA">
|
||||
|
||||
For `ref2va`, first launch the reference-conditioning capability with
|
||||
`--model-variant ref2va`, then provide conditions with role `reference`.
|
||||
Image, video, and audio references can be combined. Material tags in the
|
||||
prompt use the one-based order for each modality.
|
||||
|
||||
An image condition here is semantic reference material rather than a
|
||||
pixel-aligned first frame. Use the FL2VA tab when animating a screenshot from
|
||||
that exact starting composition.
|
||||
|
||||
```bash Command
|
||||
curl -sS -X POST http://127.0.0.1:30010/v1/videos \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "MiniMaxAI/MiniMax-H3",
|
||||
"prompt": "Use <Picture 1> as the visual subject and <Audio 1> as the sound reference, with coherent natural motion.",
|
||||
"seconds": 5,
|
||||
"task": "ref2va",
|
||||
"conditions": [
|
||||
{
|
||||
"type": "image",
|
||||
"uri": "file:///data/minimax-h3/reference.png",
|
||||
"role": "reference"
|
||||
},
|
||||
{
|
||||
"type": "audio",
|
||||
"uri": "file:///data/minimax-h3/reference.mp3",
|
||||
"role": "reference"
|
||||
}
|
||||
],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "auto",
|
||||
"duration_seconds": 5.0
|
||||
},
|
||||
"num_outputs_per_prompt": 1,
|
||||
"num_inference_steps": 50,
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
"seed": 3101
|
||||
}'
|
||||
```
|
||||
|
||||
</Tab>
|
||||
|
||||
</Tabs>
|
||||
|
||||
Poll and download any conditioned request with the same job-status and
|
||||
content endpoints used in the T2VA example. Server-local `file://` URIs must
|
||||
refer to files visible inside the SGLang server environment.
|
||||
|
||||
## 5. Sampling and output controls
|
||||
|
||||
MiniMax-H3 supports more than one output per prompt. The video API accepts
|
||||
`num_outputs_per_prompt` (or OpenAI-compatible `n`) from 1 through 10. Offline
|
||||
generation accepts `--num-outputs-per-prompt N`; `--num-outputs N` is the short
|
||||
alias. A scalar seed is expanded deterministically as `seed + output_index`, so
|
||||
the outputs do not reuse the same noise.
|
||||
|
||||
Same-prompt fan-out reuses text conditioning. On the verified 2× RTX 5090
|
||||
recipe, a 5-step two-output request completed in 155.39 seconds versus 78.11
|
||||
seconds for one output, while producing two distinct valid MP4 files. The
|
||||
independent denoise and decode passes remain sequential on this 32 GB profile
|
||||
to keep peak memory bounded; the grouped path adds essentially no orchestration
|
||||
overhead. Use server replicas when lower wall-clock latency for many variants
|
||||
matters more than per-server memory efficiency.
|
||||
|
||||
For example, set `"num_outputs_per_prompt": 2` in any request above. After the
|
||||
job completes, download both outputs by selecting each zero-based variant:
|
||||
|
||||
```bash Command
|
||||
video_id="<completed-job-id>"
|
||||
for variant in 0 1; do
|
||||
curl -sS -L \
|
||||
"http://127.0.0.1:30010/v1/videos/${video_id}/content?variant=${variant}" \
|
||||
-o "minimax-h3-${variant}.mp4"
|
||||
done
|
||||
```
|
||||
|
||||
### Choose a quality profile
|
||||
|
||||
`quality` is a request-scoped sampling parameter. One resident server can
|
||||
switch between all four profiles; an approximate request mounts its audited
|
||||
Cache-DiT policy at the batch boundary, and a later `lossless` request removes
|
||||
the hooks before denoising.
|
||||
|
||||
Start the validated server once:
|
||||
|
||||
```bash Command
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant fl2va \
|
||||
--num-gpus 4 \
|
||||
--tp-size 1 \
|
||||
--sp-degree 4 \
|
||||
--ulysses-degree 4 \
|
||||
--ring-degree 1 \
|
||||
--performance-mode speed \
|
||||
--use-fsdp-inference false \
|
||||
--enable-torch-compile false \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
Then choose a request tag:
|
||||
|
||||
<Tabs>
|
||||
|
||||
<Tab title="lossless">
|
||||
|
||||
Native denoising with no feature-cache approximation. This is the default.
|
||||
|
||||
```json Request field
|
||||
{
|
||||
"quality": "lossless"
|
||||
}
|
||||
```
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="high">
|
||||
|
||||
The least aggressive approximate profile. Use it when output should stay
|
||||
closest to the same-seed lossless trajectory.
|
||||
|
||||
```json Request field
|
||||
{
|
||||
"quality": "high"
|
||||
}
|
||||
```
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="medium">
|
||||
|
||||
The balanced profile: substantially lower latency with a larger change from
|
||||
the same-seed lossless output.
|
||||
|
||||
```json Request field
|
||||
{
|
||||
"quality": "medium"
|
||||
}
|
||||
```
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="low">
|
||||
|
||||
The fastest validated profile and the largest visual deviation. Use it for
|
||||
latency-sensitive previews and high-throughput generation.
|
||||
|
||||
```json Request field
|
||||
{
|
||||
"quality": "low"
|
||||
}
|
||||
```
|
||||
|
||||
</Tab>
|
||||
|
||||
</Tabs>
|
||||
|
||||
The measured trade-off is:
|
||||
|
||||
| `quality` | Mean inference latency | Speedup | SSIM vs lossless | PSNR vs lossless | Expected trade-off |
|
||||
| --- | ---: | ---: | ---: | ---: | --- |
|
||||
| `lossless` | 75.10 s | 1.00× | 1.000 | exact | Native reference path |
|
||||
| `high` | 53.70 s | 1.40× | 0.931 | 28.16 dB | Smallest same-seed visual change |
|
||||
| `medium` | 30.23 s | 2.48× | 0.818 | 20.40 dB | Balanced latency and visual deviation |
|
||||
| `low` | 25.81 s | 2.91× | 0.794 | 19.25 dB | Largest deviation; fastest preview path |
|
||||
|
||||
These numbers use 1344×768, 124-frame, 24 fps T2VA with 50 inference steps,
|
||||
video flow shift 12, audio flow shift 3, and three fixed prompt/seed pairs on
|
||||
4×H200. The prompts cover a quiet detailed scene, fast multi-subject action,
|
||||
and a moving close-up portrait. `inference_time_s` is averaged across the three
|
||||
prompts; the quiet-scene point is itself the mean of two repeats.
|
||||
|
||||
SSIM and PSNR compare decoded, frame-aligned output with the `lossless` result
|
||||
for the same prompt and seed. They measure trajectory deviation, not absolute
|
||||
perceptual quality: an approximate profile can produce a different but still
|
||||
plausible realization. The profiles also change the joint audio-video denoise
|
||||
trajectory, while these two metrics cover video only.
|
||||
|
||||
Approximate profiles currently accept only the exact workload and 4×H200
|
||||
deployment above; other hardware, task modes, request shapes, step counts, or
|
||||
flow shifts fail before denoising. Offline generation uses the same profile
|
||||
name, for example `sglang generate --quality medium`.
|
||||
|
||||
<Note>
|
||||
`quality` selects a model sampling profile and can change generated content.
|
||||
`output_quality` controls only output-file compression; it is a separate field.
|
||||
</Note>
|
||||
|
||||
For manually tuned Cache-DiT experiments outside that validated profile, omit
|
||||
the request `quality` field and set the process-wide environment controls
|
||||
directly. An explicit `quality: lossless` request overrides those controls and
|
||||
restores native denoising:
|
||||
|
||||
```bash Command
|
||||
SGLANG_CACHE_DIT_ENABLED=true \
|
||||
SGLANG_CACHE_DIT_FN=1 \
|
||||
SGLANG_CACHE_DIT_BN=0 \
|
||||
SGLANG_CACHE_DIT_WARMUP=4 \
|
||||
SGLANG_CACHE_DIT_RDT=0.12 \
|
||||
SGLANG_CACHE_DIT_MC=2 \
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant ref2va \
|
||||
--num-gpus 8 \
|
||||
--ulysses-degree 8 \
|
||||
--performance-mode speed \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
<Warning>
|
||||
Cache-DiT skips selected block computation and is approximate. It cannot be
|
||||
combined with FSDP inference or DiT layerwise offload. Breakable CUDA graph
|
||||
execution takes precedence and leaves Cache-DiT disabled. Tune the cache
|
||||
thresholds only after comparing both video and audio quality on the target
|
||||
task profile. A real B200 request has completed, but the named profiles above
|
||||
remain fail-closed to the audited 4×H200 workload.
|
||||
</Warning>
|
||||
|
||||
## 6. Runtime feature recipes
|
||||
|
||||
<Tabs>
|
||||
|
||||
<Tab title="Lossless runtime">
|
||||
|
||||
The recommended `speed` launch already combines resident components with
|
||||
Ulysses sequence parallelism. Validation status below applies only to the
|
||||
listed hardware and topology; it is not inherited by a similar GPU family.
|
||||
|
||||
| Feature | Validation status | Notes |
|
||||
| --- | --- | --- |
|
||||
| Ulysses sequence parallelism | Verified: 8× B200, 4× H200, 4× H100, and Ulysses1/2/4/8 on MI300X and MI355X | Use `--ulysses-degree`; Ring is not compatible with H3's packed multi-segment attention. |
|
||||
| Tensor parallelism | Verified: B200 TP2 + Ulysses4; H100 TP2 + Ulysses2 and TP4 + Ulysses1 | `--tp-size` may be combined with Ulysses when the TP-local head count remains divisible by the Ulysses degree. On 4×H100, TP2 + Ulysses2 is the measured speed default. |
|
||||
| FSDP inference | Verified: 4× B200 and 4× H100 + Ulysses4 | Preserves H3's mixed BF16/FP32 parameter policy. B200 completed the exact eager comparison; H100 completed consecutive real requests at about 57 GB peak memory per GPU. |
|
||||
| Resident components | Verified: B200, H200, 4×H100 with TP, and 1/2/4/8× MI300X and MI355X | This is the recommended single-request latency path when the complete workload fits. |
|
||||
| CPU and layerwise offload | Verified: 2× RTX 5090 TP2 | The measured lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. This status applies only to the listed topology. |
|
||||
| Breakable CUDA graph | Verified: B200 Ref2VA, opt-in | Matching eager output was observed for the captured signature, without a measured speedup. Re-capture for other shapes and reference sets. |
|
||||
| `torch.compile` | Measured: H200, opt-in | Steady-state benefit was below measurement noise, while startup increased and numerical output changed. Do not use it for consistency ground truth. |
|
||||
|
||||
The verified parallel, placement, and matching-signature BCG paths keep the
|
||||
BF16/FP32 weights and denoising math. `torch.compile` is the exception called
|
||||
out above. Always use the eager BF16/FP32 launch when producing CI consistency
|
||||
ground truth.
|
||||
|
||||
For the validated 1344×768 Ref2VA profile, use a 5504-row text bucket so both
|
||||
the server warmup and reference-conditioned requests share the captured
|
||||
signature:
|
||||
|
||||
```bash Command
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant ref2va \
|
||||
--num-gpus 8 \
|
||||
--ulysses-degree 8 \
|
||||
--performance-mode speed \
|
||||
--enable-breakable-cuda-graph true \
|
||||
--warmup-resolutions 1344x768 \
|
||||
--bcg-text-buckets 5504 \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
BCG is lossless for a matching captured signature, but capture reserves extra
|
||||
GPU memory. Re-measure the live H3 text length before reusing this bucket for a
|
||||
different task profile, reference set, resolution, or prompt template.
|
||||
|
||||
</Tab>
|
||||
|
||||
<Tab title="Online quantization">
|
||||
|
||||
On the verified 8× B200 topology, quantize the BF16 transformer at server load:
|
||||
|
||||
```bash Command
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant ref2va \
|
||||
--num-gpus 8 \
|
||||
--ulysses-degree 8 \
|
||||
--performance-mode speed \
|
||||
--quantization fp8 \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
H3 automatically keeps its video/audio patch projections, timestep MLP, and
|
||||
final video/audio heads in FP32. All other linear layers have stable full
|
||||
module prefixes, so additional layers can be kept unquantized:
|
||||
|
||||
```bash Command
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant ref2va \
|
||||
--num-gpus 8 \
|
||||
--ulysses-degree 8 \
|
||||
--quantization fp8 \
|
||||
--quantization-ignored-layers blocks.0.attn token_refiner \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
<Warning>
|
||||
Online FP8 is approximate and is not a consistency ground-truth mode. It can
|
||||
be combined with Cache-DiT, but the two approximations compound. Validate
|
||||
visual quality, audio quality, memory use, and latency on the target workload.
|
||||
The picker exposes this option only on the B200 and B300 topologies used for
|
||||
real H3 validation runs.
|
||||
</Warning>
|
||||
|
||||
</Tab>
|
||||
|
||||
</Tabs>
|
||||
|
||||
## 7. Configuration notes
|
||||
|
||||
- MiniMax-H3 produces the canonical 24 fps output; request duration is expressed through `target.duration_seconds`.
|
||||
- `target.duration_seconds` must be between 4 and 15 seconds, inclusive. The command picker defaults to the verified 5-second profile.
|
||||
- Use a 768-pixel short edge for the released quality profile. The aligned output dimensions are derived from `target.aspect_ratio`.
|
||||
- `flow_shift` controls video diffusion and `audio_flow_shift` controls audio diffusion.
|
||||
- V2V uses `task: "ref2va"` with a `video` or `video_audio` reference; it is served by the `Ref2VA` partition and is not a separate public task value.
|
||||
- `conditions[].start_time_seconds` selects a non-negative offset for a video reference. Its visual and audio streams are always sought together.
|
||||
- Ref2VA condition order is semantic and must match the one-based material tags in the prompt. For Ref2VA, `target.aspect_ratio: "auto"` resolves to the model's 16:9 fallback rather than inheriting a reference asset's geometry.
|
||||
- The distilled pipeline uses a single denoising branch, so CFG parallelism does not apply. Do not enable it: `--enable-cfg-parallel true` or `--cfg-parallel-size` greater than 1 is rejected instead of duplicating the positive branch. Explicitly disabling CFG, or setting its size to 1, remains a valid no-op.
|
||||
- The released visual VAE quality recipe uses overlapping tiled decode. SGLang keeps that recipe by default and distributes complete tiles across the decode group; this changes scheduling, not the computation inside each tile.
|
||||
- H3 rejects `--vae-config.parallel-decode-mode spatial` and `spatial_shard`: validation found output mismatches. Use the default released tiled recipe.
|
||||
- Keep the default `--encoder-parallel auto`. With the server’s default `batching_max_size` of 1, single-node H100/H200/B200/B300 recipes with peer-to-peer access fold the Qwen text encoder over otherwise idle Ulysses ranks. This is separate from DiT tensor parallelism. A pure-TP recipe already shards the encoder over its TP group and does not add a world fold.
|
||||
- For throughput-oriented serving, select **DP (batched throughput)**. The picker pairs `--encoder-parallel dp` with an editable `--batching-max-size` greater than 1; compatible requests are distributed across ranks, while every rank keeps a full encoder replica. Encoder DP requires TP1 and DiT DP1, so it is disabled for the H100 TP2 + Ulysses2 and RTX 5090 TP2 recipes. It provides no benefit for a batch of one and is not bitwise-identical to the folded deployment.
|
||||
- Use explicit **Fold** to prioritize single-request latency and encoder memory on a measured high-bandwidth single-node topology. Use **Replicate** as the compatibility path when folding or encoder DP is unsuitable.
|
||||
- `--use-fsdp-inference true` shards only the DiT. MiniMax-H3 preserves the original FP32 dtype of its patch, time, and output projections during FSDP all-gather, so this path does not trade numerical correctness for memory. On 4×H100, prefer TP2 + Ulysses2 for speed; use FSDP as an explicit capacity policy rather than assuming it is faster.
|
||||
- `speed` keeps model components resident, `auto` applies the model-aware 120 GiB residency threshold, and `memory` enables the memory-saving placement policy. Explicit `--layerwise-offload-components` overrides that placement list. DiT residency/prefetch knobs are scoped to the DiT; the text encoder and video VAE decoder use one-layer prefetch and zero residency, while the H3 video VAE encoder stays resident. When `memory` is combined with explicit FSDP, H3 instead keeps the sharded DiT on GPU and layerwise-offloads the text encoder and executable VAE decoder blocks. Use `speed` only after confirming that the complete target workload fits.
|
||||
- Breakable CUDA graph execution is an explicit opt-in, not part of the recommended `speed` preset. It requires `--enable-breakable-cuda-graph`, every served size in `--warmup-resolutions`, and `--bcg-text-buckets` that cover the live H3 condition sequence. The validated 1344×768 Ref2VA recipe uses 5504; other task profiles and reference sets may need a different value. It preserves eager output for matching captured signatures, but graph capture consumes additional GPU memory and may provide little latency benefit when Ulysses attention and collectives dominate, so benchmark it on the target topology before enabling it.
|
||||
|
||||
## 8. Benchmarks
|
||||
|
||||
The picker exposes resident and FSDP profiles on NVIDIA datacenter GPUs. GPU
|
||||
counts are properties of the selected recipes, not a claim that every platform
|
||||
requires that many GPUs. The detailed tables below report performance only for
|
||||
the configurations with collected measurements:
|
||||
|
||||
| Hardware | Default resident recipe | Other profile or topology |
|
||||
| --- | --- | --- |
|
||||
| B300 | 8× Ulysses8 resident | 8× FSDP + Ulysses8; the 8-GPU sweep is not a minimum-GPU claim. |
|
||||
| B200 | 8× Ulysses8 resident | 4× FSDP + Ulysses4 |
|
||||
| H200 | 4× Ulysses4 resident | 4× FSDP + Ulysses4 |
|
||||
| H100 | 4× TP2 + Ulysses2 resident | 4× TP4 + Ulysses1; 4× FSDP + Ulysses4 |
|
||||
| MI300X / MI355X | 8× Ulysses8 resident | 1×, 2×, and 4× scaling runs |
|
||||
| RTX 5090 | 2× TP2 + layerwise offload | — |
|
||||
|
||||
### B300 precision and encoder placement
|
||||
|
||||
A 12-configuration sweep on a single 8× B300 host, covering both checkpoint
|
||||
partitions, both transformer precisions, and all three text-encoder
|
||||
placements. It answers one question — *how long does one request take, and how
|
||||
much memory does it need*.
|
||||
|
||||
### What was measured
|
||||
|
||||
**Hardware.** 8× NVIDIA B300 SXM6, single node.
|
||||
|
||||
**Model.** `MiniMaxAI/MiniMax-H3`, both released weight partitions.
|
||||
|
||||
**Serve command.** Exactly the recipe the picker emits for B300, plus the one
|
||||
or two overlay flags under test:
|
||||
|
||||
```bash Command
|
||||
sglang serve \
|
||||
--model-path MiniMaxAI/MiniMax-H3 \
|
||||
--model-variant fl2va \
|
||||
--num-gpus 8 \
|
||||
--ulysses-degree 8 \
|
||||
--performance-mode speed \
|
||||
--host 0.0.0.0 \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
The swept axes are `--model-variant` (`fl2va` / `ref2va`), `--quantization`
|
||||
(unset for BF16 / `fp8`), and `--encoder-parallel` (`auto` / `fold` /
|
||||
`replicate`). Nothing else differs between the 12 servers.
|
||||
|
||||
This is a single-request latency sweep (`batching_max_size: 1`), so encoder DP
|
||||
is intentionally excluded: it cannot distribute a batch of one. Use the
|
||||
picker’s **DP (batched throughput)** option for a multi-request throughput
|
||||
deployment; the table below does not claim a measured H3 DP speedup.
|
||||
|
||||
**Driver.**
|
||||
|
||||
```bash Command
|
||||
python3 -m sglang.multimodal_gen.benchmarks.bench_serving \
|
||||
--host 127.0.0.1 --port 30010 \
|
||||
--model MiniMaxAI/MiniMax-H3 \
|
||||
--dataset vbench --task text-to-video \
|
||||
--num-prompts 1 --max-concurrency 1 \
|
||||
--warmup-requests 1 --warmup-inference-steps 50 \
|
||||
--extra-body '{"task":"t2va","conditions":[],"target":{"short_edge":768,"aspect_ratio":"16:9","duration_seconds":5.0},"seconds":5,"flow_shift":12.0,"audio_flow_shift":3.0}'
|
||||
```
|
||||
|
||||
**Workload**
|
||||
|
||||
| Property | Value |
|
||||
| --- | --- |
|
||||
| Output duration | 5.167 s |
|
||||
| Resolution | 1344×768 |
|
||||
| Frames | 124 @ 24 fps |
|
||||
| Denoising steps | 50 |
|
||||
| `flow_shift` / `audio_flow_shift` | 12.0 / 3.0 |
|
||||
| Requests in flight | 1 (`--max-concurrency 1`, server at `batching_max_size: 1`) |
|
||||
| Requests measured | 1 per cell, after 1 warmup request |
|
||||
|
||||
### Results
|
||||
|
||||
| Weights | Precision | Encoder | Load | Warmup | Latency | Peak/GPU |
|
||||
| --- | --- | --- | ---: | ---: | ---: | ---: |
|
||||
| FL2VA | BF16 | auto | 118.1 s | 29.65 s | **19.04 s** | 83,578 MB |
|
||||
| FL2VA | BF16 | fold | 114.0 s | 28.72 s | **19.04 s** | 83,578 MB |
|
||||
| FL2VA | BF16 | replicate | 116.0 s | 28.33 s | **19.04 s** | 124,158 MB |
|
||||
| FL2VA | FP8 | auto | 116.0 s | 27.16 s | **18.03 s** | 51,926 MB |
|
||||
| FL2VA | FP8 | fold | 116.0 s | 25.99 s | **18.04 s** | 51,926 MB |
|
||||
| FL2VA | FP8 | replicate | 118.0 s | 27.97 s | **18.04 s** | 92,506 MB |
|
||||
| Ref2VA | BF16 | auto | 114.0 s | 38.69 s | **29.12 s** | 83,968 MB |
|
||||
| Ref2VA | BF16 | fold | 118.0 s | 36.58 s | **29.13 s** | 83,968 MB |
|
||||
| Ref2VA | BF16 | replicate | 116.0 s | 35.17 s | **29.13 s** | 124,490 MB |
|
||||
| Ref2VA | FP8 | auto | 124.0 s | 34.30 s | **27.12 s** | 52,816 MB |
|
||||
| Ref2VA | FP8 | fold | 112.0 s | 34.44 s | **27.12 s** | 52,816 MB |
|
||||
| Ref2VA | FP8 | replicate | 116.0 s | 33.42 s | **27.12 s** | 93,396 MB |
|
||||
|
||||
### H100 topology comparison
|
||||
|
||||
The same four-card H100 host completed three lossless placements. TP2 with
|
||||
Ulysses2 was the fastest; TP4 used the least memory:
|
||||
|
||||
| Topology | Pipeline latency | Peak/GPU |
|
||||
| --- | ---: | ---: |
|
||||
| TP2 + Ulysses2 | 13.25 s | 66.04 GB |
|
||||
| FSDP + Ulysses4 | 13.36 s | 57.01 GB |
|
||||
| TP4 + Ulysses1 | 13.86 s | 49.80 GB |
|
||||
|
||||
### RTX 5090 capacity run
|
||||
|
||||
The verified two-card RTX 5090 host used TP2 with layerwise offload. The full
|
||||
50-step, 1344×768, 5-second request completed in 559.67 seconds: 525.05
|
||||
seconds of denoising and 33.61 seconds of decoding, with a 26.3 GiB sampled
|
||||
peak per GPU.
|
||||
|
||||
| DiT settings | 5-step denoise | Inference | Peak/GPU | Result |
|
||||
| --- | ---: | ---: | ---: | --- |
|
||||
| prefetch 1, resident 20 | 43.48 s | 78.11 s | 26.3 GiB | Selected recipe |
|
||||
| prefetch 2, resident 20 | 43.37 s | 78.06 s | 27.5 GiB | No measurable gain |
|
||||
| Ulysses2, prefetch 2, resident 10 | Did not reach warmup | — | — | Rejected |
|
||||
|
||||
### AMD Instinct task and scaling runs
|
||||
|
||||
The AMD recipes keep the released BF16/FP32 precision policy and use AITER
|
||||
packed attention. The picker emits the fastest measured topology, 8 GPUs with
|
||||
Ulysses degree 8. All runs below completed full H.264/AAC decoding and
|
||||
representative-frame inspection.
|
||||
|
||||
| Hardware | Task | Denoise | Decode | Peak/GPU |
|
||||
| --- | --- | ---: | ---: | ---: |
|
||||
| MI355X | T2VA | 55.2907 s | 9.5344 s | 97,444 MB |
|
||||
| MI355X | FL2VA | 53.7978 s | 9.4477 s | 96,922 MB |
|
||||
| MI355X | Ref2VA | 41.3812 s | 6.8247 s | 94,518 MB |
|
||||
| MI300X | T2VA | 167.4878 s | 25.3244 s | 97,272 MB |
|
||||
| MI300X | FL2VA | 150.2311 s | 12.5684 s | 96,750 MB |
|
||||
| MI300X | Ref2VA | 107.6232 s | 11.3768 s | 94,268 MB |
|
||||
|
||||
The task matrix used 8 GPUs and 50 denoising steps. The scaling matrix uses
|
||||
one 1344×768, 209-frame T2VA request and changes only the GPU count and
|
||||
matching Ulysses degree:
|
||||
|
||||
| Hardware | GPUs | Denoise | Decode | Peak/GPU |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| MI355X | 8 | 55.2907 s | 9.5344 s | 97,444 MB |
|
||||
| MI355X | 4 | 104.2294 s | 11.1824 s | 103,350 MB |
|
||||
| MI355X | 2 | 223.0246 s | 15.5330 s | 115,250 MB |
|
||||
| MI355X | 1 | 288.7968 s | 24.0472 s | 137,676 MB |
|
||||
| MI300X | 8 | 167.4878 s | 25.3244 s | 97,272 MB |
|
||||
| MI300X | 4 | 297.3727 s | 26.5067 s | 103,436 MB |
|
||||
| MI300X | 2 | 585.5401 s | 29.4909 s | 115,010 MB |
|
||||
| MI300X | 1 | 978.0886 s | 36.0142 s | 137,626 MB |
|
||||
|
||||
For a measured lower-count AMD deployment, set both `--num-gpus` and
|
||||
`--ulysses-degree` to 4, 2, or 1. AITER packed attention matched segment-wise
|
||||
BF16 SDPA at cosine similarity `0.9999991655` on MI355X and `0.9999991059` on
|
||||
MI300X.
|
||||
@@ -92,6 +92,12 @@ Video models denoise a bounded latent video sequence for each request. Use these
|
||||
href="/cookbook/diffusion/MOVA/MOVA"
|
||||
img="/cards/logos/mova.png"
|
||||
/>
|
||||
<Card
|
||||
title="MiniMax-H3"
|
||||
mode="card"
|
||||
href="/cookbook/diffusion/MiniMax/MiniMax-H3"
|
||||
img="/cards/logos/minimax.png"
|
||||
/>
|
||||
</CardGroup>
|
||||
|
||||
## Realtime / World Models
|
||||
|
||||
@@ -1488,6 +1488,13 @@
|
||||
"cookbook/diffusion/MOVA/MOVA"
|
||||
]
|
||||
},
|
||||
{
|
||||
"group": "MiniMax",
|
||||
"tag": "NEW",
|
||||
"pages": [
|
||||
"cookbook/diffusion/MiniMax/MiniMax-H3"
|
||||
]
|
||||
},
|
||||
{
|
||||
"group": "LingBot World",
|
||||
"pages": [
|
||||
|
||||
@@ -75,15 +75,17 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
|
||||
### Model and runtime
|
||||
|
||||
- `--model-path {MODEL}`: model path or Hugging Face model ID
|
||||
- `--model-variant {NAME}`: semantic checkpoint variant to load when one model repository contains multiple weight partitions. The pipeline maps this stable name to the repository layout before loading; for example, MiniMax-H3 accepts `fl2va` and `ref2va`. This is a server/load-time choice, unlike a request's `task`.
|
||||
- `--model-subfolder {PATH}`: advanced direct override for a component subfolder inside the model repository. Prefer `--model-variant` when the pipeline exposes semantic routing. If both are supplied, they must resolve to the same weight partition.
|
||||
- `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter
|
||||
- `--lora-merge-mode {auto|merge|dynamic}`: choose how LoRA is applied. `auto` statically merges regular weights and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather peaks.
|
||||
- `--num-gpus {N}`: number of GPUs to use
|
||||
- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default and keeps safe offload defaults, using FSDP only for validated DiT-offload replacement paths; `speed` also enables `--enable-torch-compile` by default unless you explicitly disable it. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes.
|
||||
- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default and keeps safe offload defaults, using FSDP only for validated DiT-offload replacement paths; `speed` also enables `--enable-torch-compile` unless the model-specific deployment config opts out or you explicitly disable it. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes.
|
||||
- `--tp-size {N}`: tensor parallelism size, mainly for encoders
|
||||
- `--sp-degree {N}`: sequence parallelism size
|
||||
- `--ulysses-degree {N}` and `--ring-degree {N}`: USP parallelism controls
|
||||
- `--enable-cfg-parallel {true|false}`: enable or explicitly disable CFG parallelism
|
||||
- `--encoder-parallel {auto|fold|dp|replicate}`: how the text/image encoders use the GPUs the DiT replica leaves idle during encoding. `auto` (the default for `generate`) TP-folds an encoder wide enough to pay for the per-layer all-reduce and replicates the rest; `fold` forces the shard whenever the dims allow it; `dp` splits a batched encode across ranks (needs `--batching-max-size > 1` to engage, and is the `serve` default); `replicate` encodes redundantly on every rank. `fold` and `replicate` are bitwise-identical to single-GPU encoding. See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel).
|
||||
- `--encoder-parallel {auto|fold|dp|replicate}`: how the text/image encoders use the GPUs the DiT replica leaves idle during encoding. `auto` (the default for both `generate` and `serve`) TP-folds an encoder wide enough to pay for the per-layer all-reduce, selects DP for a server batch when it can engage, and otherwise replicates; `fold` forces the shard whenever the dims allow it; `dp` splits a batched encode across ranks and needs `--batching-max-size > 1` to engage; `replicate` encodes redundantly on every rank. `fold` and `replicate` are bitwise-identical to single-GPU encoding. See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel).
|
||||
- `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic
|
||||
- `--enable-torch-compile {true|false}`: compile native diffusion hot paths. When no warmup mode is configured, this also enables server warmup so first real requests do not pay compile latency.
|
||||
- `--offload-during-compile {true|false}`: when compile warmup is active, temporarily layerwise-offload DiT weights and move resident non-DiT components off-device so `max-autotune` fits on tighter-memory GPUs; the configured serving residency is restored before real traffic. Skipped under existing layerwise offload, Cache-DiT, or FSDP.
|
||||
@@ -102,6 +104,8 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
|
||||
- `--prompt {PROMPT}` and `--negative-prompt {PROMPT}`
|
||||
- `--image-path {PATH} [{PATH} ...]`: input image(s) for image-to-video or image-to-image generation
|
||||
- `--num-inference-steps {STEPS}` and `--seed {SEED}`
|
||||
- `--num-outputs-per-prompt {N}` / `--num-outputs {N}`: generate multiple outputs for each prompt. A scalar seed expands as `seed + output_index`.
|
||||
- `--quality {PROFILE}`: select a model-owned request quality/performance profile. Supported names and deployment constraints are model-specific.
|
||||
- `--height {HEIGHT}`, `--width {WIDTH}`, `--num-frames {N}`, `--fps {FPS}`
|
||||
- `--output-path {PATH}`, `--output-file-name {NAME}`, `--save-output`, `--return-frames`
|
||||
|
||||
@@ -178,7 +182,7 @@ sglang generate \
|
||||
HTTP server-only arguments are ignored by `sglang generate`.
|
||||
</Note>
|
||||
|
||||
For diffusers pipelines, Cache-DiT can be enabled with `SGLANG_CACHE_DIT_ENABLED=true` or `--cache-dit-config`. See [Cache-DiT](../cache_dit).
|
||||
For supported pipelines, Cache-DiT can be enabled with `SGLANG_CACHE_DIT_ENABLED=true` or `--cache-dit-config`. See [Cache-DiT](../cache_dit).
|
||||
|
||||
For supported image pipelines, breakable CUDA graph can be enabled with `--enable-breakable-cuda-graph`, but you must declare every served resolution in `--warmup-resolutions` so warmup captures matching graph signatures.
|
||||
|
||||
|
||||
@@ -547,6 +547,10 @@ SGLang Diffusion x Cache-DiT supports almost all models originally supported in
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Hunyuan</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>HunyuanVideo</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>MiniMax</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>MiniMax-H3 (T2VA, FL2VA, and Ref2VA)</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
@@ -562,7 +566,8 @@ SGLang Diffusion x Cache-DiT supports almost all models originally supported in
|
||||
|
||||
- **SGLang-native pipelines**: Distributed Cache-DiT paths exist for supported pipelines. Hybrid SP+TP configurations add communication and cache coordination overhead, so validate them on the target model and hardware before using them as production defaults.
|
||||
- **SCM minimum steps**: SCM requires >= 8 inference steps to be effective
|
||||
- **Model support**: Only models registered in Cache-DiT's BlockAdapterRegister are supported
|
||||
- **Model support**: The model must be registered in Cache-DiT's
|
||||
`BlockAdapterRegister` or have an SGLang custom block adapter.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
|
||||
@@ -115,6 +115,12 @@ Rows are grouped when a family shares the same runtime path or optimization supp
|
||||
<td>Video-audio, 360p / 720p; local MOVA detector aliases are also supported.</td>
|
||||
<td><span className="sgd-muted">No dedicated optimization listed</span></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>MiniMax-H3</td>
|
||||
<td><div className="sgd-id-list"><code>MiniMaxAI/MiniMax-H3</code></div></td>
|
||||
<td>T2VA / FL2VA / Ref2VA, 768p at 24 fps with synchronized audio</td>
|
||||
<td><span className="sgd-chip">Cache-DiT</span><span className="sgd-chip">Online FP8</span></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Wan2.1 Fun</td>
|
||||
<td><div className="sgd-id-list"><code>weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers</code></div></td>
|
||||
@@ -548,6 +554,21 @@ Optimization columns are abbreviated to keep the matrix readable:
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>MiniMax-H3 (T2VA / FL2VA / Ref2VA image, audio, video/V2V)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>MiniMaxAI/MiniMax-H3</code></td>
|
||||
<td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768p · 24 fps</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3 (one/two-stage/TI2V/HQ)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Lightricks/LTX-2.3</code></td>
|
||||
|
||||
@@ -134,6 +134,11 @@ sglang generate \
|
||||
--save-output
|
||||
```
|
||||
|
||||
MiniMax-H3 supports this path while preserving its required FP32 patch,
|
||||
timestep, and output projections. See the
|
||||
[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#6-runtime-feature-recipes)
|
||||
for its distributed serving recipe.
|
||||
|
||||
### MXFP4 Online Quantization
|
||||
|
||||
MXFP4 provides aggressive 4-bit compression with online quantization. **Note: Requires ROCm and MI350+ (gfx95x) GPU.**
|
||||
|
||||
@@ -99,7 +99,10 @@ const selectionSpace = (config) => {
|
||||
|
||||
const walk = (dir) => readdirSync(dir, { withFileTypes: true }).flatMap((e) =>
|
||||
e.isDirectory() ? walk(join(dir, e.name))
|
||||
: (e.name.endsWith(".jsx") && !e.name.includes("benchmark") ? [join(dir, e.name)] : []));
|
||||
: (e.name.endsWith(".jsx")
|
||||
&& !e.name.includes("benchmark")
|
||||
&& e.name !== "popular-models.jsx"
|
||||
? [join(dir, e.name)] : []));
|
||||
|
||||
for (const path of walk(CONFIGS)) {
|
||||
const where = relative(join(SNIPPETS, ".."), path);
|
||||
@@ -182,6 +185,14 @@ for (const path of walk(CONFIGS)) {
|
||||
}
|
||||
}
|
||||
}
|
||||
if (typeof config.curl === "function") {
|
||||
probe((sel) => {
|
||||
const out = config.curl(sel, null);
|
||||
if (typeof out !== "string") {
|
||||
throw new Error(`curl returned ${typeof out}, expected a string`);
|
||||
}
|
||||
}, "curl");
|
||||
}
|
||||
}
|
||||
|
||||
if (failures.length) {
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
// vendor picks the selector group: blackwell | hopper | amd.
|
||||
// `multiNodeDockerFlags: string[]` (either source) adds
|
||||
// `docker run` flags the platform's fabric needs
|
||||
// groupHardware optional — set false to show one flat hardware row
|
||||
// variants/quantizations/strategies/nodesOptions LEGACY 4-dim option lists,
|
||||
// used when `matchDims` is absent (nodesOptions id is
|
||||
// `single` or `multi-N` → --nnodes N)
|
||||
@@ -40,7 +41,9 @@
|
||||
// whose verification round is open rather than absent.
|
||||
// modelNames HF slug lookup, `hw|variant|quant` then `variant|quant`
|
||||
// placeholders {{KEY}} → {target: 'command'|'curl', label, default?}
|
||||
// curl cURL template (uses {{MODEL_NAME}} + placeholders)
|
||||
// curl cURL template (uses {{MODEL_NAME}} + placeholders), or
|
||||
// `(selection, cell) => template` when the request payload
|
||||
// depends on a custom match/overlay dimension
|
||||
// benchmarkCommands optional — powers the "⚡ Reproduce" modal (speed +
|
||||
// per-eval accuracy templates)
|
||||
// defaultAccuracy optional — per-variant accuracy merged under cell.accuracy
|
||||
@@ -56,8 +59,14 @@
|
||||
// dockerImages optional — `docker run` image, keyed by
|
||||
// `hw|quant|strategy` then `hw|quant` then `hw`;
|
||||
// falls back to `lmsysorg/sglang:dev`
|
||||
// dockerMounts optional — additional `-v` mount specs
|
||||
// dockerRunCommand optional — command placed after the image and before
|
||||
// generated server flags; string or `(selection) => string`
|
||||
// runModes optional — command output tabs to show (`python` and/or
|
||||
// `docker`); defaults to both, in that order
|
||||
// `docker`), as an array or `(selection) => array`;
|
||||
// defaults to both, in that order
|
||||
// showPlaygroundLink optional — false hides the "Open the Playground" footer
|
||||
// for cookbooks that only expose the deployment matrix
|
||||
// github optional — "Submit verified cell" issue-template overrides
|
||||
// playgroundFeatures optional — consumed by _playground.jsx (see its header)
|
||||
//
|
||||
@@ -689,6 +698,9 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
const di = config.dockerImages || {};
|
||||
const image = di[`${sel.hw}|${sel.quant}|${sel.strategy}`]
|
||||
|| di[`${sel.hw}|${sel.quant}`] || di[sel.hw] || "lmsysorg/sglang:dev";
|
||||
const dockerRunCommand = typeof config.dockerRunCommand === "function"
|
||||
? config.dockerRunCommand(sel)
|
||||
: (config.dockerRunCommand || "sglang serve");
|
||||
const portFlag = flags.find((x) => x.split(/[\s=]/)[0] === "--port");
|
||||
const servePort = portFlag ? portFlag.slice("--port".length).trim() : "{{PORT}}";
|
||||
const vendorOf = (hwId) => {
|
||||
@@ -728,13 +740,14 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
multinode ? " --network host" : ` -p ${servePort}:${servePort}`,
|
||||
...(multinode ? fabricFlagsOf(sel.hw).map((f) => " " + f) : []),
|
||||
" -v ~/.cache/huggingface:/root/.cache/huggingface",
|
||||
...(config.dockerMounts || []).map((mount) => ` -v ${mount}`),
|
||||
// HF token only for gated checkpoints — configs that declare an HF_TOKEN placeholder.
|
||||
...(config.placeholders && config.placeholders.HF_TOKEN
|
||||
? [` --env "HF_TOKEN={{HF_TOKEN}}"`] : []),
|
||||
...cellEnv.map((e) => ` --env ${e}`),
|
||||
" --ipc=host",
|
||||
` ${image}`,
|
||||
" sglang serve",
|
||||
` ${dockerRunCommand}`,
|
||||
...flags.map((f) => " " + f),
|
||||
];
|
||||
cmd = dockerLines.join(" \\\n");
|
||||
@@ -1030,6 +1043,9 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
.map((hw) => ({ id: hw.id, label: hw.label, subtitle: hw.vram }));
|
||||
if (items.length) groups.push({ label: vendor.toUpperCase(), items });
|
||||
}
|
||||
if (config.groupHardware === false) {
|
||||
return [{ label: null, items: groups.flatMap((group) => group.items) }];
|
||||
}
|
||||
return groups;
|
||||
};
|
||||
|
||||
@@ -1161,8 +1177,17 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
const [benchConc, setBenchConc] = useState(null);
|
||||
const [benchAcc, setBenchAcc] = useState(null);
|
||||
const [benchCopied, setBenchCopied] = useState(null);
|
||||
const runModes = config.runModes || ["python", "docker"];
|
||||
const configuredRunModes = typeof config.runModes === "function"
|
||||
? config.runModes(sel)
|
||||
: config.runModes;
|
||||
const runModes = configuredRunModes || ["python", "docker"];
|
||||
const [runMode, setRunMode] = useState(runModes[0]); // "python" | "docker"
|
||||
const hasRunMode = runModes.includes(runMode);
|
||||
const fallbackRunMode = runModes[0];
|
||||
const activeRunMode = hasRunMode ? runMode : fallbackRunMode;
|
||||
useEffect(() => {
|
||||
if (!hasRunMode) setRunMode(fallbackRunMode);
|
||||
}, [hasRunMode, fallbackRunMode]);
|
||||
useEffect(() => { if (modal === "env") setEnvDraft(env); }, [modal, env]);
|
||||
|
||||
// Live --mamba-full-memory-ratio from the ratio calculator (K3 pages):
|
||||
@@ -1194,7 +1219,7 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
else flags.push(line);
|
||||
return { ...cell, flags };
|
||||
})();
|
||||
const command = renderCommand(cellWithRatio, sel, env, runMode);
|
||||
const command = renderCommand(cellWithRatio, sel, env, activeRunMode);
|
||||
// Speculative-decoding hint on the EFFECTIVE flags — speculation can arrive via
|
||||
// the Spec Decode overlay as well as the cell. SGLang resets
|
||||
// --max-running-requests to 48 when spec is on and it's unset; verified for both
|
||||
@@ -1263,7 +1288,9 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
return out;
|
||||
};
|
||||
const modelName = resolveModelName(sel);
|
||||
const curlText = interpolate(config.curl || "", env, modelName);
|
||||
const curlTemplate =
|
||||
typeof config.curl === "function" ? config.curl(sel, cell) : config.curl;
|
||||
const curlText = interpolate(curlTemplate || "", env, modelName);
|
||||
const hwGroups = buildHardwareGroups();
|
||||
const benchEntry = benchmarks ? findBenchmark(benchmarks, sel) : null;
|
||||
|
||||
@@ -1385,8 +1412,8 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
<div style={s.cardColumn}>
|
||||
<div style={{ ...s.title, marginBottom: "2px" }}>Hardware Platform</div>
|
||||
{hwGroups.map((g) => (
|
||||
<div key={g.label} style={s.vendorRow}>
|
||||
<div style={s.vendorLabel}>{g.label}</div>
|
||||
<div key={g.label || "hardware"} style={s.vendorRow}>
|
||||
{g.label && <div style={s.vendorLabel}>{g.label}</div>}
|
||||
<div style={s.itemsGrid(maxHwCols)}>
|
||||
{g.items.map((item) => renderButton(item, "hw", sel.hw))}
|
||||
{Array.from({ length: maxHwCols - g.items.length }).map((_, i) => (
|
||||
@@ -1431,13 +1458,13 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
key={mode}
|
||||
style={{
|
||||
...(index === runModes.length - 1
|
||||
? s.runModeChipLast(runMode === mode)
|
||||
: s.runModeChip(runMode === mode)),
|
||||
? s.runModeChipLast(activeRunMode === mode)
|
||||
: s.runModeChip(activeRunMode === mode)),
|
||||
...(runModes.length === 1 ? { borderRadius: 7 } : {}),
|
||||
}}
|
||||
onClick={() => setRunMode(mode)}
|
||||
role="tab"
|
||||
aria-selected={runMode === mode}
|
||||
aria-selected={activeRunMode === mode}
|
||||
>
|
||||
{mode === "docker" ? "Docker" : "Python"}
|
||||
</span>
|
||||
@@ -1473,38 +1500,40 @@ export const Deployment = ({ config, benchmarks }) => {
|
||||
|
||||
{/* Playground link — scrollIntoView, not an href, so the hash (which
|
||||
carries the selection) isn't overwritten. */}
|
||||
<div
|
||||
style={{
|
||||
padding: "6px 12px",
|
||||
fontSize: "12px",
|
||||
color: isDark ? "#9ca3af" : "#6b7280",
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
gap: "6px",
|
||||
}}
|
||||
>
|
||||
<span>Need to go beyond the verified matrix?</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
const el = document.getElementById("playground");
|
||||
if (el) el.scrollIntoView({ behavior: "smooth", block: "start" });
|
||||
}}
|
||||
{config.showPlaygroundLink !== false && (
|
||||
<div
|
||||
style={{
|
||||
background: "transparent",
|
||||
border: "none",
|
||||
padding: 0,
|
||||
color: isDark ? "#FDBA74" : "#C2410C",
|
||||
cursor: "pointer",
|
||||
padding: "6px 12px",
|
||||
fontSize: "12px",
|
||||
fontWeight: 600,
|
||||
textDecoration: "underline",
|
||||
textUnderlineOffset: "2px",
|
||||
color: isDark ? "#9ca3af" : "#6b7280",
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
gap: "6px",
|
||||
}}
|
||||
>
|
||||
Open the Playground →
|
||||
</button>
|
||||
</div>
|
||||
<span>Need to go beyond the verified matrix?</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
const el = document.getElementById("playground");
|
||||
if (el) el.scrollIntoView({ behavior: "smooth", block: "start" });
|
||||
}}
|
||||
style={{
|
||||
background: "transparent",
|
||||
border: "none",
|
||||
padding: 0,
|
||||
color: isDark ? "#FDBA74" : "#C2410C",
|
||||
cursor: "pointer",
|
||||
fontSize: "12px",
|
||||
fontWeight: 600,
|
||||
textDecoration: "underline",
|
||||
textUnderlineOffset: "2px",
|
||||
}}
|
||||
>
|
||||
Open the Playground →
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* cURL modal */}
|
||||
{modal === "curl" && (
|
||||
|
||||
@@ -1390,6 +1390,9 @@ export const Playground = ({ config }) => {
|
||||
const di = config.dockerImages || {};
|
||||
const image = di[`${sel.hw}|${sel.quant}|${sel.strategy}`]
|
||||
|| di[`${sel.hw}|${sel.quant}`] || di[sel.hw] || "lmsysorg/sglang:dev";
|
||||
const dockerRunCommand = typeof config.dockerRunCommand === "function"
|
||||
? config.dockerRunCommand(sel)
|
||||
: (config.dockerRunCommand || "sglang serve");
|
||||
const portFlag = f.find((x) => x.split(/[\s=]/)[0] === "--port");
|
||||
const servePort = portFlag ? portFlag.slice("--port".length).trim() : "{{PORT}}";
|
||||
// Mirrors `multiNodeDockerFlags` on the _deployment.jsx HARDWARE_CATALOG
|
||||
@@ -1406,11 +1409,12 @@ export const Playground = ({ config }) => {
|
||||
(multinode || pdMode) ? " --network host" : ` -p ${servePort}:${servePort}`,
|
||||
...(multinode ? fabricFlags.map((x) => " " + x) : []),
|
||||
" -v ~/.cache/huggingface:/root/.cache/huggingface",
|
||||
...(config.dockerMounts || []).map((mount) => ` -v ${mount}`),
|
||||
` --env "HF_TOKEN={{HF_TOKEN}}"`,
|
||||
...cellEnv.map((e) => ` --env ${e}`),
|
||||
" --ipc=host",
|
||||
` ${image}`,
|
||||
" sglang serve",
|
||||
` ${dockerRunCommand}`,
|
||||
...f.map((x) => " " + x),
|
||||
];
|
||||
cmd = dockerLines.join(" \\\n");
|
||||
|
||||
@@ -0,0 +1,618 @@
|
||||
// MiniMax-H3 diffusion deployment matrix. Consumed by _deployment.jsx.
|
||||
//
|
||||
// The mode, quantization, and encoder choices are deployment overlays because
|
||||
// they do not change which base hardware topology fits. Request sampling
|
||||
// controls remain in the generated cURL instead of being mixed into this
|
||||
// deployment matrix.
|
||||
// Hardware/profile cells remain deliberately small and carry an honest
|
||||
// verification state for the exact platform, rather than inheriting a result
|
||||
// measured on a different GPU.
|
||||
|
||||
|
||||
export const config = {
|
||||
modelName: "MiniMax-H3",
|
||||
|
||||
supportedHardware: [
|
||||
"b200",
|
||||
"b300",
|
||||
"h200",
|
||||
"h100",
|
||||
"mi300x",
|
||||
"mi355x",
|
||||
"rtx5090",
|
||||
],
|
||||
hardware: [
|
||||
{ id: "rtx5090", label: "RTX 5090", vram: "32GB", vendor: "consumer" },
|
||||
],
|
||||
groupHardware: false,
|
||||
|
||||
matchDims: [
|
||||
{
|
||||
id: "profile",
|
||||
title: "Deployment Profile",
|
||||
showWhen: (s) => ["b200", "b300", "h200", "h100"].includes(s.hw),
|
||||
options: [
|
||||
{ id: "resident", label: "Resident" },
|
||||
{
|
||||
id: "fsdp",
|
||||
label: "FSDP sharded",
|
||||
showWhen: (s) =>
|
||||
["b200", "b300", "h200", "h100"].includes(s.hw),
|
||||
},
|
||||
{
|
||||
id: "offload",
|
||||
label: "Layerwise offload",
|
||||
showWhen: (s) => s.hw === "rtx5090",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
|
||||
overlayDims: [
|
||||
{
|
||||
id: "weights",
|
||||
title: "Checkpoint Weights",
|
||||
default: "fl2va",
|
||||
options: [
|
||||
{
|
||||
id: "fl2va",
|
||||
label: "FL2VA (First-and-Last-Frame-to-Video-and-Audio)",
|
||||
flags: ["--model-variant fl2va"],
|
||||
},
|
||||
{
|
||||
id: "ref2va",
|
||||
label: "Ref2VA (Reference-to-Video-and-Audio)",
|
||||
flags: ["--model-variant ref2va"],
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
id: "mode",
|
||||
title: "Request Mode",
|
||||
default: "t2va",
|
||||
options: [
|
||||
{
|
||||
id: "t2va",
|
||||
label: "Text only",
|
||||
showWhen: (s) => s.weights === "fl2va",
|
||||
},
|
||||
{
|
||||
id: "i2va",
|
||||
label: "First frame",
|
||||
showWhen: (s) => s.weights === "fl2va",
|
||||
},
|
||||
{
|
||||
id: "l2va",
|
||||
label: "Last frame",
|
||||
showWhen: (s) => s.weights === "fl2va",
|
||||
},
|
||||
{
|
||||
id: "fl2va",
|
||||
label: "First + last frames",
|
||||
showWhen: (s) => s.weights === "fl2va",
|
||||
},
|
||||
{
|
||||
id: "ref_image",
|
||||
label: "Image reference",
|
||||
showWhen: (s) => s.weights === "ref2va",
|
||||
},
|
||||
{
|
||||
id: "ref_image_audio",
|
||||
label: "Image + audio",
|
||||
showWhen: (s) => s.weights === "ref2va",
|
||||
},
|
||||
{
|
||||
id: "v2v",
|
||||
label: "Video reference",
|
||||
showWhen: (s) => s.weights === "ref2va",
|
||||
},
|
||||
{
|
||||
id: "video_audio",
|
||||
label: "Video + soundtrack",
|
||||
showWhen: (s) => s.weights === "ref2va",
|
||||
},
|
||||
{
|
||||
id: "audio_only",
|
||||
label: "Audio reference",
|
||||
showWhen: (s) => s.weights === "ref2va",
|
||||
},
|
||||
{
|
||||
id: "mixed_ref",
|
||||
label: "Mixed references",
|
||||
showWhen: (s) => s.weights === "ref2va",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
id: "quant",
|
||||
title: "Online Quantization",
|
||||
default: "bf16",
|
||||
showWhen: (s) => ["b200", "b300"].includes(s.hw),
|
||||
options: [
|
||||
{ id: "bf16", label: "Off — Native BF16/FP32" },
|
||||
{
|
||||
id: "fp8",
|
||||
label: "FP8 — Approximate",
|
||||
showWhen: (s) => ["b200", "b300"].includes(s.hw),
|
||||
disabled: (s) => s.profile !== "resident",
|
||||
disableReason:
|
||||
"The documented FP8 operating point keeps the transformer resident; FSDP combinations have not been validated.",
|
||||
flags: ["--quantization fp8"],
|
||||
hints: [
|
||||
"Online FP8 is approximate. Validate both video and audio quality;",
|
||||
"verified B200 and B300 runs reduced memory; re-benchmark latency on the target workload.",
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
id: "encoder",
|
||||
title: "Text Encoder Parallel",
|
||||
default: "auto",
|
||||
options: [
|
||||
{
|
||||
id: "auto",
|
||||
label: "Auto (recommended)",
|
||||
hints: [
|
||||
"Auto uses folding for the single-request recipes below and can",
|
||||
"select data parallel encoding for a compatible TP1 request batch.",
|
||||
],
|
||||
},
|
||||
{
|
||||
id: "fold",
|
||||
label: "Fold (single-request)",
|
||||
flags: ["--encoder-parallel fold"],
|
||||
hints: [
|
||||
"Fold shards the resident Qwen3-VL encoder across the replica and is",
|
||||
"best suited to single-node GPUs with fast peer-to-peer links.",
|
||||
],
|
||||
},
|
||||
{
|
||||
id: "dp",
|
||||
label: "DP (batched throughput)",
|
||||
disabled: (s) =>
|
||||
s.hw === "rtx5090" ||
|
||||
(s.hw === "h100" && s.profile === "resident"),
|
||||
disableReason:
|
||||
"Encoder DP requires TP1 and DiT DP1; this verified recipe uses TP2.",
|
||||
flags: [
|
||||
"--encoder-parallel dp",
|
||||
"--batching-max-size {{BATCHING_MAX_SIZE}}",
|
||||
],
|
||||
hints: [
|
||||
"DP distributes a compatible multi-request text batch across ranks;",
|
||||
"it does not improve a batch of one and replicates encoder weights.",
|
||||
],
|
||||
},
|
||||
{
|
||||
id: "replicate",
|
||||
label: "Replicate (compatibility)",
|
||||
flags: ["--encoder-parallel replicate"],
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
|
||||
modelNames: {
|
||||
default: "MiniMaxAI/MiniMax-H3",
|
||||
},
|
||||
|
||||
placeholders: {
|
||||
HOST_IP: {
|
||||
target: "command",
|
||||
label: "Bind host",
|
||||
default: "0.0.0.0",
|
||||
},
|
||||
PORT: {
|
||||
target: "command",
|
||||
label: "Bind port",
|
||||
default: "30010",
|
||||
},
|
||||
HF_TOKEN: {
|
||||
target: "command",
|
||||
label: "HF token (Docker)",
|
||||
default: "<your-hf-token>",
|
||||
},
|
||||
MEDIA_DIR: {
|
||||
target: "command",
|
||||
label: "Host media directory (Docker)",
|
||||
default: "/data/minimax-h3",
|
||||
},
|
||||
CURL_HOST: {
|
||||
target: "curl",
|
||||
label: "Server host",
|
||||
default: "localhost",
|
||||
},
|
||||
CURL_PORT: {
|
||||
target: "curl",
|
||||
label: "Server port",
|
||||
default: "30010",
|
||||
},
|
||||
NUM_OUTPUTS: {
|
||||
target: "curl",
|
||||
label: "Outputs per prompt (1-10)",
|
||||
default: "1",
|
||||
},
|
||||
BATCHING_MAX_SIZE: {
|
||||
target: "command",
|
||||
label: "Maximum request batch size",
|
||||
default: "2",
|
||||
},
|
||||
DURATION_SECONDS: {
|
||||
target: "curl",
|
||||
label: "Duration (seconds, 4-15)",
|
||||
default: "5",
|
||||
},
|
||||
FIRST_FRAME: {
|
||||
target: "curl",
|
||||
label: "FL2VA first frame URI",
|
||||
default: "file:///data/minimax-h3/first-frame.png",
|
||||
},
|
||||
LAST_FRAME: {
|
||||
target: "curl",
|
||||
label: "FL2VA last frame URI",
|
||||
default: "file:///data/minimax-h3/last-frame.png",
|
||||
},
|
||||
INPUT_VIDEO: {
|
||||
target: "curl",
|
||||
label: "First video URI",
|
||||
default: "file:///data/minimax-h3/video-1.mp4",
|
||||
},
|
||||
INPUT_VIDEO_START_SECONDS: {
|
||||
target: "curl",
|
||||
label: "First video start (seconds)",
|
||||
default: "0",
|
||||
},
|
||||
SECOND_INPUT_VIDEO: {
|
||||
target: "curl",
|
||||
label: "Second video URI (mixed ref)",
|
||||
default: "file:///data/minimax-h3/video-2.mp4",
|
||||
},
|
||||
SECOND_INPUT_VIDEO_START_SECONDS: {
|
||||
target: "curl",
|
||||
label: "Second video start (seconds)",
|
||||
default: "0",
|
||||
},
|
||||
REFERENCE_IMAGE: {
|
||||
target: "curl",
|
||||
label: "First reference image URI",
|
||||
default: "file:///data/minimax-h3/reference-1.png",
|
||||
},
|
||||
SECOND_REFERENCE_IMAGE: {
|
||||
target: "curl",
|
||||
label: "Second reference image URI",
|
||||
default: "file:///data/minimax-h3/reference-2.png",
|
||||
},
|
||||
REFERENCE_AUDIO: {
|
||||
target: "curl",
|
||||
label: "First reference audio URI",
|
||||
default: "file:///data/minimax-h3/reference-1.mp3",
|
||||
},
|
||||
SECOND_REFERENCE_AUDIO: {
|
||||
target: "curl",
|
||||
label: "Second reference audio URI",
|
||||
default: "file:///data/minimax-h3/reference-2.mp3",
|
||||
},
|
||||
},
|
||||
|
||||
curl: (s) => {
|
||||
const request = {
|
||||
model: "{{MODEL_NAME}}",
|
||||
prompt:
|
||||
"Night-vision bedroom footage: while the owner sleeps, three cats burst in playing tiny brass instruments at full volume, freeze, then march out as if nothing happened.",
|
||||
seconds: "{{DURATION_SECONDS}}",
|
||||
task: "t2va",
|
||||
conditions: [],
|
||||
target: {
|
||||
short_edge: 768,
|
||||
aspect_ratio: "16:9",
|
||||
duration_seconds: "{{DURATION_SECONDS}}",
|
||||
},
|
||||
num_outputs_per_prompt: "{{NUM_OUTPUTS}}",
|
||||
num_inference_steps: 50,
|
||||
flow_shift: 12.0,
|
||||
audio_flow_shift: 3.0,
|
||||
seed: 1101,
|
||||
};
|
||||
const imageReference = (uri) => ({
|
||||
type: "image",
|
||||
uri,
|
||||
role: "reference",
|
||||
});
|
||||
const audioReference = (uri) => ({
|
||||
type: "audio",
|
||||
uri,
|
||||
role: "reference",
|
||||
});
|
||||
const videoReference = (uri, start, type = "video") => ({
|
||||
type,
|
||||
uri,
|
||||
role: "reference",
|
||||
start_time_seconds: start,
|
||||
});
|
||||
|
||||
if (["i2va", "l2va", "fl2va"].includes(s.mode)) {
|
||||
request.task = "fl2va";
|
||||
request.prompt =
|
||||
"Continue naturally between the supplied endpoint frame or frames, with synchronized ambient sound.";
|
||||
request.target.aspect_ratio = "auto";
|
||||
request.seed = 2101;
|
||||
request.conditions = [];
|
||||
if (s.mode !== "l2va") {
|
||||
request.conditions.push({
|
||||
type: "image",
|
||||
uri: "{{FIRST_FRAME}}",
|
||||
role: "keyframe",
|
||||
frame_index: 0,
|
||||
});
|
||||
}
|
||||
if (s.mode !== "i2va") {
|
||||
request.conditions.push({
|
||||
type: "image",
|
||||
uri: "{{LAST_FRAME}}",
|
||||
role: "keyframe",
|
||||
frame_index: -1,
|
||||
});
|
||||
}
|
||||
} else if (s.mode === "ref_image") {
|
||||
request.task = "ref2va";
|
||||
request.prompt = "Use <Picture 1> as the visual subject and style reference.";
|
||||
request.target.aspect_ratio = "auto";
|
||||
request.conditions = [imageReference("{{REFERENCE_IMAGE}}")];
|
||||
request.seed = 3101;
|
||||
} else if (s.mode === "ref_image_audio") {
|
||||
request.task = "ref2va";
|
||||
request.prompt =
|
||||
"Use <Picture 1> as the visual subject and <Audio 1> as the sound reference.";
|
||||
request.target.aspect_ratio = "auto";
|
||||
request.conditions = [
|
||||
imageReference("{{REFERENCE_IMAGE}}"),
|
||||
audioReference("{{REFERENCE_AUDIO}}"),
|
||||
];
|
||||
request.seed = 3102;
|
||||
} else if (s.mode === "v2v" || s.mode === "video_audio") {
|
||||
request.task = "ref2va";
|
||||
request.prompt =
|
||||
s.mode === "video_audio"
|
||||
? "Follow <Video 1> and its required <Audio 1> soundtrack with coherent synchronized motion."
|
||||
: "Follow the appearance and motion of <Video 1>; use its soundtrack when present.";
|
||||
request.conditions = [
|
||||
videoReference(
|
||||
"{{INPUT_VIDEO}}",
|
||||
"{{INPUT_VIDEO_START_SECONDS}}",
|
||||
s.mode === "video_audio" ? "video_audio" : "video",
|
||||
),
|
||||
];
|
||||
request.seed = s.mode === "video_audio" ? 4102 : 4101;
|
||||
} else if (s.mode === "audio_only") {
|
||||
request.task = "ref2va";
|
||||
request.prompt = "Build a coherent visual scene around <Audio 1>.";
|
||||
request.conditions = [audioReference("{{REFERENCE_AUDIO}}")];
|
||||
request.seed = 3103;
|
||||
} else if (s.mode === "mixed_ref") {
|
||||
request.task = "ref2va";
|
||||
request.prompt =
|
||||
"Combine <Picture 1>, <Picture 2>, <Audio 1>, <Audio 2>, <Video 1>, and <Video 2> in their one-based modality order.";
|
||||
request.conditions = [
|
||||
imageReference("{{REFERENCE_IMAGE}}"),
|
||||
imageReference("{{SECOND_REFERENCE_IMAGE}}"),
|
||||
audioReference("{{REFERENCE_AUDIO}}"),
|
||||
audioReference("{{SECOND_REFERENCE_AUDIO}}"),
|
||||
videoReference("{{INPUT_VIDEO}}", "{{INPUT_VIDEO_START_SECONDS}}"),
|
||||
videoReference(
|
||||
"{{SECOND_INPUT_VIDEO}}",
|
||||
"{{SECOND_INPUT_VIDEO_START_SECONDS}}",
|
||||
),
|
||||
];
|
||||
request.seed = 3104;
|
||||
}
|
||||
|
||||
const body = JSON.stringify(request, null, 2).replace(
|
||||
/"{{(NUM_OUTPUTS|DURATION_SECONDS|INPUT_VIDEO_START_SECONDS|SECOND_INPUT_VIDEO_START_SECONDS)}}"/g,
|
||||
"{{$1}}",
|
||||
);
|
||||
return `curl -sS -X POST http://{{CURL_HOST}}:{{CURL_PORT}}/v1/videos \\
|
||||
-H 'Content-Type: application/json' \\
|
||||
-d '${body}'`;
|
||||
},
|
||||
|
||||
dockerMounts: ["{{MEDIA_DIR}}:/data/minimax-h3:ro"],
|
||||
|
||||
dockerRunCommand: (s) =>
|
||||
["mi300x", "mi355x"].includes(s.hw)
|
||||
? `bash -lc 'python -m pip install -e "/sgl-workspace/sglang/python[diffusion_hip]" && exec sglang serve "$@"' --`
|
||||
: `bash -lc 'python -m pip install -e "/sgl-workspace/sglang/python[diffusion]" && exec sglang serve "$@"' --`,
|
||||
|
||||
// Publish AMD Docker only after an H3-capable ROCm image has been validated.
|
||||
runModes: (s) =>
|
||||
["mi300x", "mi355x"].includes(s.hw)
|
||||
? ["python"]
|
||||
: ["python", "docker"],
|
||||
|
||||
dockerImages: {
|
||||
b200: "lmsysorg/sglang:dev",
|
||||
b300: "lmsysorg/sglang:dev",
|
||||
h200: "lmsysorg/sglang:dev",
|
||||
h100: "lmsysorg/sglang:dev",
|
||||
},
|
||||
|
||||
showPlaygroundLink: false,
|
||||
|
||||
cells: [
|
||||
{
|
||||
match: { hw: "b200", profile: "resident" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 8",
|
||||
"--ulysses-degree 8",
|
||||
"--performance-mode speed",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
},
|
||||
{
|
||||
match: { hw: "b300", profile: "resident" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 8",
|
||||
"--ulysses-degree 8",
|
||||
"--performance-mode speed",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"This is the B300 topology used for the documented benchmark sweep, not a claimed minimum GPU count.",
|
||||
},
|
||||
{
|
||||
match: { hw: "h200", profile: "resident" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 4",
|
||||
"--ulysses-degree 4",
|
||||
"--performance-mode speed",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
},
|
||||
{
|
||||
match: { hw: "b300", profile: "fsdp" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 8",
|
||||
"--ulysses-degree 8",
|
||||
"--performance-mode speed",
|
||||
"--use-fsdp-inference true",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"FSDP reduces resident DiT memory but adds per-block parameter collectives. Prefer Resident when the full pipeline fits.",
|
||||
},
|
||||
{
|
||||
match: { hw: "h200", profile: "fsdp" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 4",
|
||||
"--ulysses-degree 4",
|
||||
"--performance-mode speed",
|
||||
"--use-fsdp-inference true",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"FSDP reduces resident DiT memory but adds per-block parameter collectives. Prefer Resident when the full pipeline fits.",
|
||||
},
|
||||
{
|
||||
match: { hw: "b200", profile: "fsdp" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 4",
|
||||
"--ulysses-degree 4",
|
||||
"--performance-mode speed",
|
||||
"--use-fsdp-inference true",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"The 4-GPU FSDP path is lossless but slower than the 8-GPU resident recipe.",
|
||||
},
|
||||
{
|
||||
match: { hw: "h100", profile: "resident" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 4",
|
||||
"--tp-size 2",
|
||||
"--ulysses-degree 2",
|
||||
"--performance-mode speed",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"Fastest measured 4× H100 80 GB topology. TP4 + Ulysses1 lowers peak memory at a small latency cost.",
|
||||
},
|
||||
{
|
||||
match: { hw: "h100", profile: "fsdp" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 4",
|
||||
"--ulysses-degree 4",
|
||||
"--performance-mode speed",
|
||||
"--use-fsdp-inference true",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"Capacity path on 4× H100 80 GB. Prefer the resident TP2 + Ulysses2 profile for latency.",
|
||||
},
|
||||
{
|
||||
match: { hw: "mi300x", profile: "resident" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
env: ["SGLANG_USE_AITER=1"],
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 8",
|
||||
"--ulysses-degree 8",
|
||||
"--performance-mode speed",
|
||||
"--attention-backend aiter",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"Validated on 1×, 2×, 4×, and 8× MI300X with BF16 and AITER packed attention. The picker emits the fastest measured 8-GPU topology; set --num-gpus and --ulysses-degree to the same lower count for a measured capacity recipe.",
|
||||
},
|
||||
{
|
||||
match: { hw: "mi355x", profile: "resident" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
env: ["SGLANG_USE_AITER=1"],
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 8",
|
||||
"--ulysses-degree 8",
|
||||
"--performance-mode speed",
|
||||
"--attention-backend aiter",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"Validated on 1×, 2×, 4×, and 8× MI355X with BF16 and AITER packed attention. The picker emits the fastest measured 8-GPU topology; set --num-gpus and --ulysses-degree to the same lower count for a measured capacity recipe.",
|
||||
},
|
||||
{
|
||||
match: { hw: "rtx5090", profile: "offload" },
|
||||
nnodes: 1,
|
||||
verified: true,
|
||||
flags: [
|
||||
"--model-path {{MODEL_NAME}}",
|
||||
"--num-gpus 2",
|
||||
"--tp-size 2",
|
||||
"--ulysses-degree 1",
|
||||
"--performance-mode memory",
|
||||
"--layerwise-offload-components dit,text_encoder,vae",
|
||||
"--dit-offload-prefetch-size 1",
|
||||
"--dit-layerwise-resident-layers 20",
|
||||
"--enable-torch-compile false",
|
||||
"--host {{HOST_IP}}",
|
||||
"--port {{PORT}}",
|
||||
],
|
||||
warn:
|
||||
"Validated lossless BF16/FP32 recipe on 2× RTX 5090 (32 GB each) with a 384 GiB-class host. TP2 avoids the full per-rank DiT replication observed with Ulysses2 on PCIe.",
|
||||
},
|
||||
],
|
||||
};
|
||||
@@ -6,6 +6,8 @@
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <sgl_kernel/impl/norm.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
|
||||
#include <cstdint>
|
||||
@@ -42,7 +44,8 @@ constexpr uint32_t active_mask() {
|
||||
}
|
||||
}
|
||||
|
||||
SGL_DEVICE float load_cache_value(const float* ptr, int64_t idx) {
|
||||
template <typename CacheDType>
|
||||
SGL_DEVICE CacheDType load_cache_value(const CacheDType* ptr, int64_t idx) {
|
||||
#ifdef USE_ROCM
|
||||
return ptr[idx];
|
||||
#else
|
||||
@@ -50,7 +53,15 @@ SGL_DEVICE float load_cache_value(const float* ptr, int64_t idx) {
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, int64_t kRopeDim, bool kIsNeox, bool kUsePDL, typename DType, typename IdType>
|
||||
template <
|
||||
int64_t kHeadDim,
|
||||
int64_t kRopeDim,
|
||||
bool kIsNeox,
|
||||
bool kUsePDL,
|
||||
typename DType,
|
||||
typename CacheDType,
|
||||
bool kRoundNormBeforeRope,
|
||||
typename IdType>
|
||||
__global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__ params) {
|
||||
using namespace device;
|
||||
|
||||
@@ -63,14 +74,17 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
|
||||
constexpr uint32_t kRotaryLanes = kRopeDim / kElemsPerThread;
|
||||
constexpr uint32_t kHalfRotaryLanes = kRotaryLanes / 2;
|
||||
constexpr uint32_t kActiveMask = active_mask<kRotaryLanes>();
|
||||
constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float);
|
||||
constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(CacheDType);
|
||||
|
||||
static_assert(kElemsPerThread % 2 == 0, "Each lane must own an even number of elements");
|
||||
static_assert(kRopeDim > 0 && kRopeDim <= kHeadDim, "Invalid rope dimension");
|
||||
static_assert(kRopeDim % kElemsPerThread == 0, "rope_dim must align with per-lane vector width");
|
||||
static_assert(
|
||||
!kIsNeox || (kRotaryLanes >= 2 && ((kRotaryLanes & (kRotaryLanes - 1)) == 0)),
|
||||
"NeoX fused qknorm+rope requires rotary lane count to be a power of 2");
|
||||
!kIsNeox || (kRotaryLanes >= 2 && kRotaryLanes % 2 == 0),
|
||||
"NeoX fused qknorm+rope requires an even rotary lane count");
|
||||
static_assert(
|
||||
!kRoundNormBeforeRope || std::is_same_v<DType, CacheDType>,
|
||||
"Rounded QKNorm+RoPE requires cache and activation dtypes to match");
|
||||
|
||||
using Packed = packed_t<DType>;
|
||||
using Storage = AlignedVector<Packed, kVecSize>;
|
||||
@@ -98,6 +112,53 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
|
||||
auto input_vec = load_as<Storage>(input, lane_id);
|
||||
const auto weight_vec = load_as<Storage>(weight_ptr, lane_id);
|
||||
|
||||
if constexpr (kRoundNormBeforeRope) {
|
||||
auto output_vec = norm::apply_norm_warp<kHeadDim>(input_vec, weight_vec, eps);
|
||||
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
|
||||
const auto cos_ptr = static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
|
||||
const auto sin_ptr = cos_ptr + kRopeDim / 2;
|
||||
|
||||
if constexpr (kIsNeox) {
|
||||
if (lane_id < kRotaryLanes) {
|
||||
const auto partner_lane =
|
||||
lane_id < kHalfRotaryLanes ? lane_id + kHalfRotaryLanes : lane_id - kHalfRotaryLanes;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecSize; ++j) {
|
||||
auto partner_vec = output_vec[j];
|
||||
auto partner_bits = reinterpret_cast<const uint32_t&>(partner_vec);
|
||||
partner_bits = __shfl_sync(kActiveMask, partner_bits, partner_lane);
|
||||
reinterpret_cast<uint32_t&>(partner_vec) = partner_bits;
|
||||
auto& values = unpack(output_vec[j]);
|
||||
const auto& partner_values = unpack(partner_vec);
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < 2; ++i) {
|
||||
const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + 2 * j + i;
|
||||
const auto cos = load_cache_value(cos_ptr, half_idx);
|
||||
const auto sin = load_cache_value(sin_ptr, half_idx);
|
||||
values[i] = lane_id < kHalfRotaryLanes ? values[i] * cos - partner_values[i] * sin
|
||||
: values[i] * cos + partner_values[i] * sin;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (lane_id < kRotaryLanes) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kVecSize; ++j) {
|
||||
auto& values = unpack(output_vec[j]);
|
||||
const auto half_idx = lane_id * kElemsPerThread / 2 + j;
|
||||
const auto cos = load_cache_value(cos_ptr, half_idx);
|
||||
const auto sin = load_cache_value(sin_ptr, half_idx);
|
||||
const auto x = values[0];
|
||||
const auto y = values[1];
|
||||
values[0] = x * cos - y * sin;
|
||||
values[1] = y * cos + x * sin;
|
||||
}
|
||||
}
|
||||
}
|
||||
store_as<Storage>(const_cast<void*>(input), output_vec, lane_id);
|
||||
continue;
|
||||
}
|
||||
|
||||
float elems[kElemsPerThread];
|
||||
float sum_of_squares = 0.0f;
|
||||
|
||||
@@ -122,27 +183,28 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
|
||||
if constexpr (kIsNeox) {
|
||||
if (lane_id < kRotaryLanes) {
|
||||
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
|
||||
const auto cos_ptr = static_cast<const float*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
|
||||
const auto cos_ptr =
|
||||
static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
|
||||
const auto sin_ptr = cos_ptr + kRopeDim / 2;
|
||||
const auto partner_lane = lane_id < kHalfRotaryLanes ? lane_id + kHalfRotaryLanes : lane_id - kHalfRotaryLanes;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemsPerThread; ++i) {
|
||||
float swapped = __shfl_xor_sync(kActiveMask, elems[i], kHalfRotaryLanes);
|
||||
float swapped = __shfl_sync(kActiveMask, elems[i], partner_lane);
|
||||
if (lane_id < kHalfRotaryLanes) {
|
||||
swapped = -swapped;
|
||||
}
|
||||
int dim_idx = static_cast<int>(lane_id * kElemsPerThread + i);
|
||||
dim_idx = (dim_idx * 2) % kRopeDim;
|
||||
const int half_idx = dim_idx / 2;
|
||||
const float cos = load_cache_value(cos_ptr, half_idx);
|
||||
const float sin = load_cache_value(sin_ptr, half_idx);
|
||||
const auto half_idx = (lane_id % kHalfRotaryLanes) * kElemsPerThread + i;
|
||||
const float cos = cast<fp32_t>(load_cache_value(cos_ptr, half_idx));
|
||||
const float sin = cast<fp32_t>(load_cache_value(sin_ptr, half_idx));
|
||||
elems[i] = elems[i] * cos + swapped * sin;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (lane_id < kRotaryLanes) {
|
||||
const auto pos = static_cast<int64_t>(static_cast<const IdType*>(positions)[token_id]);
|
||||
const auto cos_ptr = static_cast<const float*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
|
||||
const auto cos_ptr =
|
||||
static_cast<const CacheDType*>(pointer::offset(cos_sin_cache_ptr, pos * kCosSinStrideBytes));
|
||||
const auto sin_ptr = cos_ptr + kRopeDim / 2;
|
||||
|
||||
#pragma unroll
|
||||
@@ -150,8 +212,8 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
|
||||
const float x = elems[i];
|
||||
const float y = elems[i + 1];
|
||||
const int half_idx = static_cast<int>(lane_id * kElemsPerThread + i) / 2;
|
||||
const float cos = load_cache_value(cos_ptr, half_idx);
|
||||
const float sin = load_cache_value(sin_ptr, half_idx);
|
||||
const float cos = cast<fp32_t>(load_cache_value(cos_ptr, half_idx));
|
||||
const float sin = cast<fp32_t>(load_cache_value(sin_ptr, half_idx));
|
||||
elems[i] = x * cos - y * sin;
|
||||
elems[i + 1] = y * cos + x * sin;
|
||||
}
|
||||
@@ -168,11 +230,19 @@ __global__ void fused_qknorm_rope_warp(const QKNormRopeParams __grid_constant__
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
template <int64_t kHeadDim, int64_t kRopeDim, bool kIsNeox, bool kUsePDL, typename DType>
|
||||
template <
|
||||
int64_t kHeadDim,
|
||||
int64_t kRopeDim,
|
||||
bool kIsNeox,
|
||||
bool kUsePDL,
|
||||
typename DType,
|
||||
typename CacheDType,
|
||||
bool kRoundNormBeforeRope>
|
||||
struct QKNormRopeKernel {
|
||||
static_assert(kHeadDim <= 256, "Only head_dim <= 256 is supported");
|
||||
template <typename IdType>
|
||||
static constexpr auto kernel = fused_qknorm_rope_warp<kHeadDim, kRopeDim, kIsNeox, kUsePDL, DType, IdType>;
|
||||
static constexpr auto kernel =
|
||||
fused_qknorm_rope_warp<kHeadDim, kRopeDim, kIsNeox, kUsePDL, DType, CacheDType, kRoundNormBeforeRope, IdType>;
|
||||
|
||||
static void
|
||||
run(const tvm::ffi::TensorView q,
|
||||
@@ -201,7 +271,7 @@ struct QKNormRopeKernel {
|
||||
TensorMatcher({N, Q, D}).with_strides({Dq, Dd, 1}).with_dtype<DType>().with_device(device).verify(q);
|
||||
TensorMatcher({N, K, D}).with_strides({Dk, Dd, 1}).with_dtype<DType>().with_device(device).verify(k);
|
||||
TensorMatcher({D}).with_dtype<DType>().with_device(device).verify(q_weight).verify(k_weight);
|
||||
TensorMatcher({-1, R}).with_dtype<float>().with_device(device).verify(cos_sin_cache);
|
||||
TensorMatcher({-1, R}).with_dtype<CacheDType>().with_device(device).verify(cos_sin_cache);
|
||||
TensorMatcher({N}).with_dtype<int32_t, int64_t>(id_type).with_device(device).verify(positions);
|
||||
|
||||
const auto num_tokens = static_cast<uint32_t>(N.unwrap());
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
// CUDA fast path for the Ulysses sequence-parallel output head merge.
|
||||
//
|
||||
// usp_merge_heads:
|
||||
// x [W, S, B, h_local, D] (contiguous, the output all-to-all result)
|
||||
// -> out [B, S, W, h_local, D] (contiguous)
|
||||
// Replaces `x.permute(2, 1, 0, 3, 4).contiguous()` on the head_dim=2
|
||||
// output path of `_usp_output_all_to_all`.
|
||||
//
|
||||
// A pure copy (no arithmetic), so it is bit-exact with the eager permute by
|
||||
// construction. It exists because ATen's generic permute-copy reaches well
|
||||
// under half of HBM bandwidth on the packed-DiT shapes, while a single pass
|
||||
// with coalesced vectorized stores runs near roofline.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/tensor.h> // For host dtype helpers and TensorView metadata
|
||||
#include <sgl_kernel/utils.h> // For RuntimeCheck and div_ceil
|
||||
|
||||
#include <sgl_kernel/type.cuh> // For CUDA dtype aliases
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
#include <sgl_kernel/vec.cuh> // For device::AlignedVector
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace sglang_usp_relayout {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int64_t kMaxGrid = 65535;
|
||||
|
||||
inline const char* data_ptr(const tvm::ffi::TensorView& t) {
|
||||
return static_cast<const char*>(t.data_ptr()) + t.byte_offset();
|
||||
}
|
||||
|
||||
inline char* mutable_data_ptr(const tvm::ffi::TensorView& t) {
|
||||
return static_cast<char*>(t.data_ptr()) + t.byte_offset();
|
||||
}
|
||||
|
||||
inline bool aligned16(const void* p) {
|
||||
return (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
|
||||
}
|
||||
|
||||
inline int64_t numel(const tvm::ffi::TensorView& t) {
|
||||
int64_t n = 1;
|
||||
for (int i = 0; i < t.ndim(); ++i) {
|
||||
n *= t.size(i);
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
inline int64_t grid_for(int64_t total) {
|
||||
int64_t grid = host::div_ceil(total, static_cast<int64_t>(kBlockSize));
|
||||
if (grid < 1) {
|
||||
grid = 1;
|
||||
}
|
||||
if (grid > kMaxGrid) {
|
||||
grid = kMaxGrid;
|
||||
}
|
||||
return grid;
|
||||
}
|
||||
|
||||
inline bool is_dense_contiguous(const tvm::ffi::TensorView& t) {
|
||||
int64_t expected = 1;
|
||||
for (int i = t.ndim() - 1; i >= 0; --i) {
|
||||
if (t.size(i) == 1) {
|
||||
continue;
|
||||
}
|
||||
if (t.stride(i) != expected) {
|
||||
return false;
|
||||
}
|
||||
expected *= t.size(i);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void check_dtype(const tvm::ffi::TensorView& t) {
|
||||
host::RuntimeCheck(host::is_type<T>(t.dtype()), "unexpected dtype for usp_merge_heads tensor");
|
||||
}
|
||||
|
||||
// out[b, s, w, h, c] = x[w, s, b, h, c]
|
||||
template <typename T, int kVec>
|
||||
__global__ void usp_merge_heads_vec_kernel(
|
||||
T* __restrict__ out,
|
||||
const T* __restrict__ x,
|
||||
int64_t n_vec,
|
||||
int64_t d_vec, // D / kVec
|
||||
int64_t h_local,
|
||||
int64_t batch,
|
||||
int64_t seq,
|
||||
int64_t world) {
|
||||
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
|
||||
for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < n_vec; i += stride) {
|
||||
int64_t rest = i;
|
||||
const int64_t c_vec = rest % d_vec;
|
||||
rest /= d_vec;
|
||||
const int64_t h = rest % h_local;
|
||||
rest /= h_local;
|
||||
const int64_t w = rest % world;
|
||||
rest /= world;
|
||||
const int64_t s = rest % seq;
|
||||
const int64_t b = rest / seq;
|
||||
|
||||
const int64_t src_vec = ((((w * seq + s) * batch + b) * h_local) + h) * d_vec + c_vec;
|
||||
device::AlignedVector<T, kVec> val;
|
||||
val.load(x, src_vec);
|
||||
val.store(out, i);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void usp_merge_heads_scalar_kernel(
|
||||
T* __restrict__ out,
|
||||
const T* __restrict__ x,
|
||||
int64_t total,
|
||||
int64_t head_dim,
|
||||
int64_t h_local,
|
||||
int64_t batch,
|
||||
int64_t seq,
|
||||
int64_t world) {
|
||||
const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
|
||||
for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < total; i += stride) {
|
||||
int64_t rest = i;
|
||||
const int64_t c = rest % head_dim;
|
||||
rest /= head_dim;
|
||||
const int64_t h = rest % h_local;
|
||||
rest /= h_local;
|
||||
const int64_t w = rest % world;
|
||||
rest /= world;
|
||||
const int64_t s = rest % seq;
|
||||
const int64_t b = rest / seq;
|
||||
|
||||
out[i] = x[((((w * seq + s) * batch + b) * h_local) + h) * head_dim + c];
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename T>
|
||||
struct UspMergeHeadsKernel {
|
||||
static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) {
|
||||
check_dtype<T>(out);
|
||||
check_dtype<T>(x);
|
||||
host::RuntimeCheck(x.ndim() == 5, "x must be [W, S, B, h_local, D]");
|
||||
host::RuntimeCheck(out.ndim() == 5, "out must be [B, S, W, h_local, D]");
|
||||
for (auto* t : {&x, &out}) {
|
||||
host::RuntimeCheck(t->device().device_type == kDLCUDA, "usp_merge_heads tensors must be CUDA");
|
||||
host::RuntimeCheck(is_dense_contiguous(*t), "usp_merge_heads tensors must be contiguous");
|
||||
}
|
||||
const int64_t world = x.size(0);
|
||||
const int64_t seq = x.size(1);
|
||||
const int64_t batch = x.size(2);
|
||||
const int64_t h_local = x.size(3);
|
||||
const int64_t head_dim = x.size(4);
|
||||
host::RuntimeCheck(
|
||||
out.size(0) == batch && out.size(1) == seq && out.size(2) == world && out.size(3) == h_local &&
|
||||
out.size(4) == head_dim,
|
||||
"out must be the [B, S, W, h_local, D] permutation of x");
|
||||
|
||||
const int64_t total = numel(x);
|
||||
if (total == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
T* out_ptr = reinterpret_cast<T*>(mutable_data_ptr(out));
|
||||
const T* x_ptr = reinterpret_cast<const T*>(data_ptr(x));
|
||||
|
||||
constexpr int kVec = 16 / sizeof(T);
|
||||
const bool vec_ok = (head_dim % kVec == 0) && aligned16(out_ptr) && aligned16(x_ptr);
|
||||
if (vec_ok) {
|
||||
const int64_t n_vec = total / kVec;
|
||||
host::LaunchKernel(static_cast<uint32_t>(grid_for(n_vec)), kBlockSize, out.device())(
|
||||
usp_merge_heads_vec_kernel<T, kVec>, out_ptr, x_ptr, n_vec, head_dim / kVec, h_local, batch, seq, world);
|
||||
} else {
|
||||
host::LaunchKernel(static_cast<uint32_t>(grid_for(total)), kBlockSize, out.device())(
|
||||
usp_merge_heads_scalar_kernel<T>, out_ptr, x_ptr, total, head_dim, h_local, batch, seq, world);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang_usp_relayout
|
||||
@@ -55,7 +55,13 @@ struct ActivationParams {
|
||||
uint32_t expert_step;
|
||||
};
|
||||
|
||||
template <typename T, ActivationKind kAct, bool kUsePDL, bool kFilterExpert>
|
||||
template <
|
||||
typename T,
|
||||
ActivationKind kAct,
|
||||
bool kUsePDL,
|
||||
bool kFilterExpert,
|
||||
bool kRoundActivation = false,
|
||||
bool kReuseInput = false>
|
||||
__global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams params) {
|
||||
using namespace device;
|
||||
constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
|
||||
@@ -70,7 +76,7 @@ __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams para
|
||||
}
|
||||
const auto offset = tid % num_vecs;
|
||||
const auto input_offset = token_id * (num_vecs * 2) + offset;
|
||||
const auto output_offset = tid;
|
||||
const auto output_offset = kReuseInput ? input_offset : tid;
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
const auto gate = device::load_as<vec_t>(params.input, input_offset);
|
||||
const auto up = device::load_as<vec_t>(params.input, input_offset + num_vecs);
|
||||
@@ -79,9 +85,18 @@ __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams para
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
const float gate_f32 = device::cast<fp32_t>(gate[i]);
|
||||
const float up_f32 = device::cast<fp32_t>(up[i]);
|
||||
out[i] = device::cast<T>(apply_activation_f32<kAct>(gate_f32) * up_f32);
|
||||
if constexpr (kRoundActivation) {
|
||||
const T activated = device::cast<T>(apply_activation_f32<kAct>(gate_f32));
|
||||
out[i] = device::cast<T>(device::cast<fp32_t>(activated) * up_f32);
|
||||
} else {
|
||||
out[i] = device::cast<T>(apply_activation_f32<kAct>(gate_f32) * up_f32);
|
||||
}
|
||||
}
|
||||
if constexpr (kReuseInput) {
|
||||
device::store_as<vec_t>(const_cast<void*>(params.input), out, output_offset);
|
||||
} else {
|
||||
device::store_as<vec_t>(params.out, out, output_offset);
|
||||
}
|
||||
device::store_as<vec_t>(params.out, out, output_offset);
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
@@ -117,26 +132,28 @@ struct ActivationKernel {
|
||||
using kernel_fn_t = decltype(&act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false>);
|
||||
using unary_kernel_fn_t = decltype(&act_kernel<T, ActivationKind::kReLU2, kUsePDL>);
|
||||
|
||||
template <ActivationKind kAct, bool kFilterExpert>
|
||||
static constexpr kernel_fn_t activation_kernel = act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert>;
|
||||
template <ActivationKind kAct, bool kFilterExpert, bool kRoundActivation = false, bool kReuseInput = false>
|
||||
static constexpr kernel_fn_t activation_kernel =
|
||||
act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert, kRoundActivation, kReuseInput>;
|
||||
|
||||
static_assert(device::kMaxVecBytes % sizeof(T) == 0, "unsupported data type");
|
||||
|
||||
template <bool kFilterExpert>
|
||||
template <bool kFilterExpert, bool kRoundActivation = false, bool kReuseInput = false>
|
||||
static kernel_fn_t select_kernel(const std::string& type) {
|
||||
using namespace host;
|
||||
if (type == "silu") {
|
||||
return activation_kernel<ActivationKind::kSiLU, kFilterExpert>;
|
||||
return activation_kernel<ActivationKind::kSiLU, kFilterExpert, kRoundActivation, kReuseInput>;
|
||||
} else if (type == "gelu") {
|
||||
return activation_kernel<ActivationKind::kGELU, kFilterExpert>;
|
||||
return activation_kernel<ActivationKind::kGELU, kFilterExpert, kRoundActivation, kReuseInput>;
|
||||
} else if (type == "gelu_tanh") {
|
||||
return activation_kernel<ActivationKind::kGELUTanh, kFilterExpert>;
|
||||
return activation_kernel<ActivationKind::kGELUTanh, kFilterExpert, kRoundActivation, kReuseInput>;
|
||||
} else {
|
||||
Panic("unsupported activation type: ", type);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
template <bool kRoundActivation = false, bool kReuseInput = false>
|
||||
static void launch(
|
||||
const tvm::ffi::TensorView& input,
|
||||
const tvm::ffi::TensorView& out,
|
||||
@@ -151,10 +168,11 @@ struct ActivationKernel {
|
||||
auto device_ = SymbolicDevice{};
|
||||
device_.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, D_out}) //
|
||||
.with_dtype<T>()
|
||||
.with_device(device_)
|
||||
.verify(out);
|
||||
if constexpr (kReuseInput) {
|
||||
TensorMatcher({N, D_out}).with_strides({D_in, 1}).with_dtype<T>().with_device(device_).verify(out);
|
||||
} else {
|
||||
TensorMatcher({N, D_out}).with_dtype<T>().with_device(device_).verify(out);
|
||||
}
|
||||
TensorMatcher({N, D_in}) //
|
||||
.with_dtype<T>()
|
||||
.with_device(device_)
|
||||
@@ -166,13 +184,16 @@ struct ActivationKernel {
|
||||
if (num_tokens == 0) return;
|
||||
RuntimeCheck(hidden_size * 2 == D_in.unwrap(), "invalid activation dimension");
|
||||
RuntimeCheck(hidden_size % kVecSize == 0, "hidden size must be divisible by vector size");
|
||||
if constexpr (kReuseInput) {
|
||||
RuntimeCheck(input.data_ptr() == out.data_ptr(), "in-place activation output must alias input");
|
||||
}
|
||||
// only get once to avoid overhead
|
||||
const auto num_total_items = num_tokens * (hidden_size / kVecSize);
|
||||
RuntimeCheck(num_total_items <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
|
||||
const auto num_blocks = div_ceil(static_cast<uint32_t>(num_total_items), kBlockSize);
|
||||
const auto params = ActivationParams{
|
||||
.input = input.data_ptr(),
|
||||
.out = out.data_ptr(),
|
||||
.out = kReuseInput ? nullptr : out.data_ptr(),
|
||||
.hidden_dim = hidden_size,
|
||||
.num_tokens = num_tokens,
|
||||
.expert_ids = expert_ids,
|
||||
@@ -180,10 +201,10 @@ struct ActivationKernel {
|
||||
};
|
||||
if (expert_ids != nullptr) {
|
||||
RuntimeCheck(expert_step > 0, "expert_step must be positive");
|
||||
const auto kernel = select_kernel<true>(type);
|
||||
const auto kernel = select_kernel<true, kRoundActivation, kReuseInput>(type);
|
||||
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
|
||||
} else {
|
||||
const auto kernel = select_kernel<false>(type);
|
||||
const auto kernel = select_kernel<false, kRoundActivation, kReuseInput>(type);
|
||||
LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
}
|
||||
@@ -192,6 +213,16 @@ struct ActivationKernel {
|
||||
launch(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
|
||||
}
|
||||
|
||||
static void
|
||||
run_activation_with_rounding(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
|
||||
launch<true>(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
|
||||
}
|
||||
|
||||
static void run_activation_with_rounding_input_inplace(
|
||||
const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
|
||||
launch<true, true>(input, out, type, /*expert_ids=*/nullptr, /*expert_step=*/1);
|
||||
}
|
||||
|
||||
static void run_activation_filtered(
|
||||
const tvm::ffi::TensorView input,
|
||||
const tvm::ffi::TensorView out,
|
||||
|
||||
@@ -29,15 +29,26 @@ def _fast_math_flags() -> list[str]:
|
||||
|
||||
|
||||
@cache_once
|
||||
def activation_module(dtype: torch.dtype) -> Module:
|
||||
def activation_module(dtype: torch.dtype, *, fast_math: bool = True) -> Module:
|
||||
fast_math_flags = _fast_math_flags()
|
||||
if not fast_math and not fast_math_flags:
|
||||
return activation_module(dtype)
|
||||
args = make_cpp_args(dtype, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"activation",
|
||||
"activation" if fast_math else "rounded_activation",
|
||||
*args,
|
||||
cuda_files=["elementwise/activation.cuh"],
|
||||
extra_cuda_cflags=_fast_math_flags(),
|
||||
extra_cuda_cflags=fast_math_flags if fast_math else [],
|
||||
cuda_wrappers=[
|
||||
("run_activation", f"ActivationKernel<{args}>::run_activation"),
|
||||
(
|
||||
"run_activation_with_rounding",
|
||||
f"ActivationKernel<{args}>::run_activation_with_rounding",
|
||||
),
|
||||
(
|
||||
"run_activation_with_rounding_input_inplace",
|
||||
f"ActivationKernel<{args}>::run_activation_with_rounding_input_inplace",
|
||||
),
|
||||
(
|
||||
"run_activation_filtered",
|
||||
f"ActivationKernel<{args}>::run_activation_filtered",
|
||||
@@ -65,6 +76,28 @@ def _run_activation_inplace(
|
||||
module.run_activation(input_2d, out_2d, op_name)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["out"])
|
||||
def _run_activation_with_rounding_inplace(
|
||||
op_name: str, input: torch.Tensor, out: torch.Tensor
|
||||
) -> None:
|
||||
hidden_size = input.shape[-1] // 2
|
||||
# Fast-math changes FP16 SiLU at eager rounding boundaries on SM90.
|
||||
module = activation_module(input.dtype, fast_math=False)
|
||||
input_2d = input.view(-1, hidden_size * 2)
|
||||
out_2d = out.view(-1, hidden_size)
|
||||
module.run_activation_with_rounding(input_2d, out_2d, op_name)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["input"])
|
||||
def _run_silu_and_mul_with_rounding_inplace(input: torch.Tensor) -> None:
|
||||
hidden_size = input.shape[-1] // 2
|
||||
module = activation_module(input.dtype, fast_math=False)
|
||||
input_2d = input.view(-1, hidden_size * 2)
|
||||
module.run_activation_with_rounding_input_inplace(
|
||||
input_2d, input_2d[:, :hidden_size], "silu"
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["out"])
|
||||
def _run_activation_filtered_inplace(
|
||||
op_name: str,
|
||||
@@ -150,6 +183,23 @@ def silu_and_mul(
|
||||
return run_activation("silu", input, out, expert_ids, expert_step)
|
||||
|
||||
|
||||
def silu_and_mul_with_activation_rounding(
|
||||
input: torch.Tensor,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
hidden_size = input.shape[-1] // 2
|
||||
if out is None:
|
||||
out = input.new_empty(*input.shape[:-1], hidden_size)
|
||||
_run_activation_with_rounding_inplace("silu", input, out)
|
||||
return out
|
||||
|
||||
|
||||
def silu_and_mul_with_activation_rounding_(input: torch.Tensor) -> torch.Tensor:
|
||||
hidden_size = input.shape[-1] // 2
|
||||
_run_silu_and_mul_with_rounding_inplace(input)
|
||||
return input[..., :hidden_size]
|
||||
|
||||
|
||||
def gelu_and_mul(
|
||||
input: torch.Tensor,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
|
||||
@@ -26,8 +26,18 @@ def _jit_qknorm_rope_module(
|
||||
rope_dim: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
cache_dtype: torch.dtype,
|
||||
round_norm_before_rope: bool,
|
||||
) -> Module:
|
||||
args = make_cpp_args(head_dim, rope_dim, is_neox, is_arch_support_pdl(), dtype)
|
||||
args = make_cpp_args(
|
||||
head_dim,
|
||||
rope_dim,
|
||||
is_neox,
|
||||
is_arch_support_pdl(),
|
||||
dtype,
|
||||
cache_dtype,
|
||||
round_norm_before_rope,
|
||||
)
|
||||
return load_jit(
|
||||
"qknorm_rope",
|
||||
*args,
|
||||
@@ -43,6 +53,8 @@ def can_use_fused_inplace_qknorm_rope(
|
||||
rope_dim: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
cache_dtype: torch.dtype = torch.float32,
|
||||
round_norm_before_rope: bool = False,
|
||||
) -> bool:
|
||||
if head_dim not in (64, 128, 256):
|
||||
logger.warning(f"Unsupported head_dim={head_dim} for JIT fused QKNorm+RoPE")
|
||||
@@ -62,15 +74,29 @@ def can_use_fused_inplace_qknorm_rope(
|
||||
return False
|
||||
if is_neox:
|
||||
rotary_lanes = rope_dim // elems_per_thread
|
||||
if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
|
||||
if rotary_lanes < 2 or rotary_lanes % 2:
|
||||
logger.warning(
|
||||
"rope_dim=%s yields invalid rotary_lanes=%s for neox fused QKNorm+RoPE; rotary lane count must be a power of 2",
|
||||
"rope_dim=%s yields invalid rotary_lanes=%s for neox fused QKNorm+RoPE; rotary lane count must be even",
|
||||
rope_dim,
|
||||
rotary_lanes,
|
||||
)
|
||||
return False
|
||||
if round_norm_before_rope and cache_dtype != dtype:
|
||||
logger.warning(
|
||||
"Exact fused QKNorm+RoPE requires cache dtype %s to match activation dtype %s",
|
||||
cache_dtype,
|
||||
dtype,
|
||||
)
|
||||
return False
|
||||
try:
|
||||
_jit_qknorm_rope_module(head_dim, rope_dim, is_neox, dtype)
|
||||
_jit_qknorm_rope_module(
|
||||
head_dim,
|
||||
rope_dim,
|
||||
is_neox,
|
||||
dtype,
|
||||
cache_dtype,
|
||||
round_norm_before_rope,
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load JIT fused QKNorm+RoPE kernel: {e}")
|
||||
@@ -90,8 +116,16 @@ def fused_inplace_qknorm_rope(
|
||||
eps: float = 1e-6,
|
||||
head_dim: int = 0,
|
||||
rope_dim: int = 0,
|
||||
round_norm_before_rope: bool = False,
|
||||
) -> None:
|
||||
head_dim = head_dim or q.size(-1)
|
||||
rope_dim = rope_dim or cos_sin_cache.size(-1)
|
||||
module = _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, q.dtype)
|
||||
module = _jit_qknorm_rope_module(
|
||||
head_dim,
|
||||
rope_dim,
|
||||
is_neox,
|
||||
q.dtype,
|
||||
cos_sin_cache.dtype,
|
||||
round_norm_before_rope,
|
||||
)
|
||||
module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _round_bf16_to_fp32(value):
|
||||
# force the eager BF16 kernel boundary so Triton cannot contract the next add
|
||||
bits = value.to(tl.int32, bitcast=True)
|
||||
rounding_bias = 0x7FFF + ((bits >> 16) & 1)
|
||||
rounded_bits = (bits + rounding_bias) & -65536
|
||||
return rounded_bits.to(tl.float32, bitcast=True)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _indexed_scale_shift_bf16_kernel(
|
||||
output_ptr,
|
||||
x_ptr,
|
||||
shift_ptr,
|
||||
scale_ptr,
|
||||
indices_ptr,
|
||||
hidden_size,
|
||||
stride_x_row,
|
||||
stride_shift_row,
|
||||
stride_scale_row,
|
||||
stride_indices,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
columns = tl.arange(0, BLOCK_N)
|
||||
mask = columns < hidden_size
|
||||
index = tl.load(indices_ptr + row * stride_indices)
|
||||
|
||||
x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
shift = tl.load(
|
||||
shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
scale = tl.load(
|
||||
scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
|
||||
one_plus_scale = _round_bf16_to_fp32(1.0 + scale)
|
||||
scaled = _round_bf16_to_fp32(x * one_plus_scale)
|
||||
tl.store(
|
||||
output_ptr + row * stride_x_row + columns,
|
||||
scaled + shift,
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _indexed_gate_bf16_kernel(
|
||||
output_ptr,
|
||||
x_ptr,
|
||||
gate_ptr,
|
||||
other_ptr,
|
||||
indices_ptr,
|
||||
hidden_size,
|
||||
stride_x_row,
|
||||
stride_gate_row,
|
||||
stride_other_row,
|
||||
stride_indices,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
columns = tl.arange(0, BLOCK_N)
|
||||
mask = columns < hidden_size
|
||||
index = tl.load(indices_ptr + row * stride_indices)
|
||||
|
||||
x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
gate = tl.load(
|
||||
gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
other = tl.load(
|
||||
other_ptr + row * stride_other_row + columns, mask=mask, other=0.0
|
||||
).to(tl.float32)
|
||||
|
||||
gated = _round_bf16_to_fp32(gate * other)
|
||||
tl.store(
|
||||
output_ptr + row * stride_x_row + columns,
|
||||
x + gated,
|
||||
mask=mask,
|
||||
)
|
||||
|
||||
|
||||
def indexed_scale_shift_bf16_(
|
||||
x: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
rows, hidden_size = x.shape
|
||||
if rows == 0:
|
||||
return x
|
||||
block_n = triton.next_power_of_2(hidden_size)
|
||||
_indexed_scale_shift_bf16_kernel[(rows,)](
|
||||
x,
|
||||
x,
|
||||
shift,
|
||||
scale,
|
||||
indices,
|
||||
hidden_size,
|
||||
x.stride(0),
|
||||
shift.stride(0),
|
||||
scale.stride(0),
|
||||
indices.stride(0),
|
||||
BLOCK_N=block_n,
|
||||
num_warps=8,
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
def indexed_gate_bf16_(
|
||||
x: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
other: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
rows, hidden_size = x.shape
|
||||
if rows == 0:
|
||||
return x
|
||||
block_n = triton.next_power_of_2(hidden_size)
|
||||
_indexed_gate_bf16_kernel[(rows,)](
|
||||
x,
|
||||
x,
|
||||
gate,
|
||||
other,
|
||||
indices,
|
||||
hidden_size,
|
||||
x.stride(0),
|
||||
gate.stride(0),
|
||||
other.stride(0),
|
||||
indices.stride(0),
|
||||
BLOCK_N=block_n,
|
||||
num_warps=8,
|
||||
)
|
||||
return x
|
||||
@@ -5,6 +5,81 @@ import triton.language as tl # type: ignore
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fp32_mul_add_rn(x, scale, residual):
|
||||
"""Match separate CUDA FP32 multiply and add rounding (no FMA)."""
|
||||
return tl.inline_asm_elementwise(
|
||||
asm="""{
|
||||
.reg .f32 product;
|
||||
mul.rn.f32 product, $1, $2;
|
||||
add.rn.f32 $0, $3, product;
|
||||
}""",
|
||||
constraints="=f,f,f,f",
|
||||
args=(x, scale, residual),
|
||||
dtype=tl.float32,
|
||||
is_pure=True,
|
||||
pack=1,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_scaled_residual_add_exact_kernel(
|
||||
output_ptr,
|
||||
residual_ptr,
|
||||
x_ptr,
|
||||
scale_ptr,
|
||||
numel: tl.constexpr,
|
||||
width: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < numel
|
||||
x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32)
|
||||
scale = tl.load(scale_ptr + offsets % width, mask=mask)
|
||||
residual = tl.load(residual_ptr + offsets, mask=mask)
|
||||
output = _fp32_mul_add_rn(x, scale, residual)
|
||||
tl.store(output_ptr + offsets, output, mask=mask)
|
||||
|
||||
|
||||
def try_fused_scaled_residual_add_exact(
|
||||
residual: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
) -> torch.Tensor | None:
|
||||
"""Fuse ``residual + x * scale`` without changing eager FP32 rounding."""
|
||||
if (
|
||||
not current_platform.is_cuda()
|
||||
or torch.is_grad_enabled()
|
||||
or torch.compiler.is_compiling()
|
||||
or residual.dtype != torch.float32
|
||||
or x.dtype not in (torch.float16, torch.bfloat16)
|
||||
or scale.dtype != torch.float32
|
||||
or not residual.is_cuda
|
||||
or residual.device != x.device
|
||||
or residual.device != scale.device
|
||||
or residual.shape != x.shape
|
||||
or scale.shape != (x.shape[-1],)
|
||||
or not residual.is_contiguous()
|
||||
or not x.is_contiguous()
|
||||
or not scale.is_contiguous()
|
||||
or x.numel() == 0
|
||||
):
|
||||
return None
|
||||
|
||||
output = torch.empty_like(residual)
|
||||
block_size = 1024
|
||||
_fused_scaled_residual_add_exact_kernel[(triton.cdiv(x.numel(), block_size),)](
|
||||
output,
|
||||
residual,
|
||||
x,
|
||||
scale,
|
||||
numel=x.numel(),
|
||||
width=x.shape[-1],
|
||||
BLOCK_SIZE=block_size,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_layernorm_scale_shift_gate_select01_kernel(
|
||||
output_ptr,
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _pack_qkv_destination_major_kernel(
|
||||
output_ptr,
|
||||
q_ptr,
|
||||
k_ptr,
|
||||
v_ptr,
|
||||
total_elements,
|
||||
rows,
|
||||
local_heads,
|
||||
head_size,
|
||||
stride_q_row,
|
||||
stride_q_head,
|
||||
stride_k_row,
|
||||
stride_k_head,
|
||||
stride_v_row,
|
||||
stride_v_head,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < total_elements
|
||||
|
||||
dim = offsets % head_size
|
||||
head_slot = offsets // head_size
|
||||
local_head = head_slot % local_heads
|
||||
row_slot = head_slot // local_heads
|
||||
row = row_slot % rows
|
||||
destination = row_slot // rows
|
||||
global_head = destination * local_heads + local_head
|
||||
|
||||
q = tl.load(
|
||||
q_ptr + row * stride_q_row + global_head * stride_q_head + dim,
|
||||
mask=mask,
|
||||
)
|
||||
k = tl.load(
|
||||
k_ptr + row * stride_k_row + global_head * stride_k_head + dim,
|
||||
mask=mask,
|
||||
)
|
||||
v = tl.load(
|
||||
v_ptr + row * stride_v_row + global_head * stride_v_head + dim,
|
||||
mask=mask,
|
||||
)
|
||||
output_base = head_slot * (3 * head_size) + dim
|
||||
tl.store(output_ptr + output_base, q, mask=mask)
|
||||
tl.store(output_ptr + output_base + head_size, k, mask=mask)
|
||||
tl.store(output_ptr + output_base + 2 * head_size, v, mask=mask)
|
||||
|
||||
|
||||
def pack_qkv_destination_major(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
world_size: int,
|
||||
) -> torch.Tensor:
|
||||
rows, global_heads, head_size = q.shape
|
||||
local_heads = global_heads // world_size
|
||||
output = torch.empty(
|
||||
world_size,
|
||||
rows,
|
||||
local_heads,
|
||||
3 * head_size,
|
||||
dtype=q.dtype,
|
||||
device=q.device,
|
||||
)
|
||||
total_elements = rows * global_heads * head_size
|
||||
if total_elements == 0:
|
||||
return output
|
||||
|
||||
block_size = 1024
|
||||
_pack_qkv_destination_major_kernel[(triton.cdiv(total_elements, block_size),)](
|
||||
output,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
total_elements,
|
||||
rows,
|
||||
local_heads,
|
||||
head_size,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
k.stride(0),
|
||||
k.stride(1),
|
||||
v.stride(0),
|
||||
v.stride(1),
|
||||
BLOCK_SIZE=block_size,
|
||||
num_warps=8,
|
||||
)
|
||||
return output
|
||||
@@ -0,0 +1,83 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_usp_relayout_module(dtype: torch.dtype) -> Module:
|
||||
args = make_cpp_args(dtype)
|
||||
return load_jit(
|
||||
"diffusion_usp_relayout",
|
||||
*args,
|
||||
cuda_files=["diffusion/usp_relayout.cuh"],
|
||||
cuda_wrappers=[
|
||||
(
|
||||
"usp_merge_heads",
|
||||
"sglang_usp_relayout::" f"UspMergeHeadsKernel<{args}>::run",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _fake_merge_heads(x: torch.Tensor) -> torch.Tensor:
|
||||
world, seq, batch, h_local, head_dim = x.shape
|
||||
return x.new_empty((batch, seq, world, h_local, head_dim))
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
op_name="diffusion_usp_merge_heads",
|
||||
mutates_args=[],
|
||||
fake_impl=_fake_merge_heads,
|
||||
)
|
||||
def _usp_merge_heads_custom_op(x: torch.Tensor) -> torch.Tensor:
|
||||
world, seq, batch, h_local, head_dim = x.shape
|
||||
out = x.new_empty((batch, seq, world, h_local, head_dim))
|
||||
module = _jit_usp_relayout_module(x.dtype)
|
||||
module.usp_merge_heads(out, x)
|
||||
return out
|
||||
|
||||
|
||||
def can_use_usp_merge_heads(x: torch.Tensor) -> bool:
|
||||
return (
|
||||
isinstance(x, torch.Tensor)
|
||||
and torch.version.hip is None
|
||||
and x.is_cuda
|
||||
and x.dtype in _SUPPORTED_DTYPES
|
||||
and x.dim() == 5
|
||||
and x.numel() > 0
|
||||
and x.is_contiguous()
|
||||
)
|
||||
|
||||
|
||||
def _usp_merge_heads_cuda(x: torch.Tensor) -> torch.Tensor:
|
||||
"""[W, S, B, h_local, D] -> [B, S, W, h_local, D] contiguous.
|
||||
|
||||
Bit-exact single-pass replacement for
|
||||
``x.permute(2, 1, 0, 3, 4).contiguous()`` on the Ulysses output path.
|
||||
"""
|
||||
if not can_use_usp_merge_heads(x):
|
||||
raise RuntimeError("unsupported input for usp_merge_heads CUDA")
|
||||
return _usp_merge_heads_custom_op(x)
|
||||
|
||||
|
||||
def usp_merge_heads(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Merge Ulysses output heads with an exact eager fallback.
|
||||
|
||||
The backend selection lives here so callers only express the layout
|
||||
transformation. Unsupported devices, layouts, and compiled regions retain
|
||||
the original PyTorch operation.
|
||||
"""
|
||||
if not torch.compiler.is_compiling() and can_use_usp_merge_heads(x):
|
||||
return _usp_merge_heads_cuda(x)
|
||||
return x.permute(2, 1, 0, 3, 4).contiguous()
|
||||
@@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus
|
||||
## Key Features
|
||||
|
||||
SGLang Diffusion has the following features:
|
||||
- Broad model support: Wan, FastWan, FLUX, Qwen-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3, LingBot World, SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
|
||||
- Broad model support: Wan, FastWan, FLUX, Qwen-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3, MiniMax-H3, LingBot World, SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
|
||||
- Fast inference speed: empowered by optimized `sgl-kernel` kernels, scheduler/runtime improvements, caching acceleration, and native diffusion hot-path optimizations
|
||||
- Ease of use: OpenAI-compatible api, CLI, and python sdk support
|
||||
- Multi-platform support:
|
||||
|
||||
@@ -111,6 +111,21 @@ def _infer_slo_base_time_ms_from_warmups(
|
||||
return float(np.median(candidates_ms)) if candidates_ms else None
|
||||
|
||||
|
||||
def _parse_extra_body(raw: Optional[str]) -> Dict[str, Any]:
|
||||
"""Parses --extra-body, which is merged over every generated payload."""
|
||||
if not raw:
|
||||
return {}
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"--extra-body is not valid JSON: {exc}") from exc
|
||||
if not isinstance(parsed, dict):
|
||||
raise ValueError(
|
||||
f"--extra-body must be a JSON object, got {type(parsed).__name__}."
|
||||
)
|
||||
return parsed
|
||||
|
||||
|
||||
def _populate_slo_ms_from_warmups(
|
||||
requests_list: List[RequestFuncInput], warmup_pairs: List[tuple], args
|
||||
) -> List[RequestFuncInput]:
|
||||
@@ -496,6 +511,10 @@ async def benchmark(args):
|
||||
if args.base_url is None:
|
||||
args.base_url = NetworkAddress(args.host, args.port).to_url()
|
||||
|
||||
# Parsed before the service wait and the dataset download so a malformed
|
||||
# value fails immediately instead of after minutes of setup.
|
||||
extra_body = _parse_extra_body(args.extra_body)
|
||||
|
||||
# Wait for service
|
||||
wait_for_service(args.base_url)
|
||||
|
||||
@@ -571,6 +590,13 @@ async def benchmark(args):
|
||||
requests_list = dataset.get_requests()
|
||||
logger.info(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
|
||||
|
||||
if extra_body:
|
||||
logger.info(f"Merging --extra-body into every request: {extra_body}")
|
||||
requests_list = [
|
||||
replace(req, extra_body={**req.extra_body, **extra_body})
|
||||
for req in requests_list
|
||||
]
|
||||
|
||||
# Limit concurrency
|
||||
if args.max_concurrency is not None:
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency)
|
||||
@@ -588,10 +614,10 @@ async def benchmark(args):
|
||||
# Run warmup requests
|
||||
warmup_pairs: List[tuple] = []
|
||||
if args.warmup_requests and requests_list:
|
||||
# The server always overrides warmup requests to use
|
||||
# num_inference_steps=1 (see Req.set_as_warmup), so we match
|
||||
# that here to keep the benchmark's SLO estimation consistent.
|
||||
warmup_steps = 1
|
||||
# Defaults to 1 to match the server's own boot warmup (see
|
||||
# Req.set_as_warmup) and keep SLO estimation consistent. Raise it
|
||||
# for models that reject a 1-step schedule, such as MiniMax-H3.
|
||||
warmup_steps = args.warmup_inference_steps
|
||||
logger.info(
|
||||
f"Running {args.warmup_requests} warmup request(s) with "
|
||||
f"num_inference_steps={warmup_steps}..."
|
||||
@@ -828,6 +854,22 @@ if __name__ == "__main__":
|
||||
default=1,
|
||||
help="Number of warmup requests to run before measurement.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-inference-steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Denoise steps for warmup requests. Raise it for models that "
|
||||
"reject a 1-step schedule.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=str,
|
||||
default=None,
|
||||
help="JSON object merged over each JSON request body, for contract "
|
||||
'fields the generic payload omits (e.g. \'{"task": "t2va"}\' for '
|
||||
"MiniMax-H3). Multipart image requests forward it as an extra_body "
|
||||
"form field instead, which the server may not unpack.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-inference-steps",
|
||||
type=int,
|
||||
|
||||
@@ -12,6 +12,7 @@ from sglang.multimodal_gen.configs.models.dits.lingbot_world import (
|
||||
LingBotWorldVideoConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.dits.longlive2 import LongLive2VideoConfig
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
|
||||
from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
|
||||
from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
|
||||
from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import (
|
||||
@@ -27,6 +28,7 @@ __all__ = [
|
||||
"Ideogram4DistilledDiTConfig",
|
||||
"LingBotWorldVideoConfig",
|
||||
"LongLive2VideoConfig",
|
||||
"MiniMaxH3DiTConfig",
|
||||
"WanVideoConfig",
|
||||
"Hunyuan3DDiTConfig",
|
||||
"MOVAAudioConfig",
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
from sglang.multimodal_gen.configs.models.fsdp import is_block
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
|
||||
MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64
|
||||
MINIMAX_H3_ADALN_MODALITY_NUM = 3
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
|
||||
|
||||
lora_param_names_mapping: dict = field(default_factory=dict)
|
||||
|
||||
_supported_attention_backends: set[AttentionBackendEnum] = field(
|
||||
default_factory=lambda: {
|
||||
AttentionBackendEnum.FA,
|
||||
AttentionBackendEnum.AITER,
|
||||
AttentionBackendEnum.TORCH_SDPA,
|
||||
}
|
||||
)
|
||||
|
||||
num_layers: int = 50
|
||||
token_refiner_num_layers: int = 2
|
||||
hidden_size: int = 5376
|
||||
num_attention_heads: int = 56
|
||||
attention_head_dim: int = 128
|
||||
ffn_hidden_size: int = 14336
|
||||
latents_dim: int = 24
|
||||
audio_latents_dim: int = 32
|
||||
patch_size: tuple[int, int, int] = (1, 2, 2)
|
||||
text_dim: int = 5120
|
||||
timestep_input_dim: int = 256
|
||||
time_embed_hidden_size: int = 5376
|
||||
time_embed_dim: int = 2688
|
||||
adaln_out_features: int = 18 * 5376
|
||||
final_adaln_out_features: int = 2 * 5376
|
||||
rope_inv_freq_len: int = 16
|
||||
norm_eps: float = 1e-5
|
||||
qk_norm_eps: float = 1e-5
|
||||
final_norm_eps: float = 1e-5
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if isinstance(self.patch_size, list):
|
||||
self.patch_size = tuple(self.patch_size)
|
||||
if len(self.patch_size) != 3:
|
||||
raise ValueError(f"patch_size must have 3 values, got {self.patch_size}.")
|
||||
self.num_channels_latents = self.latents_dim
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3DiTConfig(DiTConfig):
|
||||
arch_config: MiniMaxH3DiTArchConfig = field(default_factory=MiniMaxH3DiTArchConfig)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_ADALN_MODALITY_NUM",
|
||||
"MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT",
|
||||
"MiniMaxH3DiTArchConfig",
|
||||
"MiniMaxH3DiTConfig",
|
||||
]
|
||||
@@ -21,6 +21,10 @@ from sglang.multimodal_gen.configs.models.encoders.ideogram import (
|
||||
Ideogram4TextEncoderConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
|
||||
from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
|
||||
MiniMaxH3Qwen3VLArchConfig,
|
||||
MiniMaxH3Qwen3VLConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
|
||||
from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig
|
||||
from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
|
||||
@@ -36,6 +40,8 @@ __all__ = [
|
||||
"Flux2MistralTextConfig",
|
||||
"build_flux2_text_messages",
|
||||
"LlamaConfig",
|
||||
"MiniMaxH3Qwen3VLArchConfig",
|
||||
"MiniMaxH3Qwen3VLConfig",
|
||||
"Qwen3TextConfig",
|
||||
"Qwen3VLConfig",
|
||||
"T5Config",
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Native Qwen3-VL encoder configuration for MiniMax H3."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sglang.multimodal_gen.configs.models.encoders.qwen3vl import (
|
||||
Qwen3VLArchConfig,
|
||||
Qwen3VLConfig,
|
||||
)
|
||||
|
||||
MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3Qwen3VLArchConfig(Qwen3VLArchConfig):
|
||||
"""The checkpoint is Qwen3-VL-32B, consumed at hidden_states[50]."""
|
||||
|
||||
architectures: list[str] = field(
|
||||
default_factory=lambda: ["MiniMaxH3Qwen3VLEncoder"]
|
||||
)
|
||||
hidden_size: int = 5120
|
||||
intermediate_size: int = 25600
|
||||
num_hidden_layers: int = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
|
||||
num_attention_heads: int = 64
|
||||
num_key_value_heads: int = 8
|
||||
head_dim: int = 128
|
||||
text_len: int = 262144
|
||||
hidden_state_skip_layer: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3Qwen3VLConfig(Qwen3VLConfig):
|
||||
arch_config: MiniMaxH3Qwen3VLArchConfig = field(
|
||||
default_factory=MiniMaxH3Qwen3VLArchConfig
|
||||
)
|
||||
|
||||
def post_diffusers_config_update(self) -> None:
|
||||
"""Select the in-tree extractor after loading the HF architecture."""
|
||||
|
||||
arch = self.arch_config
|
||||
arch.architectures = ["MiniMaxH3Qwen3VLEncoder"]
|
||||
arch.hidden_size = int(arch.text_config.hidden_size)
|
||||
arch.intermediate_size = int(arch.text_config.intermediate_size)
|
||||
arch.num_attention_heads = int(arch.text_config.num_attention_heads)
|
||||
arch.num_key_value_heads = int(arch.text_config.num_key_value_heads)
|
||||
arch.head_dim = int(arch.text_config.head_dim)
|
||||
arch.num_hidden_layers = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
|
||||
arch.text_config.num_hidden_layers = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
|
||||
arch.text_config.output_hidden_states = False
|
||||
arch.text_config.use_cache = False
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER",
|
||||
"MiniMaxH3Qwen3VLArchConfig",
|
||||
"MiniMaxH3Qwen3VLConfig",
|
||||
]
|
||||
@@ -3,6 +3,12 @@
|
||||
from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
|
||||
MiniMaxH3AudioVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
|
||||
MiniMaxH3VideoVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
|
||||
StableDiffusion3VAEConfig,
|
||||
)
|
||||
@@ -11,6 +17,8 @@ from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig
|
||||
__all__ = [
|
||||
"DacVAEConfig",
|
||||
"HunyuanVAEConfig",
|
||||
"MiniMaxH3AudioVAEConfig",
|
||||
"MiniMaxH3VideoVAEConfig",
|
||||
"StableDiffusion3VAEConfig",
|
||||
"WanVAEConfig",
|
||||
"Hunyuan3DVAEConfig",
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_contract import (
|
||||
validate_minimax_h3_vae_latent_stats,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3AudioVAEArchConfig(VAEArchConfig):
|
||||
sample_rate: int = 32000
|
||||
latent_channels: int = 32
|
||||
latents_mean: list[float] | None = None
|
||||
latents_std: list[float] | None = None
|
||||
output_channel: int = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3AudioVAEConfig(VAEConfig):
|
||||
arch_config: MiniMaxH3AudioVAEArchConfig = field(
|
||||
default_factory=MiniMaxH3AudioVAEArchConfig
|
||||
)
|
||||
load_encoder: bool = True
|
||||
load_decoder: bool = True
|
||||
|
||||
def post_init(self) -> None:
|
||||
validate_minimax_h3_vae_latent_stats(
|
||||
self.arch_config,
|
||||
component_name="audio_vae",
|
||||
expected_channels=32,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3AudioVAEArchConfig", "MiniMaxH3AudioVAEConfig"]
|
||||
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import math
|
||||
from typing import Protocol
|
||||
|
||||
|
||||
class MiniMaxH3LatentStatsConfig(Protocol):
|
||||
latent_channels: int
|
||||
latents_mean: list[float] | None
|
||||
latents_std: list[float] | None
|
||||
|
||||
|
||||
class MiniMaxH3VAEContractError(ValueError):
|
||||
def __init__(self, component_name: str, detail: str) -> None:
|
||||
super().__init__(f"MiniMax H3 {component_name} {detail}")
|
||||
self.component_name = component_name
|
||||
self.detail = detail
|
||||
|
||||
def __reduce__(self):
|
||||
# BaseException pickles via cls(*args); rebuild from the two ctor args
|
||||
# so the error propagates cleanly across process boundaries.
|
||||
return (type(self), (self.component_name, self.detail))
|
||||
|
||||
|
||||
def validate_minimax_h3_vae_latent_stats(
|
||||
arch_config: MiniMaxH3LatentStatsConfig,
|
||||
component_name: str,
|
||||
expected_channels: int,
|
||||
) -> None:
|
||||
if arch_config.latent_channels != expected_channels:
|
||||
raise MiniMaxH3VAEContractError(
|
||||
component_name,
|
||||
"latent_channels must be "
|
||||
f"{expected_channels}, got {arch_config.latent_channels!r}",
|
||||
)
|
||||
|
||||
for field_name, values in (
|
||||
("latents_mean", arch_config.latents_mean),
|
||||
("latents_std", arch_config.latents_std),
|
||||
):
|
||||
if values is None:
|
||||
raise MiniMaxH3VAEContractError(
|
||||
component_name,
|
||||
f"config.json missing {field_name}",
|
||||
)
|
||||
if not isinstance(values, list) or not all(
|
||||
isinstance(value, (int, float)) and not isinstance(value, bool)
|
||||
for value in values
|
||||
):
|
||||
raise MiniMaxH3VAEContractError(
|
||||
component_name,
|
||||
f"config.json {field_name} must be a list of numbers",
|
||||
)
|
||||
if len(values) != expected_channels:
|
||||
raise MiniMaxH3VAEContractError(
|
||||
component_name,
|
||||
f"config.json {field_name} must contain exactly "
|
||||
f"{expected_channels} values, got {len(values)}",
|
||||
)
|
||||
if field_name == "latents_mean" and not all(
|
||||
math.isfinite(value) for value in values
|
||||
):
|
||||
raise MiniMaxH3VAEContractError(
|
||||
component_name,
|
||||
"config.json latents_mean values must be finite",
|
||||
)
|
||||
if field_name == "latents_std" and not all(
|
||||
math.isfinite(value) and value > 0 for value in values
|
||||
):
|
||||
raise MiniMaxH3VAEContractError(
|
||||
component_name,
|
||||
"config.json latents_std values must be finite and greater than zero",
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MiniMaxH3VAEContractError",
|
||||
"validate_minimax_h3_vae_latent_stats",
|
||||
]
|
||||
@@ -0,0 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_contract import (
|
||||
validate_minimax_h3_vae_latent_stats,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3VideoVAEArchConfig(VAEArchConfig):
|
||||
latent_channels: int = 24
|
||||
latents_mean: list[float] | None = None
|
||||
latents_std: list[float] | None = None
|
||||
temporal_compression_ratio: int = 4
|
||||
spatial_compression_ratio: int = 16
|
||||
vae_clip_length: int = 17
|
||||
vae_token_drop: int = 3
|
||||
vae_encoder_tiling: int = 1
|
||||
vae_decoder_tiling: int = 1
|
||||
vae_parallel_tiling: int = 1
|
||||
vae_tile_size: int = 256
|
||||
vae_tile_overlap_min: int = 64
|
||||
vae_chunk_dim: int = -1
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3VideoVAEConfig(VAEConfig):
|
||||
arch_config: MiniMaxH3VideoVAEArchConfig = field(
|
||||
default_factory=MiniMaxH3VideoVAEArchConfig
|
||||
)
|
||||
load_encoder: bool = True
|
||||
load_decoder: bool = True
|
||||
use_tiling: bool = True
|
||||
use_parallel_tiling: bool = True
|
||||
# The released checkpoint's quality contract uses overlapping latent
|
||||
# tiles. Parallel tiling distributes whole tiles without changing that
|
||||
# recipe. Spatial-shard decode is rejected because validation found output
|
||||
# mismatches on H3.
|
||||
parallel_decode_mode: str = "tiled"
|
||||
|
||||
def resolved_parallel_decode_mode(self) -> str:
|
||||
if self.parallel_decode_mode == "auto":
|
||||
return "tiled"
|
||||
if self.parallel_decode_mode in ("spatial", "spatial_shard"):
|
||||
raise ValueError(
|
||||
"MiniMax H3 rejects spatial-shard VAE decode because it failed "
|
||||
"the released quality contract; use tiled"
|
||||
)
|
||||
if self.parallel_decode_mode == "tiled":
|
||||
return "tiled"
|
||||
if self.parallel_decode_mode == "patch":
|
||||
raise ValueError("MiniMax H3 does not support patch VAE decode; use tiled")
|
||||
raise ValueError(
|
||||
f"unsupported MiniMax H3 VAE parallel decode mode "
|
||||
f"{self.parallel_decode_mode!r}"
|
||||
)
|
||||
|
||||
def post_init(self) -> None:
|
||||
self.resolved_parallel_decode_mode()
|
||||
validate_minimax_h3_vae_latent_stats(
|
||||
self.arch_config,
|
||||
component_name="video_vae",
|
||||
expected_channels=24,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3VideoVAEArchConfig", "MiniMaxH3VideoVAEConfig"]
|
||||
@@ -40,6 +40,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
|
||||
LTX2PipelineConfig,
|
||||
LTX23PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
|
||||
@@ -86,4 +89,5 @@ __all__ = [
|
||||
"LTX23PipelineConfig",
|
||||
"LingBotWorldCausalDMDConfig",
|
||||
"LingBotWorldV2CausalDMDConfig",
|
||||
"MiniMaxH3PipelineConfig",
|
||||
]
|
||||
|
||||
@@ -267,6 +267,11 @@ class PipelineConfig:
|
||||
# return the model-specific config for optimal deployment setting
|
||||
return ModelDeploymentConfig()
|
||||
|
||||
def validate_server_args(self, server_args: Any) -> None:
|
||||
"""Validate model-owned constraints after server args are normalized."""
|
||||
|
||||
del server_args
|
||||
|
||||
# Wan2.2 TI2V parameters
|
||||
boundary_ratio: float | None = None
|
||||
|
||||
@@ -394,6 +399,11 @@ class PipelineConfig:
|
||||
"""
|
||||
return self.task_type in (ModelTaskType.T2I, ModelTaskType.T2V)
|
||||
|
||||
def supports_disaggregation(self) -> bool:
|
||||
"""Return whether multi-service disaggregated deployment is supported."""
|
||||
|
||||
return True
|
||||
|
||||
def supports_native_grouped_requests(self):
|
||||
"""Return whether dynamic batches should run as grouped Req lists."""
|
||||
return False
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
|
||||
from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
|
||||
MiniMaxH3Qwen3VLConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
|
||||
MiniMaxH3AudioVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
|
||||
MiniMaxH3VideoVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.base import (
|
||||
ModelTaskType,
|
||||
PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
|
||||
ModelDeploymentConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3PipelineConfig(PipelineConfig):
|
||||
"""MiniMax H3 native audio-video pipeline configuration."""
|
||||
|
||||
# Canonical H3 materials are prepared by the model-specific stages. The
|
||||
# generic TI2V image resize would both duplicate that work and overwrite
|
||||
# the already-resolved target canvas.
|
||||
skip_input_image_preprocess: bool = True
|
||||
native_only_components = (
|
||||
"text_encoder",
|
||||
"transformer",
|
||||
"video_vae",
|
||||
"audio_vae",
|
||||
)
|
||||
task_type: ModelTaskType = ModelTaskType.TI2V
|
||||
dit_config: MiniMaxH3DiTConfig = field(default_factory=MiniMaxH3DiTConfig)
|
||||
vae_config: MiniMaxH3VideoVAEConfig = field(default_factory=MiniMaxH3VideoVAEConfig)
|
||||
audio_vae_config: MiniMaxH3AudioVAEConfig = field(
|
||||
default_factory=MiniMaxH3AudioVAEConfig
|
||||
)
|
||||
dit_precision: str = "bf16"
|
||||
# The video VAE remains fp32-resident because it also encodes keyframes.
|
||||
# Decode follows the released fp16-autocast recipe unless the user
|
||||
# explicitly disables autocast.
|
||||
vae_precision: str = "fp32"
|
||||
vae_decode_precision: str = "fp16"
|
||||
audio_vae_precision: str = "fp32"
|
||||
text_encoder_configs: tuple[MiniMaxH3Qwen3VLConfig, ...] = field(
|
||||
default_factory=lambda: (MiniMaxH3Qwen3VLConfig(),)
|
||||
)
|
||||
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
|
||||
text_encoder_extra_args: list[dict] = field(default_factory=lambda: [{}])
|
||||
# The released checkpoint is CFG-distilled and has one positive branch.
|
||||
should_use_guidance: bool = False
|
||||
output_audio_sample_rate: int | None = 32000
|
||||
output_audio_channels: int | None = 2
|
||||
output_av_drift_tolerance_s: float | None = 0.25
|
||||
|
||||
def accepts_audio_input(self) -> bool:
|
||||
return True
|
||||
|
||||
def supports_disaggregation(self) -> bool:
|
||||
return False
|
||||
|
||||
@property
|
||||
def requires_audio_output(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_model_deployment_config(self) -> ModelDeploymentConfig:
|
||||
return ModelDeploymentConfig(
|
||||
speed_mode_enable_torch_compile_by_default=False,
|
||||
keep_resident_min_available_gb=120,
|
||||
keep_resident_components=("dit", "text_encoder", "vae"),
|
||||
auto_enable_cfg_parallel=False,
|
||||
supports_cfg_parallel=False,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _server_arg_value(value):
|
||||
return getattr(value, "value", value)
|
||||
|
||||
def validate_quality_deployment(self, server_args) -> None:
|
||||
"""Fail closed unless the resident server matches the measured profile."""
|
||||
|
||||
attention_backend = self._server_arg_value(server_args.attention_backend)
|
||||
attention_backend = (
|
||||
str(attention_backend).strip().lower()
|
||||
if attention_backend is not None
|
||||
else None
|
||||
)
|
||||
capability = current_platform.get_device_capability()
|
||||
capability_int = capability.to_int() if capability is not None else None
|
||||
device_name = (
|
||||
current_platform.get_device_name()
|
||||
if current_platform.is_cuda()
|
||||
else type(current_platform).__name__
|
||||
)
|
||||
model_variant = str(server_args.model_variant or "fl2va").lower()
|
||||
actual = {
|
||||
"attention_backend": attention_backend,
|
||||
"backend": self._server_arg_value(server_args.backend),
|
||||
"component_attention_backends": {},
|
||||
"enable_breakable_cuda_graph": server_args.enable_breakable_cuda_graph,
|
||||
"enable_torch_compile": server_args.enable_torch_compile,
|
||||
"is_dit_layerwise_offload_selected": (
|
||||
server_args.is_dit_layerwise_offload_selected
|
||||
),
|
||||
"model_variant": model_variant,
|
||||
"num_gpus": server_args.num_gpus,
|
||||
"performance_mode": server_args.performance_mode,
|
||||
"quantization": server_args.quantization,
|
||||
"regional_compile": server_args.regional_compile,
|
||||
"ring_degree": server_args.ring_degree,
|
||||
"sp_degree": server_args.sp_degree,
|
||||
"tp_size": server_args.tp_size,
|
||||
"ulysses_degree": server_args.ulysses_degree,
|
||||
"use_fsdp_inference": server_args.use_fsdp_inference,
|
||||
}
|
||||
actual["component_attention_backends"] = dict(
|
||||
server_args.component_attention_backends or {}
|
||||
)
|
||||
expected = {
|
||||
"attention_backend": {None, "fa"},
|
||||
"backend": {"auto", "sglang"},
|
||||
"component_attention_backends": {},
|
||||
"enable_breakable_cuda_graph": False,
|
||||
"enable_torch_compile": False,
|
||||
"is_dit_layerwise_offload_selected": False,
|
||||
"model_variant": "fl2va",
|
||||
"num_gpus": 4,
|
||||
"performance_mode": "speed",
|
||||
"quantization": None,
|
||||
"regional_compile": False,
|
||||
"ring_degree": 1,
|
||||
"sp_degree": 4,
|
||||
"tp_size": 1,
|
||||
"ulysses_degree": 4,
|
||||
"use_fsdp_inference": False,
|
||||
}
|
||||
mismatches = {
|
||||
name: {"expected": wanted, "actual": actual[name]}
|
||||
for name, wanted in expected.items()
|
||||
if (
|
||||
actual[name] not in wanted
|
||||
if isinstance(wanted, set)
|
||||
else actual[name] != wanted
|
||||
)
|
||||
}
|
||||
if (
|
||||
not current_platform.is_cuda()
|
||||
or "H200" not in device_name.upper()
|
||||
or capability_int != 90
|
||||
):
|
||||
mismatches["device"] = {
|
||||
"expected": "NVIDIA H200 (compute capability 9.0)",
|
||||
"actual": f"{device_name} (compute capability {capability_int})",
|
||||
}
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
"MiniMax-H3 approximate quality profiles are validated only for "
|
||||
f"the strict 4xH200 fl2va deployment; mismatches: {mismatches}"
|
||||
)
|
||||
|
||||
def validate_server_args(self, server_args) -> None:
|
||||
# Reject known-inexact VAE modes before any large component download.
|
||||
self.vae_config.resolved_parallel_decode_mode()
|
||||
attention_backend = self._server_arg_value(server_args.attention_backend)
|
||||
if str(attention_backend).strip().lower() == "sage_attn":
|
||||
raise ValueError(
|
||||
"MiniMax-H3 does not support SageAttention: the current packed "
|
||||
"varlen path does not preserve model output"
|
||||
)
|
||||
|
||||
def select_vae_weight_files(
|
||||
self,
|
||||
safetensors_list: list[str],
|
||||
component_model_path: str,
|
||||
component_name: str,
|
||||
vae_precision: str,
|
||||
) -> list[str]:
|
||||
if component_name == "video_vae":
|
||||
return [os.path.join(component_model_path, "source", "model.safetensors")]
|
||||
return safetensors_list
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3PipelineConfig"]
|
||||
@@ -25,6 +25,10 @@ class ModelDeploymentConfig:
|
||||
auto_enable_cfg_parallel: bool = True
|
||||
# degree 1 keeps CFG parallel disabled and leaves GPUs available for SP
|
||||
auto_cfg_parallel_degree_by_num_gpus: tuple[tuple[int, int], ...] = ()
|
||||
# Let performance_mode=speed opt into torch.compile unless the model has
|
||||
# established that the compiled path changes its numerical contract.
|
||||
speed_mode_enable_torch_compile_by_default: bool = True
|
||||
supports_cfg_parallel: bool = True
|
||||
|
||||
def get_auto_cfg_parallel_degree(self, num_gpus: int) -> int:
|
||||
for candidate_num_gpus, cfg_degree in self.auto_cfg_parallel_degree_by_num_gpus:
|
||||
|
||||
@@ -0,0 +1,305 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import msgspec
|
||||
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
|
||||
|
||||
_MINIMAX_H3_MAX_SIGNED_SEED = (1 << 63) - 1
|
||||
|
||||
|
||||
def _optional_unit_float(value: Any, field_name: str) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise ValueError(f"{field_name} must be a number")
|
||||
out = float(value)
|
||||
if out < 0.0 or out > 1.0:
|
||||
raise ValueError(f"{field_name} must be in [0, 1]")
|
||||
return out
|
||||
|
||||
|
||||
def _optional_positive_finite_float(value: Any, field_name: str) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise ValueError(f"{field_name} must be a number")
|
||||
out = float(value)
|
||||
if not math.isfinite(out) or out <= 0.0:
|
||||
raise ValueError(f"{field_name} must be a positive finite number")
|
||||
return out
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3SamplingParams(SamplingParams):
|
||||
height: int = 512
|
||||
width: int = 896
|
||||
num_inference_steps: int = 50
|
||||
num_frames: int = field(default=1, init=False)
|
||||
fps: int = field(default=24, init=False)
|
||||
negative_prompt: None = field(default=None, init=False)
|
||||
guidance_scale: float = field(default=1.0, init=False)
|
||||
guidance_scale_2: None = field(default=None, init=False)
|
||||
true_cfg_scale: None = field(default=None, init=False)
|
||||
guidance_rescale: float = field(default=0.0, init=False)
|
||||
cfg_normalization: float = field(default=0.0, init=False)
|
||||
imgvid_cond_noise_aug_for_inference: float | None = None
|
||||
audio_cond_noise_aug_for_inference: float | None = None
|
||||
task: str | None = None
|
||||
conditions: list[dict[str, Any]] | None = None
|
||||
target: dict[str, Any] | None = None
|
||||
audio_flow_shift: float | None = None
|
||||
output_mode: str | None = field(
|
||||
default=None,
|
||||
metadata={"batch_sig_exclude": True},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def video_request_extra_fields(cls) -> frozenset[str]:
|
||||
return frozenset(
|
||||
{
|
||||
"task",
|
||||
"conditions",
|
||||
"target",
|
||||
"audio_flow_shift",
|
||||
"audio_guidance_scale",
|
||||
"quality",
|
||||
"output_mode",
|
||||
"imgvid_cond_noise_aug_for_inference",
|
||||
"audio_cond_noise_aug_for_inference",
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _video_hooks():
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.video_adapter import (
|
||||
MiniMaxH3VideoModelAdapter,
|
||||
)
|
||||
|
||||
return MiniMaxH3VideoModelAdapter()
|
||||
|
||||
@classmethod
|
||||
def lower_video_request_kwargs(
|
||||
cls,
|
||||
request: Any,
|
||||
kwargs: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return cls._video_hooks().lower_video_request_kwargs(request, kwargs)
|
||||
|
||||
def prepare_video_request_for_queue(self, req: Any) -> None:
|
||||
hooks = self._video_hooks()
|
||||
hooks.validate_sampling_params(self)
|
||||
hooks.prepare_for_queue_sync(req)
|
||||
|
||||
def expand_video_request_outputs_for_queue(self, req: Any) -> list[Any]:
|
||||
"""Use the same independent-seed grouped path in serve and generate."""
|
||||
|
||||
from sglang.multimodal_gen.runtime.entrypoints.utils import (
|
||||
expand_request_outputs,
|
||||
)
|
||||
|
||||
return expand_request_outputs(req)
|
||||
|
||||
def prepare_synthetic_warmup_request_for_queue(
|
||||
self, req: Any, server_args: Any
|
||||
) -> None:
|
||||
"""Lower generic warmup into one valid native partition request.
|
||||
|
||||
This intentionally calls the existing pre-queue resolver directly
|
||||
instead of the public video admission hook: synthetic warmup disables
|
||||
file delivery, while the public H3 contract correctly requires it.
|
||||
"""
|
||||
selected_variant = getattr(server_args, "model_variant", None)
|
||||
if selected_variant is not None:
|
||||
selected_partition = str(selected_variant).strip().lower()
|
||||
else:
|
||||
selected_path = server_args.model_subfolder or server_args.model_path
|
||||
selected_partition = os.path.basename(
|
||||
os.path.normpath(str(selected_path))
|
||||
).lower()
|
||||
|
||||
if selected_partition == "ref2va":
|
||||
image_path = req.image_path
|
||||
if isinstance(image_path, list):
|
||||
if not image_path:
|
||||
raise ValueError(
|
||||
"MiniMax H3 Ref2VA synthetic warmup requires an image"
|
||||
)
|
||||
image_path = image_path[0]
|
||||
if not isinstance(image_path, str) or not image_path:
|
||||
raise ValueError("MiniMax H3 Ref2VA synthetic warmup requires an image")
|
||||
task = "ref2va"
|
||||
conditions = [
|
||||
{
|
||||
"type": "image",
|
||||
"uri": image_path,
|
||||
"role": "reference",
|
||||
}
|
||||
]
|
||||
else:
|
||||
task = "t2va"
|
||||
conditions = []
|
||||
|
||||
self.task = task
|
||||
self.conditions = conditions
|
||||
self.target = {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 5.0,
|
||||
}
|
||||
selected_seed = req.seed if isinstance(req.seed, int) else int(req.seed[0])
|
||||
req.extra.update(self.build_request_extra(_seed_override=int(selected_seed)))
|
||||
self._video_hooks().prepare_for_queue_sync(req)
|
||||
|
||||
def project_video_queued_job_fields(self, req: Any) -> dict[str, str]:
|
||||
return self._video_hooks().project_queued_job_fields(req)
|
||||
|
||||
def validate_video_final_outputs(
|
||||
self,
|
||||
output_paths: list[str],
|
||||
req: Any,
|
||||
) -> dict[str, str]:
|
||||
return self._video_hooks().validate_final_outputs_sync(output_paths, req)
|
||||
|
||||
def cleanup_video_request(self, req: Any) -> None:
|
||||
self._video_hooks().cleanup_request_sync(req)
|
||||
|
||||
def _adjust(self, server_args) -> None:
|
||||
"""Apply generic path/output adjustments without deriving time shape.
|
||||
|
||||
The generic helper normally rewrites ``num_frames`` for temporal VAE
|
||||
alignment and GPU sharding. MiniMax H3 resolves that shape from the
|
||||
canonical target during pre-queue admission, so those rewrites must
|
||||
not leak into the offline parameter object. Keep the transport
|
||||
metadata at its internal sentinel values until pre-queue populates the
|
||||
actual batch shape.
|
||||
"""
|
||||
|
||||
super()._adjust(server_args)
|
||||
self.fps = 24
|
||||
self.num_frames = 1
|
||||
|
||||
def _validate(self) -> None:
|
||||
self.fps = 24
|
||||
self.num_frames = 1
|
||||
if isinstance(self.target, Mapping):
|
||||
self.target = {
|
||||
field_name: self.target[field_name]
|
||||
for field_name in (
|
||||
"short_edge",
|
||||
"aspect_ratio",
|
||||
"duration_seconds",
|
||||
)
|
||||
if field_name in self.target
|
||||
}
|
||||
super()._validate()
|
||||
_optional_positive_finite_float(self.flow_shift, "flow_shift")
|
||||
_optional_positive_finite_float(self.audio_flow_shift, "audio_flow_shift")
|
||||
if self.enable_frame_interpolation:
|
||||
raise ValueError(
|
||||
"MiniMax H3 does not support enable_frame_interpolation: the "
|
||||
"accepted delivery contract is the canonical 24 fps output"
|
||||
)
|
||||
if self.enable_upscaling:
|
||||
raise ValueError(
|
||||
"MiniMax H3 does not support enable_upscaling: the accepted "
|
||||
"delivery contract is the resolved target canvas"
|
||||
)
|
||||
if self.enable_teacache:
|
||||
raise ValueError(
|
||||
"MiniMax H3 does not support enable_teacache: its packed "
|
||||
"video/audio denoise loop has no lossless TeaCache contract"
|
||||
)
|
||||
if self.rollout:
|
||||
raise ValueError(
|
||||
"MiniMax H3 does not support rollout: its coupled video/audio "
|
||||
"scheduler has no SchedulerRLMixin contract"
|
||||
)
|
||||
if self.return_trajectory_latents or self.return_trajectory_decoded:
|
||||
raise ValueError(
|
||||
"MiniMax H3 does not support trajectory output for its coupled "
|
||||
"video/audio denoise state"
|
||||
)
|
||||
seeds = self.seed if isinstance(self.seed, list) else [self.seed]
|
||||
for seed in seeds:
|
||||
if seed > _MINIMAX_H3_MAX_SIGNED_SEED:
|
||||
raise ValueError(
|
||||
"MiniMax H3 seed must not exceed the signed int64 maximum, "
|
||||
f"got {seed}"
|
||||
)
|
||||
if (
|
||||
isinstance(self.seed, int)
|
||||
and self.seed + self.num_outputs_per_prompt - 1
|
||||
> _MINIMAX_H3_MAX_SIGNED_SEED
|
||||
):
|
||||
raise ValueError(
|
||||
"MiniMax H3 scalar seed plus output index must fit the "
|
||||
"signed int64 upper bound"
|
||||
)
|
||||
|
||||
def build_request_extra(self, *, _seed_override: int | None = None) -> dict:
|
||||
_optional_unit_float(
|
||||
self.imgvid_cond_noise_aug_for_inference,
|
||||
"imgvid_cond_noise_aug_for_inference",
|
||||
)
|
||||
_optional_unit_float(
|
||||
self.audio_cond_noise_aug_for_inference,
|
||||
"audio_cond_noise_aug_for_inference",
|
||||
)
|
||||
extra = super().build_request_extra()
|
||||
if self.task is not None:
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.request_validation import (
|
||||
minimax_h3_validate_canonical_request,
|
||||
)
|
||||
|
||||
extra["minimax_h3_canonical_request"] = (
|
||||
minimax_h3_validate_canonical_request(
|
||||
task=self.task,
|
||||
prompt=self.prompt,
|
||||
conditions=self.conditions,
|
||||
target=self.target,
|
||||
flow_shift=self.flow_shift,
|
||||
audio_flow_shift=self.audio_flow_shift,
|
||||
seed=(
|
||||
_seed_override
|
||||
if _seed_override is not None
|
||||
else self.seed if isinstance(self.seed, int) else None
|
||||
),
|
||||
)
|
||||
)
|
||||
elif (
|
||||
self.conditions is not None
|
||||
or self.target is not None
|
||||
or self.flow_shift is not None
|
||||
or self.audio_flow_shift is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"task is required when conditions/target/flow_shift/"
|
||||
"audio_flow_shift are provided"
|
||||
)
|
||||
return extra
|
||||
|
||||
def refresh_request_extra_after_output_expansion(self, req: Any) -> None:
|
||||
"""Copy validated canonical identity with the selected scalar Req seed."""
|
||||
|
||||
selected_seed = getattr(req, "seed", None)
|
||||
if isinstance(selected_seed, int):
|
||||
canonical_key = "minimax_h3_canonical_request"
|
||||
canonical = dict(req.extra[canonical_key])
|
||||
canonical["seed"] = selected_seed
|
||||
req.extra[canonical_key] = canonical
|
||||
resolved_plan_key = "minimax_h3_resolved_plan"
|
||||
resolved_plan = req.extra.get(resolved_plan_key)
|
||||
if resolved_plan is not None:
|
||||
req.extra[resolved_plan_key] = msgspec.structs.replace(
|
||||
resolved_plan, seed=selected_seed
|
||||
)
|
||||
else:
|
||||
req.extra.update(self.build_request_extra())
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3SamplingParams"]
|
||||
@@ -123,6 +123,10 @@ class SamplingParams:
|
||||
)
|
||||
output_quality: str | None = "default"
|
||||
output_compression: int | None = None
|
||||
# Model-owned, request-scoped approximate acceleration profile. Models
|
||||
# that support it must validate the deployment and workload explicitly.
|
||||
# It intentionally participates in the dynamic-batch signature.
|
||||
quality: str = "lossless"
|
||||
|
||||
# Frame interpolation
|
||||
enable_frame_interpolation: bool = False
|
||||
@@ -328,6 +332,64 @@ class SamplingParams:
|
||||
if self.realtime_chunk_size is not None:
|
||||
req.realtime_chunk_size = self.realtime_chunk_size
|
||||
|
||||
@classmethod
|
||||
def video_request_extra_fields(cls) -> frozenset[str]:
|
||||
"""Declare model-specific multipart video fields accepted by this type."""
|
||||
|
||||
return frozenset()
|
||||
|
||||
@classmethod
|
||||
def lower_video_request_kwargs(
|
||||
cls,
|
||||
request: Any,
|
||||
kwargs: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Adapt generic video-API kwargs before constructing this params type."""
|
||||
del request
|
||||
return kwargs
|
||||
|
||||
def prepare_video_request_for_queue(self, req: Any) -> None:
|
||||
"""Resolve model-specific admission facts before a video job is queued."""
|
||||
del req
|
||||
|
||||
def expand_video_request_outputs_for_queue(self, req: Any) -> list[Any] | None:
|
||||
"""Return per-output requests when a model owns grouped execution.
|
||||
|
||||
``None`` preserves the default model-native ``num_outputs`` handling.
|
||||
Models that need the framework's independent-seed request expansion
|
||||
can opt in after their shared pre-queue work has completed.
|
||||
"""
|
||||
del req
|
||||
return None
|
||||
|
||||
def prepare_synthetic_warmup_request_for_queue(
|
||||
self, req: Any, server_args: Any
|
||||
) -> None:
|
||||
"""Resolve model-specific facts for one synthetic warmup request."""
|
||||
del req, server_args
|
||||
|
||||
def project_video_queued_job_fields(self, req: Any) -> dict[str, str]:
|
||||
"""Return model-resolved fields to publish with the queued video job."""
|
||||
del req
|
||||
return {}
|
||||
|
||||
def validate_video_final_outputs(
|
||||
self,
|
||||
output_paths: list[str],
|
||||
req: Any,
|
||||
) -> dict[str, str]:
|
||||
"""Validate final files and return truthful completion metadata."""
|
||||
del output_paths, req
|
||||
return {}
|
||||
|
||||
def cleanup_video_request(self, req: Any) -> None:
|
||||
"""Release request-scoped resources owned by the model integration."""
|
||||
del req
|
||||
|
||||
def refresh_request_extra_after_output_expansion(self, req: Any) -> None:
|
||||
"""Refresh request identity after assigning a per-output seed."""
|
||||
del req
|
||||
|
||||
def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int:
|
||||
"""Convert output_quality string to compression level."""
|
||||
if data_type == DataType.ACTION:
|
||||
@@ -346,6 +408,11 @@ class SamplingParams:
|
||||
f"prompt_path must be a txt file, got {self.prompt_path!r}"
|
||||
)
|
||||
|
||||
if not isinstance(self.quality, str) or not self.quality.strip():
|
||||
raise ValueError(
|
||||
f"quality must be a non-empty string, got {self.quality!r}"
|
||||
)
|
||||
|
||||
# These are always required to be sane regardless of pipeline.
|
||||
if (
|
||||
not isinstance(self.num_outputs_per_prompt, int)
|
||||
@@ -862,10 +929,20 @@ class SamplingParams:
|
||||
type=int,
|
||||
help="Output compression level (0-100, higher means better quality but larger file size)",
|
||||
)
|
||||
add_argument(
|
||||
"--quality",
|
||||
type=str,
|
||||
help=(
|
||||
"Select a model-owned quality/performance profile. "
|
||||
"Support and validated deployment constraints are model-specific."
|
||||
),
|
||||
)
|
||||
add_argument(
|
||||
"--num-outputs-per-prompt",
|
||||
"--num-outputs",
|
||||
dest="num_outputs_per_prompt",
|
||||
type=int,
|
||||
help="Number of outputs to generate per prompt",
|
||||
help="Number of outputs to generate per prompt (alias: --num-outputs)",
|
||||
)
|
||||
add_argument(
|
||||
"--seed",
|
||||
|
||||
@@ -37,6 +37,7 @@ from sglang.multimodal_gen.configs.pipeline_configs import (
|
||||
HunyuanConfig,
|
||||
LingBotWorldCausalDMDConfig,
|
||||
LingBotWorldV2CausalDMDConfig,
|
||||
MiniMaxH3PipelineConfig,
|
||||
WanI2V480PConfig,
|
||||
WanI2V720PConfig,
|
||||
WanT2V480PConfig,
|
||||
@@ -141,6 +142,7 @@ from sglang.multimodal_gen.configs.sample.ltx_2 import (
|
||||
LTX23HQSamplingParams,
|
||||
LTX23SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.mova import (
|
||||
MOVA_360P_SamplingParams,
|
||||
MOVA_720P_SamplingParams,
|
||||
@@ -826,6 +828,18 @@ def _register_configs():
|
||||
lambda hf_id: "mova" in hf_id.lower() and "720p" in hf_id.lower()
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=MiniMaxH3SamplingParams,
|
||||
pipeline_config_cls=MiniMaxH3PipelineConfig,
|
||||
hf_model_paths=[
|
||||
"MiniMaxAI/MiniMax-H3",
|
||||
"MiniMax/MiniMax-H3",
|
||||
],
|
||||
model_detectors=[
|
||||
lambda model_id: "minimaxh3"
|
||||
in model_id.lower().replace("-", "").replace("_", "")
|
||||
],
|
||||
)
|
||||
# FLUX
|
||||
register_configs(
|
||||
sampling_param_cls=FluxSamplingParams,
|
||||
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0
|
||||
# ==============================================================================
|
||||
"""MiniMax H3 breakable CUDA graph packed-prompt padding."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
|
||||
MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
|
||||
prompt_padding as bcg_utils,
|
||||
)
|
||||
|
||||
|
||||
def is_minimax_h3_transformer(current_model: Any, call_kwargs: dict) -> bool:
|
||||
return (
|
||||
bcg_utils.transformer_class_name_matches(current_model, "minimaxh3")
|
||||
and "prompt_embeds" in call_kwargs
|
||||
and "packed_seq_params" in call_kwargs
|
||||
and "refiner_packed_seq_params" in call_kwargs
|
||||
and "text_pos_info" in call_kwargs
|
||||
)
|
||||
|
||||
|
||||
def _position_ids(info: Any) -> torch.Tensor | None:
|
||||
if isinstance(info, dict):
|
||||
ids = info.get("position_ids")
|
||||
else:
|
||||
ids = getattr(info, "position_ids", None)
|
||||
return ids if torch.is_tensor(ids) else None
|
||||
|
||||
|
||||
def _replace_position_ids(info: Any, ids: torch.Tensor) -> dict[str, Any]:
|
||||
if isinstance(info, dict):
|
||||
return {**info, "position_ids": ids}
|
||||
# H3 currently passes dictionaries. Avoid mutating an unknown request
|
||||
# object if an alternate frontend supplies one.
|
||||
return {"position_ids": ids}
|
||||
|
||||
|
||||
def _replace_psp(
|
||||
psp: Any,
|
||||
*,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
max_seqlen_q: int,
|
||||
) -> dict[str, Any]:
|
||||
if isinstance(psp, dict):
|
||||
return {
|
||||
**psp,
|
||||
"cu_seqlens_q": cu_seqlens_q,
|
||||
"max_seqlen_q": max_seqlen_q,
|
||||
}
|
||||
return {
|
||||
"cu_seqlens_q": cu_seqlens_q,
|
||||
"max_seqlen_q": max_seqlen_q,
|
||||
}
|
||||
|
||||
|
||||
def _aligned(value: int) -> int:
|
||||
alignment = MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
|
||||
return (value + alignment - 1) // alignment * alignment
|
||||
|
||||
|
||||
def pad_minimax_h3_prompt_kwargs(
|
||||
call_kwargs: dict, current_model: Any, buckets: tuple[int, ...]
|
||||
) -> dict:
|
||||
prompt = bcg_utils.first_tensor(call_kwargs.get("prompt_embeds"))
|
||||
text_pos = _position_ids(call_kwargs.get("text_pos_info"))
|
||||
img_pos = _position_ids(call_kwargs.get("img_pos_info"))
|
||||
audio_pos = _position_ids(call_kwargs.get("audio_pos_info"))
|
||||
x = call_kwargs.get("x")
|
||||
if (
|
||||
not torch.is_tensor(prompt)
|
||||
or prompt.dim() < 2
|
||||
or not torch.is_tensor(text_pos)
|
||||
or not torch.is_tensor(img_pos)
|
||||
or not torch.is_tensor(audio_pos)
|
||||
or not torch.is_tensor(x)
|
||||
or x.dim() != 3
|
||||
):
|
||||
return call_kwargs
|
||||
|
||||
text_len = int(prompt.shape[0])
|
||||
bucket = bcg_utils.select_text_bucket(text_len, buckets)
|
||||
if bucket is None:
|
||||
return call_kwargs
|
||||
|
||||
# All used H3 rows are disjoint text, image/video, or audio rows. Derive
|
||||
# the used/media lengths from tensor shapes so padding itself does not
|
||||
# perform a GPU-to-host .item() synchronization on every denoising step.
|
||||
media_rows = int(img_pos.numel()) + int(audio_pos.numel())
|
||||
used = text_len + media_rows
|
||||
source_seq = int(x.shape[1])
|
||||
if source_seq < _aligned(used):
|
||||
return call_kwargs
|
||||
|
||||
out = dict(call_kwargs)
|
||||
# request-local row lists have prompt-dependent shapes, so keep them out
|
||||
# of bucketed BCG signatures and rebuild the rows in the eager break
|
||||
out.pop("local_embedding_layout", None)
|
||||
# Request-static H3 denoising normally carries the live refined-text
|
||||
# length as a host integer to avoid a per-step device scalar read. Host
|
||||
# integers are baked into BCG signatures, however, so different prompt
|
||||
# lengths would miss the same text bucket. Make only the BCG-padded copy a
|
||||
# scalar tensor; the eager embedding break reads its updated replay value.
|
||||
refined_len = out.get("refined_prompt_embeds_length")
|
||||
if refined_len is not None and not torch.is_tensor(refined_len):
|
||||
out["refined_prompt_embeds_length"] = torch.tensor(
|
||||
int(refined_len),
|
||||
dtype=torch.int64,
|
||||
device=prompt.device,
|
||||
)
|
||||
if text_len < bucket:
|
||||
out["prompt_embeds"] = bcg_utils.pad_tensor_dim(prompt, dim=0, target=bucket)
|
||||
# These rows exist only to stabilize the BCG input signature. The
|
||||
# model's eager embedding break trims prompt/text_pos/refiner metadata
|
||||
# back to ``text_len`` before any projection or attention, so their
|
||||
# values never enter the model.
|
||||
dummy_text_pos = torch.arange(
|
||||
used,
|
||||
used + (bucket - text_len),
|
||||
dtype=text_pos.dtype,
|
||||
device=text_pos.device,
|
||||
)
|
||||
out["text_pos_info"] = _replace_position_ids(
|
||||
out["text_pos_info"], torch.cat((text_pos.view(-1), dummy_text_pos))
|
||||
)
|
||||
|
||||
# Do not grow the main packed sequence to media_rows + bucket. Changing
|
||||
# the SP row partition changes GEMM shapes and is measurably non-bitwise
|
||||
# even though dummy rows live in an independent attention segment. A
|
||||
# capture is therefore reusable only inside the request's existing
|
||||
# 64-row packed-sequence alignment group; other groups safely miss the
|
||||
# signature and run eager.
|
||||
packed_cu = torch.tensor([0, used, source_seq], dtype=torch.int32, device=x.device)
|
||||
out["packed_seq_params"] = _replace_psp(
|
||||
out["packed_seq_params"],
|
||||
cu_seqlens_q=packed_cu,
|
||||
# FA accepts an upper bound; keeping this bucket-stable is required
|
||||
# because non-tensor values are baked into the BCG signature.
|
||||
max_seqlen_q=source_seq,
|
||||
)
|
||||
refiner_cu = torch.tensor(
|
||||
[0, text_len, bucket],
|
||||
dtype=torch.int32,
|
||||
device=prompt.device,
|
||||
)
|
||||
out["refiner_packed_seq_params"] = _replace_psp(
|
||||
out["refiner_packed_seq_params"],
|
||||
cu_seqlens_q=refiner_cu,
|
||||
max_seqlen_q=bucket,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
bcg_utils.register_prompt_padder(
|
||||
is_minimax_h3_transformer, pad_minimax_h3_prompt_kwargs
|
||||
)
|
||||
@@ -301,6 +301,7 @@ def _ensure_model_padders_registered() -> None:
|
||||
_model_padders_registered = True
|
||||
from sglang.multimodal_gen.runtime.breakable_cuda_graph.model_padders import ( # noqa: F401
|
||||
ideogram,
|
||||
minimax_h3,
|
||||
qwen_image,
|
||||
zimage,
|
||||
)
|
||||
|
||||
@@ -38,6 +38,20 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_gro
|
||||
_original_similarity = None
|
||||
|
||||
|
||||
def disable_cache_on_transformer(transformer: torch.nn.Module) -> torch.nn.Module:
|
||||
"""Remove Cache-DiT hooks so subsequent requests use the native forward."""
|
||||
|
||||
logger.info("Disabling cache-dit on %s", type(transformer).__name__)
|
||||
target = getattr(transformer, "_sglang_cache_dit_adapter", transformer)
|
||||
cache_dit.disable_cache(target)
|
||||
if target is not transformer:
|
||||
del transformer._sglang_cache_dit_adapter
|
||||
for name in ("_is_parallelized", "_parallelism_config"):
|
||||
if hasattr(transformer, name):
|
||||
delattr(transformer, name)
|
||||
return transformer
|
||||
|
||||
|
||||
def _patch_cache_dit_similarity():
|
||||
from cache_dit.caching.cache_contexts import cache_manager
|
||||
|
||||
@@ -268,6 +282,7 @@ DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS: dict[str, DualTransformerBlockAdapterSpec]
|
||||
_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern]] = {
|
||||
"ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3),
|
||||
"Krea2Transformer2DModel": ("transformer_blocks", ForwardPattern.Pattern_3),
|
||||
"MiniMaxH3DiTModel": ("blocks", ForwardPattern.Pattern_3),
|
||||
}
|
||||
|
||||
|
||||
@@ -412,6 +427,8 @@ def enable_cache_on_transformer(
|
||||
calibrator_config=calibrator_config,
|
||||
parallelism_config=None,
|
||||
)
|
||||
if custom_adapter is not None:
|
||||
transformer._sglang_cache_dit_adapter = custom_adapter
|
||||
|
||||
if parallelism_config is not None:
|
||||
context_manager = getattr(transformer, "_context_manager", None)
|
||||
|
||||
@@ -257,8 +257,17 @@ IPC_A2A = IpcA2AState()
|
||||
def ipc_a2a_ready(group) -> bool:
|
||||
"""True when the IPC transport is enabled and initialized (initializes
|
||||
lazily on the first eager call; never inside a graph capture)."""
|
||||
from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
if not envs.SGLANG_DIFFUSION_IPC_A2A or IPC_A2A.failed:
|
||||
return False
|
||||
# TP+Ulysses groups are strided in global-rank order, while this transport
|
||||
# supports the adjacent two-device topology used by TP1+U2. Reject the
|
||||
# transport consistently before lazy initialization so no rank enters IPC
|
||||
# while its peer falls back to NCCL.
|
||||
if not current_platform.is_cuda() or get_tp_world_size() > 1:
|
||||
return False
|
||||
if IPC_A2A.inited:
|
||||
return True
|
||||
if torch.cuda.is_current_stream_capturing():
|
||||
|
||||
@@ -220,6 +220,7 @@ class DiffGenerator:
|
||||
)
|
||||
|
||||
request_groups: list[list[Req]] = []
|
||||
parent_requests: list[tuple[Req, int]] = []
|
||||
image_paths_per_prompt = self._resolve_image_paths_per_prompt(
|
||||
prompts, sampling_params_orig.image_path
|
||||
)
|
||||
@@ -243,13 +244,32 @@ class DiffGenerator:
|
||||
sampling_params=sampling_params,
|
||||
external_trace_header=external_trace_header,
|
||||
)
|
||||
request_groups.append(
|
||||
expand_request_outputs(
|
||||
req,
|
||||
num_prompts=len(prompts),
|
||||
prompt_index=i,
|
||||
parent_requests.append((req, i))
|
||||
|
||||
for req, prompt_index in parent_requests:
|
||||
sampling_params = req.sampling_params
|
||||
try:
|
||||
if sampling_params.data_type == DataType.VIDEO:
|
||||
sampling_params.prepare_video_request_for_queue(req)
|
||||
request_groups.append(
|
||||
expand_request_outputs(
|
||||
req,
|
||||
num_prompts=len(prompts),
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
if sampling_params.data_type == DataType.VIDEO:
|
||||
sampling_params.cleanup_video_request(req)
|
||||
for prepared_requests in request_groups:
|
||||
if (
|
||||
prepared_requests
|
||||
and prepared_requests[0].data_type == DataType.VIDEO
|
||||
):
|
||||
prepared_requests[0].sampling_params.cleanup_video_request(
|
||||
prepared_requests[0]
|
||||
)
|
||||
raise
|
||||
|
||||
results: list[GenerationResult] = []
|
||||
total_start_time = time.perf_counter()
|
||||
@@ -285,6 +305,10 @@ class DiffGenerator:
|
||||
)
|
||||
for idx, path in enumerate(output_file_paths):
|
||||
req = requests[idx]
|
||||
if req.data_type == DataType.VIDEO:
|
||||
req.sampling_params.validate_video_final_outputs(
|
||||
[path], req
|
||||
)
|
||||
results.append(
|
||||
GenerationResult(
|
||||
**self._result_common(
|
||||
@@ -346,6 +370,11 @@ class DiffGenerator:
|
||||
|
||||
for idx in range(len(samples_out)):
|
||||
req = requests[idx]
|
||||
output_file_path = req.output_file_path(1, 0)
|
||||
if req.data_type == DataType.VIDEO and req.save_output:
|
||||
req.sampling_params.validate_video_final_outputs(
|
||||
[output_file_path], req
|
||||
)
|
||||
results.append(
|
||||
GenerationResult(
|
||||
**self._result_common(
|
||||
@@ -355,12 +384,23 @@ class DiffGenerator:
|
||||
frames=frames_out[idx],
|
||||
audio=audios_out[idx],
|
||||
prompt_index=global_output_index + idx,
|
||||
output_file_path=req.output_file_path(1, 0),
|
||||
output_file_path=output_file_path,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Generation failed: %s", e, exc_info=True)
|
||||
finally:
|
||||
if requests and requests[0].data_type == DataType.VIDEO:
|
||||
try:
|
||||
# Pre-queue resources are shared by the shallow
|
||||
# per-output Req copies, so one idempotent cleanup is
|
||||
# sufficient for the whole parent request.
|
||||
requests[0].sampling_params.cleanup_video_request(requests[0])
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to clean up model-owned video request resources",
|
||||
exc_info=True,
|
||||
)
|
||||
global_output_index += len(requests)
|
||||
|
||||
total_gen_time = time.perf_counter() - total_start_time
|
||||
|
||||
@@ -134,6 +134,9 @@ class VideoGenerationsRequest(BaseModel):
|
||||
diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend
|
||||
# Performance profiling
|
||||
perf_dump_path: Optional[str] = None
|
||||
profile: Optional[bool] = False
|
||||
num_profiled_timesteps: Optional[int] = None
|
||||
profile_all_stages: Optional[bool] = False
|
||||
|
||||
|
||||
class VideoListResponse(BaseModel):
|
||||
|
||||
@@ -341,10 +341,14 @@ async def _save_url_image_to_path(image_url: str, target_path: str) -> str:
|
||||
async def process_generation_batch(
|
||||
scheduler_client: AsyncSchedulerClient,
|
||||
batch,
|
||||
*,
|
||||
scheduler_batches=None,
|
||||
) -> tuple[list[str], OutputBatch]:
|
||||
total_start_time = time.perf_counter()
|
||||
with trace_req(batch.trace_ctx), log_generation_timer(logger, batch.prompt):
|
||||
result = await scheduler_client.forward([batch])
|
||||
result = await scheduler_client.forward(
|
||||
scheduler_batches if scheduler_batches is not None else [batch]
|
||||
)
|
||||
|
||||
if (
|
||||
result.output is None
|
||||
|
||||
@@ -83,6 +83,120 @@ def _parse_form_extra_value(value: Any) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
_MULTIPART_EXTRA_FORM_FIELDS = (
|
||||
"use_duration_template",
|
||||
"use_resolution_template",
|
||||
"use_system_prompt",
|
||||
"use_guardrails",
|
||||
"guardrails",
|
||||
"video_path",
|
||||
"video_url",
|
||||
"generate_sound",
|
||||
"sound_duration",
|
||||
"condition_frame_indexes",
|
||||
"action_mode",
|
||||
"domain_id",
|
||||
"domain_name",
|
||||
"raw_action_dim",
|
||||
"action_fps",
|
||||
"action",
|
||||
"action_view_point",
|
||||
"action_normalization",
|
||||
"condition_frame_indexes_vision",
|
||||
"condition_video_keep",
|
||||
)
|
||||
|
||||
|
||||
def _video_sampling_params_cls(server_args) -> type[SamplingParams]:
|
||||
"""Resolve the params type selected for the current server."""
|
||||
|
||||
sampling_params_cls = SamplingParams
|
||||
if server_args.pipeline_class_name:
|
||||
from sglang.multimodal_gen.registry import get_pipeline_config_classes
|
||||
|
||||
config_classes = get_pipeline_config_classes(server_args.pipeline_class_name)
|
||||
if config_classes is not None:
|
||||
_, sampling_params_cls = config_classes
|
||||
if sampling_params_cls is SamplingParams:
|
||||
from sglang.multimodal_gen.registry import get_model_info
|
||||
|
||||
model_info = get_model_info(
|
||||
server_args.model_path,
|
||||
backend=server_args.backend,
|
||||
model_id=server_args.model_id,
|
||||
)
|
||||
if model_info is not None:
|
||||
sampling_params_cls = model_info.sampling_param_cls
|
||||
return sampling_params_cls
|
||||
|
||||
|
||||
def _multipart_extra_form_keys(
|
||||
sampling_params_cls: type[SamplingParams],
|
||||
) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
*VideoGenerationsRequest.model_fields,
|
||||
*_MULTIPART_EXTRA_FORM_FIELDS,
|
||||
*sorted(sampling_params_cls.video_request_extra_fields()),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _filter_multipart_declared_fields(
|
||||
extra_from_form: Dict[str, Any],
|
||||
sampling_params_cls: type[SamplingParams],
|
||||
) -> Dict[str, Any]:
|
||||
declared = set(_multipart_extra_form_keys(sampling_params_cls))
|
||||
return {key: value for key, value in extra_from_form.items() if key in declared}
|
||||
|
||||
|
||||
def _merge_multipart_extra_form_fields(
|
||||
raw_form: Any,
|
||||
extra_from_form: Dict[str, Any],
|
||||
sampling_params_cls: type[SamplingParams],
|
||||
) -> None:
|
||||
for key in _multipart_extra_form_keys(sampling_params_cls):
|
||||
if key in raw_form and key not in extra_from_form:
|
||||
extra_from_form[key] = _parse_form_extra_value(raw_form[key])
|
||||
|
||||
|
||||
def _multipart_video_extras(
|
||||
raw_form: Any,
|
||||
*,
|
||||
extra_body: Any,
|
||||
extra_params: Any,
|
||||
sampling_params_cls: type[SamplingParams],
|
||||
) -> Dict[str, Any]:
|
||||
"""Build and validate multipart extras once for request construction."""
|
||||
|
||||
extra_from_form: Dict[str, Any] = {}
|
||||
if extra_body:
|
||||
try:
|
||||
extra_from_form = flatten_extra_params(json.loads(extra_body))
|
||||
except (json.JSONDecodeError, ValueError, TypeError) as exc:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="extra_body is not valid JSON"
|
||||
) from exc
|
||||
if extra_params:
|
||||
try:
|
||||
extra_from_form.update(
|
||||
flatten_extra_params({"extra_params": json.loads(extra_params)})
|
||||
)
|
||||
except (json.JSONDecodeError, ValueError, TypeError) as exc:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="extra_params is not valid JSON"
|
||||
) from exc
|
||||
_merge_multipart_extra_form_fields(
|
||||
raw_form,
|
||||
extra_from_form,
|
||||
sampling_params_cls,
|
||||
)
|
||||
flatten_extra_params(extra_from_form)
|
||||
return _filter_multipart_declared_fields(extra_from_form, sampling_params_cls)
|
||||
|
||||
|
||||
def _is_probably_video_source(source: Any) -> bool:
|
||||
content_type = (getattr(source, "content_type", "") or "").lower()
|
||||
if content_type.startswith("video/"):
|
||||
@@ -224,45 +338,52 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
|
||||
server_args.pipeline_config.action_stats_path
|
||||
)
|
||||
|
||||
return build_sampling_params(
|
||||
request_id,
|
||||
prompt=request.prompt,
|
||||
num_outputs_per_prompt=max(1, min(int(num_outputs), 10)),
|
||||
size=request.size,
|
||||
width=request.width,
|
||||
height=request.height,
|
||||
num_frames=num_frames,
|
||||
fps=fps,
|
||||
image_path=image_path,
|
||||
video_path=video_path,
|
||||
output_file_name=request_id,
|
||||
seed=request.seed,
|
||||
generator_device=request.generator_device,
|
||||
num_inference_steps=request.num_inference_steps,
|
||||
guidance_scale=request.guidance_scale,
|
||||
guidance_scale_2=request.guidance_scale_2,
|
||||
negative_prompt=request.negative_prompt,
|
||||
max_sequence_length=request.max_sequence_length,
|
||||
flow_shift=request.flow_shift,
|
||||
use_duration_template=_extra_value(request, "use_duration_template"),
|
||||
use_resolution_template=_extra_value(request, "use_resolution_template"),
|
||||
use_system_prompt=_extra_value(request, "use_system_prompt"),
|
||||
use_guardrails=_extra_value(request, "use_guardrails"),
|
||||
enable_teacache=request.enable_teacache,
|
||||
enable_frame_interpolation=request.enable_frame_interpolation,
|
||||
frame_interpolation_exp=request.frame_interpolation_exp,
|
||||
frame_interpolation_scale=request.frame_interpolation_scale,
|
||||
frame_interpolation_model_path=request.frame_interpolation_model_path,
|
||||
enable_upscaling=request.enable_upscaling,
|
||||
upscaling_model_path=request.upscaling_model_path,
|
||||
upscaling_scale=request.upscaling_scale,
|
||||
output_path=request.output_path,
|
||||
output_compression=request.output_compression,
|
||||
output_quality=request.output_quality,
|
||||
perf_dump_path=request.perf_dump_path,
|
||||
diffusers_kwargs=request.diffusers_kwargs,
|
||||
kwargs = {
|
||||
"prompt": request.prompt,
|
||||
"num_outputs_per_prompt": max(1, min(int(num_outputs), 10)),
|
||||
"size": request.size,
|
||||
"width": request.width,
|
||||
"height": request.height,
|
||||
"num_frames": num_frames,
|
||||
"fps": fps,
|
||||
"image_path": image_path,
|
||||
"video_path": video_path,
|
||||
"output_file_name": request_id,
|
||||
"seed": request.seed,
|
||||
"generator_device": request.generator_device,
|
||||
"num_inference_steps": request.num_inference_steps,
|
||||
"guidance_scale": request.guidance_scale,
|
||||
"guidance_scale_2": request.guidance_scale_2,
|
||||
"true_cfg_scale": request.true_cfg_scale,
|
||||
"negative_prompt": request.negative_prompt,
|
||||
"max_sequence_length": request.max_sequence_length,
|
||||
"flow_shift": request.flow_shift,
|
||||
"use_duration_template": _extra_value(request, "use_duration_template"),
|
||||
"use_resolution_template": _extra_value(request, "use_resolution_template"),
|
||||
"use_system_prompt": _extra_value(request, "use_system_prompt"),
|
||||
"use_guardrails": _extra_value(request, "use_guardrails"),
|
||||
"enable_teacache": request.enable_teacache,
|
||||
"enable_frame_interpolation": request.enable_frame_interpolation,
|
||||
"frame_interpolation_exp": request.frame_interpolation_exp,
|
||||
"frame_interpolation_scale": request.frame_interpolation_scale,
|
||||
"frame_interpolation_model_path": request.frame_interpolation_model_path,
|
||||
"enable_upscaling": request.enable_upscaling,
|
||||
"upscaling_model_path": request.upscaling_model_path,
|
||||
"upscaling_scale": request.upscaling_scale,
|
||||
"output_path": request.output_path,
|
||||
"output_compression": request.output_compression,
|
||||
"output_quality": request.output_quality,
|
||||
"perf_dump_path": request.perf_dump_path,
|
||||
"profile": request.profile,
|
||||
"num_profiled_timesteps": request.num_profiled_timesteps,
|
||||
"profile_all_stages": request.profile_all_stages,
|
||||
"diffusers_kwargs": request.diffusers_kwargs,
|
||||
**cosmos3_kwargs,
|
||||
)
|
||||
}
|
||||
|
||||
sampling_params_cls = _video_sampling_params_cls(server_args)
|
||||
kwargs = sampling_params_cls.lower_video_request_kwargs(request, kwargs)
|
||||
return build_sampling_params(request_id, **kwargs)
|
||||
|
||||
|
||||
# extract metadata which http_server needs to know
|
||||
@@ -311,6 +432,7 @@ async def _dispatch_job_async(
|
||||
job_id: str,
|
||||
batch: Req,
|
||||
*,
|
||||
scheduler_batches: list[Req] | None = None,
|
||||
temp_dirs: list[str] | None = None,
|
||||
output_persistent: bool = True,
|
||||
) -> None:
|
||||
@@ -318,9 +440,30 @@ async def _dispatch_job_async(
|
||||
|
||||
try:
|
||||
save_file_path_list, result = await process_generation_batch(
|
||||
async_scheduler_client, batch
|
||||
async_scheduler_client,
|
||||
batch,
|
||||
scheduler_batches=scheduler_batches,
|
||||
)
|
||||
save_file_path = save_file_path_list[0]
|
||||
try:
|
||||
final_media_fields = await asyncio.to_thread(
|
||||
batch.sampling_params.validate_video_final_outputs,
|
||||
save_file_path_list,
|
||||
batch,
|
||||
)
|
||||
except Exception:
|
||||
for output_path in save_file_path_list:
|
||||
try:
|
||||
os.remove(output_path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except OSError:
|
||||
logger.warning(
|
||||
"Failed to remove rejected video output %s",
|
||||
output_path,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
cloud_url = await cloud_storage.upload_and_cleanup(save_file_path)
|
||||
|
||||
@@ -343,13 +486,32 @@ async def _dispatch_job_async(
|
||||
update_fields = add_common_data_to_response(
|
||||
update_fields, request_id=job_id, result=result
|
||||
)
|
||||
update_fields.update(final_media_fields)
|
||||
await VIDEO_STORE.update_fields(job_id, update_fields)
|
||||
except Exception as e:
|
||||
logger.error(f"{e}")
|
||||
await VIDEO_STORE.update_fields(
|
||||
job_id, {"status": "failed", "error": {"message": str(e)}}
|
||||
job_id,
|
||||
{
|
||||
"status": "failed",
|
||||
"error": {"message": str(e)},
|
||||
"url": None,
|
||||
"file_path": None,
|
||||
"file_paths": None,
|
||||
"num_outputs": None,
|
||||
},
|
||||
)
|
||||
finally:
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
batch.sampling_params.cleanup_video_request,
|
||||
batch,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to clean up model-owned video request resources",
|
||||
exc_info=True,
|
||||
)
|
||||
for td in temp_dirs or []:
|
||||
shutil.rmtree(td, ignore_errors=True)
|
||||
|
||||
@@ -376,6 +538,8 @@ async def create_video(
|
||||
generator_device: Optional[str] = Form("cuda"),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
guidance_scale: Optional[float] = Form(None),
|
||||
guidance_scale_2: Optional[float] = Form(None),
|
||||
true_cfg_scale: Optional[float] = Form(None),
|
||||
num_inference_steps: Optional[int] = Form(None),
|
||||
max_sequence_length: Optional[int] = Form(None),
|
||||
flow_shift: Optional[float] = Form(None),
|
||||
@@ -398,6 +562,22 @@ async def create_video(
|
||||
|
||||
server_args = get_global_server_args()
|
||||
task_type = server_args.pipeline_config.task_type
|
||||
is_multipart = "multipart/form-data" in content_type
|
||||
raw_form: Any = None
|
||||
extra_from_form: Dict[str, Any] = {}
|
||||
|
||||
# Parse model-specific multipart metadata before creating request-owned
|
||||
# directories or saving uploads, so malformed JSON leaves no resources.
|
||||
if is_multipart:
|
||||
if not prompt:
|
||||
raise HTTPException(status_code=400, detail="prompt is required")
|
||||
raw_form = await request.form()
|
||||
extra_from_form = _multipart_video_extras(
|
||||
raw_form,
|
||||
extra_body=extra_body,
|
||||
extra_params=extra_params,
|
||||
sampling_params_cls=_video_sampling_params_cls(server_args),
|
||||
)
|
||||
|
||||
# Resolve input upload directory (may be a temp dir when saving is disabled)
|
||||
temp_dirs: list[str] = []
|
||||
@@ -411,14 +591,11 @@ async def create_video(
|
||||
# Resolve output directory
|
||||
effective_output_path = server_args.output_path
|
||||
output_persistent = True
|
||||
if "multipart/form-data" not in content_type:
|
||||
if not is_multipart:
|
||||
# JSON body may carry a per-request output_path; checked after parsing below
|
||||
pass
|
||||
|
||||
if "multipart/form-data" in content_type:
|
||||
if not prompt:
|
||||
raise HTTPException(status_code=400, detail="prompt is required")
|
||||
|
||||
if is_multipart:
|
||||
video_input_path = None
|
||||
image_sources = merge_image_input_list(input_reference, reference_url)
|
||||
if video_reference is not None:
|
||||
@@ -462,52 +639,10 @@ async def create_video(
|
||||
status_code=400, detail=f"Failed to process image source: {str(e)}"
|
||||
)
|
||||
|
||||
# Parse extra_body JSON (if provided in multipart form) to get fps/num_frames overrides
|
||||
extra_from_form: Dict[str, Any] = {}
|
||||
if extra_body:
|
||||
try:
|
||||
extra_from_form = flatten_extra_params(json.loads(extra_body))
|
||||
except Exception:
|
||||
extra_from_form = {}
|
||||
if extra_params:
|
||||
try:
|
||||
extra_from_form.update(
|
||||
flatten_extra_params({"extra_params": json.loads(extra_params)})
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def form_value(name: str, value: Any) -> Any:
|
||||
selected = value if value is not None else extra_from_form.get(name)
|
||||
return _parse_form_extra_value(selected)
|
||||
|
||||
raw_form = await request.form()
|
||||
for key in (
|
||||
"use_duration_template",
|
||||
"use_resolution_template",
|
||||
"use_system_prompt",
|
||||
"use_guardrails",
|
||||
"guardrails",
|
||||
"video_path",
|
||||
"video_url",
|
||||
"generate_sound",
|
||||
"sound_duration",
|
||||
"condition_frame_indexes",
|
||||
"action_mode",
|
||||
"domain_id",
|
||||
"domain_name",
|
||||
"raw_action_dim",
|
||||
"action_fps",
|
||||
"action",
|
||||
"action_view_point",
|
||||
"action_normalization",
|
||||
"condition_frame_indexes_vision",
|
||||
"condition_video_keep",
|
||||
):
|
||||
if key in raw_form and key not in extra_from_form:
|
||||
extra_from_form[key] = _parse_form_extra_value(raw_form[key])
|
||||
flatten_extra_params(extra_from_form)
|
||||
|
||||
request_field_names = set(VideoGenerationsRequest.model_fields)
|
||||
extra_request_fields = {
|
||||
key: value
|
||||
@@ -536,6 +671,8 @@ async def create_video(
|
||||
negative_prompt=form_value("negative_prompt", negative_prompt),
|
||||
num_inference_steps=form_value("num_inference_steps", num_inference_steps),
|
||||
guidance_scale=form_value("guidance_scale", guidance_scale),
|
||||
guidance_scale_2=form_value("guidance_scale_2", guidance_scale_2),
|
||||
true_cfg_scale=form_value("true_cfg_scale", true_cfg_scale),
|
||||
max_sequence_length=form_value("max_sequence_length", max_sequence_length),
|
||||
flow_shift=form_value("flow_shift", flow_shift),
|
||||
enable_teacache=form_value("enable_teacache", enable_teacache),
|
||||
@@ -639,30 +776,59 @@ async def create_video(
|
||||
try:
|
||||
sampling_params = _build_video_sampling_params(request_id, req)
|
||||
except (ValueError, TypeError) as e:
|
||||
for td in temp_dirs:
|
||||
shutil.rmtree(td, ignore_errors=True)
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
job = _video_job_from_sampling(request_id, req, sampling_params)
|
||||
await VIDEO_STORE.upsert(request_id, job)
|
||||
batch: Req | None = None
|
||||
scheduler_batches: list[Req] | None = None
|
||||
try:
|
||||
# Build Req for scheduler.
|
||||
trace_headers = extract_trace_headers(request.headers)
|
||||
batch = prepare_request(
|
||||
server_args=server_args,
|
||||
sampling_params=sampling_params,
|
||||
external_trace_header=trace_headers,
|
||||
)
|
||||
# Add diffusers_kwargs if provided.
|
||||
if req.diffusers_kwargs:
|
||||
batch.extra["diffusers_kwargs"] = req.diffusers_kwargs
|
||||
if "max_sequence_length" in req.diffusers_kwargs:
|
||||
batch.max_sequence_length = req.diffusers_kwargs["max_sequence_length"]
|
||||
if "flow_shift" in req.diffusers_kwargs:
|
||||
batch.flow_shift = req.diffusers_kwargs["flow_shift"]
|
||||
await asyncio.to_thread(
|
||||
sampling_params.prepare_video_request_for_queue,
|
||||
batch,
|
||||
)
|
||||
scheduler_batches = sampling_params.expand_video_request_outputs_for_queue(
|
||||
batch
|
||||
)
|
||||
job = _video_job_from_sampling(request_id, req, sampling_params)
|
||||
job.update(sampling_params.project_video_queued_job_fields(batch))
|
||||
await VIDEO_STORE.upsert(request_id, job)
|
||||
except Exception as e:
|
||||
if batch is not None:
|
||||
try:
|
||||
await asyncio.to_thread(sampling_params.cleanup_video_request, batch)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to clean up rejected video request resources",
|
||||
exc_info=True,
|
||||
)
|
||||
for td in temp_dirs:
|
||||
shutil.rmtree(td, ignore_errors=True)
|
||||
if isinstance(e, (TypeError, ValueError)):
|
||||
raise HTTPException(status_code=400, detail=str(e)) from e
|
||||
raise
|
||||
|
||||
# Build Req for scheduler
|
||||
trace_headers = extract_trace_headers(request.headers)
|
||||
batch = prepare_request(
|
||||
server_args=server_args,
|
||||
sampling_params=sampling_params,
|
||||
external_trace_header=trace_headers,
|
||||
)
|
||||
# Add diffusers_kwargs if provided
|
||||
if req.diffusers_kwargs:
|
||||
batch.extra["diffusers_kwargs"] = req.diffusers_kwargs
|
||||
if "max_sequence_length" in req.diffusers_kwargs:
|
||||
batch.max_sequence_length = req.diffusers_kwargs["max_sequence_length"]
|
||||
if "flow_shift" in req.diffusers_kwargs:
|
||||
batch.flow_shift = req.diffusers_kwargs["flow_shift"]
|
||||
assert batch is not None
|
||||
# Enqueue the job asynchronously and return immediately
|
||||
asyncio.create_task(
|
||||
_dispatch_job_async(
|
||||
request_id,
|
||||
batch,
|
||||
scheduler_batches=scheduler_batches,
|
||||
temp_dirs=temp_dirs or None,
|
||||
output_persistent=output_persistent,
|
||||
)
|
||||
@@ -717,6 +883,21 @@ async def delete_video(video_id: str = Path(...)):
|
||||
return VideoResponse(**job)
|
||||
|
||||
|
||||
def _select_video_variant_path(job: dict, variant: str | None) -> str | None:
|
||||
file_paths = job.get("file_paths")
|
||||
if file_paths:
|
||||
try:
|
||||
variant_index = 0 if variant is None else int(variant)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if 0 <= variant_index < len(file_paths):
|
||||
return file_paths[variant_index]
|
||||
return None
|
||||
if variant not in (None, "0", 0):
|
||||
return None
|
||||
return job.get("file_path")
|
||||
|
||||
|
||||
@router.get("/{video_id}/content")
|
||||
async def download_video_content(
|
||||
video_id: str = Path(...), variant: Optional[str] = Query(None)
|
||||
@@ -731,9 +912,13 @@ async def download_video_content(
|
||||
detail=f"Video has been uploaded to cloud storage. Please use the cloud URL: {job.get('url')}",
|
||||
)
|
||||
|
||||
file_path = job.get("file_path")
|
||||
if not file_path or not os.path.exists(file_path):
|
||||
file_path = _select_video_variant_path(job, variant)
|
||||
if job.get("status") not in {"completed", "failed"}:
|
||||
raise HTTPException(status_code=404, detail="Generation is still in-progress")
|
||||
if not file_path or not os.path.exists(file_path):
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Video variant {variant} not found"
|
||||
)
|
||||
|
||||
media_type = "video/mp4" # default variant
|
||||
return FileResponse(
|
||||
|
||||
@@ -341,6 +341,7 @@ def expand_request_outputs(
|
||||
req.seed = seeds[0]
|
||||
req.seeds = None
|
||||
req.generator = None
|
||||
req.sampling_params.refresh_request_extra_after_output_expansion(req)
|
||||
return [req]
|
||||
|
||||
expanded: list[Req] = []
|
||||
@@ -365,6 +366,9 @@ def expand_request_outputs(
|
||||
output_req.output_file_name = _with_output_index_suffix(
|
||||
req.output_file_name, output_index
|
||||
)
|
||||
output_req.sampling_params.refresh_request_extra_after_output_expansion(
|
||||
output_req
|
||||
)
|
||||
output_req.validate()
|
||||
expanded.append(output_req)
|
||||
|
||||
@@ -487,7 +491,7 @@ def _try_save_cuda_video_direct(
|
||||
if video.shape[0] != 3:
|
||||
return False
|
||||
|
||||
frames = (video * 255).clamp(0, 255).to(torch.uint8)
|
||||
frames = (video * 255).clamp_(0, 255).to(torch.uint8)
|
||||
frames = frames.permute(1, 2, 3, 0).contiguous()
|
||||
num_frames, height, width, _ = frames.shape
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
|
||||
@@ -15,7 +16,10 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER_GFX95
|
||||
from sglang.multimodal_gen.runtime.platforms.aiter import (
|
||||
USE_AITER_GFX95,
|
||||
USE_AITER_GFX942,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -205,3 +209,38 @@ class AITerImpl(AttentionImpl):
|
||||
return_lse=True,
|
||||
)
|
||||
return output
|
||||
|
||||
@torch.compiler.disable
|
||||
def forward_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
) -> torch.Tensor:
|
||||
del cu_seqlens_host
|
||||
if USE_AITER_GFX942:
|
||||
# The grouped-varlen ASM kernel hangs on H3's ~64K packed
|
||||
# sequences on gfx942; AITER's Triton path handles this shape.
|
||||
attention_func = importlib.import_module(
|
||||
"aiter.ops.triton.attention.mha"
|
||||
).flash_attn_varlen_func
|
||||
else:
|
||||
attention_func = aiter.flash_attn_varlen_func
|
||||
|
||||
cu_seqlens = cu_seqlens.to(device=query.device, dtype=torch.int32).contiguous()
|
||||
output = attention_func(
|
||||
q=query.contiguous(),
|
||||
k=key.contiguous(),
|
||||
v=value.contiguous(),
|
||||
cu_seqlens_q=cu_seqlens,
|
||||
cu_seqlens_k=cu_seqlens,
|
||||
max_seqlen_q=max_seqlen,
|
||||
max_seqlen_k=max_seqlen,
|
||||
softmax_scale=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
)
|
||||
return output[0] if isinstance(output, tuple) else output
|
||||
|
||||
@@ -170,6 +170,20 @@ class AttentionImpl(ABC, Generic[T]):
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError
|
||||
|
||||
def forward_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} does not implement packed varlen attention"
|
||||
)
|
||||
|
||||
|
||||
def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
|
||||
return wrap_method_with_debug_kernel_once(
|
||||
|
||||
@@ -443,3 +443,28 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
return out_tensor
|
||||
|
||||
raise ValueError(f"flash attention version {fa_ver} is not supported.")
|
||||
|
||||
def forward_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
) -> torch.Tensor:
|
||||
del cu_seqlens_host
|
||||
output = flash_attn_varlen_func(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seqlens_q=cu_seqlens,
|
||||
cu_seqlens_k=cu_seqlens,
|
||||
max_seqlen_q=max_seqlen,
|
||||
max_seqlen_k=max_seqlen,
|
||||
softmax_scale=self.softmax_scale,
|
||||
causal=self.causal,
|
||||
ver=fa_ver,
|
||||
)
|
||||
return output[0] if isinstance(output, tuple) else output
|
||||
|
||||
@@ -94,6 +94,35 @@ class SDPAImpl(AttentionImpl):
|
||||
output = output.transpose(1, 2)
|
||||
return output
|
||||
|
||||
def forward_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
) -> torch.Tensor:
|
||||
del max_seqlen
|
||||
bounds = (
|
||||
cu_seqlens_host
|
||||
if cu_seqlens_host is not None
|
||||
else tuple(int(item) for item in cu_seqlens.tolist())
|
||||
)
|
||||
output = torch.empty_like(query)
|
||||
for start, stop in zip(bounds[:-1], bounds[1:]):
|
||||
if start == stop:
|
||||
continue
|
||||
segment = self.forward(
|
||||
query[start:stop].unsqueeze(0),
|
||||
key[start:stop].unsqueeze(0),
|
||||
value[start:stop].unsqueeze(0),
|
||||
None,
|
||||
)
|
||||
output[start:stop].copy_(segment[0])
|
||||
return output
|
||||
|
||||
|
||||
class CudnnSDPABackend(SDPABackend):
|
||||
@staticmethod
|
||||
@@ -122,7 +151,7 @@ class DynamicCudnnSDPABackend(SDPABackend):
|
||||
return DynamicCudnnSDPAImpl
|
||||
|
||||
|
||||
class DynamicCudnnSDPAImpl(AttentionImpl):
|
||||
class DynamicCudnnSDPAImpl(SDPAImpl):
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
|
||||
@@ -8,6 +8,10 @@ import torch.distributed as dist
|
||||
import torch.distributed._functional_collectives as ft_c
|
||||
from torch.distributed.tensor.experimental._attention import _cp_options
|
||||
|
||||
from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
|
||||
pack_qkv_destination_major,
|
||||
)
|
||||
from sglang.kernels.ops.diffusion.usp_relayout import usp_merge_heads
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_sp_group,
|
||||
get_ulysses_parallel_rank,
|
||||
@@ -280,6 +284,47 @@ def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
|
||||
def _usp_input_all_to_all_packed_qkv(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Exchange 3D Q/K/V with one destination-major Ulysses collective."""
|
||||
world_size = get_ulysses_parallel_world_size()
|
||||
if world_size <= 1:
|
||||
return q, k, v
|
||||
|
||||
assert q.ndim == 3 and q.shape == k.shape == v.shape
|
||||
s_local, h_global, head_size = q.shape
|
||||
assert h_global % world_size == 0
|
||||
h_local = h_global // world_size
|
||||
|
||||
if (
|
||||
q.is_cuda
|
||||
and q.dtype in (torch.float16, torch.bfloat16)
|
||||
and q.dtype == k.dtype == v.dtype
|
||||
and q.stride(-1) == k.stride(-1) == v.stride(-1) == 1
|
||||
and not torch.compiler.is_compiling()
|
||||
):
|
||||
packed = pack_qkv_destination_major(q, k, v, world_size)
|
||||
else:
|
||||
packed = torch.empty(
|
||||
(world_size, s_local, h_local, 3 * head_size),
|
||||
dtype=q.dtype,
|
||||
device=q.device,
|
||||
)
|
||||
for index, tensor in enumerate((q, k, v)):
|
||||
head_shards = tensor.view(s_local, world_size, h_local, head_size).permute(
|
||||
1, 0, 2, 3
|
||||
)
|
||||
packed[..., index * head_size : (index + 1) * head_size].copy_(head_shards)
|
||||
|
||||
packed = _usp_all_to_all_single(packed)
|
||||
packed = packed.reshape(s_local * world_size, h_local, 3 * head_size)
|
||||
q, k, v = packed.split(head_size, dim=-1)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def _usp_input_all_to_all_varlen(
|
||||
x: torch.Tensor, seq_lens: list[int], head_dim: int = 1
|
||||
) -> torch.Tensor:
|
||||
@@ -419,7 +464,7 @@ def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
|
||||
x = x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, h_global, s_local, d)
|
||||
else: # head_dim == 2
|
||||
# Shape transition: [world_size, s_local, b, h_local, d] -> [b, s_local, world_size, h_local, d]
|
||||
x = x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, s_local, h_global, d)
|
||||
x = usp_merge_heads(x).reshape(b, s_local, h_global, d)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
@@ -113,9 +113,12 @@ class ComponentLoader(ABC):
|
||||
return {}
|
||||
|
||||
def should_raise_customized_load_error(
|
||||
self, _server_args: ServerArgs, _component_name: str
|
||||
self, server_args: ServerArgs, component_name: str
|
||||
) -> bool:
|
||||
return False
|
||||
native_only_components = getattr(
|
||||
server_args.pipeline_config, "native_only_components", ()
|
||||
)
|
||||
return component_name in native_only_components
|
||||
|
||||
@staticmethod
|
||||
def _is_component_set_as_layerwise_load(
|
||||
|
||||
@@ -54,7 +54,6 @@ class ImageEncoderLoader(TextEncoderLoader):
|
||||
finalize_encoder_folding(
|
||||
encoder_config,
|
||||
server_args.encoder_parallel,
|
||||
batched=server_args.batching_max_size > 1,
|
||||
)
|
||||
|
||||
# Always start with local device; load_model will adjust for offload if needed
|
||||
|
||||
+42
-3
@@ -2,7 +2,7 @@ import dataclasses
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Generator, Iterable
|
||||
from collections.abc import Callable, Generator, Iterable
|
||||
from contextlib import nullcontext
|
||||
from typing import cast
|
||||
|
||||
@@ -39,6 +39,7 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import (
|
||||
safetensors_weights_iterator,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.encoders.base import (
|
||||
TextEncoder,
|
||||
finalize_encoder_folding,
|
||||
get_folding_tp_group,
|
||||
)
|
||||
@@ -184,6 +185,7 @@ class TextEncoderLoader(ComponentLoader):
|
||||
model_name_or_path: str,
|
||||
fall_back_to_pt: bool,
|
||||
allow_patterns_overrides: list[str] | None,
|
||||
key_filter: Callable[[str], bool] | None = None,
|
||||
) -> tuple[str, list[str], bool]:
|
||||
"""Prepare weights for the model.
|
||||
|
||||
@@ -216,7 +218,10 @@ class TextEncoderLoader(ComponentLoader):
|
||||
|
||||
if use_safetensors:
|
||||
hf_weights_files = filter_duplicate_safetensors_files(
|
||||
hf_weights_files, hf_folder, index_file
|
||||
hf_weights_files,
|
||||
hf_folder,
|
||||
index_file,
|
||||
key_filter=key_filter,
|
||||
)
|
||||
else:
|
||||
hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
|
||||
@@ -237,20 +242,39 @@ class TextEncoderLoader(ComponentLoader):
|
||||
self,
|
||||
source: "Source",
|
||||
to_cpu: bool,
|
||||
key_filter: Callable[[str], bool] | None = None,
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
"""get an iterator for the model weights based on the load format."""
|
||||
source_key_filter: Callable[[str], bool] | None
|
||||
if key_filter is None:
|
||||
source_key_filter = None
|
||||
else:
|
||||
|
||||
def include_source_weight(name: str) -> bool:
|
||||
return key_filter(source.prefix + name)
|
||||
|
||||
source_key_filter = include_source_weight
|
||||
|
||||
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
|
||||
source.model_or_path,
|
||||
source.fall_back_to_pt,
|
||||
source.allow_patterns_overrides,
|
||||
key_filter=source_key_filter,
|
||||
)
|
||||
if use_safetensors:
|
||||
weights_iterator = safetensors_weights_iterator(
|
||||
hf_weights_files,
|
||||
to_cpu=to_cpu,
|
||||
key_filter=source_key_filter,
|
||||
)
|
||||
else:
|
||||
weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
|
||||
if source_key_filter is not None:
|
||||
weights_iterator = (
|
||||
(name, tensor)
|
||||
for name, tensor in weights_iterator
|
||||
if source_key_filter(name)
|
||||
)
|
||||
|
||||
# apply the prefix.
|
||||
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
|
||||
@@ -261,6 +285,10 @@ class TextEncoderLoader(ComponentLoader):
|
||||
model_path: str,
|
||||
to_cpu: bool,
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
key_filter = cast(
|
||||
Callable[[str], bool] | None,
|
||||
getattr(model, "should_materialize_checkpoint_weight", None),
|
||||
)
|
||||
primary_weights = TextEncoderLoader.Source(
|
||||
model_path,
|
||||
prefix="",
|
||||
@@ -270,6 +298,7 @@ class TextEncoderLoader(ComponentLoader):
|
||||
yield from self._get_weights_iterator(
|
||||
primary_weights,
|
||||
to_cpu,
|
||||
key_filter,
|
||||
)
|
||||
|
||||
secondary_weights = cast(
|
||||
@@ -280,6 +309,7 @@ class TextEncoderLoader(ComponentLoader):
|
||||
yield from self._get_weights_iterator(
|
||||
source,
|
||||
to_cpu,
|
||||
key_filter,
|
||||
)
|
||||
|
||||
def load_customized(
|
||||
@@ -314,11 +344,20 @@ class TextEncoderLoader(ComponentLoader):
|
||||
)
|
||||
if post_diffusers_config_update is not None:
|
||||
post_diffusers_config_update()
|
||||
model_cls, _ = ModelRegistry.resolve_model_cls(
|
||||
getattr(encoder_config, "architectures", [])
|
||||
)
|
||||
# real dims are populated now; resolve fold vs replicate
|
||||
finalize_encoder_folding(
|
||||
encoder_config,
|
||||
server_args.encoder_parallel,
|
||||
batched=server_args.batching_max_size > 1,
|
||||
prefer_dp=(
|
||||
server_args.batching_max_size > 1
|
||||
and (server_args.tp_size or 1) == 1
|
||||
and (server_args.dp_size or 1) == 1
|
||||
and issubclass(model_cls, TextEncoder)
|
||||
and model_cls.supports_dp_encode
|
||||
),
|
||||
)
|
||||
encoder_dtype = server_args.pipeline_config.text_encoder_precisions[
|
||||
encoder_index
|
||||
|
||||
+17
-11
@@ -28,6 +28,20 @@ _is_npu = is_npu()
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _warn_if_expected_param_dtype_missing(
|
||||
model: torch.nn.Module, expected_dtype: torch.dtype | None
|
||||
) -> None:
|
||||
if expected_dtype is None:
|
||||
return
|
||||
param_dtypes = {param.dtype for param in model.parameters()}
|
||||
if expected_dtype not in param_dtypes:
|
||||
logger.warning(
|
||||
"Model parameter dtypes do not include expected param dtype, %s vs %s",
|
||||
param_dtypes,
|
||||
expected_dtype,
|
||||
)
|
||||
|
||||
|
||||
def _server_args_for_transformer_component(
|
||||
server_args: ServerArgs, component_name: str
|
||||
) -> ServerArgs:
|
||||
@@ -89,7 +103,8 @@ class TransformerLoader(ComponentLoader):
|
||||
# Don't let a quantized load quietly fall back to the unquantized native
|
||||
# model. That would drop the requested precision and bury the real error.
|
||||
return (
|
||||
component_server_args.transformer_weights_path is not None
|
||||
super().should_raise_customized_load_error(server_args, component_name)
|
||||
or component_server_args.transformer_weights_path is not None
|
||||
or component_server_args.quantization is not None
|
||||
)
|
||||
|
||||
@@ -185,15 +200,6 @@ class TransformerLoader(ComponentLoader):
|
||||
for post_load_hook in quant_spec.post_load_hooks:
|
||||
post_load_hook(model)
|
||||
|
||||
# considering the existent of mixed-precision models (e.g., nunchaku)
|
||||
if (
|
||||
next(model.parameters()).dtype != quant_spec.param_dtype
|
||||
and quant_spec.param_dtype
|
||||
):
|
||||
logger.warning(
|
||||
"Model dtype does not match expected param dtype, %s vs %s",
|
||||
next(model.parameters()).dtype,
|
||||
quant_spec.param_dtype,
|
||||
)
|
||||
_warn_if_expected_param_dtype_missing(model, quant_spec.param_dtype)
|
||||
|
||||
return model
|
||||
|
||||
@@ -142,10 +142,12 @@ class VAELoader(ComponentLoader):
|
||||
should_offload = self.should_offload(server_args)
|
||||
target_device = self.target_device(should_offload)
|
||||
|
||||
# Check for auto_map first (custom VAE classes)
|
||||
native_only = component_name in getattr(
|
||||
server_args.pipeline_config, "native_only_components", ()
|
||||
)
|
||||
auto_map = config.get("auto_map", {})
|
||||
auto_model_map = auto_map.get("AutoModel")
|
||||
if auto_model_map:
|
||||
if auto_model_map and not native_only:
|
||||
module_path, cls_name = auto_model_map.rsplit(".", 1)
|
||||
custom_module_file = os.path.join(component_model_path, f"{module_path}.py")
|
||||
spec = importlib.util.spec_from_file_location("_custom", custom_module_file)
|
||||
@@ -191,16 +193,18 @@ class VAELoader(ComponentLoader):
|
||||
for sf_path in safetensors_list:
|
||||
loaded.update(safetensors_load_file(sf_path))
|
||||
_backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
|
||||
vae.load_state_dict(loaded, strict=False)
|
||||
strict_load = native_only
|
||||
vae.load_state_dict(loaded, strict=strict_load)
|
||||
|
||||
state_keys = set(vae.state_dict().keys())
|
||||
loaded_keys = set(loaded.keys())
|
||||
missing_keys = sorted(state_keys - loaded_keys)
|
||||
unexpected_keys = sorted(loaded_keys - state_keys)
|
||||
if missing_keys:
|
||||
logger.warning("VAE missing keys: %s", missing_keys)
|
||||
if unexpected_keys:
|
||||
logger.warning("VAE unexpected keys: %s", unexpected_keys)
|
||||
if not strict_load:
|
||||
state_keys = set(vae.state_dict().keys())
|
||||
loaded_keys = set(loaded.keys())
|
||||
missing_keys = sorted(state_keys - loaded_keys)
|
||||
unexpected_keys = sorted(loaded_keys - state_keys)
|
||||
if missing_keys:
|
||||
logger.warning("VAE missing keys: %s", missing_keys)
|
||||
if unexpected_keys:
|
||||
logger.warning("VAE unexpected keys: %s", unexpected_keys)
|
||||
|
||||
if _should_use_channels_last_3d(server_args, component_name):
|
||||
n = _convert_conv3d_weights_to_channels_last_3d(vae)
|
||||
|
||||
@@ -20,6 +20,7 @@ from torch.distributed.fsdp import (
|
||||
FSDPModule,
|
||||
MixedPrecisionPolicy,
|
||||
fully_shard,
|
||||
register_fsdp_forward_method,
|
||||
)
|
||||
from torch.nn.modules.module import _IncompatibleKeys
|
||||
|
||||
@@ -204,7 +205,8 @@ def maybe_load_fsdp_model(
|
||||
Args:
|
||||
param_dtype: Data type for model parameters, also used for:
|
||||
- Model initialization context (set_default_torch_dtype)
|
||||
- FSDP mixed precision policy
|
||||
- FSDP mixed precision policy unless the model preserves mixed
|
||||
original parameter dtypes
|
||||
- Weight loading and casting
|
||||
reduce_dtype: Data type for gradient reduction in FSDP mixed precision.
|
||||
strict: If True, enforce strict state dict loading (all keys must match).
|
||||
@@ -215,8 +217,19 @@ def maybe_load_fsdp_model(
|
||||
|
||||
# 1. prepare for loading
|
||||
default_torch_dtype = param_dtype if param_dtype else torch.bfloat16
|
||||
# Some native models deliberately mix FP32 projections with lower-precision
|
||||
# blocks. FSDP must all-gather those parameters in their original dtypes;
|
||||
# the thread-local compute dtype below remains the requested default.
|
||||
fsdp_param_dtype = (
|
||||
None
|
||||
if fsdp_inference and getattr(model_cls, "_fsdp_mixed_dtype_params", False)
|
||||
else default_torch_dtype
|
||||
)
|
||||
mp_policy = MixedPrecisionPolicy(
|
||||
default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
|
||||
param_dtype=fsdp_param_dtype,
|
||||
reduce_dtype=reduce_dtype,
|
||||
output_dtype=output_dtype,
|
||||
cast_forward_inputs=False,
|
||||
)
|
||||
|
||||
set_mixed_precision_policy(
|
||||
@@ -279,6 +292,8 @@ def maybe_load_fsdp_model(
|
||||
fsdp_shard_conditions=getattr(model, "_fsdp_shard_conditions", None),
|
||||
pin_cpu_memory=pin_cpu_memory,
|
||||
)
|
||||
if callable(getattr(model, "refine_prompt_embeds", None)):
|
||||
register_fsdp_forward_method(model, "refine_prompt_embeds")
|
||||
|
||||
param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
|
||||
|
||||
|
||||
@@ -601,7 +601,12 @@ def _resolve_quant_config(
|
||||
# in source dtype and are quantized in
|
||||
# process_weights_after_loading.
|
||||
quant_cls = get_quantization_config(server_args.quantization)
|
||||
return quant_cls()
|
||||
quant_kwargs = {}
|
||||
if server_args.quantization in {"fp8", "mxfp4"}:
|
||||
quant_kwargs["ignored_layers"] = getattr(
|
||||
server_args, "quantization_ignored_layers", None
|
||||
)
|
||||
return quant_cls(**quant_kwargs)
|
||||
|
||||
quant_config = get_quant_config(hf_config, component_model_path)
|
||||
if quant_config is None and server_args.transformer_weights_path:
|
||||
|
||||
@@ -65,7 +65,10 @@ def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
|
||||
# So, we use the index_file to
|
||||
# look up which safetensors files should be used.
|
||||
def filter_duplicate_safetensors_files(
|
||||
hf_weights_files: list[str], hf_folder: str, index_file: str
|
||||
hf_weights_files: list[str],
|
||||
hf_folder: str,
|
||||
index_file: str,
|
||||
key_filter: Callable[[str], bool] | None = None,
|
||||
) -> list[str]:
|
||||
# model.safetensors.index.json is a mapping from keys in the
|
||||
# torch state_dict to safetensors file holding that weight.
|
||||
@@ -79,6 +82,9 @@ def filter_duplicate_safetensors_files(
|
||||
weight_map = json.load(f)["weight_map"]
|
||||
weight_files_in_index = set()
|
||||
for weight_name in weight_map:
|
||||
# remove only shards whose indexed tensors are all filtered
|
||||
if key_filter is not None and not key_filter(weight_name):
|
||||
continue
|
||||
weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
|
||||
# Filter out any fields that are not found in the index file.
|
||||
hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
|
||||
|
||||
@@ -107,6 +107,21 @@ class _ExpandedOutputParts:
|
||||
trajectory_decoded_parts: list[list[torch.Tensor]] | None = None
|
||||
|
||||
|
||||
def _worker_cpu_intra_op_threads(num_gpus: int) -> int | None:
|
||||
"""CPU intra-op thread budget for one of `num_gpus` co-located workers.
|
||||
|
||||
torch defaults the intra-op pool to every host core in every worker, so
|
||||
co-located workers oversubscribe the host num_gpus-fold and any CPU op
|
||||
past the ~32k-element parallel grain pays pool wakeup contention instead
|
||||
of microseconds (measured 500x on request-static packed layouts). An
|
||||
explicit OMP_NUM_THREADS keeps deployer intent (returns None).
|
||||
"""
|
||||
if "OMP_NUM_THREADS" in os.environ:
|
||||
return None
|
||||
cpu_count = os.cpu_count() or 1
|
||||
return max(1, min(16, cpu_count // max(1, num_gpus)))
|
||||
|
||||
|
||||
class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
"""
|
||||
A worker that executes the model on a single GPU.
|
||||
@@ -198,6 +213,9 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
def init_device_and_model(self) -> None:
|
||||
"""Initialize the device and load the model."""
|
||||
torch.get_device_module().set_device(self.local_rank)
|
||||
intra_op_threads = _worker_cpu_intra_op_threads(self.server_args.num_gpus)
|
||||
if intra_op_threads is not None:
|
||||
torch.set_num_threads(intra_op_threads)
|
||||
# Set environment variables for distributed initialization
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = str(self.master_port)
|
||||
|
||||
@@ -67,8 +67,12 @@ class LayerwiseOffloadManager:
|
||||
)
|
||||
self.copy_stream = torch.get_device_module().Stream()
|
||||
|
||||
# ``named_parameters()`` is relative to ``model``, just like the path in
|
||||
# ``layers_attr_str``. Anchor the match so a manager for top-level
|
||||
# ``blocks`` cannot also capture an unrelated nested list such as
|
||||
# ``token_refiner.blocks`` whose forward hooks run at a different time.
|
||||
self._layer_name_re = re.compile(
|
||||
rf"(^|\.){re.escape(layers_attr_str)}\.(\d+)(\.|$)"
|
||||
rf"^{re.escape(layers_attr_str)}\.(?P<layer_idx>\d+)(\.|$)"
|
||||
)
|
||||
|
||||
# layer_idx -> {dtype: consolidated_pinned_cpu_tensor}
|
||||
@@ -99,7 +103,7 @@ class LayerwiseOffloadManager:
|
||||
if not m:
|
||||
return None
|
||||
try:
|
||||
return int(m.group(2))
|
||||
return int(m.group("layer_idx"))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@@ -612,6 +616,10 @@ class LayerwiseOffloadableModuleMixin:
|
||||
self.layerwise_offload_managers = []
|
||||
named_modules = dict(self.named_modules())
|
||||
configured_layer_names = []
|
||||
# These legacy tuning knobs are explicitly DiT-scoped. Auxiliary
|
||||
# components still support layerwise streaming, but their layers run
|
||||
# once per component use and get no reuse benefit from DiT residency.
|
||||
dit_tuning_enabled = self.layerwise_offload_dit_group_enabled
|
||||
for layer_name in self.layer_names:
|
||||
module_list = named_modules.get(layer_name)
|
||||
if not isinstance(module_list, (torch.nn.ModuleList, torch.nn.Sequential)):
|
||||
@@ -620,14 +628,17 @@ class LayerwiseOffloadableModuleMixin:
|
||||
continue
|
||||
|
||||
num_layers = len(module_list)
|
||||
if server_args.dit_offload_prefetch_size < 1.0:
|
||||
prefetch_size = 1 + int(
|
||||
round(server_args.dit_offload_prefetch_size * (num_layers - 1))
|
||||
)
|
||||
prefetch_value = (
|
||||
server_args.dit_offload_prefetch_size if dit_tuning_enabled else 0.0
|
||||
)
|
||||
if prefetch_value < 1.0:
|
||||
prefetch_size = 1 + int(round(prefetch_value * (num_layers - 1)))
|
||||
else:
|
||||
prefetch_size = int(server_args.dit_offload_prefetch_size)
|
||||
prefetch_size = int(prefetch_value)
|
||||
|
||||
resident_value = server_args.dit_layerwise_resident_layers
|
||||
resident_value = (
|
||||
server_args.dit_layerwise_resident_layers if dit_tuning_enabled else 0.0
|
||||
)
|
||||
if resident_value <= 0:
|
||||
resident_layers = 0
|
||||
elif resident_value < 1.0:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -123,11 +123,11 @@ def encoder_dp_worthwhile(
|
||||
|
||||
|
||||
def finalize_encoder_folding(
|
||||
config: EncoderConfig, policy: str = "auto", batched: bool = False
|
||||
config: EncoderConfig, policy: str = "auto", prefer_dp: bool = False
|
||||
) -> None:
|
||||
"""resolve fold-vs-replicate once real dims are known (post update_model_arch,
|
||||
pre construction); folding shards the weights, so it rules out dp for the
|
||||
lifetime of the loaded model. `batched` is the batching ceiling being > 1."""
|
||||
lifetime of the loaded model. `prefer_dp` means the runtime can engage dp."""
|
||||
if config.parallel_folding_mode is None:
|
||||
return
|
||||
group = get_folding_tp_group(config)
|
||||
@@ -138,7 +138,7 @@ def finalize_encoder_folding(
|
||||
# a batched encode prefers dp (one all_gather) over folding (an
|
||||
# all_reduce per layer), so leave a dp-capable encoder unsharded
|
||||
keep = (
|
||||
not (batched and encoder_dp_capable(config))
|
||||
not (prefer_dp and encoder_dp_capable(config))
|
||||
and encoder_folding_worthwhile(config, group.world_size)
|
||||
and group_has_measured_topology(group)
|
||||
)
|
||||
|
||||
@@ -0,0 +1,193 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Native, TP-foldable Qwen3-VL layer-50 encoder for MiniMax H3."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
|
||||
from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
|
||||
MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER,
|
||||
MiniMaxH3Qwen3VLConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
|
||||
from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
|
||||
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel
|
||||
|
||||
MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
|
||||
_LAYER_WEIGHT_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.")
|
||||
|
||||
|
||||
def _is_unconsumed_checkpoint_weight(name: str) -> bool:
|
||||
"""Weights intentionally absent from the layer-50 feature extractor."""
|
||||
|
||||
if name == "lm_head.weight" or name.startswith("model.language_model.norm."):
|
||||
return True
|
||||
match = _LAYER_WEIGHT_RE.match(name)
|
||||
return bool(match and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER)
|
||||
|
||||
|
||||
class MiniMaxH3Qwen3VLEncoder(TextEncoder):
|
||||
"""Qwen3-VL-32B multimodal backbone ending at hidden_states[50].
|
||||
|
||||
The component loader builds and loads this module under the encoder-folding
|
||||
TP group. A TP=1/SP=8 DiT deployment therefore shards the encoder over all
|
||||
eight otherwise-idle ranks during encoding.
|
||||
"""
|
||||
|
||||
supports_dp_encode = True
|
||||
|
||||
@staticmethod
|
||||
def should_materialize_checkpoint_weight(name: str) -> bool:
|
||||
return (
|
||||
"rotary_emb.inv_freq" not in name
|
||||
and not _is_unconsumed_checkpoint_weight(name)
|
||||
)
|
||||
|
||||
def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None:
|
||||
super().__init__(config)
|
||||
arch = config.arch_config
|
||||
selected_layer = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER
|
||||
if int(arch.text_config.num_hidden_layers) != selected_layer:
|
||||
raise ValueError(
|
||||
"MiniMax H3 Qwen3-VL config must be trimmed to "
|
||||
f"{selected_layer} language layers before construction"
|
||||
)
|
||||
self.model = Qwen3VLModel(arch, use_tensor_parallel=True)
|
||||
# H3 consumes the unnormalized output immediately after layer 49.
|
||||
self.model.language_model.norm = nn.Identity()
|
||||
self.image_token_id = int(arch.image_token_id)
|
||||
self.video_token_id = int(arch.video_token_id)
|
||||
self.selected_lm_layer = selected_layer
|
||||
self.hidden_dim = MINIMAX_H3_QWEN3VL_HIDDEN_DIM
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.parameters()).device
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
position_ids: torch.Tensor | None = None,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
output_attentions: bool | None = None,
|
||||
output_hidden_states: bool | None = None,
|
||||
**kwargs: Any,
|
||||
) -> BaseEncoderOutput:
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
position_ids=position_ids,
|
||||
attention_mask=attention_mask,
|
||||
inputs_embeds=inputs_embeds,
|
||||
output_attentions=False,
|
||||
output_hidden_states=False,
|
||||
return_dict=True,
|
||||
use_cache=False,
|
||||
**kwargs,
|
||||
)
|
||||
return BaseEncoderOutput(last_hidden_state=outputs.last_hidden_state)
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
*,
|
||||
pixel_values: torch.Tensor | None = None,
|
||||
image_grid_thw: torch.Tensor | None = None,
|
||||
pixel_values_videos: torch.Tensor | None = None,
|
||||
video_grid_thw: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if input_ids.dim() != 1:
|
||||
raise ValueError(f"input_ids must be 1-D, got {list(input_ids.shape)}")
|
||||
if (pixel_values is None) != (image_grid_thw is None):
|
||||
raise ValueError("pixel_values and image_grid_thw must be given together")
|
||||
if (pixel_values_videos is None) != (video_grid_thw is None):
|
||||
raise ValueError(
|
||||
"pixel_values_videos and video_grid_thw must be given together"
|
||||
)
|
||||
|
||||
host_ids = input_ids.to(device="cpu", dtype=torch.long)[None]
|
||||
host_image_grid_thw = (
|
||||
image_grid_thw.to(device="cpu", dtype=torch.long)
|
||||
if image_grid_thw is not None
|
||||
else None
|
||||
)
|
||||
host_video_grid_thw = (
|
||||
video_grid_thw.to(device="cpu", dtype=torch.long)
|
||||
if video_grid_thw is not None
|
||||
else None
|
||||
)
|
||||
position_ids = None
|
||||
if host_image_grid_thw is not None or host_video_grid_thw is not None:
|
||||
position_ids, _ = self.model.get_rope_index(
|
||||
host_ids,
|
||||
host_image_grid_thw,
|
||||
host_video_grid_thw,
|
||||
attention_mask=torch.ones_like(host_ids),
|
||||
)
|
||||
ids = host_ids.to(self.device)
|
||||
call_kwargs: dict[str, Any] = {
|
||||
"input_ids": ids,
|
||||
"attention_mask": torch.ones_like(ids),
|
||||
"output_attentions": False,
|
||||
"output_hidden_states": False,
|
||||
"return_dict": True,
|
||||
"use_cache": False,
|
||||
}
|
||||
if position_ids is not None:
|
||||
call_kwargs["position_ids"] = position_ids.to(self.device)
|
||||
if pixel_values is not None:
|
||||
call_kwargs["pixel_values"] = pixel_values.to(self.device, torch.bfloat16)
|
||||
call_kwargs["image_grid_thw"] = host_image_grid_thw
|
||||
if pixel_values_videos is not None:
|
||||
call_kwargs["pixel_values_videos"] = pixel_values_videos.to(
|
||||
self.device, torch.bfloat16
|
||||
)
|
||||
call_kwargs["video_grid_thw"] = host_video_grid_thw
|
||||
|
||||
hidden = self.model(**call_kwargs).last_hidden_state[0].to(torch.bfloat16)
|
||||
expected_shape = [int(ids.shape[1]), self.hidden_dim]
|
||||
if list(hidden.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"unexpected hidden shape {list(hidden.shape)}, "
|
||||
f"expected {expected_shape}"
|
||||
)
|
||||
return hidden
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
params = dict(self.named_parameters(remove_duplicate=False))
|
||||
loaded: set[str] = set()
|
||||
for name, loaded_weight in weights:
|
||||
if not self.should_materialize_checkpoint_weight(name):
|
||||
continue
|
||||
param = params.get(name)
|
||||
if param is None:
|
||||
raise KeyError(
|
||||
f"Unexpected MiniMax H3 Qwen3-VL checkpoint weight: {name}"
|
||||
)
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
try:
|
||||
weight_loader(param, loaded_weight.to(param.dtype))
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"Failed to load MiniMax H3 Qwen3-VL weight "
|
||||
f"{name!r}: checkpoint={tuple(loaded_weight.shape)}, "
|
||||
f"parameter={tuple(param.shape)}"
|
||||
) from exc
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
|
||||
|
||||
EntryClass = MiniMaxH3Qwen3VLEncoder
|
||||
|
||||
__all__ = ["MiniMaxH3Qwen3VLEncoder"]
|
||||
@@ -4,9 +4,10 @@ from transformers import (
|
||||
Cache,
|
||||
DynamicCache,
|
||||
)
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
|
||||
from transformers.utils import TransformersKwargs, is_torchdynamo_compiling
|
||||
from transformers.utils.generic import is_flash_attention_requested
|
||||
from transformers.vision_utils import get_vision_cu_seqlens, get_vision_position_ids
|
||||
|
||||
from sglang.multimodal_gen.configs.models.encoders.qwen3vl import Qwen3VLConfig
|
||||
from sglang.multimodal_gen.runtime.distributed import (
|
||||
@@ -208,7 +209,11 @@ class Qwen3VLTextAttention(nn.Module):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layer_idx = layer_idx
|
||||
self.head_dim = config.hidden_size // config.num_attention_heads
|
||||
self.head_dim = (
|
||||
int(config.head_dim)
|
||||
if getattr(config, "head_dim", None) is not None
|
||||
else config.hidden_size // config.num_attention_heads
|
||||
)
|
||||
self.total_num_heads = config.num_attention_heads
|
||||
self.total_num_key_value_heads = config.num_key_value_heads
|
||||
tp_size = _tp_world_size() if use_tensor_parallel else 1
|
||||
@@ -582,14 +587,6 @@ class Qwen3VLTextModel(nn.Module):
|
||||
else:
|
||||
text_position_ids = position_ids[0]
|
||||
|
||||
attention_mask = create_causal_mask(
|
||||
config=self.config,
|
||||
inputs_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
past_key_values=past_key_values,
|
||||
position_ids=text_position_ids,
|
||||
)
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
# create position embeddings to be shared across the decoder layers
|
||||
@@ -651,7 +648,8 @@ class Qwen3VLTextModel(nn.Module):
|
||||
):
|
||||
visual_pos_masks = visual_pos_masks.to(hidden_states.device)
|
||||
visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype)
|
||||
local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds
|
||||
local_this = hidden_states[visual_pos_masks, :]
|
||||
local_this.add_(visual_embeds)
|
||||
hidden_states[visual_pos_masks, :] = local_this
|
||||
return hidden_states
|
||||
|
||||
@@ -664,10 +662,13 @@ class Qwen3VLModel(nn.Module):
|
||||
config: Qwen3VLConfig
|
||||
_no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
|
||||
|
||||
def __init__(self, config):
|
||||
def __init__(self, config, *, use_tensor_parallel: bool = False):
|
||||
super().__init__()
|
||||
self.visual = Qwen3VLVisionModel._from_config(config.vision_config)
|
||||
self.language_model = Qwen3VLTextModel(config.text_config)
|
||||
self.language_model = Qwen3VLTextModel(
|
||||
config.text_config,
|
||||
use_tensor_parallel=use_tensor_parallel,
|
||||
)
|
||||
self.rope_deltas = None # cache rope_deltas here
|
||||
self.config = config
|
||||
|
||||
@@ -868,6 +869,25 @@ class Qwen3VLModel(nn.Module):
|
||||
# Same implementation as for images
|
||||
return self.get_image_features(pixel_values_videos, video_grid_thw)
|
||||
|
||||
def _get_flat_visual_features(
|
||||
self,
|
||||
pixel_values: torch.FloatTensor,
|
||||
grid_thw: Optional[torch.LongTensor],
|
||||
):
|
||||
pixel_values = pixel_values.type(self.visual.dtype)
|
||||
vision_kwargs = {}
|
||||
if grid_thw is not None and grid_thw.device.type == "cpu":
|
||||
if not is_flash_attention_requested(self.visual.config):
|
||||
vision_kwargs = {
|
||||
"position_ids": get_vision_position_ids(
|
||||
grid_thw, self.visual.spatial_merge_size
|
||||
).to(pixel_values.device),
|
||||
"cu_seqlens": get_vision_cu_seqlens(grid_thw),
|
||||
}
|
||||
grid_thw = grid_thw.to(pixel_values.device)
|
||||
visual_out = self.visual(pixel_values, grid_thw=grid_thw, **vision_kwargs)
|
||||
return visual_out.pooler_output, visual_out.deepstack_features
|
||||
|
||||
def get_image_features(
|
||||
self,
|
||||
pixel_values: torch.FloatTensor,
|
||||
@@ -882,10 +902,9 @@ class Qwen3VLModel(nn.Module):
|
||||
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
|
||||
The temporal, height and width of feature shape of each image in LLM.
|
||||
"""
|
||||
pixel_values = pixel_values.type(self.visual.dtype)
|
||||
visual_out = self.visual(pixel_values, grid_thw=image_grid_thw)
|
||||
image_embeds = visual_out.pooler_output
|
||||
deepstack_image_embeds = visual_out.deepstack_features
|
||||
image_embeds, deepstack_image_embeds = self._get_flat_visual_features(
|
||||
pixel_values, image_grid_thw
|
||||
)
|
||||
split_sizes = (
|
||||
image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2
|
||||
).tolist()
|
||||
@@ -996,35 +1015,40 @@ class Qwen3VLModel(nn.Module):
|
||||
return_dict if return_dict is not None else self.config.use_return_dict
|
||||
)
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds_owned = inputs_embeds is None
|
||||
if inputs_embeds_owned:
|
||||
inputs_embeds = self.get_input_embeddings()(input_ids)
|
||||
|
||||
image_mask = None
|
||||
video_mask = None
|
||||
|
||||
if pixel_values is not None:
|
||||
image_embeds, deepstack_image_embeds = self.get_image_features( # long
|
||||
image_embeds, deepstack_image_embeds = self._get_flat_visual_features(
|
||||
pixel_values, image_grid_thw
|
||||
)
|
||||
image_embeds = torch.cat(image_embeds, dim=0).to(
|
||||
inputs_embeds.device, inputs_embeds.dtype
|
||||
)
|
||||
image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
|
||||
image_mask, _ = self.get_placeholder_mask(
|
||||
input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
|
||||
)
|
||||
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
|
||||
if inputs_embeds_owned:
|
||||
inputs_embeds.masked_scatter_(image_mask, image_embeds)
|
||||
else:
|
||||
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
|
||||
inputs_embeds_owned = True
|
||||
|
||||
if pixel_values_videos is not None:
|
||||
video_embeds, deepstack_video_embeds = self.get_video_features(
|
||||
video_embeds, deepstack_video_embeds = self._get_flat_visual_features(
|
||||
pixel_values_videos, video_grid_thw
|
||||
)
|
||||
video_embeds = torch.cat(video_embeds, dim=0).to(
|
||||
inputs_embeds.device, inputs_embeds.dtype
|
||||
)
|
||||
video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
|
||||
_, video_mask = self.get_placeholder_mask(
|
||||
input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
|
||||
)
|
||||
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
|
||||
if inputs_embeds_owned:
|
||||
inputs_embeds.masked_scatter_(video_mask, video_embeds)
|
||||
else:
|
||||
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
|
||||
inputs_embeds_owned = True
|
||||
|
||||
visual_pos_masks = None
|
||||
deepstack_visual_embeds = None
|
||||
@@ -1040,8 +1064,8 @@ class Qwen3VLModel(nn.Module):
|
||||
deepstack_image_embeds, deepstack_video_embeds
|
||||
):
|
||||
embed_joint = img_embed.new_zeros(
|
||||
visual_pos_masks.sum(), img_embed.shape[-1]
|
||||
).to(img_embed.device)
|
||||
img_embed.shape[0] + vid_embed.shape[0], img_embed.shape[-1]
|
||||
)
|
||||
embed_joint[image_mask_joint, :] = img_embed
|
||||
embed_joint[video_mask_joint, :] = vid_embed
|
||||
deepstack_visual_embeds.append(embed_joint)
|
||||
|
||||
+214
@@ -0,0 +1,214 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _require_finite_tensor(tensor: torch.Tensor, name: str) -> None:
|
||||
if not bool(torch.isfinite(tensor).all().item()):
|
||||
raise ValueError(f"{name} must be finite")
|
||||
|
||||
|
||||
def _validate_unit_timestep(timestep: torch.Tensor, name: str) -> None:
|
||||
if not isinstance(timestep, torch.Tensor):
|
||||
raise ValueError(f"{name} must be a torch.Tensor")
|
||||
if not torch.is_floating_point(timestep):
|
||||
raise ValueError(f"{name} must be a floating point tensor")
|
||||
_require_finite_tensor(timestep, name)
|
||||
out_of_range = (timestep < 0) | (timestep > 1)
|
||||
if bool(out_of_range.any().item()):
|
||||
raise ValueError(f"{name} must be in [0, 1]")
|
||||
|
||||
|
||||
def _validate_sigma(value: float, name: str) -> float:
|
||||
sigma = float(value)
|
||||
if not math.isfinite(sigma):
|
||||
raise ValueError(f"{name} must be finite")
|
||||
if sigma < 0.0:
|
||||
raise ValueError(f"{name} must be non-negative")
|
||||
return sigma
|
||||
|
||||
|
||||
def _validate_timestep_sigma_pair(
|
||||
timestep: torch.Tensor,
|
||||
sigma_curr: float,
|
||||
name: str,
|
||||
) -> float:
|
||||
_validate_unit_timestep(timestep, f"{name}_timestep")
|
||||
sigma = _validate_sigma(sigma_curr, f"{name}_sigma_curr")
|
||||
expected = 1.0 - timestep.detach().to(dtype=torch.float32)
|
||||
actual = torch.full_like(expected, sigma)
|
||||
if not torch.allclose(actual, expected, rtol=1e-5, atol=1e-5):
|
||||
raise ValueError(f"{name}_sigma_curr must equal 1 - {name}_timestep")
|
||||
return sigma
|
||||
|
||||
|
||||
def minimax_h3_rf_v_to_x0(
|
||||
xt: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
if xt.shape != v.shape:
|
||||
raise ValueError(f"xt and v shapes must match, got {xt.shape} vs {v.shape}")
|
||||
if not torch.is_floating_point(xt):
|
||||
raise ValueError("xt must be a floating point tensor")
|
||||
if not torch.is_floating_point(v):
|
||||
raise ValueError("v must be a floating point tensor")
|
||||
_require_finite_tensor(xt, "xt")
|
||||
_require_finite_tensor(v, "v")
|
||||
_validate_unit_timestep(timestep, "timestep")
|
||||
x0 = _minimax_h3_rf_v_to_x0(xt, v, timestep)
|
||||
_require_finite_tensor(x0, "x0")
|
||||
return x0
|
||||
|
||||
|
||||
def _minimax_h3_rf_v_to_x0(
|
||||
xt: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
cond_t = timestep.to(device=xt.device, dtype=xt.dtype)
|
||||
while cond_t.ndim < xt.ndim:
|
||||
cond_t = cond_t.unsqueeze(-1)
|
||||
sigma_t = 1 - cond_t
|
||||
return xt + sigma_t * v
|
||||
|
||||
|
||||
def minimax_h3_euler_eta0_step(
|
||||
state: torch.Tensor,
|
||||
denoised: torch.Tensor,
|
||||
*,
|
||||
sigma_curr: float,
|
||||
sigma_next: float,
|
||||
) -> torch.Tensor:
|
||||
if state.shape != denoised.shape:
|
||||
raise ValueError(
|
||||
f"state and denoised shapes must match, got {state.shape} vs "
|
||||
f"{denoised.shape}"
|
||||
)
|
||||
if not torch.is_floating_point(state):
|
||||
raise ValueError("state must be a floating point tensor")
|
||||
if not torch.is_floating_point(denoised):
|
||||
raise ValueError("denoised must be a floating point tensor")
|
||||
_require_finite_tensor(state, "state")
|
||||
_require_finite_tensor(denoised, "denoised")
|
||||
sigma_curr = _validate_sigma(sigma_curr, "sigma_curr")
|
||||
sigma_next = _validate_sigma(sigma_next, "sigma_next")
|
||||
if sigma_curr == 0.0 and sigma_next != 0.0:
|
||||
raise ValueError("sigma_next must be 0 when sigma_curr is 0")
|
||||
out = _minimax_h3_euler_eta0_step(
|
||||
state,
|
||||
denoised,
|
||||
sigma_curr=sigma_curr,
|
||||
sigma_next=sigma_next,
|
||||
)
|
||||
_require_finite_tensor(out, "euler_eta0_step output")
|
||||
return out
|
||||
|
||||
|
||||
def _minimax_h3_euler_eta0_step(
|
||||
state: torch.Tensor,
|
||||
denoised: torch.Tensor,
|
||||
*,
|
||||
sigma_curr: float,
|
||||
sigma_next: float,
|
||||
sigma_ratio: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if sigma_curr == 0.0:
|
||||
return state
|
||||
compute_dtype = torch.float32
|
||||
if state.dtype not in (torch.float16, torch.bfloat16):
|
||||
compute_dtype = state.dtype
|
||||
if sigma_ratio is None:
|
||||
sigma_curr_t = state.new_tensor(sigma_curr, dtype=compute_dtype)
|
||||
sigma_next_t = state.new_tensor(sigma_next, dtype=compute_dtype)
|
||||
ratio = sigma_next_t / sigma_curr_t
|
||||
else:
|
||||
ratio = sigma_ratio.to(device=state.device, dtype=compute_dtype)
|
||||
out = ratio * state.to(dtype=compute_dtype) + (1.0 - ratio) * denoised.to(
|
||||
dtype=compute_dtype
|
||||
)
|
||||
return out.to(dtype=state.dtype)
|
||||
|
||||
|
||||
class MiniMaxH3EulerAncestralEta0SchedulerAdapter:
|
||||
def __init__(self, **config: Any) -> None:
|
||||
if config:
|
||||
raise ValueError(
|
||||
f"{type(self).__name__} does not accept config fields: "
|
||||
f"{sorted(config)}"
|
||||
)
|
||||
|
||||
def set_shift(self, _flow_shift: float) -> None:
|
||||
"""Ignore flow shift, matching the previous loader-specific path."""
|
||||
|
||||
def step_denoising(
|
||||
self,
|
||||
*,
|
||||
input_visual_latent: torch.Tensor,
|
||||
input_audio_latent: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
noise_pred_visual: torch.Tensor,
|
||||
noise_pred_audio: torch.Tensor,
|
||||
sigma_curr: float,
|
||||
sigma_next: float,
|
||||
video_timestep: torch.Tensor | None = None,
|
||||
audio_timestep: torch.Tensor | None = None,
|
||||
video_sigma_curr: float | None = None,
|
||||
video_sigma_next: float | None = None,
|
||||
audio_sigma_curr: float | None = None,
|
||||
audio_sigma_next: float | None = None,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
visual_timestep = timestep if video_timestep is None else video_timestep
|
||||
audio_timestep = timestep if audio_timestep is None else audio_timestep
|
||||
visual_sigma_curr = sigma_curr if video_sigma_curr is None else video_sigma_curr
|
||||
visual_sigma_next = sigma_next if video_sigma_next is None else video_sigma_next
|
||||
audio_sigma_curr = sigma_curr if audio_sigma_curr is None else audio_sigma_curr
|
||||
audio_sigma_next = sigma_next if audio_sigma_next is None else audio_sigma_next
|
||||
visual_sigma_curr = _validate_timestep_sigma_pair(
|
||||
visual_timestep,
|
||||
visual_sigma_curr,
|
||||
"video",
|
||||
)
|
||||
audio_sigma_curr = _validate_timestep_sigma_pair(
|
||||
audio_timestep,
|
||||
audio_sigma_curr,
|
||||
"audio",
|
||||
)
|
||||
|
||||
denoised_visual = minimax_h3_rf_v_to_x0(
|
||||
input_visual_latent,
|
||||
noise_pred_visual,
|
||||
visual_timestep,
|
||||
)
|
||||
denoised_audio = minimax_h3_rf_v_to_x0(
|
||||
input_audio_latent,
|
||||
noise_pred_audio,
|
||||
audio_timestep,
|
||||
)
|
||||
return {
|
||||
"output_visual_latent": minimax_h3_euler_eta0_step(
|
||||
input_visual_latent,
|
||||
denoised_visual,
|
||||
sigma_curr=visual_sigma_curr,
|
||||
sigma_next=visual_sigma_next,
|
||||
),
|
||||
"output_audio_latent": minimax_h3_euler_eta0_step(
|
||||
input_audio_latent,
|
||||
denoised_audio,
|
||||
sigma_curr=audio_sigma_curr,
|
||||
sigma_next=audio_sigma_next,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
EntryClass = MiniMaxH3EulerAncestralEta0SchedulerAdapter
|
||||
|
||||
__all__ = [
|
||||
"MiniMaxH3EulerAncestralEta0SchedulerAdapter",
|
||||
"minimax_h3_euler_eta0_step",
|
||||
"minimax_h3_rf_v_to_x0",
|
||||
]
|
||||
@@ -0,0 +1,118 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_audio import (
|
||||
MiniMaxH3AudioVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
|
||||
MiniMaxH3VideoVAEConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.vaes.minimax_h3_audio_vae import (
|
||||
DacAudioVAE,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.vaes.minimax_h3_video_vae import (
|
||||
AutoencoderKLLegacy,
|
||||
)
|
||||
|
||||
|
||||
class MiniMaxH3VideoVAE(AutoencoderKLLegacy, LayerwiseOffloadableModuleMixin):
|
||||
layerwise_offload_dit_group_enabled = False
|
||||
# EncoderFCN3D indexes its down containers instead of calling them, so they
|
||||
# cannot host layerwise hooks. Keep the small encoder resident.
|
||||
layer_names = ["decoder.transformer_blocks"]
|
||||
|
||||
def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None:
|
||||
arch = config.arch_config
|
||||
parallel_decode_mode = config.resolved_parallel_decode_mode()
|
||||
use_tiled_decode = config.use_tiling and parallel_decode_mode == "tiled"
|
||||
super().__init__(
|
||||
in_channels=3,
|
||||
out_ch=3,
|
||||
ch=128,
|
||||
embed_dim=24,
|
||||
z_channels=24,
|
||||
use_3d_conv=True,
|
||||
zq_ch_encoder=None,
|
||||
zq_ch_decoder=None,
|
||||
num_res_blocks=2,
|
||||
num_res_blocks_decoder=None,
|
||||
ch_mult=[1, 2, 2, 4, 4, 8],
|
||||
space_down=[2, 2, 2, 2, 1, 1],
|
||||
space_up=[1, 2, 2, 2, 2, 1],
|
||||
time_down=[1, 2, 2, 1, 1, 1],
|
||||
time_up=None,
|
||||
padding_mode="reflect",
|
||||
padding_mode_t=None,
|
||||
use_t_isolated_gn=True,
|
||||
causal_encoder=True,
|
||||
causal_decoder=False,
|
||||
use_vit_decoder=True,
|
||||
vit_decoder_kwargs={
|
||||
"dim_head": 64,
|
||||
"ffn_activation_fn": "silu",
|
||||
"ffn_use_gated": True,
|
||||
"heads": 32,
|
||||
"norm_affine": True,
|
||||
"norm_type": "rms_norm",
|
||||
"num_layers": 36,
|
||||
"qk_norm_affine": False,
|
||||
"qk_norm_type": "rms_norm",
|
||||
"rope_dim_ratio": 0.75,
|
||||
"rope_theta": 100.0,
|
||||
},
|
||||
shift_factor=0.0,
|
||||
scaling_factor=1.0,
|
||||
pixel_norm_type="imagenet",
|
||||
clip_length=arch.vae_clip_length,
|
||||
token_drop=arch.vae_token_drop,
|
||||
encoder_tiling=bool(arch.vae_encoder_tiling),
|
||||
decoder_tiling=use_tiled_decode,
|
||||
parallel_tiling=use_tiled_decode
|
||||
and config.use_parallel_decode
|
||||
and config.use_parallel_tiling
|
||||
and bool(arch.vae_parallel_tiling),
|
||||
tile_size=int(arch.vae_tile_size),
|
||||
tile_overlap_min=int(arch.vae_tile_overlap_min),
|
||||
encoder_parallel=False,
|
||||
decoder_parallel=False,
|
||||
chunk_dim=int(arch.vae_chunk_dim),
|
||||
)
|
||||
self.sglang_config = config
|
||||
self.use_parallel_decode = config.use_parallel_decode
|
||||
self.parallel_decode_mode = parallel_decode_mode
|
||||
|
||||
def prepare_decoder_autocast_weights(self, dtype) -> int:
|
||||
return self.decoder.prepare_autocast_linear_weights(dtype)
|
||||
|
||||
|
||||
class MiniMaxH3AudioVAE(DacAudioVAE, LayerwiseOffloadableModuleMixin):
|
||||
layerwise_offload_dit_group_enabled = False
|
||||
# BigVGAN stores each executable upsampler inside a one-element ModuleList.
|
||||
# The outer ``decoder.ups`` containers are indexed but never called, so hooks
|
||||
# must target the inner lists whose ConvTranspose1d modules run forward.
|
||||
layer_names = [
|
||||
"encoder.block",
|
||||
*(f"decoder.ups.{index}" for index in range(7)),
|
||||
"decoder.resblocks",
|
||||
]
|
||||
|
||||
def __init__(self, config: MiniMaxH3AudioVAEConfig) -> None:
|
||||
super().__init__(
|
||||
encoder_dim=64,
|
||||
encoder_rates=[2, 4, 4, 5, 5],
|
||||
latent_dim=2048,
|
||||
decoder_dim=1024,
|
||||
decoder_rates=[5, 5, 2, 2, 2, 2, 2],
|
||||
sample_rate=32000,
|
||||
vae_latent_channels=32,
|
||||
attn_proj=True,
|
||||
decoder_type="bigvgan",
|
||||
)
|
||||
self.config = config
|
||||
|
||||
|
||||
EntryClass = [MiniMaxH3VideoVAE, MiniMaxH3AudioVAE]
|
||||
|
||||
__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"]
|
||||
@@ -0,0 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .audio_vae import DacAudioVAE
|
||||
|
||||
__all__ = ["DacAudioVAE"]
|
||||
@@ -0,0 +1,177 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
if "sinc" in dir(torch):
|
||||
sinc = torch.sinc
|
||||
else:
|
||||
# This code is adopted from adefossez's julius.core.sinc under the MIT License
|
||||
# https://adefossez.github.io/julius/julius/core.html
|
||||
def sinc(x: torch.Tensor):
|
||||
"""
|
||||
Implementation of sinc, i.e. sin(pi * x) / (pi * x)
|
||||
__Warning__: Different to julius.sinc, the input is multiplied by `pi`!
|
||||
"""
|
||||
return torch.where(
|
||||
x == 0,
|
||||
torch.tensor(1.0, device=x.device, dtype=x.dtype),
|
||||
torch.sin(math.pi * x) / math.pi / x,
|
||||
)
|
||||
|
||||
|
||||
# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
|
||||
# https://adefossez.github.io/julius/julius/lowpass.html
|
||||
def kaiser_sinc_filter1d(
|
||||
cutoff, half_width, kernel_size
|
||||
): # return filter [1,1,kernel_size]
|
||||
even = kernel_size % 2 == 0
|
||||
half_size = kernel_size // 2
|
||||
|
||||
# For kaiser window
|
||||
delta_f = 4 * half_width
|
||||
A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
|
||||
if A > 50.0:
|
||||
beta = 0.1102 * (A - 8.7)
|
||||
elif A >= 21.0:
|
||||
beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
|
||||
else:
|
||||
beta = 0.0
|
||||
window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
|
||||
|
||||
# ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
|
||||
if even:
|
||||
time = torch.arange(-half_size, half_size) + 0.5
|
||||
else:
|
||||
time = torch.arange(kernel_size) - half_size
|
||||
if cutoff == 0:
|
||||
filter_ = torch.zeros_like(time)
|
||||
else:
|
||||
filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
|
||||
"""
|
||||
Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal.
|
||||
"""
|
||||
filter_ /= filter_.sum()
|
||||
filter = filter_.view(1, 1, kernel_size)
|
||||
|
||||
return filter
|
||||
|
||||
|
||||
class LowPassFilter1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
cutoff=0.5,
|
||||
half_width=0.6,
|
||||
stride: int = 1,
|
||||
padding: bool = True,
|
||||
padding_mode: str = "replicate",
|
||||
kernel_size: int = 12,
|
||||
):
|
||||
"""
|
||||
kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible.
|
||||
"""
|
||||
super().__init__()
|
||||
if cutoff < -0.0:
|
||||
raise ValueError("Minimum cutoff must be larger than zero.")
|
||||
if cutoff > 0.5:
|
||||
raise ValueError("A cutoff above 0.5 does not make sense.")
|
||||
self.kernel_size = kernel_size
|
||||
self.even = kernel_size % 2 == 0
|
||||
self.pad_left = kernel_size // 2 - int(self.even)
|
||||
self.pad_right = kernel_size // 2
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.padding_mode = padding_mode
|
||||
filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
|
||||
self.register_buffer("filter", filter)
|
||||
|
||||
# Input [B, C, T]
|
||||
def forward(self, x):
|
||||
_, C, _ = x.shape
|
||||
|
||||
if self.padding:
|
||||
x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
|
||||
out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class UpSample1d(nn.Module):
|
||||
def __init__(self, ratio=2, kernel_size=None):
|
||||
super().__init__()
|
||||
self.ratio = ratio
|
||||
self.kernel_size = (
|
||||
int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
|
||||
)
|
||||
self.stride = ratio
|
||||
self.pad = self.kernel_size // ratio - 1
|
||||
self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
|
||||
self.pad_right = (
|
||||
self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
|
||||
)
|
||||
filter = kaiser_sinc_filter1d(
|
||||
cutoff=0.5 / ratio,
|
||||
half_width=0.6 / ratio,
|
||||
kernel_size=self.kernel_size,
|
||||
)
|
||||
self.register_buffer("filter", filter)
|
||||
|
||||
def forward(self, x):
|
||||
_, C, _ = x.shape
|
||||
|
||||
x = F.pad(x, (self.pad, self.pad), mode="replicate")
|
||||
x = F.conv_transpose1d(
|
||||
x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
|
||||
)
|
||||
x.mul_(self.ratio)
|
||||
x = x[..., self.pad_left : -self.pad_right]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class DownSample1d(nn.Module):
|
||||
def __init__(self, ratio=2, kernel_size=None):
|
||||
super().__init__()
|
||||
self.ratio = ratio
|
||||
self.kernel_size = (
|
||||
int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
|
||||
)
|
||||
self.lowpass = LowPassFilter1d(
|
||||
cutoff=0.5 / ratio,
|
||||
half_width=0.6 / ratio,
|
||||
stride=ratio,
|
||||
kernel_size=self.kernel_size,
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
xx = self.lowpass(x)
|
||||
|
||||
return xx
|
||||
|
||||
|
||||
class Activation1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
activation,
|
||||
up_ratio: int = 2,
|
||||
down_ratio: int = 2,
|
||||
up_kernel_size: int = 12,
|
||||
down_kernel_size: int = 12,
|
||||
):
|
||||
super().__init__()
|
||||
self.up_ratio = up_ratio
|
||||
self.down_ratio = down_ratio
|
||||
self.act = activation
|
||||
self.upsample = UpSample1d(up_ratio, up_kernel_size)
|
||||
self.downsample = DownSample1d(down_ratio, down_kernel_size)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.upsample(x)
|
||||
x = self.act(x)
|
||||
x = self.downsample(x)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,307 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# DAC-lineage audio VAE: waveform encoder + BigVGAN decoder (inference-only bundle).
|
||||
import math
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
from torch.nn.utils.parametrizations import weight_norm
|
||||
|
||||
from .bigvgan import AttrDict, BigVGAN
|
||||
|
||||
|
||||
class GeGluMlp(nn.Module):
|
||||
def __init__(self, in_features, hidden_features):
|
||||
super().__init__()
|
||||
self.norm = nn.LayerNorm(in_features)
|
||||
self.act = nn.GELU(approximate="tanh")
|
||||
self.w0 = nn.Linear(in_features, hidden_features)
|
||||
self.w1 = nn.Linear(in_features, hidden_features)
|
||||
self.w2 = nn.Linear(hidden_features, in_features)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm(x)
|
||||
x = self.act(self.w0(x)).mul_(self.w1(x))
|
||||
x = self.w2(x)
|
||||
return x
|
||||
|
||||
|
||||
class CausalAttention(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, num_heads):
|
||||
super().__init__()
|
||||
if in_dim > out_dim:
|
||||
# assert in_dim // num_heads == out_dim
|
||||
self.head_dim = in_dim // num_heads
|
||||
self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False)
|
||||
self.q_bias = nn.Parameter(torch.zeros(in_dim))
|
||||
self.v_bias = nn.Parameter(torch.zeros(in_dim))
|
||||
self.register_buffer("zero_k_bias", torch.zeros(in_dim))
|
||||
else:
|
||||
# assert out_dim // num_heads == in_dim
|
||||
self.head_dim = out_dim // num_heads
|
||||
self.qkv = nn.Linear(in_dim, out_dim * 3, bias=False)
|
||||
self.q_bias = nn.Parameter(torch.zeros(out_dim))
|
||||
self.v_bias = nn.Parameter(torch.zeros(out_dim))
|
||||
self.register_buffer("zero_k_bias", torch.zeros(out_dim))
|
||||
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.scale = self.head_dim**-0.5
|
||||
self.proj = nn.Linear(out_dim, out_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
B, N, C = x.shape
|
||||
qkv = F.linear(
|
||||
input=x,
|
||||
weight=self.qkv.weight,
|
||||
bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)),
|
||||
)
|
||||
q, k, v = (
|
||||
qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
|
||||
.permute(2, 0, 3, 1, 4)
|
||||
.unbind(0)
|
||||
)
|
||||
|
||||
x = scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True
|
||||
)
|
||||
|
||||
if self.in_dim > self.out_dim:
|
||||
x = torch.mean(x, dim=1)
|
||||
if self.in_dim // self.num_heads != self.out_dim:
|
||||
x = nn.functional.adaptive_avg_pool1d(x, self.out_dim)
|
||||
else:
|
||||
x = x.transpose(1, 2).reshape(B, N, -1)
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class AttnProjection(nn.Module):
|
||||
def __init__(
|
||||
self, in_dim, out_dim, num_heads, norm_layer=nn.LayerNorm, mlp_ratio=2
|
||||
):
|
||||
super().__init__()
|
||||
assert out_dim % in_dim == 0 or in_dim % out_dim == 0
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
self.norm1 = norm_layer(in_dim)
|
||||
self.attn = CausalAttention(in_dim, out_dim, num_heads)
|
||||
self.proj = nn.Linear(in_dim, out_dim)
|
||||
self.norm3 = norm_layer(in_dim)
|
||||
|
||||
self.norm2 = norm_layer(out_dim)
|
||||
hidden_dim = int(out_dim * mlp_ratio)
|
||||
self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim)
|
||||
# self.mlp = FeedForward(out_dim, out_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x)))
|
||||
return self.mlp(self.norm2(x)).add_(x)
|
||||
|
||||
|
||||
def WNConv1d(*args, **kwargs):
|
||||
return weight_norm(nn.Conv1d(*args, **kwargs))
|
||||
|
||||
|
||||
@torch.jit.script
|
||||
def snake(x, alpha):
|
||||
shape = x.shape
|
||||
x = x.reshape(shape[0], shape[1], -1)
|
||||
x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
|
||||
x = x.reshape(shape)
|
||||
return x
|
||||
|
||||
|
||||
class Snake1d(nn.Module):
|
||||
def __init__(self, channels):
|
||||
super().__init__()
|
||||
self.alpha = nn.Parameter(torch.ones(1, channels, 1))
|
||||
|
||||
def forward(self, x):
|
||||
return snake(x, self.alpha)
|
||||
|
||||
|
||||
class ResidualUnit(nn.Module):
|
||||
def __init__(self, dim: int = 16, dilation: int = 1):
|
||||
super().__init__()
|
||||
pad = ((7 - 1) * dilation) // 2
|
||||
self.block = nn.Sequential(
|
||||
Snake1d(dim),
|
||||
WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
|
||||
Snake1d(dim),
|
||||
WNConv1d(dim, dim, kernel_size=1),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
y = self.block(x)
|
||||
pad = (x.shape[-1] - y.shape[-1]) // 2
|
||||
if pad > 0:
|
||||
x = x[..., pad:-pad]
|
||||
return x + y
|
||||
|
||||
|
||||
class EncoderBlock(nn.Module):
|
||||
def __init__(self, dim: int = 16, stride: int = 1):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
ResidualUnit(dim // 2, dilation=1),
|
||||
ResidualUnit(dim // 2, dilation=3),
|
||||
ResidualUnit(dim // 2, dilation=9),
|
||||
Snake1d(dim // 2),
|
||||
WNConv1d(
|
||||
dim // 2,
|
||||
dim,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
padding=math.ceil(stride / 2),
|
||||
),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int = 64,
|
||||
strides: list = [2, 4, 8, 8],
|
||||
d_latent: int = 64,
|
||||
):
|
||||
super().__init__()
|
||||
# Create first convolution
|
||||
self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
|
||||
|
||||
# Create EncoderBlocks that double channels as they downsample by `stride`
|
||||
for stride in strides:
|
||||
d_model *= 2
|
||||
self.block += [EncoderBlock(d_model, stride=stride)]
|
||||
|
||||
# Create last convolution
|
||||
self.block += [
|
||||
Snake1d(d_model),
|
||||
WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
|
||||
]
|
||||
|
||||
# Wrap black into nn.Sequential
|
||||
self.block = nn.Sequential(*self.block)
|
||||
self.enc_dim = d_model
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class DacAudioVAE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
encoder_dim: int = 64,
|
||||
encoder_rates: List[int] = [2, 4, 8, 8],
|
||||
latent_dim: int = None,
|
||||
decoder_dim: int = 1536,
|
||||
decoder_rates: List[int] = [8, 8, 4, 2],
|
||||
sample_rate: int = 44100,
|
||||
vae_latent_channels: int = 64,
|
||||
attn_proj: bool = False,
|
||||
decoder_type: str = "bigvgan",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.encoder_dim = encoder_dim
|
||||
self.encoder_rates = encoder_rates
|
||||
self.decoder_dim = decoder_dim
|
||||
self.decoder_rates = decoder_rates
|
||||
self.sample_rate = sample_rate
|
||||
self.attn_proj = attn_proj
|
||||
self.decoder_type = decoder_type
|
||||
|
||||
if latent_dim is None:
|
||||
latent_dim = encoder_dim * (2 ** len(encoder_rates))
|
||||
|
||||
self.latent_dim = latent_dim
|
||||
|
||||
self.hop_length = np.prod(encoder_rates)
|
||||
self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
|
||||
|
||||
if latent_dim % vae_latent_channels == 0:
|
||||
self.attn_proj_dim = vae_latent_channels
|
||||
else:
|
||||
# smallest power of two >= vae_latent_channels
|
||||
self.attn_proj_dim = 2 ** int(np.ceil(np.log2(vae_latent_channels)))
|
||||
|
||||
self.mean_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1)
|
||||
self.logs_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1)
|
||||
|
||||
self.dec_in_proj = nn.Conv1d(vae_latent_channels, latent_dim, 1)
|
||||
|
||||
if self.decoder_type == "bigvgan":
|
||||
if sample_rate == 16000:
|
||||
bigvgan_conf = {
|
||||
"resblock": "1",
|
||||
"num_mels": latent_dim,
|
||||
"upsample_rates": [5, 5, 2, 2, 2, 2],
|
||||
"upsample_kernel_sizes": [9, 9, 4, 4, 4, 4],
|
||||
"upsample_initial_channel": decoder_dim,
|
||||
"resblock_kernel_sizes": [3, 7, 11],
|
||||
"resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
|
||||
"use_tanh_at_final": False,
|
||||
"use_bias_at_final": False,
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": True,
|
||||
}
|
||||
elif sample_rate == 32000:
|
||||
bigvgan_conf = {
|
||||
"resblock": "1",
|
||||
"num_mels": latent_dim,
|
||||
"upsample_rates": [5, 5, 2, 2, 2, 2, 2],
|
||||
"upsample_kernel_sizes": [9, 9, 4, 4, 4, 4, 4],
|
||||
"upsample_initial_channel": decoder_dim,
|
||||
"resblock_kernel_sizes": [3, 7, 11],
|
||||
"resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
|
||||
"use_tanh_at_final": False,
|
||||
"use_bias_at_final": False,
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": True,
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Invalid sample_rate: {sample_rate}")
|
||||
|
||||
h = AttrDict(**bigvgan_conf)
|
||||
self.decoder = BigVGAN(h)
|
||||
else:
|
||||
raise ValueError(f"Invalid decoder type: {self.decoder_type}")
|
||||
|
||||
if self.attn_proj:
|
||||
self.pre_block = AttnProjection(latent_dim, self.attn_proj_dim, num_heads=8)
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
|
||||
def preprocess(self, audio_data, sample_rate):
|
||||
if sample_rate is None:
|
||||
sample_rate = self.sample_rate
|
||||
|
||||
length = audio_data.shape[-1]
|
||||
right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
|
||||
if right_pad:
|
||||
audio_data = nn.functional.pad(audio_data, (0, right_pad))
|
||||
|
||||
return audio_data
|
||||
|
||||
def decode(self, z: torch.Tensor):
|
||||
"""Decode given latent codes and return audio data
|
||||
|
||||
Parameters
|
||||
----------
|
||||
z : Tensor[B x D x T]
|
||||
Continuous latent representation
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor[B x 1 x length]
|
||||
Decoded audio data.
|
||||
"""
|
||||
z = self.dec_in_proj(z)
|
||||
return self.decoder(z)
|
||||
@@ -0,0 +1,255 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import Conv1d, ConvTranspose1d, Parameter
|
||||
from torch.nn.utils.parametrizations import weight_norm
|
||||
|
||||
from .alias_free import Activation1d
|
||||
|
||||
|
||||
def get_padding(kernel_size, dilation=1):
|
||||
return int((kernel_size * dilation - dilation) / 2)
|
||||
|
||||
|
||||
# Adapted from https://github.com/EdwardDixon/snake under the MIT license.
|
||||
@torch.jit.script
|
||||
def snakebeta(x, alpha, beta):
|
||||
shape = x.shape
|
||||
x = x.reshape(shape[0], shape[1], -1)
|
||||
x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
|
||||
x = x.reshape(shape)
|
||||
return x
|
||||
|
||||
|
||||
class SnakeBeta(nn.Module):
|
||||
def __init__(
|
||||
self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
|
||||
):
|
||||
super(SnakeBeta, self).__init__()
|
||||
self.in_features = in_features
|
||||
|
||||
self.alpha_logscale = alpha_logscale
|
||||
if self.alpha_logscale:
|
||||
self.alpha = Parameter(torch.zeros(in_features) * alpha)
|
||||
self.beta = Parameter(torch.zeros(in_features) * alpha)
|
||||
else:
|
||||
self.alpha = Parameter(torch.ones(in_features) * alpha)
|
||||
self.beta = Parameter(torch.ones(in_features) * alpha)
|
||||
|
||||
self.alpha.requires_grad = alpha_trainable
|
||||
self.beta.requires_grad = alpha_trainable
|
||||
self.no_div_by_zero = 0.000000001
|
||||
|
||||
def forward(self, x):
|
||||
alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
|
||||
beta = self.beta.unsqueeze(0).unsqueeze(-1)
|
||||
if self.alpha_logscale:
|
||||
alpha = torch.exp(alpha)
|
||||
beta = torch.exp(beta)
|
||||
x = snakebeta(x, alpha, beta)
|
||||
return x
|
||||
|
||||
|
||||
class AttrDict(dict):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(AttrDict, self).__init__(*args, **kwargs)
|
||||
self.__dict__ = self
|
||||
|
||||
|
||||
class AMPBlock1(torch.nn.Module):
|
||||
"""
|
||||
AMPBlock applies SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
|
||||
AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
|
||||
|
||||
Args:
|
||||
h (AttrDict): Hyperparameters.
|
||||
channels (int): Number of convolution channels.
|
||||
kernel_size (int): Size of the convolution kernel. Default is 3.
|
||||
dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
|
||||
activation (str): Activation function type. Must be 'snakebeta'.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
h: AttrDict,
|
||||
channels: int,
|
||||
kernel_size: int = 3,
|
||||
dilation: tuple = (1, 3, 5),
|
||||
activation: str = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.h = h
|
||||
|
||||
self.convs1 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
dilation=d,
|
||||
padding=get_padding(kernel_size, d),
|
||||
)
|
||||
)
|
||||
for d in dilation
|
||||
]
|
||||
)
|
||||
|
||||
self.convs2 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
)
|
||||
for _ in range(len(dilation))
|
||||
]
|
||||
)
|
||||
|
||||
self.num_layers = len(self.convs1) + len(
|
||||
self.convs2
|
||||
) # Total number of conv layers
|
||||
|
||||
if activation == "snakebeta":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(
|
||||
activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)
|
||||
)
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"activation incorrectly specified. check the config file and look for 'activation'."
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
activation_iter = iter(self.activations)
|
||||
for c1, c2 in zip(self.convs1, self.convs2):
|
||||
a1 = next(activation_iter)
|
||||
a2 = next(activation_iter)
|
||||
xt = a1(x)
|
||||
xt = c1(xt)
|
||||
xt = a2(xt)
|
||||
xt = c2(xt)
|
||||
x = xt.add_(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class BigVGAN(torch.nn.Module):
|
||||
"""
|
||||
BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
|
||||
|
||||
Args:
|
||||
h (AttrDict): Hyperparameters.
|
||||
"""
|
||||
|
||||
def __init__(self, h: AttrDict):
|
||||
super().__init__()
|
||||
self.h = h
|
||||
|
||||
self.num_kernels = len(h.resblock_kernel_sizes)
|
||||
self.num_upsamples = len(h.upsample_rates)
|
||||
|
||||
# Pre-conv
|
||||
self.conv_pre = weight_norm(
|
||||
Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
|
||||
)
|
||||
|
||||
# Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
|
||||
if h.resblock == "1":
|
||||
resblock_class = AMPBlock1
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}"
|
||||
)
|
||||
|
||||
# Transposed conv-based upsamplers. does not apply anti-aliasing
|
||||
self.ups = nn.ModuleList()
|
||||
for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
|
||||
self.ups.append(
|
||||
nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
ConvTranspose1d(
|
||||
h.upsample_initial_channel // (2**i),
|
||||
h.upsample_initial_channel // (2 ** (i + 1)),
|
||||
k,
|
||||
u,
|
||||
padding=(k - u) // 2,
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
|
||||
self.resblocks = nn.ModuleList()
|
||||
for i in range(len(self.ups)):
|
||||
ch = h.upsample_initial_channel // (2 ** (i + 1))
|
||||
for j, (k, d) in enumerate(
|
||||
zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
|
||||
):
|
||||
self.resblocks.append(
|
||||
resblock_class(h, ch, k, d, activation=h.activation)
|
||||
)
|
||||
|
||||
# Post-conv
|
||||
if h.activation != "snakebeta":
|
||||
raise NotImplementedError(
|
||||
"activation incorrectly specified. check the config file and look for 'activation'."
|
||||
)
|
||||
activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale)
|
||||
|
||||
self.activation_post = Activation1d(activation=activation_post)
|
||||
|
||||
# Whether to use bias for the final conv_post. Default to True for backward compatibility
|
||||
self.use_bias_at_final = h.get("use_bias_at_final", True)
|
||||
self.conv_post = weight_norm(
|
||||
Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
|
||||
)
|
||||
|
||||
# Final tanh activation. Defaults to True for backward compatibility
|
||||
self.use_tanh_at_final = h.get("use_tanh_at_final", True)
|
||||
|
||||
def forward(self, x):
|
||||
# Pre-conv
|
||||
x = self.conv_pre(x)
|
||||
|
||||
for i in range(self.num_upsamples):
|
||||
# Upsampling
|
||||
for i_up in range(len(self.ups[i])):
|
||||
x = self.ups[i][i_up](x)
|
||||
# AMP blocks
|
||||
xs = None
|
||||
for j in range(self.num_kernels):
|
||||
if xs is None:
|
||||
xs = self.resblocks[i * self.num_kernels + j](x)
|
||||
else:
|
||||
xs += self.resblocks[i * self.num_kernels + j](x)
|
||||
x = xs.div_(self.num_kernels)
|
||||
|
||||
# Post-conv
|
||||
x = self.activation_post(x)
|
||||
x = self.conv_post(x)
|
||||
# Final tanh activation
|
||||
if self.use_tanh_at_final:
|
||||
x.tanh_()
|
||||
else:
|
||||
x.clamp_(min=-1.0, max=1.0) # Bound the output to [-1, 1]
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .klvae import AutoencoderKLLegacy
|
||||
|
||||
__all__ = ["AutoencoderKLLegacy"]
|
||||
@@ -0,0 +1,174 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Attention module for the MiniMax H3 visual VAE (inference-only bundle).
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
from diffusers.utils import logging
|
||||
|
||||
from .flash import flash_attn
|
||||
from .vit_utils import _env_flag, apply_rotary_pos_emb_qk
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def _vit_norm_input(module, hidden_states):
|
||||
if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1"):
|
||||
return hidden_states.float()
|
||||
weight = getattr(module, "weight", None)
|
||||
return hidden_states.to(getattr(weight, "dtype", hidden_states.dtype))
|
||||
|
||||
|
||||
def _apply_qk_norm(module, hidden_states):
|
||||
if (
|
||||
_env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1")
|
||||
and isinstance(module, (nn.LayerNorm, nn.RMSNorm))
|
||||
and getattr(module, "weight", None) is None
|
||||
and getattr(module, "bias", None) is None
|
||||
and hidden_states.is_cuda
|
||||
and hidden_states.dtype in (torch.float16, torch.bfloat16)
|
||||
and not torch.is_grad_enabled()
|
||||
and not torch.compiler.is_compiling()
|
||||
):
|
||||
# CUDA LayerNorm/RMSNorm accumulates half/bfloat16 inputs in FP32.
|
||||
# With no affine parameters its half output is bit-identical to the
|
||||
# released FP32-norm-then-cast recipe, without two full-tensor casts.
|
||||
with torch.autocast("cuda", enabled=False):
|
||||
return module(hidden_states)
|
||||
return module(_vit_norm_input(module, hidden_states)).to(hidden_states.dtype)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
heads,
|
||||
dim_head,
|
||||
embed_dim: Optional[int] = None,
|
||||
qk_norm_type: Optional[str] = None,
|
||||
qk_norm_affine: bool = False,
|
||||
bias: bool = True,
|
||||
out_bias: Optional[bool] = None,
|
||||
eps: float = 1e-5,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim_head = dim_head
|
||||
self.heads = heads
|
||||
self.attn_inner_dim = dim_head * heads
|
||||
self.embed_dim = embed_dim if embed_dim is not None else self.attn_inner_dim
|
||||
|
||||
out_bias = out_bias if out_bias is not None else bias
|
||||
|
||||
if qk_norm_type is None:
|
||||
self.norm_q = None
|
||||
self.norm_k = None
|
||||
elif qk_norm_type == "layer_norm":
|
||||
self.norm_q = nn.LayerNorm(
|
||||
dim_head, eps=eps, elementwise_affine=qk_norm_affine
|
||||
)
|
||||
self.norm_k = nn.LayerNorm(
|
||||
dim_head, eps=eps, elementwise_affine=qk_norm_affine
|
||||
)
|
||||
elif qk_norm_type == "rms_norm":
|
||||
self.norm_q = nn.RMSNorm(
|
||||
dim_head, eps=eps, elementwise_affine=qk_norm_affine
|
||||
)
|
||||
self.norm_k = nn.RMSNorm(
|
||||
dim_head, eps=eps, elementwise_affine=qk_norm_affine
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"unknown qk_norm_type: {qk_norm_type}. Should be None,'layer_norm','rms_norm'"
|
||||
)
|
||||
|
||||
self.to_qkv = nn.Linear(self.embed_dim, self.attn_inner_dim * 3, bias=bias)
|
||||
|
||||
self.to_out = nn.Linear(self.attn_inner_dim, self.embed_dim, bias=out_bias)
|
||||
|
||||
if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0):
|
||||
logger.warning(f"Unused kwargs: {kwargs}")
|
||||
|
||||
def _perform_attention(self, query, key, value, pack_info):
|
||||
cu_seqlens = pack_info.get("cu_seqlens", None)
|
||||
mask_mod = pack_info.get("mask_mod", None)
|
||||
block_sparse = pack_info.get("block_sparse", None)
|
||||
valid_seq_len = pack_info.get("valid_seq_len", None)
|
||||
|
||||
if cu_seqlens is not None:
|
||||
raise NotImplementedError(
|
||||
"varlen attention is not supported in this inference-only bundle"
|
||||
)
|
||||
|
||||
padded_seq_len = query.shape[1]
|
||||
if valid_seq_len is not None:
|
||||
valid_seq_len = int(valid_seq_len)
|
||||
if not 0 < valid_seq_len <= padded_seq_len:
|
||||
raise ValueError(
|
||||
"valid_seq_len must be in (0, padded_seq_len], got "
|
||||
f"{valid_seq_len} for padded_seq_len={padded_seq_len}"
|
||||
)
|
||||
query = query[:, :valid_seq_len]
|
||||
key = key[:, :valid_seq_len]
|
||||
value = value[:, :valid_seq_len]
|
||||
|
||||
if mask_mod is not None:
|
||||
hidden_states = flash_attn(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
mask_mod=mask_mod,
|
||||
block_sparse=block_sparse,
|
||||
)
|
||||
else:
|
||||
hidden_states = flash_attn(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
)
|
||||
|
||||
if valid_seq_len is not None and valid_seq_len < padded_seq_len:
|
||||
hidden_states = torch.cat(
|
||||
[
|
||||
hidden_states,
|
||||
hidden_states.new_zeros(
|
||||
hidden_states.shape[0],
|
||||
padded_seq_len - valid_seq_len,
|
||||
hidden_states.shape[2],
|
||||
hidden_states.shape[3],
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
return hidden_states
|
||||
|
||||
def perform_attention(self, query, key, value, pack_info={}):
|
||||
return self._perform_attention(query, key, value, pack_info)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
rotary_pos_emb: Optional[torch.Tensor] = None,
|
||||
pack_info: dict = {},
|
||||
) -> torch.Tensor:
|
||||
batch_size, seq_len, _ = hidden_states.shape
|
||||
|
||||
qkv = self.to_qkv(hidden_states)
|
||||
qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head)
|
||||
query, key, value = torch.chunk(qkv, 3, dim=-1)
|
||||
|
||||
if self.norm_q is not None:
|
||||
query = _apply_qk_norm(self.norm_q, query)
|
||||
if self.norm_k is not None:
|
||||
key = _apply_qk_norm(self.norm_k, key)
|
||||
|
||||
if rotary_pos_emb is not None:
|
||||
query, key = apply_rotary_pos_emb_qk(query, key, rotary_pos_emb)
|
||||
|
||||
hidden_states = self.perform_attention(query, key, value, pack_info)
|
||||
|
||||
hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
|
||||
hidden_states = self.to_out(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,281 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Transformer building blocks for the MiniMax H3 visual VAE ViT decoder.
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.utils import logging
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
|
||||
from sglang.kernels.ops.activation.activation import (
|
||||
silu_and_mul_with_activation_rounding,
|
||||
)
|
||||
from sglang.kernels.ops.diffusion.triton.scale_shift import (
|
||||
try_fused_scaled_residual_add_exact,
|
||||
)
|
||||
|
||||
from .attention import Attention
|
||||
from .vit_utils import _env_flag, _vit_torch_compile_kwargs
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def _vit_norm_input(module, hidden_states):
|
||||
if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1"):
|
||||
return hidden_states.float()
|
||||
return hidden_states.to(getattr(module.weight, "dtype", hidden_states.dtype))
|
||||
|
||||
|
||||
def _scaled_residual_add(residual, x, scale):
|
||||
fused = try_fused_scaled_residual_add_exact(residual, x, scale)
|
||||
return residual + x * scale if fused is None else fused
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
dim_out: Optional[int] = None,
|
||||
mult: int = 4,
|
||||
activation_fn: str = "silu",
|
||||
bias: bool = True,
|
||||
use_gated: bool = True,
|
||||
glu_balanced: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
ratio = 2 / 3 if (use_gated and glu_balanced) else 1
|
||||
inner_dim = round(dim * mult * ratio)
|
||||
dim_out = dim_out if dim_out is not None else dim
|
||||
self.use_gated = use_gated
|
||||
|
||||
if use_gated:
|
||||
self.w1 = nn.Linear(dim, inner_dim * 2, bias=bias)
|
||||
else:
|
||||
self.w1 = nn.Linear(dim, inner_dim, bias=bias)
|
||||
|
||||
if activation_fn == "silu":
|
||||
self.act_fn = nn.SiLU()
|
||||
elif activation_fn == "gelu":
|
||||
self.act_fn = nn.GELU()
|
||||
elif activation_fn == "gelu-approximate":
|
||||
self.act_fn = nn.GELU(approximate="tanh")
|
||||
else:
|
||||
raise ValueError(f"Unsupported activation function: {activation_fn}")
|
||||
|
||||
self.w2 = nn.Linear(inner_dim, dim_out, bias=bias)
|
||||
self._compile_forward_enabled = _env_flag(
|
||||
"MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE", "0"
|
||||
)
|
||||
self._compile_forward_fatal = _env_flag(
|
||||
"MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE_FATAL", "0"
|
||||
)
|
||||
self._compiled_forward = None
|
||||
|
||||
def _forward_impl(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.w1(hidden_states)
|
||||
|
||||
if self.use_gated:
|
||||
if (
|
||||
isinstance(self.act_fn, nn.SiLU)
|
||||
and hidden_states.is_cuda
|
||||
and hidden_states.dtype in (torch.float16, torch.bfloat16)
|
||||
and hidden_states.is_contiguous()
|
||||
and hidden_states.shape[-1] % 32 == 0
|
||||
):
|
||||
hidden_states = silu_and_mul_with_activation_rounding(hidden_states)
|
||||
else:
|
||||
gate, hidden_states = hidden_states.chunk(2, dim=-1)
|
||||
hidden_states = self.act_fn(gate).mul_(hidden_states)
|
||||
else:
|
||||
hidden_states = self.act_fn(hidden_states)
|
||||
|
||||
hidden_states = self.w2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
def _get_forward_impl(self):
|
||||
if not self._compile_forward_enabled:
|
||||
return self._forward_impl
|
||||
if self._compiled_forward is not None:
|
||||
return self._compiled_forward
|
||||
if not hasattr(torch, "compile"):
|
||||
message = (
|
||||
"torch.compile is unavailable; falling back to eager ViT FeedForward"
|
||||
)
|
||||
if self._compile_forward_fatal:
|
||||
raise RuntimeError(message)
|
||||
logger.warning(f"[ViTFeedForward] {message}")
|
||||
self._compile_forward_enabled = False
|
||||
return self._forward_impl
|
||||
|
||||
kwargs = _vit_torch_compile_kwargs(
|
||||
"MINIMAX_H3_VAE_DECODER_VIT_FF_TORCH_COMPILE"
|
||||
)
|
||||
try:
|
||||
self._compiled_forward = torch.compile(self._forward_impl, **kwargs)
|
||||
logger.info(f"[ViTFeedForward] torch.compile enabled kwargs={kwargs}")
|
||||
except Exception as exc:
|
||||
if self._compile_forward_fatal:
|
||||
raise
|
||||
logger.warning(
|
||||
f"[ViTFeedForward] torch.compile setup failed: {type(exc).__name__}: {exc}; "
|
||||
"falling back to eager"
|
||||
)
|
||||
self._compile_forward_enabled = False
|
||||
self._compiled_forward = None
|
||||
return self._forward_impl
|
||||
return self._compiled_forward
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
forward_impl = self._get_forward_impl()
|
||||
try:
|
||||
return forward_impl(hidden_states)
|
||||
except Exception as exc:
|
||||
if (
|
||||
self._compile_forward_enabled
|
||||
and self._compiled_forward is not None
|
||||
and forward_impl is self._compiled_forward
|
||||
and not self._compile_forward_fatal
|
||||
):
|
||||
logger.warning(
|
||||
f"[ViTFeedForward] compiled forward failed: {type(exc).__name__}: {exc}; "
|
||||
"disabling compile and retrying eager"
|
||||
)
|
||||
self._compile_forward_enabled = False
|
||||
self._compiled_forward = None
|
||||
return self._forward_impl(hidden_states)
|
||||
raise
|
||||
|
||||
|
||||
class RotaryEmbeddingND(nn.Module):
|
||||
def __init__(self, dim, rotary_base=10000, n_dim=3, use_angle=False):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.n_dim = n_dim
|
||||
|
||||
if dim % (2 * n_dim) != 0:
|
||||
raise ValueError(
|
||||
f"head_dim {dim} must be divisible by 2 * n_dim {2 * n_dim}"
|
||||
)
|
||||
|
||||
if use_angle:
|
||||
self.angle_scale = 2.0 * math.pi
|
||||
else:
|
||||
self.angle_scale = 1.0
|
||||
|
||||
inv_freq = 1 / rotary_base ** torch.arange(
|
||||
0, 1, 2 * n_dim / dim, dtype=torch.float32
|
||||
)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def forward(self, img_ids):
|
||||
B, N, D = img_ids.shape
|
||||
if D != self.n_dim:
|
||||
raise ValueError(f"Expected {self.n_dim} dimensions, got {D}")
|
||||
|
||||
with torch.autocast("cuda", enabled=False):
|
||||
angles = (
|
||||
self.angle_scale
|
||||
* img_ids[:, :, :, None]
|
||||
* self.inv_freq.to(img_ids.device)[None, None, None, :]
|
||||
)
|
||||
angles = angles.flatten(2, 3)
|
||||
angles = angles.tile(2)
|
||||
angles = angles.unsqueeze(2)
|
||||
|
||||
cos = torch.cos(angles)
|
||||
sin = torch.sin(angles)
|
||||
|
||||
return cos.to(dtype=img_ids.dtype), sin.to(dtype=img_ids.dtype)
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
heads: int,
|
||||
dim_head: int,
|
||||
embed_dim: Optional[int] = None,
|
||||
ffn_glu_balanced: bool = False,
|
||||
norm_type: str = "layer_norm",
|
||||
norm_affine: bool = True,
|
||||
qk_norm_type: str = "rms_norm",
|
||||
qk_norm_affine: bool = False,
|
||||
ffn_activation_fn: str = "silu",
|
||||
ffn_use_gated: bool = True,
|
||||
use_scale: bool = True,
|
||||
bias: bool = True,
|
||||
eps: float = 1e-5,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
dim = embed_dim if embed_dim is not None else dim_head * heads
|
||||
self.use_scale = use_scale
|
||||
|
||||
if norm_type == "layer_norm":
|
||||
norm_class = nn.LayerNorm
|
||||
elif norm_type == "rms_norm":
|
||||
norm_class = nn.RMSNorm
|
||||
else:
|
||||
raise ValueError(f"unknown norm_type {norm_type}")
|
||||
|
||||
self.norm1 = norm_class(
|
||||
dim,
|
||||
elementwise_affine=norm_affine,
|
||||
eps=eps,
|
||||
)
|
||||
self.attn = Attention(
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
embed_dim=dim,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qk_norm_affine=qk_norm_affine,
|
||||
bias=bias,
|
||||
eps=eps,
|
||||
**kwargs,
|
||||
)
|
||||
if use_scale:
|
||||
self.scale1 = nn.Parameter(torch.zeros(dim))
|
||||
|
||||
self.norm2 = norm_class(
|
||||
dim,
|
||||
elementwise_affine=norm_affine,
|
||||
eps=eps,
|
||||
)
|
||||
self.ff = FeedForward(
|
||||
dim=dim,
|
||||
activation_fn=ffn_activation_fn,
|
||||
bias=bias,
|
||||
use_gated=ffn_use_gated,
|
||||
glu_balanced=ffn_glu_balanced,
|
||||
)
|
||||
if use_scale:
|
||||
self.scale2 = nn.Parameter(torch.zeros(dim))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
rotary_pos_emb: Optional[torch.FloatTensor] = None,
|
||||
pack_info: dict = {},
|
||||
):
|
||||
norm_hidden_states = self.norm1(_vit_norm_input(self.norm1, hidden_states)).to(
|
||||
hidden_states.dtype
|
||||
)
|
||||
attn_output = self.attn(norm_hidden_states, rotary_pos_emb, pack_info)
|
||||
if self.use_scale:
|
||||
hidden_states = _scaled_residual_add(
|
||||
hidden_states, attn_output, self.scale1
|
||||
)
|
||||
else:
|
||||
hidden_states = hidden_states + attn_output
|
||||
|
||||
norm_hidden_states = self.norm2(_vit_norm_input(self.norm2, hidden_states)).to(
|
||||
hidden_states.dtype
|
||||
)
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
if self.use_scale:
|
||||
hidden_states = _scaled_residual_add(hidden_states, ff_output, self.scale2)
|
||||
else:
|
||||
hidden_states = hidden_states + ff_output
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,83 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# 3D convolution for the MiniMax H3 visual VAE.
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class BaseConv3d(nn.Conv3d):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
):
|
||||
super().__init__(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode,
|
||||
)
|
||||
padding_mode = "constant" if padding_mode == "zeros" else padding_mode
|
||||
padding_mode_t = "constant" if padding_mode_t == "zeros" else padding_mode_t
|
||||
self.pad_mode = padding_mode
|
||||
self.pad_mode_t = padding_mode_t or ("constant" if causal else "replicate")
|
||||
self.causal = causal
|
||||
|
||||
def _apply_temporal_padding(self, x):
|
||||
B, C, D, H, W = x.shape
|
||||
if D > 1:
|
||||
pad_size = (
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
self.padding[0] * 2 if self.causal else self.padding[0],
|
||||
0 if self.causal else self.padding[0],
|
||||
)
|
||||
return F.pad(x, pad_size, mode=self.pad_mode_t)
|
||||
else:
|
||||
if self.pad_mode_t == "constant":
|
||||
assert self.causal, "Zeros padding is only supported for causal mode"
|
||||
return F.pad(
|
||||
x,
|
||||
(0, 0, 0, 0, self.kernel_size[0] - 1, 0),
|
||||
mode="constant",
|
||||
)
|
||||
else:
|
||||
return x.expand(-1, -1, self.kernel_size[0], -1, -1)
|
||||
|
||||
def _apply_padding(self, x):
|
||||
if sum(self.padding) == 0:
|
||||
return x
|
||||
|
||||
x = F.pad(
|
||||
x,
|
||||
(self.padding[2], self.padding[2], self.padding[1], self.padding[1], 0, 0),
|
||||
mode=self.pad_mode,
|
||||
)
|
||||
|
||||
x = self._apply_temporal_padding(x)
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
if sum(self.padding) == 0:
|
||||
return super().forward(x)
|
||||
|
||||
x = self._apply_padding(x)
|
||||
return F.conv3d(
|
||||
x,
|
||||
self.weight,
|
||||
self.bias,
|
||||
stride=self.stride,
|
||||
padding=0,
|
||||
dilation=self.dilation,
|
||||
)
|
||||
@@ -0,0 +1,190 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Torch-native attention implemented with PyTorch SDPA instead of FA4/CUTLASS.
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
_BLOCK_CAUSAL_MASK_MOD_CACHE = {}
|
||||
|
||||
|
||||
def _auto_sdpa_backend_name() -> str | None:
|
||||
"""Return the ROCm-only correctness fallback for H3 video-VAE SDPA."""
|
||||
if torch.version.hip is None:
|
||||
return None
|
||||
|
||||
from sglang.srt.utils import is_gfx95_supported
|
||||
|
||||
# Fused ROCm SDPA corrupts the dense ViT decode on gfx950. Keep every
|
||||
# non-gfx950 platform, including CUDA, on PyTorch's unchanged auto path.
|
||||
return "math" if is_gfx95_supported() else None
|
||||
|
||||
|
||||
_AUTO_SDPA_BACKEND = _auto_sdpa_backend_name()
|
||||
|
||||
|
||||
def _as_bool_mask(mask, *, device):
|
||||
if not isinstance(mask, torch.Tensor):
|
||||
mask = torch.as_tensor(mask, device=device)
|
||||
return mask.to(device=device, dtype=torch.bool)
|
||||
|
||||
|
||||
def _ensure_nonempty_rows(mask):
|
||||
if mask.numel() == 0 or mask.shape[-1] == 0:
|
||||
return mask
|
||||
empty = ~mask.any(dim=-1)
|
||||
mask[..., 0] |= empty
|
||||
return mask
|
||||
|
||||
|
||||
def _sdpa_kernel_context():
|
||||
backend_name = os.environ.get("MINIMAX_H3_TORCH_SDPA_BACKEND", "auto").lower()
|
||||
if backend_name in {"", "auto", "default"}:
|
||||
backend_name = _AUTO_SDPA_BACKEND
|
||||
if backend_name is None:
|
||||
return nullcontext()
|
||||
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
backends = {
|
||||
"math": SDPBackend.MATH,
|
||||
"flash": SDPBackend.FLASH_ATTENTION,
|
||||
"flash_attention": SDPBackend.FLASH_ATTENTION,
|
||||
"efficient": SDPBackend.EFFICIENT_ATTENTION,
|
||||
"mem_efficient": SDPBackend.EFFICIENT_ATTENTION,
|
||||
"cudnn": SDPBackend.CUDNN_ATTENTION,
|
||||
"cudnn_attention": SDPBackend.CUDNN_ATTENTION,
|
||||
}
|
||||
if backend_name not in backends:
|
||||
raise ValueError(
|
||||
"MINIMAX_H3_TORCH_SDPA_BACKEND must be one of "
|
||||
f"{sorted([*backends, 'auto', 'default'])}, got {backend_name!r}"
|
||||
)
|
||||
return sdpa_kernel(backends=[backends[backend_name]])
|
||||
|
||||
|
||||
def _sdpa_attention(query, key, value, causal=False, attn_mask=None):
|
||||
# query/key/value arrive as [B, S, H, D]; PyTorch SDPA expects
|
||||
# [B, H, S, D].
|
||||
q = query.transpose(1, 2)
|
||||
k = key.transpose(1, 2)
|
||||
v = value.transpose(1, 2)
|
||||
if attn_mask is not None and attn_mask.dim() == 3:
|
||||
attn_mask = attn_mask.unsqueeze(0)
|
||||
with _sdpa_kernel_context():
|
||||
out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=0.0,
|
||||
is_causal=causal,
|
||||
)
|
||||
return out.transpose(1, 2).nan_to_num(0.0)
|
||||
|
||||
|
||||
def _mask_mod_to_dense(mask_mod, batch, heads, q_len, kv_len, device, aux_tensors=None):
|
||||
q_idx = torch.arange(q_len, device=device).view(q_len, 1)
|
||||
kv_idx = torch.arange(kv_len, device=device).view(1, kv_len)
|
||||
dense = torch.empty((batch, heads, q_len, kv_len), dtype=torch.bool, device=device)
|
||||
for b in range(batch):
|
||||
b_idx = torch.tensor(b, device=device)
|
||||
for h in range(heads):
|
||||
h_idx = torch.tensor(h, device=device)
|
||||
mask = mask_mod(b_idx, h_idx, q_idx, kv_idx, None, aux_tensors)
|
||||
dense[b, h] = _as_bool_mask(mask, device=device)
|
||||
return _ensure_nonempty_rows(dense)
|
||||
|
||||
|
||||
#########################################################
|
||||
# Block causal attention
|
||||
#########################################################
|
||||
|
||||
|
||||
def make_block_causal_mask_mod(num_tokens, block_size, num_special=0, suffix=False):
|
||||
if num_tokens < 0:
|
||||
raise ValueError(f"num_tokens must be non-negative, got {num_tokens}")
|
||||
if block_size <= 0:
|
||||
raise ValueError(f"block_size must be positive, got {block_size}")
|
||||
if num_special < 0:
|
||||
raise ValueError(f"num_special must be non-negative, got {num_special}")
|
||||
|
||||
cache_key = (num_tokens, block_size, num_special, suffix)
|
||||
if cache_key in _BLOCK_CAUSAL_MASK_MOD_CACHE:
|
||||
return _BLOCK_CAUSAL_MASK_MOD_CACHE[cache_key]
|
||||
|
||||
if suffix:
|
||||
|
||||
def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors):
|
||||
del b, h, seqlen_info, aux_tensors
|
||||
q_is_special = q_idx >= num_tokens
|
||||
kv_is_special = kv_idx >= num_tokens
|
||||
return (
|
||||
q_is_special
|
||||
| kv_is_special
|
||||
| (q_idx // block_size >= kv_idx // block_size)
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors):
|
||||
del b, h, seqlen_info, aux_tensors
|
||||
q_is_special = q_idx < num_special
|
||||
kv_is_special = kv_idx < num_special
|
||||
q_block_idx = (q_idx - num_special) // block_size
|
||||
kv_block_idx = (kv_idx - num_special) // block_size
|
||||
return q_is_special | kv_is_special | (q_block_idx >= kv_block_idx)
|
||||
|
||||
mask_mod.block_sparse_cache_key = (
|
||||
"block_causal",
|
||||
num_tokens,
|
||||
block_size,
|
||||
num_special,
|
||||
suffix,
|
||||
)
|
||||
_BLOCK_CAUSAL_MASK_MOD_CACHE[cache_key] = mask_mod
|
||||
return mask_mod
|
||||
|
||||
|
||||
#########################################################
|
||||
# Public entry point
|
||||
#########################################################
|
||||
|
||||
|
||||
@torch.compiler.disable
|
||||
def flash_attn(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
causal: bool = False,
|
||||
mask_mod=None,
|
||||
block_sparse=None,
|
||||
aux_tensors=None,
|
||||
) -> torch.Tensor:
|
||||
use_masked = mask_mod is not None or block_sparse is not None
|
||||
|
||||
if block_sparse is not None and mask_mod is None:
|
||||
raise ValueError("block_sparse requires mask_mod")
|
||||
if causal and mask_mod is not None:
|
||||
raise ValueError(
|
||||
"causal must be encoded in mask_mod when using masked attention"
|
||||
)
|
||||
if aux_tensors is not None and not use_masked:
|
||||
raise ValueError("aux_tensors is only supported with masked attention")
|
||||
|
||||
if use_masked:
|
||||
batch, q_len, heads, _ = query.shape
|
||||
kv_len = key.shape[1]
|
||||
dense_mask = _mask_mod_to_dense(
|
||||
mask_mod,
|
||||
batch,
|
||||
heads,
|
||||
q_len,
|
||||
kv_len,
|
||||
query.device,
|
||||
aux_tensors=aux_tensors,
|
||||
)
|
||||
return _sdpa_attention(query, key, value, attn_mask=dense_mask)
|
||||
|
||||
return _sdpa_attention(query, key, value, causal=causal)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,283 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Torch-native normalization for the MiniMax H3 visual VAE.
|
||||
import math
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .conv import BaseConv3d
|
||||
|
||||
|
||||
def _validate_activation(activation):
|
||||
valid_activations = {"identity", "silu", "relu"}
|
||||
if activation not in valid_activations:
|
||||
raise ValueError(
|
||||
f"Unsupported activation: {activation}. Supported: {valid_activations}"
|
||||
)
|
||||
|
||||
|
||||
def _apply_activation(x, activation):
|
||||
_validate_activation(activation)
|
||||
if activation == "identity":
|
||||
return x
|
||||
if activation == "silu":
|
||||
return F.silu(x)
|
||||
return F.relu(x)
|
||||
|
||||
|
||||
def _merge_time_to_batch(x):
|
||||
batch, channels, depth, height, width = x.shape
|
||||
return (
|
||||
x.permute(0, 2, 1, 3, 4)
|
||||
.contiguous()
|
||||
.view(batch * depth, channels, 1, height, width)
|
||||
)
|
||||
|
||||
|
||||
def _split_time_from_batch(x, batch):
|
||||
batch_depth, channels, _, height, width = x.shape
|
||||
depth = batch_depth // batch
|
||||
return (
|
||||
x.view(batch, depth, channels, height, width)
|
||||
.permute(0, 2, 1, 3, 4)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
|
||||
def fused_group_norm(x, num_groups, weight, bias, eps=1e-5, activation="silu"):
|
||||
out = F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps)
|
||||
return _apply_activation(out, activation)
|
||||
|
||||
|
||||
def fused_spatial_norm(
|
||||
f,
|
||||
num_groups,
|
||||
norm_weight,
|
||||
norm_bias,
|
||||
dynamic_scale,
|
||||
dynamic_bias,
|
||||
eps=1e-5,
|
||||
activation="silu",
|
||||
):
|
||||
norm_f = F.group_norm(
|
||||
f,
|
||||
num_groups,
|
||||
weight=norm_weight,
|
||||
bias=norm_bias,
|
||||
eps=eps,
|
||||
)
|
||||
out = norm_f * dynamic_scale + dynamic_bias
|
||||
return _apply_activation(out, activation)
|
||||
|
||||
|
||||
class DummyAffine(torch.nn.Module):
|
||||
def __init__(self, num_channels, affine=True):
|
||||
super().__init__()
|
||||
if affine:
|
||||
self.weight = torch.nn.Parameter(torch.ones(num_channels))
|
||||
self.bias = torch.nn.Parameter(torch.zeros(num_channels))
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
def forward(self, input):
|
||||
if self.weight is None:
|
||||
return input
|
||||
shape = [1, -1] + [1] * (input.dim() - 2)
|
||||
return input * self.weight.view(*shape) + self.bias.view(*shape)
|
||||
|
||||
|
||||
class FusedGroupNorm3D(torch.nn.Module):
|
||||
"""Compatibility wrapper implemented with native PyTorch ops."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_groups,
|
||||
num_channels,
|
||||
eps=1e-5,
|
||||
affine=True,
|
||||
activation="silu",
|
||||
cond_channels=None,
|
||||
use_t_isolated_gn=False,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
):
|
||||
super().__init__()
|
||||
_validate_activation(activation)
|
||||
self.num_groups = num_groups
|
||||
self.num_channels = num_channels
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
self.activation = activation
|
||||
self.use_t_isolated_gn = use_t_isolated_gn
|
||||
|
||||
if cond_channels is not None:
|
||||
self.use_spatial_affine = True
|
||||
self.norm_layer = DummyAffine(num_channels, affine=affine)
|
||||
self.conv_y = BaseConv3d(
|
||||
cond_channels,
|
||||
num_channels,
|
||||
kernel_size=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
self.conv_b = BaseConv3d(
|
||||
cond_channels,
|
||||
num_channels,
|
||||
kernel_size=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
else:
|
||||
self.use_spatial_affine = False
|
||||
if self.affine:
|
||||
self.weight = torch.nn.Parameter(torch.ones(num_channels))
|
||||
self.bias = torch.nn.Parameter(torch.zeros(num_channels))
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
def forward(self, f, cond=None):
|
||||
need_reshape = self.use_t_isolated_gn and f.dim() == 5
|
||||
batch = f.shape[0] if need_reshape else None
|
||||
f_size = f.shape[-3:]
|
||||
if need_reshape:
|
||||
f = _merge_time_to_batch(f)
|
||||
|
||||
if self.use_spatial_affine:
|
||||
scale = self.conv_y(cond)
|
||||
bias = self.conv_b(cond)
|
||||
if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:
|
||||
scale = F.interpolate(scale, size=f_size, mode="nearest")
|
||||
bias = F.interpolate(bias, size=f_size, mode="nearest")
|
||||
if need_reshape:
|
||||
scale = _merge_time_to_batch(scale)
|
||||
bias = _merge_time_to_batch(bias)
|
||||
out = fused_spatial_norm(
|
||||
f,
|
||||
self.num_groups,
|
||||
self.norm_layer.weight,
|
||||
self.norm_layer.bias,
|
||||
scale,
|
||||
bias,
|
||||
self.eps,
|
||||
self.activation,
|
||||
)
|
||||
else:
|
||||
if cond is not None:
|
||||
raise NotImplementedError("Dynamic affine is not defined")
|
||||
weight = self.weight if self.affine else None
|
||||
bias = self.bias if self.affine else None
|
||||
out = fused_group_norm(
|
||||
f, self.num_groups, weight, bias, self.eps, self.activation
|
||||
)
|
||||
|
||||
if need_reshape:
|
||||
out = _split_time_from_batch(out, batch)
|
||||
return out
|
||||
|
||||
|
||||
class TemporalIsolatedGroupNorm(nn.GroupNorm):
|
||||
def forward(self, input):
|
||||
if input.dim() == 5:
|
||||
batch = input.shape[0]
|
||||
input = _merge_time_to_batch(input)
|
||||
output = super().forward(input)
|
||||
return _split_time_from_batch(output, batch)
|
||||
return super().forward(input)
|
||||
|
||||
|
||||
class SpatialNorm3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
f_channels,
|
||||
zq_channels,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
use_t_isolated_gn=False,
|
||||
):
|
||||
super().__init__()
|
||||
norm_cls = TemporalIsolatedGroupNorm if use_t_isolated_gn else nn.GroupNorm
|
||||
self.norm_layer = norm_cls(
|
||||
num_groups=32, num_channels=f_channels, eps=1e-6, affine=True
|
||||
)
|
||||
|
||||
self.conv_y = BaseConv3d(
|
||||
zq_channels,
|
||||
f_channels,
|
||||
kernel_size=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
self.conv_b = BaseConv3d(
|
||||
zq_channels,
|
||||
f_channels,
|
||||
kernel_size=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
def forward(self, f, zq):
|
||||
f_size = f.shape[-3:]
|
||||
norm_f = self.norm_layer(f)
|
||||
scale = self.conv_y(zq)
|
||||
bias = self.conv_b(zq)
|
||||
|
||||
if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:
|
||||
scale = F.interpolate(scale, size=f_size, mode="nearest")
|
||||
bias = F.interpolate(bias, size=f_size, mode="nearest")
|
||||
|
||||
return norm_f * scale + bias
|
||||
|
||||
|
||||
def get_spatial_norm_3d(
|
||||
num_channels,
|
||||
cond_channels,
|
||||
*,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
use_t_isolated_gn=False,
|
||||
):
|
||||
if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":
|
||||
return FusedGroupNorm3D(
|
||||
num_groups=32,
|
||||
num_channels=num_channels,
|
||||
eps=1e-6,
|
||||
affine=True,
|
||||
cond_channels=cond_channels,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
return SpatialNorm3D(
|
||||
num_channels,
|
||||
cond_channels,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
)
|
||||
|
||||
|
||||
def get_group_norm_3d(num_channels, use_t_isolated_gn=False):
|
||||
if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":
|
||||
return FusedGroupNorm3D(
|
||||
num_groups=32,
|
||||
num_channels=num_channels,
|
||||
eps=1e-6,
|
||||
affine=True,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
)
|
||||
|
||||
norm_cls = TemporalIsolatedGroupNorm if use_t_isolated_gn else nn.GroupNorm
|
||||
return norm_cls(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Tensor pre/post-processing for the MiniMax H3 visual VAE.
|
||||
import math
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers.utils import logging
|
||||
from einops import rearrange
|
||||
from torchvision.transforms import Normalize
|
||||
|
||||
NORM_CONFIGS = {
|
||||
"imagenet": {
|
||||
"mean": (0.485, 0.456, 0.406),
|
||||
"std": (0.229, 0.224, 0.225),
|
||||
},
|
||||
"simple": {
|
||||
"mean": (0.5, 0.5, 0.5),
|
||||
"std": (0.5, 0.5, 0.5),
|
||||
},
|
||||
"raw": {
|
||||
"mean": (0.0, 0.0, 0.0),
|
||||
"std": (1.0, 1.0, 1.0),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_norm_constants(
|
||||
norm_type: str = "imagenet",
|
||||
) -> Tuple[Tuple[float, ...], Tuple[float, ...]]:
|
||||
if norm_type not in NORM_CONFIGS:
|
||||
raise ValueError(
|
||||
f"Unknown norm_type: {norm_type}. Must be one of {list(NORM_CONFIGS.keys())}"
|
||||
)
|
||||
config = NORM_CONFIGS[norm_type]
|
||||
return config["mean"], config["std"]
|
||||
|
||||
|
||||
def get_normalize_transform(
|
||||
norm_type: str = "imagenet", *, inplace: bool = False
|
||||
) -> Normalize:
|
||||
mean, std = get_norm_constants(norm_type)
|
||||
return Normalize(mean, std, inplace=inplace)
|
||||
|
||||
|
||||
def get_denormalize_transform(norm_type: str = "imagenet") -> Normalize:
|
||||
mean, std = get_norm_constants(norm_type)
|
||||
inv_mean = tuple(-m / s for m, s in zip(mean, std))
|
||||
inv_std = tuple(1.0 / s for s in std)
|
||||
return Normalize(inv_mean, inv_std)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class VAEProcessor:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
vae_ratio,
|
||||
vae_ratio_t,
|
||||
clip_length,
|
||||
frame_overlap,
|
||||
token_overlap,
|
||||
tokens_chunk_size,
|
||||
isolated_last_frame,
|
||||
latent_patch_size,
|
||||
crop_mode,
|
||||
pixel_norm_type="imagenet",
|
||||
transform=None,
|
||||
transform_rev=None,
|
||||
use_3d_conv=False,
|
||||
):
|
||||
self.vae_ratio = vae_ratio
|
||||
self.vae_ratio_t = vae_ratio_t
|
||||
self.clip_length = clip_length
|
||||
self.frame_overlap = frame_overlap
|
||||
self.token_overlap = token_overlap
|
||||
self.tokens_chunk_size = tokens_chunk_size
|
||||
self.isolated_last_frame = isolated_last_frame
|
||||
self.latent_patch_size = latent_patch_size
|
||||
self.crop_mode = crop_mode
|
||||
self.transform = transform or get_normalize_transform(pixel_norm_type)
|
||||
self._runtime_owned_transform = (
|
||||
get_normalize_transform(pixel_norm_type, inplace=True)
|
||||
if transform is None
|
||||
else None
|
||||
)
|
||||
self.transform_rev = transform_rev or get_denormalize_transform(pixel_norm_type)
|
||||
self.use_3d_conv = use_3d_conv
|
||||
|
||||
def _ensure_list(self, data):
|
||||
return data if isinstance(data, list) else [data]
|
||||
|
||||
def _align_to_total_patch_size(self, h, w):
|
||||
total_patch_size = self.latent_patch_size * self.vae_ratio
|
||||
new_h = (h // total_patch_size) * total_patch_size
|
||||
new_w = (w // total_patch_size) * total_patch_size
|
||||
return new_h, new_w
|
||||
|
||||
def _crop_to_align(self, tensor, new_h, new_w, is_video=False):
|
||||
if is_video:
|
||||
_, _, _, h, w = tensor.shape
|
||||
else:
|
||||
_, _, h, w = tensor.shape
|
||||
|
||||
if self.crop_mode == "center":
|
||||
top = (h - new_h) // 2
|
||||
left = (w - new_w) // 2
|
||||
else:
|
||||
top = 0
|
||||
left = 0
|
||||
|
||||
if is_video:
|
||||
return tensor[:, :, :, top : top + new_h, left : left + new_w]
|
||||
else:
|
||||
return tensor[:, :, top : top + new_h, left : left + new_w]
|
||||
|
||||
def _align_target_token(self, T, mode):
|
||||
intra_tail = self.clip_length % self.vae_ratio_t
|
||||
min_frames = intra_tail or self.vae_ratio_t
|
||||
full_chunks = T // self.clip_length
|
||||
remainder = T % self.clip_length
|
||||
|
||||
if remainder == 0:
|
||||
return max(T, min_frames)
|
||||
|
||||
if mode == "pad":
|
||||
aligned_r = (
|
||||
math.ceil((remainder - intra_tail) / self.vae_ratio_t)
|
||||
* self.vae_ratio_t
|
||||
+ intra_tail
|
||||
)
|
||||
if aligned_r > self.clip_length:
|
||||
return (full_chunks + 1) * self.clip_length + intra_tail
|
||||
return full_chunks * self.clip_length + aligned_r
|
||||
else: # trim
|
||||
k = (remainder - intra_tail) // self.vae_ratio_t
|
||||
if k >= 0:
|
||||
target = (
|
||||
full_chunks * self.clip_length + k * self.vae_ratio_t + intra_tail
|
||||
)
|
||||
return max(target, min_frames)
|
||||
elif full_chunks > 0:
|
||||
return full_chunks * self.clip_length
|
||||
else:
|
||||
return min_frames
|
||||
|
||||
def _align_target(self, T, mode, granularity):
|
||||
if granularity == "chunk":
|
||||
step = self.clip_length
|
||||
tail = self.frame_overlap
|
||||
if self.isolated_last_frame:
|
||||
tail += 1
|
||||
|
||||
k = math.ceil((T - tail) / step) if mode == "pad" else (T - tail) // step
|
||||
return max(k, 1) * step + tail
|
||||
|
||||
isolated_extra = 1 if self.isolated_last_frame else 0
|
||||
return self._align_target_token(T - isolated_extra, mode) + isolated_extra
|
||||
|
||||
def align_video_length(self, video_length, mode="pad", granularity="chunk"):
|
||||
target = self._align_target(video_length, mode, granularity)
|
||||
delta = target - video_length
|
||||
if delta > 0 and mode == "trim":
|
||||
raise ValueError(
|
||||
f"Cannot trim {video_length} frames to valid length {target}: "
|
||||
f"not enough frames (granularity={granularity})"
|
||||
)
|
||||
return delta
|
||||
|
||||
def align_video_length_2pass(self, video_length):
|
||||
"""Return the leading/trailing frame pads and trailing latent drop.
|
||||
|
||||
This is the continuation-prefix (2-pass) alignment. The caller temporarily disables the model's normal token
|
||||
drop and keeps these mirrored processor fields at zero.
|
||||
"""
|
||||
if self.isolated_last_frame:
|
||||
raise ValueError(
|
||||
"align_video_length_2pass does not support isolated_last_frame"
|
||||
)
|
||||
if self.token_overlap != 0 or self.frame_overlap != 0:
|
||||
raise ValueError("align_video_length_2pass requires token_drop=0 alignment")
|
||||
|
||||
leading = self.align_video_length(video_length, mode="pad", granularity="token")
|
||||
token_aligned = video_length + leading
|
||||
trailing = self.align_video_length(
|
||||
token_aligned, mode="pad", granularity="chunk"
|
||||
)
|
||||
|
||||
if trailing > 0:
|
||||
intra_tail = self.clip_length % self.vae_ratio_t
|
||||
full_chunks = token_aligned // self.clip_length
|
||||
remainder = token_aligned % self.clip_length
|
||||
real_tokens = full_chunks * self.tokens_chunk_size
|
||||
if remainder > 0:
|
||||
real_tokens += (remainder - intra_tail) // self.vae_ratio_t + 1
|
||||
drop_tokens = self.get_latent_length(token_aligned + trailing) - real_tokens
|
||||
else:
|
||||
drop_tokens = 0
|
||||
|
||||
return leading, trailing, drop_tokens
|
||||
|
||||
def get_suitable_video_length(self, video_length, verbose=False):
|
||||
used_frame_length = video_length + self.align_video_length(
|
||||
video_length, mode="trim", granularity="chunk"
|
||||
)
|
||||
if verbose:
|
||||
logger.info(
|
||||
f"Pick first {used_frame_length} frames from {video_length}-frame video"
|
||||
)
|
||||
return used_frame_length
|
||||
|
||||
def get_latent_length(self, video_length):
|
||||
tail_frame = self.frame_overlap
|
||||
tail_token = self.token_overlap
|
||||
if self.isolated_last_frame:
|
||||
tail_frame += 1
|
||||
tail_token += 1
|
||||
|
||||
video_length = self.get_suitable_video_length(video_length)
|
||||
latent_length = (
|
||||
int((video_length - tail_frame) // self.clip_length)
|
||||
* self.tokens_chunk_size
|
||||
+ tail_token
|
||||
)
|
||||
return latent_length
|
||||
|
||||
def transform_tensor(self, tensor, *, runtime_owned=False):
|
||||
B, T = None, None
|
||||
if tensor.ndim == 5:
|
||||
if tensor.shape[2] == 3:
|
||||
tensor = tensor.transpose(1, 2)
|
||||
B, _, T, _, _ = tensor.shape
|
||||
tensor = rearrange(tensor, "b c t h w -> (b t) c h w")
|
||||
elif tensor.ndim == 4:
|
||||
if tensor.shape[0] == 3:
|
||||
tensor = tensor.transpose(0, 1)
|
||||
elif tensor.ndim == 3:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
else:
|
||||
raise ValueError(f"Unsupported tensor shape: {tensor.shape}")
|
||||
|
||||
transform = (
|
||||
self._runtime_owned_transform
|
||||
if runtime_owned and self._runtime_owned_transform is not None
|
||||
else self.transform
|
||||
)
|
||||
tensor = transform(tensor)
|
||||
|
||||
if B is not None and T is not None:
|
||||
tensor = rearrange(tensor, "(b t) c h w -> b c t h w", b=B, t=T)
|
||||
|
||||
return tensor.contiguous()
|
||||
|
||||
def revert_tensor(self, tensor):
|
||||
B, T = None, None
|
||||
if self.use_3d_conv:
|
||||
tensor = tensor.unsqueeze(2) if tensor.ndim == 4 else tensor
|
||||
B, _, T, _, _ = tensor.shape
|
||||
tensor = rearrange(tensor, "b c t h w -> (b t) c h w")
|
||||
tensor_rev = self.transform_rev(tensor).clamp_(0, 1)
|
||||
if B is not None:
|
||||
tensor_rev = rearrange(tensor_rev, "(b t) c h w -> b c t h w", b=B, t=T)
|
||||
return tensor_rev.contiguous()
|
||||
|
||||
@staticmethod
|
||||
def convert_numpy_to_tensor(numpy_array, device=None):
|
||||
if isinstance(numpy_array, list):
|
||||
numpy_array = np.stack(numpy_array, axis=0)
|
||||
tensor = torch.from_numpy(numpy_array)
|
||||
# Keep decoded uint8 pixels compact across the host-to-device copy.
|
||||
# Casting the full video on CPU quadruples both the temporary host
|
||||
# allocation and transfer volume for no loss of information.
|
||||
if device is not None:
|
||||
tensor = tensor.to(device)
|
||||
tensor = tensor.permute(0, 3, 1, 2)
|
||||
return tensor.to(torch.float32).div_(255.0)
|
||||
@@ -0,0 +1,276 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# 3D causal CNN encoder for the MiniMax H3 visual VAE (inference-only bundle).
|
||||
import os
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .conv import BaseConv3d
|
||||
from .norm import get_group_norm_3d, get_spatial_norm_3d
|
||||
|
||||
# ============================================================================
|
||||
# 3D CNN Components
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def norm_silu(x, norm, cond=None):
|
||||
if cond is None:
|
||||
return F.silu(norm(x), inplace=True)
|
||||
else:
|
||||
return F.silu(norm(x, cond), inplace=True)
|
||||
|
||||
|
||||
class Downsample3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
time_stride=1,
|
||||
space_stride=2,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
):
|
||||
super().__init__()
|
||||
self.time_stride = time_stride
|
||||
self.space_stride = space_stride
|
||||
|
||||
assert time_stride in [1, 2]
|
||||
assert space_stride in [1, 2, 3]
|
||||
|
||||
self.conv = BaseConv3d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
padding=(1, 0, 0),
|
||||
stride=(time_stride, space_stride, space_stride),
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
self.causal = self.conv.causal
|
||||
self.pad_mode = self.conv.pad_mode
|
||||
|
||||
def forward(self, x):
|
||||
if self.space_stride == 2:
|
||||
pad = (0, 1, 0, 1, 0, 0)
|
||||
x = F.pad(x, pad, mode=self.pad_mode)
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class ResnetBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels=None,
|
||||
zq_ch=None,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
use_t_isolated_gn=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.use_fused_norm = (
|
||||
os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true"
|
||||
)
|
||||
|
||||
if zq_ch is None:
|
||||
self.norm1 = get_group_norm_3d(
|
||||
in_channels, use_t_isolated_gn=use_t_isolated_gn
|
||||
)
|
||||
self.norm2 = get_group_norm_3d(
|
||||
out_channels, use_t_isolated_gn=use_t_isolated_gn
|
||||
)
|
||||
else:
|
||||
self.norm1 = get_spatial_norm_3d(
|
||||
in_channels,
|
||||
zq_ch,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
)
|
||||
self.norm2 = get_spatial_norm_3d(
|
||||
out_channels,
|
||||
zq_ch,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
)
|
||||
|
||||
self.conv1 = BaseConv3d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
self.conv2 = BaseConv3d(
|
||||
out_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
self.nin_shortcut = BaseConv3d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
def forward(self, x, zq=None):
|
||||
h = x
|
||||
|
||||
if self.use_fused_norm:
|
||||
h = self.norm1(h, zq)
|
||||
else:
|
||||
h = norm_silu(h, self.norm1, zq)
|
||||
|
||||
h = self.conv1(h)
|
||||
|
||||
if self.use_fused_norm:
|
||||
h = self.norm2(h, zq)
|
||||
else:
|
||||
h = norm_silu(h, self.norm2, zq)
|
||||
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return h.add_(x)
|
||||
|
||||
|
||||
class EncoderFCN3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
ch,
|
||||
ch_mult,
|
||||
space_down,
|
||||
time_down,
|
||||
num_res_blocks,
|
||||
in_channels,
|
||||
z_channels,
|
||||
double_z=False,
|
||||
zq_ch=None,
|
||||
padding_mode="zeros",
|
||||
padding_mode_t=None,
|
||||
causal=True,
|
||||
use_t_isolated_gn=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_levels = len(ch_mult)
|
||||
|
||||
if isinstance(num_res_blocks, int):
|
||||
self.num_res_blocks = [num_res_blocks] * self.num_levels
|
||||
else:
|
||||
self.num_res_blocks = num_res_blocks
|
||||
|
||||
self.space_down_factors = space_down
|
||||
self.time_down_factors = time_down
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.use_fused_norm = (
|
||||
os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true"
|
||||
)
|
||||
|
||||
block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
|
||||
block_in = [block_mid[0]] + block_mid[:-1]
|
||||
block_out = block_mid
|
||||
|
||||
conv_kwargs = dict(
|
||||
padding_mode=padding_mode,
|
||||
padding_mode_t=padding_mode_t,
|
||||
causal=causal,
|
||||
)
|
||||
|
||||
self.conv_in = BaseConv3d(
|
||||
in_channels, block_in[0], kernel_size=3, padding=1, **conv_kwargs
|
||||
)
|
||||
|
||||
self.down = nn.ModuleList()
|
||||
for i_level in range(self.num_levels):
|
||||
down = nn.Module()
|
||||
|
||||
down.block = nn.ModuleList()
|
||||
for i in range(self.num_res_blocks[i_level]):
|
||||
down.block.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=block_in[i_level] if i == 0 else block_mid[i_level],
|
||||
out_channels=block_mid[i_level],
|
||||
zq_ch=zq_ch,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
**conv_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
if space_down[i_level] * time_down[i_level] > 1:
|
||||
down.downsample = Downsample3D(
|
||||
block_mid[i_level],
|
||||
block_out[i_level],
|
||||
time_stride=time_down[i_level],
|
||||
space_stride=space_down[i_level],
|
||||
**conv_kwargs,
|
||||
)
|
||||
else:
|
||||
if block_out[i_level] != block_mid[i_level]:
|
||||
down.downsample = BaseConv3d(
|
||||
block_mid[i_level],
|
||||
block_out[i_level],
|
||||
kernel_size=1,
|
||||
**conv_kwargs,
|
||||
)
|
||||
|
||||
self.down.append(down)
|
||||
|
||||
if zq_ch is None:
|
||||
self.norm_out = get_group_norm_3d(
|
||||
block_out[-1], use_t_isolated_gn=use_t_isolated_gn
|
||||
)
|
||||
else:
|
||||
self.norm_out = get_spatial_norm_3d(
|
||||
block_out[-1],
|
||||
zq_ch,
|
||||
use_t_isolated_gn=use_t_isolated_gn,
|
||||
**conv_kwargs,
|
||||
)
|
||||
|
||||
self.conv_out = BaseConv3d(
|
||||
block_out[-1],
|
||||
2 * z_channels if double_z else z_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
**conv_kwargs,
|
||||
)
|
||||
|
||||
def forward(self, x, zq=None):
|
||||
h = self.conv_in(x)
|
||||
for i_level in range(self.num_levels):
|
||||
for i_block in range(self.num_res_blocks[i_level]):
|
||||
h = self.down[i_level].block[i_block](h, zq)
|
||||
if hasattr(self.down[i_level], "downsample"):
|
||||
h = self.down[i_level].downsample(h)
|
||||
|
||||
if self.use_fused_norm:
|
||||
h = self.norm_out(h, zq)
|
||||
else:
|
||||
h = norm_silu(h, self.norm_out, zq)
|
||||
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
@@ -0,0 +1,374 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# ViT3D decoder for the MiniMax H3 visual VAE (inference-only bundle).
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils import logging
|
||||
|
||||
from .base_module import RotaryEmbeddingND, TransformerBlock
|
||||
from .flash import make_block_causal_mask_mod
|
||||
from .vit_utils import create_token_ids, prepare_rotary_pos_emb
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def _linear_with_module_dtype(linear, tensor, out_dtype=None):
|
||||
weight = getattr(linear, "weight", None)
|
||||
target_dtype = getattr(weight, "dtype", tensor.dtype)
|
||||
output = linear(tensor.to(target_dtype))
|
||||
if out_dtype is not None and output.dtype != out_dtype:
|
||||
output = output.to(out_dtype)
|
||||
return output
|
||||
|
||||
|
||||
def _pack_tensors_3d(tensors, patch_size, patch_size_t):
|
||||
batch_size, num_channels_tensors, temporal, height, width = tensors.shape
|
||||
|
||||
tensors = tensors.view(
|
||||
batch_size,
|
||||
num_channels_tensors,
|
||||
temporal // patch_size_t,
|
||||
patch_size_t,
|
||||
height // patch_size,
|
||||
patch_size,
|
||||
width // patch_size,
|
||||
patch_size,
|
||||
)
|
||||
tensors = tensors.permute(0, 2, 4, 6, 1, 3, 5, 7)
|
||||
tensors = tensors.reshape(
|
||||
batch_size,
|
||||
(temporal // patch_size_t) * (height // patch_size) * (width // patch_size),
|
||||
num_channels_tensors * patch_size_t * patch_size * patch_size,
|
||||
)
|
||||
return tensors
|
||||
|
||||
|
||||
def _unpack_tensors_3d(tensors, patch_size, patch_size_t, temporal, height, width):
|
||||
batch_size, num_patches, channels = tensors.shape
|
||||
num_channels_tensors = channels // (patch_size_t * patch_size * patch_size)
|
||||
|
||||
tensors = tensors.view(
|
||||
batch_size,
|
||||
temporal // patch_size_t,
|
||||
height // patch_size,
|
||||
width // patch_size,
|
||||
num_channels_tensors,
|
||||
patch_size_t,
|
||||
patch_size,
|
||||
patch_size,
|
||||
)
|
||||
tensors = tensors.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
|
||||
tensors = tensors.reshape(batch_size, num_channels_tensors, temporal, height, width)
|
||||
return tensors
|
||||
|
||||
|
||||
class ViTBase(ModelMixin, ConfigMixin):
|
||||
"""Base class for ViT Encoder and Decoder with common functionality."""
|
||||
|
||||
_no_split_modules = ["TransformerBlock"]
|
||||
|
||||
def _init_weights(self):
|
||||
def basic_init(m):
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
self.apply(basic_init)
|
||||
|
||||
def init_mask_config(self, dim, is_3d=False):
|
||||
self._mask_dim = dim
|
||||
self._mask_is_3d = is_3d
|
||||
self.register_buffer("mask_token", torch.zeros(1, 1, dim))
|
||||
|
||||
def set_mask_config(self, mask_config):
|
||||
self.mask_prob = mask_config.get("mask_prob", 0.0)
|
||||
self.mask_enabled = self.mask_prob > 0
|
||||
self.mask_style = mask_config.get("mask_style", "replace")
|
||||
if self.mask_enabled and self.mask_style == "drop" and self.mask_prob < 1.0:
|
||||
logger.warning("mask_style='drop' with mask_prob < 1.0")
|
||||
if self._mask_is_3d:
|
||||
self.temporal_scale_range = mask_config.get(
|
||||
"temporal_scale_range", (0.3, 0.5)
|
||||
)
|
||||
self.spatial_scale_range = mask_config.get(
|
||||
"spatial_scale_range", (0.1, 0.25)
|
||||
)
|
||||
self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.75)
|
||||
self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.95)
|
||||
else:
|
||||
self.spatial_scale_range = mask_config.get(
|
||||
"spatial_scale_range", (0.15, 0.15)
|
||||
)
|
||||
self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.5)
|
||||
self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.75)
|
||||
self.aspect_ratio_range = mask_config.get("aspect_ratio_range", (0.75, 1.5))
|
||||
self.max_retries = mask_config.get("max_retries", 100)
|
||||
if (
|
||||
self.mask_enabled
|
||||
and self.mask_style == "drop"
|
||||
and getattr(self, "t_causal", False)
|
||||
):
|
||||
logger.warning("mask_style='drop' with t_causal may cause issues")
|
||||
if self.mask_enabled and "mask_token" in self._buffers:
|
||||
del self._buffers["mask_token"]
|
||||
self.mask_token = nn.Parameter(torch.randn(1, 1, self._mask_dim) * 0.02)
|
||||
|
||||
def init_suffix_tokens(self, dim, num_register_tokens, has_cls_token=True):
|
||||
self.num_register_tokens = num_register_tokens
|
||||
if num_register_tokens > 0:
|
||||
self.register_tokens = nn.Parameter(
|
||||
torch.randn(1, num_register_tokens, dim) * 0.02
|
||||
)
|
||||
else:
|
||||
self.register_tokens = None
|
||||
if has_cls_token:
|
||||
self.cls_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02)
|
||||
|
||||
def apply_mask_preprocess(self, hidden_states, img_ids, patch_dims, num_suffix):
|
||||
if self.training and self.mask_enabled:
|
||||
raise NotImplementedError(
|
||||
"mask modeling is not supported in this inference-only bundle"
|
||||
)
|
||||
return hidden_states, img_ids
|
||||
|
||||
def forward_transformer_blocks(self, hidden_states, rotary_pos_emb, pack_info=None):
|
||||
if pack_info is None:
|
||||
pack_info = {}
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(hidden_states, rotary_pos_emb, pack_info)
|
||||
return hidden_states
|
||||
|
||||
def apply_mask_postprocess(self, hidden_states, num_patches):
|
||||
if self.training and self.mask_enabled and self.mask_style == "drop":
|
||||
raise NotImplementedError(
|
||||
"mask modeling is not supported in this inference-only bundle"
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class ViT3DDecoder(ViTBase):
|
||||
"""Vision Transformer Video Decoder using TransformerBlock."""
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: int = 16,
|
||||
patch_size_t: int = 4,
|
||||
t_causal: bool = False,
|
||||
in_channels: int = 16,
|
||||
out_channels: int = 3,
|
||||
num_layers: int = 24,
|
||||
heads: int = 16,
|
||||
dim_head: int = 64,
|
||||
norm_type: str = "layer_norm",
|
||||
norm_affine: bool = True,
|
||||
qk_norm_type: str = None,
|
||||
qk_norm_affine: bool = False,
|
||||
ffn_activation_fn: str = "gelu",
|
||||
ffn_use_gated: bool = False,
|
||||
rope_theta: float = 100.0,
|
||||
rope_dim_ratio: float = 1.0,
|
||||
bias: bool = True,
|
||||
eps: float = 1e-5,
|
||||
num_register_tokens: int = 4,
|
||||
mask_config: dict = {},
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
dim = heads * dim_head
|
||||
rope_apply_dim = int(dim_head * rope_dim_ratio)
|
||||
|
||||
self.pos_embed = RotaryEmbeddingND(
|
||||
rope_apply_dim, rope_theta, n_dim=3, use_angle=True
|
||||
)
|
||||
|
||||
self.x_embedder = nn.Linear(in_channels, dim)
|
||||
|
||||
self.init_suffix_tokens(dim, num_register_tokens, has_cls_token=False)
|
||||
|
||||
self.t_causal = t_causal
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
TransformerBlock(
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
norm_type=norm_type,
|
||||
norm_affine=norm_affine,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qk_norm_affine=qk_norm_affine,
|
||||
ffn_activation_fn=ffn_activation_fn,
|
||||
ffn_use_gated=ffn_use_gated,
|
||||
bias=bias,
|
||||
eps=eps,
|
||||
**kwargs,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.norm_out = nn.LayerNorm(dim, elementwise_affine=norm_affine, eps=eps)
|
||||
patch_dim = out_channels * patch_size_t * patch_size * patch_size
|
||||
self.proj_out = nn.Linear(dim, patch_dim)
|
||||
|
||||
self.init_mask_config(dim, is_3d=True)
|
||||
self.set_mask_config(mask_config)
|
||||
|
||||
self._rotary_pos_emb_cache = None
|
||||
self._autocast_linear_dtype = None
|
||||
|
||||
if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0):
|
||||
logger.warning(f"Unused kwargs: {kwargs}")
|
||||
|
||||
def _apply(self, fn, recurse=True):
|
||||
result = super()._apply(fn, recurse=recurse)
|
||||
self._rotary_pos_emb_cache = None
|
||||
self._autocast_linear_dtype = None
|
||||
return result
|
||||
|
||||
def prepare_autocast_linear_weights(self, dtype: torch.dtype) -> int:
|
||||
"""Keep decoder-block linear weights in their autocast compute dtype.
|
||||
|
||||
PyTorch autocast does not cache casts for these frozen parameters, so
|
||||
tiled decode otherwise converts every FP32 weight and bias once per
|
||||
block invocation. Persisting the rounded values is numerically
|
||||
equivalent to the per-call autocast conversion. The embedding and
|
||||
output projections stay FP32 because their calls explicitly disable
|
||||
autocast.
|
||||
"""
|
||||
|
||||
if dtype not in (torch.float16, torch.bfloat16):
|
||||
raise ValueError(
|
||||
"MiniMax H3 decoder autocast weights require fp16 or bf16, "
|
||||
f"got {dtype}"
|
||||
)
|
||||
if self._autocast_linear_dtype == dtype:
|
||||
return 0
|
||||
|
||||
converted = 0
|
||||
for block in self.transformer_blocks:
|
||||
for linear in (
|
||||
block.attn.to_qkv,
|
||||
block.attn.to_out,
|
||||
block.ff.w1,
|
||||
block.ff.w2,
|
||||
):
|
||||
if linear.weight.dtype != dtype:
|
||||
linear.to(dtype=dtype)
|
||||
converted += 1
|
||||
self._autocast_linear_dtype = dtype
|
||||
return converted
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
B, C, latent_T, latent_H, latent_W = x.shape
|
||||
patch_size = self.config.patch_size
|
||||
patch_size_t = self.config.patch_size_t
|
||||
num_suffix = 1 + self.num_register_tokens
|
||||
|
||||
hidden_states = _pack_tensors_3d(x, 1, 1)
|
||||
latent_size = (latent_T, latent_H, latent_W)
|
||||
|
||||
with torch.autocast("cuda", enabled=False):
|
||||
hidden_states = _linear_with_module_dtype(
|
||||
self.x_embedder, hidden_states, hidden_states.dtype
|
||||
)
|
||||
|
||||
num_patches = hidden_states.shape[1]
|
||||
|
||||
tokens = [hidden_states]
|
||||
|
||||
if self.register_tokens is not None:
|
||||
register_tokens = self.register_tokens.expand(B, -1, -1)
|
||||
tokens.append(register_tokens)
|
||||
|
||||
cls_token = torch.zeros_like(hidden_states[:, 0:1, :])
|
||||
tokens.append(cls_token)
|
||||
hidden_states = torch.cat(tokens, dim=1)
|
||||
|
||||
patch_dims = [latent_T, latent_H, latent_W]
|
||||
rotary_dtype = (
|
||||
torch.get_autocast_dtype("cuda")
|
||||
if x.is_cuda and torch.is_autocast_enabled("cuda")
|
||||
else hidden_states.dtype
|
||||
)
|
||||
cache_enabled = (
|
||||
not self.training
|
||||
and not self.mask_enabled
|
||||
and not torch.compiler.is_compiling()
|
||||
)
|
||||
cache_key = (
|
||||
B,
|
||||
latent_T,
|
||||
latent_H,
|
||||
latent_W,
|
||||
num_suffix,
|
||||
x.device,
|
||||
x.dtype,
|
||||
rotary_dtype,
|
||||
)
|
||||
cache_record = self._rotary_pos_emb_cache if cache_enabled else None
|
||||
cache_hit = cache_record is not None and cache_record[0] == cache_key
|
||||
if cache_hit:
|
||||
img_ids = cache_record[1]
|
||||
else:
|
||||
img_ids = create_token_ids(latent_size, x.device, x.dtype).expand(B, -1, -1)
|
||||
suffix_ids = torch.zeros(
|
||||
(B, num_suffix, 3), device=x.device, dtype=img_ids.dtype
|
||||
)
|
||||
img_ids = torch.cat([img_ids, suffix_ids], dim=1)
|
||||
|
||||
hidden_states, img_ids = self.apply_mask_preprocess(
|
||||
hidden_states, img_ids, patch_dims, num_suffix
|
||||
)
|
||||
cache_img_ids = img_ids
|
||||
|
||||
pack_info = {}
|
||||
if self.t_causal:
|
||||
spatial_size = latent_H * latent_W
|
||||
mask_mod = make_block_causal_mask_mod(
|
||||
num_tokens=num_patches,
|
||||
block_size=spatial_size,
|
||||
suffix=True,
|
||||
)
|
||||
pack_info["mask_mod"] = mask_mod
|
||||
|
||||
if cache_hit:
|
||||
rotary_pos_emb = cache_record[2]
|
||||
else:
|
||||
rotary_pos_emb = prepare_rotary_pos_emb(
|
||||
self.pos_embed(img_ids),
|
||||
dtype=rotary_dtype,
|
||||
)
|
||||
if cache_enabled:
|
||||
self._rotary_pos_emb_cache = (
|
||||
cache_key,
|
||||
cache_img_ids,
|
||||
rotary_pos_emb,
|
||||
)
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(hidden_states, rotary_pos_emb, pack_info)
|
||||
|
||||
hidden_states = self.norm_out(hidden_states)
|
||||
|
||||
hidden_states = self.apply_mask_postprocess(hidden_states, num_patches)
|
||||
|
||||
with torch.autocast("cuda", enabled=False):
|
||||
output = _linear_with_module_dtype(
|
||||
self.proj_out, hidden_states, hidden_states.dtype
|
||||
)
|
||||
|
||||
output = output[:, :num_patches, :]
|
||||
|
||||
video_t = latent_size[0] * patch_size_t
|
||||
video_h = latent_size[1] * patch_size
|
||||
video_w = latent_size[2] * patch_size
|
||||
output = _unpack_tensors_3d(
|
||||
output, patch_size, patch_size_t, video_t, video_h, video_w
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,255 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# ViT runtime helpers for the MiniMax H3 visual VAE.
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from diffusers.utils import logging
|
||||
|
||||
|
||||
def _env_flag(name, default="0"):
|
||||
value = os.environ.get(name, default)
|
||||
return str(value).strip().lower() in ("1", "true", "yes", "on")
|
||||
|
||||
|
||||
def _env_optional_bool(name, default=""):
|
||||
value = str(os.environ.get(name, default)).strip().lower()
|
||||
if value in ("", "default", "auto", "none", "unset"):
|
||||
return None
|
||||
return value not in ("0", "false", "no", "off", "disabled")
|
||||
|
||||
|
||||
def _vit_torch_compile_kwargs(prefix):
|
||||
kwargs = {}
|
||||
backend = os.environ.get(f"{prefix}_BACKEND", "inductor").strip()
|
||||
mode = os.environ.get(f"{prefix}_MODE", "reduce-overhead").strip()
|
||||
if backend and backend.lower() not in ("default", "none"):
|
||||
kwargs["backend"] = backend
|
||||
if mode and mode.lower() not in ("default", "none"):
|
||||
kwargs["mode"] = mode
|
||||
kwargs["fullgraph"] = _env_flag(f"{prefix}_FULLGRAPH", "0")
|
||||
dynamic = _env_optional_bool(f"{prefix}_DYNAMIC")
|
||||
if dynamic is not None:
|
||||
kwargs["dynamic"] = dynamic
|
||||
return kwargs
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def create_token_ids(
|
||||
patch_dims, device, dtype, id_type="length_normalized", flatten=True
|
||||
):
|
||||
coords_list = []
|
||||
|
||||
if isinstance(id_type, str):
|
||||
id_type_list = [id_type] * len(patch_dims)
|
||||
elif isinstance(id_type, list):
|
||||
id_type_list = id_type
|
||||
if len(id_type_list) != len(patch_dims):
|
||||
raise ValueError("id_type list must match patch_dims")
|
||||
else:
|
||||
raise ValueError("id_type must be a string or a list")
|
||||
|
||||
if "area_normalized" in id_type_list or id_type == "area_normalized":
|
||||
raise NotImplementedError(
|
||||
"area_normalized id_type is not supported in this inference-only bundle"
|
||||
)
|
||||
|
||||
for _dim_size, _id_type in zip(patch_dims, id_type_list):
|
||||
if isinstance(_dim_size, torch.Tensor):
|
||||
coords_list.append(_dim_size.to(device=device, dtype=dtype))
|
||||
continue
|
||||
|
||||
if _id_type == "length_normalized":
|
||||
coords = torch.arange(0.5, _dim_size, dtype=dtype, device=device)
|
||||
coords = coords / _dim_size
|
||||
coords = 2.0 * coords - 1.0
|
||||
else:
|
||||
coords = torch.arange(_dim_size, dtype=dtype, device=device)
|
||||
|
||||
coords_list.append(coords)
|
||||
|
||||
coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1)
|
||||
if flatten:
|
||||
coords = coords.flatten(0, len(patch_dims) - 1)
|
||||
|
||||
return coords.unsqueeze(0)
|
||||
|
||||
|
||||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
x1, x2 = torch.chunk(x, 2, dim=-1)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def _apply_rotary_pos_emb_impl(
|
||||
t: torch.Tensor, rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor]
|
||||
) -> torch.Tensor:
|
||||
cos, sin = rotary_pos_emb[:2]
|
||||
|
||||
if cos.dim() != 4:
|
||||
raise ValueError(f"cos must be [B, N, 1, D], got {cos.shape}")
|
||||
|
||||
cos = cos.to(t.dtype)
|
||||
sin = sin.to(t.dtype)
|
||||
|
||||
rot_dim = cos.shape[-1]
|
||||
t_dim = t.shape[-1]
|
||||
|
||||
if rot_dim < t_dim:
|
||||
t_rot, t_pass = t[..., :rot_dim], t[..., rot_dim:]
|
||||
scaled = t_rot * cos
|
||||
scaled.add_(_rotate_half(t_rot) * sin)
|
||||
t_rot = scaled
|
||||
t = torch.cat((t_rot, t_pass), dim=-1)
|
||||
else:
|
||||
scaled = t * cos
|
||||
scaled.add_(_rotate_half(t) * sin)
|
||||
t = scaled
|
||||
|
||||
return t
|
||||
|
||||
|
||||
def prepare_rotary_pos_emb(
|
||||
rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor],
|
||||
*,
|
||||
dtype: torch.dtype,
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
"""Prebuild the native Q/K rotary cache once per ViT decoder forward."""
|
||||
cos, sin = rotary_pos_emb
|
||||
if (
|
||||
not cos.is_cuda
|
||||
or dtype not in (torch.float16, torch.bfloat16)
|
||||
or cos.shape != sin.shape
|
||||
or cos.dim() != 4
|
||||
or cos.shape[0] != 1
|
||||
or cos.shape[2] != 1
|
||||
or cos.shape[-1] % 2
|
||||
or _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0")
|
||||
):
|
||||
return cos, sin
|
||||
|
||||
cos = cos.to(dtype=dtype)
|
||||
sin = sin.to(dtype=dtype)
|
||||
half = cos.shape[-1] // 2
|
||||
# RotaryEmbeddingND repeats each half. The native kernel consumes the
|
||||
# compact NeoX cache [cos_half | sin_half].
|
||||
cache = torch.cat(
|
||||
(cos[0, :, 0, :half], sin[0, :, 0, :half]),
|
||||
dim=-1,
|
||||
).contiguous()
|
||||
positions = torch.arange(
|
||||
cos.shape[1],
|
||||
dtype=torch.long,
|
||||
device=cos.device,
|
||||
)
|
||||
return cos, sin, cache, positions
|
||||
|
||||
|
||||
_COMPILED_APPLY_ROTARY_POS_EMB = None
|
||||
_APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = False
|
||||
|
||||
|
||||
def _get_apply_rotary_pos_emb_impl():
|
||||
global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED
|
||||
if _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED or not _env_flag(
|
||||
"MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0"
|
||||
):
|
||||
return _apply_rotary_pos_emb_impl
|
||||
if _COMPILED_APPLY_ROTARY_POS_EMB is not None:
|
||||
return _COMPILED_APPLY_ROTARY_POS_EMB
|
||||
if not hasattr(torch, "compile"):
|
||||
message = (
|
||||
"torch.compile is unavailable; falling back to eager ViT rotary embedding"
|
||||
)
|
||||
if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"):
|
||||
raise RuntimeError(message)
|
||||
logger.warning(f"[ViTRope] {message}")
|
||||
_APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True
|
||||
return _apply_rotary_pos_emb_impl
|
||||
|
||||
kwargs = _vit_torch_compile_kwargs("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE")
|
||||
try:
|
||||
_COMPILED_APPLY_ROTARY_POS_EMB = torch.compile(
|
||||
_apply_rotary_pos_emb_impl, **kwargs
|
||||
)
|
||||
logger.info(f"[ViTRope] torch.compile enabled kwargs={kwargs}")
|
||||
except Exception as exc:
|
||||
if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"):
|
||||
raise
|
||||
logger.warning(
|
||||
f"[ViTRope] torch.compile setup failed: {type(exc).__name__}: {exc}; "
|
||||
"falling back to eager"
|
||||
)
|
||||
_APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True
|
||||
_COMPILED_APPLY_ROTARY_POS_EMB = None
|
||||
return _apply_rotary_pos_emb_impl
|
||||
return _COMPILED_APPLY_ROTARY_POS_EMB
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(
|
||||
t: torch.Tensor, rotary_pos_emb: Sequence[torch.Tensor]
|
||||
) -> torch.Tensor:
|
||||
global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED
|
||||
fn = _get_apply_rotary_pos_emb_impl()
|
||||
try:
|
||||
return fn(t, rotary_pos_emb)
|
||||
except Exception as exc:
|
||||
if fn is _COMPILED_APPLY_ROTARY_POS_EMB and not _env_flag(
|
||||
"MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"
|
||||
):
|
||||
logger.warning(
|
||||
f"[ViTRope] compiled call failed: {type(exc).__name__}: {exc}; "
|
||||
"disabling compile and retrying eager"
|
||||
)
|
||||
_APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True
|
||||
_COMPILED_APPLY_ROTARY_POS_EMB = None
|
||||
return _apply_rotary_pos_emb_impl(t, rotary_pos_emb)
|
||||
raise
|
||||
|
||||
|
||||
def apply_rotary_pos_emb_qk(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
rotary_pos_emb: Sequence[torch.Tensor],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Apply the exact native NeoX rotary kernel to Q/K together when possible."""
|
||||
if (
|
||||
len(rotary_pos_emb) == 4
|
||||
and query.is_cuda
|
||||
and query.shape == key.shape
|
||||
and query.dtype == key.dtype
|
||||
and query.dtype in (torch.float16, torch.bfloat16)
|
||||
and query.dim() == 4
|
||||
and query.shape[0] == 1
|
||||
and not torch.compiler.is_compiling()
|
||||
):
|
||||
_, _, cache, positions = rotary_pos_emb
|
||||
if (
|
||||
cache.is_cuda
|
||||
and cache.dtype == query.dtype
|
||||
and cache.dim() == 2
|
||||
and cache.shape[0] == query.shape[1]
|
||||
and cache.shape[1] <= query.shape[-1]
|
||||
and positions.is_cuda
|
||||
and positions.shape == (query.shape[1],)
|
||||
):
|
||||
from sgl_kernel import rotary_embedding
|
||||
|
||||
query = query.contiguous()
|
||||
key = key.contiguous()
|
||||
rotary_embedding(
|
||||
positions,
|
||||
query.view(query.shape[1], -1),
|
||||
key.view(key.shape[1], -1),
|
||||
query.shape[-1],
|
||||
cache,
|
||||
True,
|
||||
)
|
||||
return query, key
|
||||
|
||||
return (
|
||||
apply_rotary_pos_emb(query, rotary_pos_emb),
|
||||
apply_rotary_pos_emb(key, rotary_pos_emb),
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages import InputValidationStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3 import (
|
||||
MiniMaxH3AudioEncodingStage,
|
||||
MiniMaxH3DecodingStage,
|
||||
MiniMaxH3DenoisingStage,
|
||||
MiniMaxH3LatentPreparationStage,
|
||||
MiniMaxH3TextEncodingStage,
|
||||
MiniMaxH3TimestepPreparationStage,
|
||||
MiniMaxH3VisualEncodingStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
|
||||
MiniMaxH3PartitionAdmissionStage,
|
||||
MiniMaxH3ReleaseMetadata,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
|
||||
|
||||
class MiniMaxH3Pipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
pipeline_name = "MiniMaxH3Pipeline"
|
||||
default_model_subfolder = "FL2VA"
|
||||
is_video_pipeline = True
|
||||
pipeline_config_cls = MiniMaxH3PipelineConfig
|
||||
sampling_params_cls = MiniMaxH3SamplingParams
|
||||
_required_config_modules = [
|
||||
"processor",
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"video_vae",
|
||||
"audio_vae",
|
||||
# scheduler intentionally absent: model_index carries scheduler=null;
|
||||
# per-modality sigma schedules are generated in TimestepPreparation
|
||||
# from the task profile, and the loop scheduler math lives in
|
||||
# scheduling_minimax_h3_euler_ancestral (stages accept scheduler=None).
|
||||
"transformer",
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def model_subfolder_for_variant(variant: str) -> str:
|
||||
if not isinstance(variant, str) or not variant.strip():
|
||||
raise ValueError("MiniMax H3 model variant must be a non-empty string")
|
||||
normalized = variant.strip().lower()
|
||||
subfolders = {
|
||||
"fl2va": "FL2VA",
|
||||
"ref2va": "Ref2VA",
|
||||
}
|
||||
try:
|
||||
return subfolders[normalized]
|
||||
except KeyError as exc:
|
||||
raise ValueError(
|
||||
f"unsupported MiniMax H3 model variant {variant!r}; "
|
||||
f"supported: {sorted(subfolders)!r}"
|
||||
) from exc
|
||||
|
||||
def _load_config(self):
|
||||
model_variant = self.server_args.model_variant
|
||||
if model_variant is not None:
|
||||
semantic_subfolder = self.model_subfolder_for_variant(model_variant)
|
||||
explicit_subfolder = self.server_args.model_subfolder
|
||||
if (
|
||||
explicit_subfolder is not None
|
||||
and explicit_subfolder.strip().lower() != semantic_subfolder.lower()
|
||||
):
|
||||
raise ValueError(
|
||||
"MiniMax H3 --model-variant and --model-subfolder select "
|
||||
f"different weight partitions: variant={model_variant!r} maps to "
|
||||
f"{semantic_subfolder!r}, model_subfolder="
|
||||
f"{explicit_subfolder!r}"
|
||||
)
|
||||
self.server_args.model_subfolder = semantic_subfolder
|
||||
model_index = super()._load_config()
|
||||
self.release_metadata = MiniMaxH3ReleaseMetadata.from_model_index(model_index)
|
||||
if (
|
||||
model_variant is not None
|
||||
and self.release_metadata.partition != model_variant.strip().lower()
|
||||
):
|
||||
raise ValueError(
|
||||
"MiniMax H3 loaded checkpoint partition does not match "
|
||||
f"--model-variant {model_variant!r}"
|
||||
)
|
||||
return model_index
|
||||
|
||||
def validate_disagg_role(self, role: RoleType) -> None:
|
||||
if role != RoleType.MONOLITHIC:
|
||||
raise ValueError(
|
||||
"MiniMaxH3Pipeline only supports monolithic deployment; "
|
||||
f"disaggregation role {role.value!r} is not supported"
|
||||
)
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
|
||||
# Per-model sigma override from model_index.json; contract tests
|
||||
# construct the pipeline without model_path, hence the guard.
|
||||
release_metadata = getattr(self, "release_metadata", None)
|
||||
sigma_shift_scales = (
|
||||
release_metadata.sigma_shift_scales
|
||||
if release_metadata is not None
|
||||
else None
|
||||
)
|
||||
self.add_stage(InputValidationStage())
|
||||
if release_metadata is not None:
|
||||
self.add_stage(MiniMaxH3PartitionAdmissionStage(release_metadata))
|
||||
self.add_stage(
|
||||
MiniMaxH3TextEncodingStage(
|
||||
text_encoder=self.get_module("text_encoder"),
|
||||
tokenizer=self.get_module("tokenizer"),
|
||||
processor=self.get_module("processor"),
|
||||
)
|
||||
)
|
||||
self.add_stage(
|
||||
MiniMaxH3VisualEncodingStage(
|
||||
video_vae=self.get_module("video_vae"),
|
||||
vae_arch_config=server_args.pipeline_config.vae_config.arch_config,
|
||||
)
|
||||
)
|
||||
self.add_stage(
|
||||
MiniMaxH3AudioEncodingStage(
|
||||
audio_vae=self.get_module("audio_vae"),
|
||||
vae_arch_config=server_args.pipeline_config.audio_vae_config.arch_config,
|
||||
)
|
||||
)
|
||||
self.add_stage(MiniMaxH3LatentPreparationStage())
|
||||
self.add_stage(
|
||||
MiniMaxH3TimestepPreparationStage(
|
||||
sigma_shift_scales=sigma_shift_scales,
|
||||
)
|
||||
)
|
||||
self.add_stage(
|
||||
MiniMaxH3DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
pipeline=self,
|
||||
)
|
||||
)
|
||||
self.add_stage(
|
||||
MiniMaxH3DecodingStage(
|
||||
video_vae=self.get_module("video_vae"),
|
||||
audio_vae=self.get_module("audio_vae"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
EntryClass = MiniMaxH3Pipeline
|
||||
@@ -82,6 +82,7 @@ class ComposedPipelineBase(ABC):
|
||||
|
||||
# the name of the pipeline it associated with, in diffusers
|
||||
pipeline_name: str
|
||||
default_model_subfolder: str | None = None
|
||||
|
||||
def is_lora_effective(self):
|
||||
return False
|
||||
@@ -172,7 +173,32 @@ class ComposedPipelineBase(ABC):
|
||||
self.modules[module_name] = module
|
||||
|
||||
def _load_config(self) -> dict[str, Any]:
|
||||
model_path = maybe_download_model(self.model_path, force_diffusers_model=True)
|
||||
model_subfolder = self.server_args.model_subfolder
|
||||
if model_subfolder is None and not os.path.isfile(
|
||||
os.path.join(self.model_path, "model_index.json")
|
||||
):
|
||||
model_subfolder = self.default_model_subfolder
|
||||
|
||||
if model_subfolder is None:
|
||||
model_path = maybe_download_model(
|
||||
self.model_path, force_diffusers_model=True
|
||||
)
|
||||
else:
|
||||
model_subfolder = os.path.normpath(model_subfolder)
|
||||
if (
|
||||
os.path.isabs(model_subfolder)
|
||||
or model_subfolder == ".."
|
||||
or model_subfolder.startswith(f"..{os.sep}")
|
||||
):
|
||||
raise ValueError(
|
||||
f"model_subfolder must stay inside the model repository: {model_subfolder!r}"
|
||||
)
|
||||
model_root = maybe_download_model(
|
||||
self.model_path,
|
||||
allow_patterns=[f"{model_subfolder}/**"],
|
||||
)
|
||||
model_path = os.path.join(model_root, model_subfolder)
|
||||
|
||||
self.model_path = model_path
|
||||
logger.info("Model path: %s", model_path)
|
||||
config = verify_model_config_and_directory(model_path)
|
||||
@@ -444,13 +470,26 @@ class ComposedPipelineBase(ABC):
|
||||
component_load_specs: list[ComponentLoadSpec] = []
|
||||
|
||||
# enqueue only real weight loads (e.g., scheduler, tokenizer is excluded); skipped/provided modules keep old handling
|
||||
for index, (
|
||||
module_name,
|
||||
(
|
||||
transformers_or_diffusers,
|
||||
architecture,
|
||||
),
|
||||
) in enumerate(model_index.items()):
|
||||
for index, (module_name, component_spec) in enumerate(model_index.items()):
|
||||
# Diffusers uses JSON null for unavailable optional components.
|
||||
# Check before unpacking the normal [library, architecture] pair.
|
||||
if component_spec is None:
|
||||
logger.warning(
|
||||
"Module %s in model_index.json has null value, removing from required_config_modules",
|
||||
module_name,
|
||||
)
|
||||
if module_name in self.required_config_modules:
|
||||
self.required_config_modules.remove(module_name)
|
||||
continue
|
||||
if (
|
||||
not isinstance(component_spec, (list, tuple))
|
||||
or len(component_spec) != 2
|
||||
):
|
||||
raise ValueError(
|
||||
f"Module {module_name!r} in model_index.json must be null or "
|
||||
f"a [library, architecture] pair, got {component_spec!r}"
|
||||
)
|
||||
transformers_or_diffusers, architecture = component_spec
|
||||
if transformers_or_diffusers is None:
|
||||
logger.warning(
|
||||
"Module %s in model_index.json has null value, removing from required_config_modules",
|
||||
|
||||
@@ -161,13 +161,23 @@ class DecodingStage(PipelineStage):
|
||||
def scale_and_shift(self, latents: torch.Tensor, server_args):
|
||||
return scale_and_shift_latents(latents, server_args, self.vae)
|
||||
|
||||
def _get_vae_decode_fn(self, vae, server_args: ServerArgs):
|
||||
def _get_vae_decode_fn(
|
||||
self,
|
||||
vae,
|
||||
server_args: ServerArgs,
|
||||
*,
|
||||
decode_fn=None,
|
||||
compiled_callable: ActiveTargetCompiledCallable | None = None,
|
||||
):
|
||||
decode_fn = decode_fn or vae.decode
|
||||
if not server_args.enable_torch_compile or not isinstance(vae, nn.Module):
|
||||
return vae.decode
|
||||
return decode_fn
|
||||
|
||||
compiled_callable = compiled_callable or self._compiled_vae_decode
|
||||
|
||||
will_compile = (
|
||||
self._compiled_vae_decode.target_id != id(vae)
|
||||
or self._compiled_vae_decode.compiled_module is None
|
||||
compiled_callable.target_id != id(vae)
|
||||
or compiled_callable.compiled_module is None
|
||||
)
|
||||
if current_platform.is_npu():
|
||||
compile_kwargs = build_torch_compile_kwargs(mode=None)
|
||||
@@ -183,8 +193,8 @@ class DecodingStage(PipelineStage):
|
||||
if will_compile:
|
||||
logger.info("Compiling VAE decode with mode: %s", mode)
|
||||
|
||||
return self._compiled_vae_decode.get_or_compile(
|
||||
vae, vae.decode, compile_kwargs=compile_kwargs
|
||||
return compiled_callable.get_or_compile(
|
||||
vae, decode_fn, compile_kwargs=compile_kwargs
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
@@ -30,6 +30,7 @@ class StageDedupMixin:
|
||||
deduplicated_output_fields: ClassVar[tuple[str, ...]] = ()
|
||||
deduplicated_tensor_tree_output_fields: ClassVar[tuple[str, ...]] = ()
|
||||
deduplicated_deepcopy_output_fields: ClassVar[tuple[str, ...]] = ()
|
||||
deduplicated_extra_output_keys: ClassVar[tuple[str, ...]] = ()
|
||||
deduplicated_extra_tensor_tree_output_keys: ClassVar[tuple[str, ...]] = ()
|
||||
|
||||
def run_grouped_requests(
|
||||
@@ -60,6 +61,7 @@ class StageDedupMixin:
|
||||
cls.deduplicated_output_fields
|
||||
or cls.deduplicated_tensor_tree_output_fields
|
||||
or cls.deduplicated_deepcopy_output_fields
|
||||
or cls.deduplicated_extra_output_keys
|
||||
or cls.deduplicated_extra_tensor_tree_output_keys
|
||||
)
|
||||
|
||||
@@ -109,8 +111,8 @@ class StageDedupMixin:
|
||||
tensor references, which is the low-overhead path for read-only outputs
|
||||
such as embeddings. Tensor-tree fields recursively clone tensors.
|
||||
Deepcopy fields are for mutable request-local runtime objects, such as
|
||||
scheduler instances. Extra keys clone selected ``Req.extra`` entries
|
||||
without replacing the destination extra dict.
|
||||
scheduler instances. Extra output keys follow the same shallow-copy
|
||||
contract, while extra tensor-tree keys recursively clone tensors.
|
||||
"""
|
||||
for field in self.deduplicated_output_fields:
|
||||
setattr(dst, field, self.copy_stage_output(getattr(src, field)))
|
||||
@@ -118,6 +120,9 @@ class StageDedupMixin:
|
||||
setattr(dst, field, self.clone_tensor_tree(getattr(src, field)))
|
||||
for field in self.deduplicated_deepcopy_output_fields:
|
||||
setattr(dst, field, deepcopy(getattr(src, field)))
|
||||
for key in self.deduplicated_extra_output_keys:
|
||||
if key in src.extra:
|
||||
dst.extra[key] = self.copy_stage_output(src.extra[key])
|
||||
for key in self.deduplicated_extra_tensor_tree_output_keys:
|
||||
if key in src.extra:
|
||||
dst.extra[key] = self.clone_tensor_tree(src.extra[key])
|
||||
|
||||
@@ -76,6 +76,9 @@ from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_c
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
is_fsdp_managed_module,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
is_layerwise_offloaded_module,
|
||||
@@ -227,7 +230,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
num_attention_heads = (
|
||||
self.server_args.pipeline_config.dit_config.num_attention_heads
|
||||
)
|
||||
attn_head_size = hidden_size // num_attention_heads
|
||||
attn_head_size = getattr(
|
||||
self.server_args.pipeline_config.dit_config,
|
||||
"attention_head_dim",
|
||||
hidden_size // num_attention_heads,
|
||||
)
|
||||
|
||||
# torch compile
|
||||
# list of offloaded dit modules if torch compile is enabled. cleared after compile and warmup
|
||||
@@ -341,10 +348,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
not args.enable_torch_compile
|
||||
or not args.offload_during_compile
|
||||
or not args.warmup
|
||||
# a subclass with its own forward would never run the restore
|
||||
or type(self).forward is not DenoisingStage.forward
|
||||
or not self._owns_compile_warmup_lifecycle()
|
||||
or args.use_fsdp_inference
|
||||
or envs.SGLANG_CACHE_DIT_ENABLED
|
||||
or self._cache_dit_requested()
|
||||
or not isinstance(module, LayerwiseOffloadableModuleMixin)
|
||||
or is_layerwise_offloaded_module(module)
|
||||
):
|
||||
@@ -353,6 +359,15 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
if is_layerwise_offloaded_module(module):
|
||||
self._offloaded_dit_modules_for_compile.append(module)
|
||||
|
||||
def _owns_compile_warmup_lifecycle(self) -> bool:
|
||||
"""Whether ``forward`` enters ``_offload_for_torch_compile_warmup``.
|
||||
|
||||
Custom denoising loops opt in explicitly after wiring the same restore
|
||||
lifecycle. This keeps the safety guard without silently disabling the
|
||||
optimization solely because a model overrides ``forward``.
|
||||
"""
|
||||
return type(self).forward is DenoisingStage.forward
|
||||
|
||||
def _move_resident_components_for_warmup(self) -> list[torch.nn.Module]:
|
||||
"""Move resident non-DiT components off-device while the warmup
|
||||
denoising (the compile/autotune peak) runs; forward() moves them back."""
|
||||
@@ -365,6 +380,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
if (
|
||||
isinstance(module, torch.nn.Module)
|
||||
and id(module) not in dit_ids
|
||||
and not is_fsdp_managed_module(module)
|
||||
and not is_layerwise_offloaded_module(module)
|
||||
):
|
||||
param = next(module.parameters(), None)
|
||||
@@ -386,7 +402,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
self.server_args, "enable_torch_compile", False
|
||||
) or not isinstance(module, nn.Module):
|
||||
return
|
||||
if envs.SGLANG_CACHE_DIT_ENABLED and not self._cache_dit_enabled:
|
||||
if self._cache_dit_requested() and not self._cache_dit_enabled:
|
||||
logger.debug("Deferring torch.compile until cache-dit is enabled")
|
||||
return
|
||||
if self._torch_compile_registry.is_compiled(module):
|
||||
@@ -434,6 +450,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
def _cache_dit_dual_model_name(self) -> str:
|
||||
return "wan2.2"
|
||||
|
||||
def _cache_dit_requested(self) -> bool:
|
||||
return envs.SGLANG_CACHE_DIT_ENABLED
|
||||
|
||||
def _cache_dit_secondary_uses_primary_config(self) -> bool:
|
||||
return False
|
||||
|
||||
@@ -604,7 +623,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
|
||||
# Keep cache-dit disabled for ordinary warmup, but allow torch.compile
|
||||
# warmup to mount cache-dit before Dynamo traces the transformer.
|
||||
if not envs.SGLANG_CACHE_DIT_ENABLED:
|
||||
if not self._cache_dit_requested():
|
||||
return
|
||||
if batch.is_warmup and not getattr(
|
||||
self.server_args, "enable_torch_compile", False
|
||||
@@ -692,9 +711,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
logger.info(
|
||||
"cache-dit enabled on transformer (steps=%d, Fn=%d, Bn=%d, rdt=%.3f)",
|
||||
primary_num_steps,
|
||||
envs.SGLANG_CACHE_DIT_FN,
|
||||
envs.SGLANG_CACHE_DIT_BN,
|
||||
envs.SGLANG_CACHE_DIT_RDT,
|
||||
primary_config.Fn_compute_blocks,
|
||||
primary_config.Bn_compute_blocks,
|
||||
primary_config.residual_diff_threshold,
|
||||
)
|
||||
|
||||
self._cache_dit_enabled = True
|
||||
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""MiniMax H3-specific pipeline stages."""
|
||||
|
||||
from .stages.audio_encoding import MiniMaxH3AudioEncodingStage
|
||||
from .stages.decoding import MiniMaxH3DecodingStage
|
||||
from .stages.denoising import MiniMaxH3DenoisingStage
|
||||
from .stages.latent_preparation import MiniMaxH3LatentPreparationStage
|
||||
from .stages.text_encoding import MiniMaxH3TextEncodingStage
|
||||
from .stages.timestep_preparation import MiniMaxH3TimestepPreparationStage
|
||||
from .stages.visual_encoding import MiniMaxH3VisualEncodingStage
|
||||
|
||||
__all__ = [
|
||||
"MiniMaxH3AudioEncodingStage",
|
||||
"MiniMaxH3DecodingStage",
|
||||
"MiniMaxH3DenoisingStage",
|
||||
"MiniMaxH3LatentPreparationStage",
|
||||
"MiniMaxH3TextEncodingStage",
|
||||
"MiniMaxH3TimestepPreparationStage",
|
||||
"MiniMaxH3VisualEncodingStage",
|
||||
]
|
||||
+278
@@ -0,0 +1,278 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 keyframe target-canvas preparation.
|
||||
|
||||
Geometry behavior:
|
||||
- auto-aspect canvases delegate to the shared adaptive v2 shape resolver;
|
||||
- cover-crop: aspect-preserving max-scale LANCZOS resize + center crop,
|
||||
upscaling refused unless explicitly allowed.
|
||||
|
||||
Both the Qwen presentation (pixel_values) and the visual-condition tokenizer consume
|
||||
the SAME prepared canvas image, so preparation
|
||||
is cached per request in batch.extra.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
|
||||
MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
|
||||
)
|
||||
|
||||
MINIMAX_H3_CANVAS_MULTIPLE = 32
|
||||
MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY = "minimax_h3_prepared_keyframes"
|
||||
|
||||
|
||||
def minimax_h3_cover_crop_plan(
|
||||
*,
|
||||
source_width: int,
|
||||
source_height: int,
|
||||
target_width: int,
|
||||
target_height: int,
|
||||
allow_upscale: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""Deterministic aspect-preserving cover-crop transform."""
|
||||
if source_width <= 0 or source_height <= 0:
|
||||
raise ValueError("cover_crop requires positive source width/height")
|
||||
scale = max(
|
||||
target_width / float(source_width), target_height / float(source_height)
|
||||
)
|
||||
if scale > 1.0 and not allow_upscale:
|
||||
raise ValueError(
|
||||
"target_canvas cover_crop would upscale the source; set "
|
||||
f"allow_upscale=true (source={source_width}x{source_height}, "
|
||||
f"target={target_width}x{target_height})"
|
||||
)
|
||||
resized_width = max(target_width, int(round(source_width * scale)))
|
||||
resized_height = max(target_height, int(round(source_height * scale)))
|
||||
left = max(0, (resized_width - target_width) // 2)
|
||||
top = max(0, (resized_height - target_height) // 2)
|
||||
return {
|
||||
"scale": scale,
|
||||
"resized_size": (resized_width, resized_height),
|
||||
"crop_box": (left, top, left + target_width, top + target_height),
|
||||
}
|
||||
|
||||
|
||||
def minimax_h3_prepare_keyframe_canvas(
|
||||
image: Any,
|
||||
*,
|
||||
target_width: int,
|
||||
target_height: int,
|
||||
allow_upscale: bool = False,
|
||||
) -> Any:
|
||||
"""Prepare a PIL image onto the target canvas.
|
||||
|
||||
Identity (no resample) when the image already IS the canvas.
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
image = image.convert("RGB")
|
||||
if image.size == (target_width, target_height):
|
||||
return image
|
||||
plan = minimax_h3_cover_crop_plan(
|
||||
source_width=image.size[0],
|
||||
source_height=image.size[1],
|
||||
target_width=target_width,
|
||||
target_height=target_height,
|
||||
allow_upscale=allow_upscale,
|
||||
)
|
||||
resized = image.resize(plan["resized_size"], Image.Resampling.LANCZOS)
|
||||
return resized.crop(plan["crop_box"])
|
||||
|
||||
|
||||
def minimax_h3_stretch_keyframe_canvas(
|
||||
image: Any,
|
||||
*,
|
||||
target_width: int,
|
||||
target_height: int,
|
||||
) -> Any:
|
||||
"""Stretch the FL first frame directly onto the resolved target canvas."""
|
||||
|
||||
from PIL import Image
|
||||
|
||||
image = image.convert("RGB")
|
||||
if image.size == (target_width, target_height):
|
||||
return image
|
||||
return image.resize((target_width, target_height), Image.Resampling.LANCZOS)
|
||||
|
||||
|
||||
def _keyframe_materials(plan: Any) -> list[Any]:
|
||||
return [m for m in plan.materials if m.material_chain == "image.target_canvas"]
|
||||
|
||||
|
||||
def _keyframe_canvas_size(shape: Any) -> tuple[int, int]:
|
||||
geometry = str(shape["geometry"])
|
||||
if geometry != "resolved_v2":
|
||||
raise ValueError(
|
||||
"fl2va keyframe preparation requires pre-queue resolved_v2 "
|
||||
f"geometry, got {geometry!r}"
|
||||
)
|
||||
return int(shape["width"]), int(shape["height"])
|
||||
|
||||
|
||||
def _validate_keyframe_materials(plan: Any, keyframes: list[Any]) -> tuple[int, ...]:
|
||||
if str(plan.task) != "fl2va":
|
||||
raise ValueError("keyframe target-canvas materials require plan.task='fl2va'")
|
||||
semantic_indices = tuple(material.frame_index for material in keyframes)
|
||||
if semantic_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
|
||||
raise ValueError(
|
||||
"fl2va keyframes must use one of the ordered frame_index signatures "
|
||||
f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {semantic_indices!r}"
|
||||
)
|
||||
frame_count = plan.shape.get("frame_count")
|
||||
if isinstance(frame_count, bool) or not isinstance(frame_count, int):
|
||||
raise ValueError("fl2va keyframe preparation requires an integer frame_count")
|
||||
if frame_count <= 1:
|
||||
raise ValueError("fl2va keyframe preparation requires frame_count > 1")
|
||||
expected_pixels = tuple(
|
||||
frame_count - 1 if index == -1 else index for index in semantic_indices
|
||||
)
|
||||
resolved_pixels = tuple(material.resolved_frame_index for material in keyframes)
|
||||
if resolved_pixels != expected_pixels:
|
||||
raise ValueError(
|
||||
"fl2va keyframe resolved_frame_index values disagree with semantic "
|
||||
f"anchors: expected {expected_pixels!r}, got {resolved_pixels!r}"
|
||||
)
|
||||
return semantic_indices
|
||||
|
||||
|
||||
def minimax_h3_prepared_keyframes(batch: Any, plan: Any) -> dict[str, Any]:
|
||||
"""Resolve + prepare one or two fl2va keyframes once per request.
|
||||
|
||||
The target canvas is shared across keyframes and must already be frozen by
|
||||
the pre-queue probe/resolve hook.
|
||||
Top-level ``image`` / ``canvas_width`` / ``canvas_height`` keys mirror the
|
||||
first-keyframe payload for compatibility; per-keyframe entries live under
|
||||
``images``.
|
||||
"""
|
||||
keyframes = _keyframe_materials(plan)
|
||||
semantic_indices = _validate_keyframe_materials(plan, keyframes)
|
||||
cached = batch.extra.get(MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY)
|
||||
if cached is not None:
|
||||
cached_indices = tuple(cached.get("semantic_frame_indices") or ())
|
||||
cached_images = cached.get("images") or ()
|
||||
if cached_indices != semantic_indices or len(cached_images) != len(keyframes):
|
||||
raise ValueError(
|
||||
"cached fl2va keyframe preparation disagrees with the resolved plan"
|
||||
)
|
||||
return cached
|
||||
|
||||
canvas_w, canvas_h = _keyframe_canvas_size(plan.shape)
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.prequeue import (
|
||||
MINIMAX_H3_PROBE_FACTS_EXTRA_KEY,
|
||||
MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY,
|
||||
)
|
||||
|
||||
probe_facts = batch.extra.get(MINIMAX_H3_PROBE_FACTS_EXTRA_KEY)
|
||||
material_shapes = batch.extra.get(MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY)
|
||||
for material in keyframes:
|
||||
condition_index = int(material.condition_index)
|
||||
facts = (
|
||||
probe_facts.get(condition_index) if isinstance(probe_facts, dict) else None
|
||||
)
|
||||
material_shape = (
|
||||
material_shapes.get(condition_index)
|
||||
if isinstance(material_shapes, dict)
|
||||
else None
|
||||
)
|
||||
if not isinstance(facts, dict) or not isinstance(material_shape, dict):
|
||||
raise ValueError(
|
||||
"fl2va keyframe preparation requires cached pre-queue probe and "
|
||||
f"shape facts for conditions[{condition_index}]"
|
||||
)
|
||||
if (
|
||||
int(material_shape.get("width") or 0),
|
||||
int(material_shape.get("height") or 0),
|
||||
) != (canvas_w, canvas_h):
|
||||
raise ValueError(
|
||||
"fl2va keyframe material shape disagrees with the resolved target: "
|
||||
f"condition={condition_index}, material="
|
||||
f"{material_shape.get('width')}x{material_shape.get('height')}, "
|
||||
f"target={canvas_w}x{canvas_h}"
|
||||
)
|
||||
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
|
||||
minimax_h3_localize_material_uri,
|
||||
)
|
||||
|
||||
entries: list[dict[str, Any]] = []
|
||||
for keyframe_index, material in enumerate(keyframes):
|
||||
image_path = minimax_h3_localize_material_uri(
|
||||
batch,
|
||||
material.uri,
|
||||
condition_type=material.condition_type,
|
||||
condition_index=int(material.condition_index),
|
||||
)
|
||||
with Image.open(image_path) as source_image:
|
||||
image = ImageOps.exif_transpose(source_image)
|
||||
# A request's first semantic keyframe is its geometry anchor, including
|
||||
# the single-image last-frame-only signature [-1]. Only the second image in the
|
||||
# two-keyframe FL extension is a follower and receives cover-crop.
|
||||
prepared_image = (
|
||||
minimax_h3_stretch_keyframe_canvas(
|
||||
image,
|
||||
target_width=canvas_w,
|
||||
target_height=canvas_h,
|
||||
)
|
||||
if keyframe_index == 0
|
||||
else minimax_h3_prepare_keyframe_canvas(
|
||||
image,
|
||||
target_width=canvas_w,
|
||||
target_height=canvas_h,
|
||||
allow_upscale=True,
|
||||
)
|
||||
)
|
||||
entries.append(
|
||||
{
|
||||
"image": prepared_image,
|
||||
"canvas_width": canvas_w,
|
||||
"canvas_height": canvas_h,
|
||||
"condition_index": int(material.condition_index),
|
||||
"frame_index": (
|
||||
None if material.frame_index is None else int(material.frame_index)
|
||||
),
|
||||
"resolved_frame_index": (
|
||||
None
|
||||
if material.resolved_frame_index is None
|
||||
else int(material.resolved_frame_index)
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
payload = {
|
||||
"image": entries[0]["image"],
|
||||
"canvas_width": canvas_w,
|
||||
"canvas_height": canvas_h,
|
||||
"images": entries,
|
||||
"semantic_frame_indices": [
|
||||
int(item["frame_index"])
|
||||
for item in entries
|
||||
if item.get("frame_index") is not None
|
||||
],
|
||||
"pixel_frame_indices": [
|
||||
int(item["resolved_frame_index"])
|
||||
for item in entries
|
||||
if item.get("resolved_frame_index") is not None
|
||||
],
|
||||
"frame_count": (
|
||||
int(plan.shape["frame_count"])
|
||||
if plan.shape.get("frame_count") is not None
|
||||
else None
|
||||
),
|
||||
}
|
||||
batch.extra[MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY] = payload
|
||||
return payload
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_CANVAS_MULTIPLE",
|
||||
"MINIMAX_H3_PREPARED_KEYFRAMES_EXTRA_KEY",
|
||||
"minimax_h3_cover_crop_plan",
|
||||
"minimax_h3_prepare_keyframe_canvas",
|
||||
"minimax_h3_prepared_keyframes",
|
||||
"minimax_h3_stretch_keyframe_canvas",
|
||||
]
|
||||
+194
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 visual/audio condition-noise augmentation.
|
||||
|
||||
The request's condition timestep is applied to both the tensor value and the
|
||||
DiT timestep. Tokenizer artifacts remain clean
|
||||
and reusable; this module materializes the fixed noised anchors immediately
|
||||
before the denoise loop.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
|
||||
minimax_h3_patchify_video_latent,
|
||||
)
|
||||
|
||||
# Channel-major packed audio rows always carry a stereo layout.
|
||||
MINIMAX_H3_AUDIO_COND_CHANNELS = 2
|
||||
|
||||
|
||||
def minimax_h3_imgvid_cond_noise_aug_rows(
|
||||
clean_rows: torch.Tensor,
|
||||
*,
|
||||
condition_shapes: Sequence[Sequence[int]],
|
||||
target_latent_t: int,
|
||||
imgvid_cond_num_frames: int,
|
||||
seed: int,
|
||||
noise_aug: float,
|
||||
) -> torch.Tensor:
|
||||
"""Apply the imgvid-condition RF noise recipe to packed clean rows.
|
||||
|
||||
``condition_shapes`` contains ``(latent_t, latent_h, latent_w)`` in packed
|
||||
visual-condition order. A new CPU generator with the same row seed is
|
||||
created for every condition. Under the dependent-noise policy, each draw
|
||||
uses the target temporal length plus the template's imgvid-condition frame
|
||||
count, then slices the prefix matching the current condition.
|
||||
"""
|
||||
|
||||
noise_aug = float(noise_aug)
|
||||
if not 0.0 <= noise_aug <= 1.0:
|
||||
raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}")
|
||||
if noise_aug == 1.0:
|
||||
return clean_rows
|
||||
if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 96:
|
||||
raise ValueError(
|
||||
"clean imgvid condition rows must have shape [n, 96], got "
|
||||
f"{list(clean_rows.shape)}"
|
||||
)
|
||||
|
||||
target_latent_t = int(target_latent_t)
|
||||
imgvid_cond_num_frames = int(imgvid_cond_num_frames)
|
||||
if target_latent_t <= 0:
|
||||
raise ValueError(f"target_latent_t must be positive, got {target_latent_t}")
|
||||
if imgvid_cond_num_frames <= 0:
|
||||
raise ValueError(
|
||||
"imgvid_cond_num_frames must be positive when condition rows exist, "
|
||||
f"got {imgvid_cond_num_frames}"
|
||||
)
|
||||
|
||||
parsed_shapes: list[tuple[int, int, int]] = []
|
||||
expected_rows = 0
|
||||
for raw_shape in condition_shapes:
|
||||
if len(raw_shape) != 3:
|
||||
raise ValueError(
|
||||
"each imgvid condition shape must be (latent_t, latent_h, latent_w), "
|
||||
f"got {list(raw_shape)}"
|
||||
)
|
||||
latent_t, latent_h, latent_w = (int(value) for value in raw_shape)
|
||||
if latent_t <= 0 or latent_h <= 0 or latent_w <= 0:
|
||||
raise ValueError(
|
||||
f"imgvid condition shape must be positive, got {list(raw_shape)}"
|
||||
)
|
||||
if latent_h % 2 or latent_w % 2:
|
||||
raise ValueError(
|
||||
"imgvid condition spatial dimensions must be divisible by 2, "
|
||||
f"got {(latent_t, latent_h, latent_w)}"
|
||||
)
|
||||
parsed_shapes.append((latent_t, latent_h, latent_w))
|
||||
expected_rows += latent_t * (latent_h // 2) * (latent_w // 2)
|
||||
if not parsed_shapes:
|
||||
raise ValueError("condition_shapes must not be empty")
|
||||
if int(clean_rows.shape[0]) != expected_rows:
|
||||
raise ValueError(
|
||||
f"clean imgvid condition rows {int(clean_rows.shape[0])} != "
|
||||
f"shape-derived rows {expected_rows}"
|
||||
)
|
||||
|
||||
out: list[torch.Tensor] = []
|
||||
row_offset = 0
|
||||
timestep = torch.tensor(noise_aug, dtype=torch.float32, device=clean_rows.device)
|
||||
for latent_t, latent_h, latent_w in parsed_shapes:
|
||||
full_t = target_latent_t + imgvid_cond_num_frames
|
||||
if full_t < latent_t:
|
||||
raise ValueError(
|
||||
f"condition latent_t {latent_t} exceeds the noise draw "
|
||||
f"length {full_t}"
|
||||
)
|
||||
generator = torch.Generator(device="cpu").manual_seed(int(seed))
|
||||
noise = torch.randn(
|
||||
1,
|
||||
24,
|
||||
full_t,
|
||||
latent_h,
|
||||
latent_w,
|
||||
generator=generator,
|
||||
dtype=torch.float32,
|
||||
device="cpu",
|
||||
)[:, :, :latent_t]
|
||||
noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]).to(
|
||||
device=clean_rows.device, dtype=torch.float32
|
||||
)
|
||||
row_count = int(noise_rows.shape[0])
|
||||
clean_part = clean_rows[row_offset : row_offset + row_count].to(torch.float32)
|
||||
out.append(timestep * clean_part + (1.0 - timestep) * noise_rows)
|
||||
row_offset += row_count
|
||||
return (out[0] if len(out) == 1 else torch.cat(out, dim=0)).contiguous()
|
||||
|
||||
|
||||
def minimax_h3_audio_cond_noise_aug_rows(
|
||||
clean_rows: torch.Tensor,
|
||||
*,
|
||||
condition_audio_t: Sequence[int],
|
||||
seed: int,
|
||||
noise_aug: float,
|
||||
) -> torch.Tensor:
|
||||
"""Apply the audio-condition RF noise recipe to packed clean rows.
|
||||
|
||||
``condition_audio_t`` contains the latent T of each audio-bearing
|
||||
condition in canonical request order. Noise is drawn per condition
|
||||
element, with a fresh CPU generator seeded with ``seed + 1`` for every
|
||||
element. Consequently each condition restarts the
|
||||
same RNG stream; concatenating the rows and drawing once would be
|
||||
numerically different for ordered multi-reference requests.
|
||||
|
||||
The mix is intentionally evaluated on CPU in fp32 before the packed rows
|
||||
are transferred to the DiT device.
|
||||
"""
|
||||
|
||||
noise_aug = float(noise_aug)
|
||||
if not 0.0 <= noise_aug <= 1.0:
|
||||
raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}")
|
||||
if noise_aug == 1.0:
|
||||
return clean_rows
|
||||
if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 32:
|
||||
raise ValueError(
|
||||
"clean audio condition rows must have shape [n, 32], got "
|
||||
f"{list(clean_rows.shape)}"
|
||||
)
|
||||
|
||||
audio_channels = MINIMAX_H3_AUDIO_COND_CHANNELS
|
||||
parsed_audio_t = [int(value) for value in condition_audio_t]
|
||||
if not parsed_audio_t:
|
||||
raise ValueError("condition_audio_t must not be empty")
|
||||
if any(value <= 0 for value in parsed_audio_t):
|
||||
raise ValueError(
|
||||
f"condition audio latent lengths must be positive, got {parsed_audio_t}"
|
||||
)
|
||||
expected_rows = audio_channels * sum(parsed_audio_t)
|
||||
if int(clean_rows.shape[0]) != expected_rows:
|
||||
raise ValueError(
|
||||
f"clean audio condition rows {int(clean_rows.shape[0])} != "
|
||||
f"shape-derived rows {expected_rows}"
|
||||
)
|
||||
|
||||
out: list[torch.Tensor] = []
|
||||
row_offset = 0
|
||||
timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu")
|
||||
for audio_t in parsed_audio_t:
|
||||
row_count = audio_channels * audio_t
|
||||
clean_part = (
|
||||
clean_rows[row_offset : row_offset + row_count]
|
||||
.detach()
|
||||
.to(device="cpu", dtype=torch.float32)
|
||||
)
|
||||
generator = torch.Generator(device="cpu").manual_seed(int(seed) + 1)
|
||||
noise = torch.randn(
|
||||
clean_part.shape,
|
||||
generator=generator,
|
||||
dtype=torch.float32,
|
||||
device="cpu",
|
||||
)
|
||||
out.append(timestep * clean_part + (1.0 - timestep) * noise)
|
||||
row_offset += row_count
|
||||
rows = out[0] if len(out) == 1 else torch.cat(out, dim=0)
|
||||
return rows.to(device=clean_rows.device, dtype=torch.float32).contiguous()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"minimax_h3_audio_cond_noise_aug_rows",
|
||||
"minimax_h3_imgvid_cond_noise_aug_rows",
|
||||
]
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
# Direct-encode text embeddings: {"positive":
|
||||
# {"hidden_states": Tensor[text_len, 5120] bf16 cpu, "text_len": int}}
|
||||
MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY = "minimax_h3_text_embeddings"
|
||||
# Direct keyframe encode: {"rows": Tensor[n_rows, 96] fp32 cpu,
|
||||
# "latent_h": int, "latent_w": int, "canvas_height": int,
|
||||
# "canvas_width": int, "keyframes": [...],
|
||||
# "semantic_frame_indices": [...], "pixel_frame_indices": [...]}
|
||||
MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY = "minimax_h3_keyframe_cond_rows"
|
||||
# Direct sigma schedules: {"video": [float], "audio": [float]}
|
||||
MINIMAX_H3_SIGMAS_EXTRA_KEY = "minimax_h3_sigmas"
|
||||
# Direct denoise state: {"initial_video_rows", "initial_audio_rows",
|
||||
# "latent_t", "latent_h", "latent_w", "audio_t"}
|
||||
MINIMAX_H3_DENOISE_STATE_EXTRA_KEY = "minimax_h3_denoise_state"
|
||||
# ref2va direct reference encodes.
|
||||
MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY = "minimax_h3_reference_image_rows"
|
||||
MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY = "minimax_h3_reference_audio_rows"
|
||||
MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY = "minimax_h3_reference_video_rows"
|
||||
MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY = "minimax_h3_prepared_reference_video"
|
||||
|
||||
MINIMAX_H3_SUPPORTED_FPS = 24
|
||||
MINIMAX_H3_MIN_DURATION_SECONDS = 4.0
|
||||
MINIMAX_H3_MAX_DURATION_SECONDS = 15.0
|
||||
|
||||
# The distilled checkpoint has exactly one positive denoise branch.
|
||||
MINIMAX_H3_DEFAULT_BRANCHES: tuple = ({"name": "cond_1"},)
|
||||
|
||||
# Audited 4xH200 T2VA profiles. The tuple is
|
||||
# (warmup steps, residual-difference threshold, max consecutive cached steps).
|
||||
MINIMAX_H3_QUALITY_PROFILES: dict[str, tuple[int, float, int] | None] = {
|
||||
"lossless": None,
|
||||
"high": (4, 0.04, 1),
|
||||
"medium": (4, 0.12, 3),
|
||||
"low": (4, 0.24, 3),
|
||||
}
|
||||
+519
@@ -0,0 +1,519 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 cfg-distilled full denoise loop.
|
||||
|
||||
Per step, the positive presentation is forwarded exactly once. Video and audio
|
||||
target rows chain through the Euler-eta0 update while visual and audio condition
|
||||
rows stay pinned to their noised step-0 anchors.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import AbstractContextManager, nullcontext
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
|
||||
MINIMAX_H3_ADALN_MODALITY_NUM,
|
||||
)
|
||||
|
||||
MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
|
||||
# ref2va audio reference anchor timestep
|
||||
MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0
|
||||
# Packed row widths: video rows are [1,2,2]-patchified 24-channel latents
|
||||
# (24 * 1 * 2 * 2 = 96); audio rows carry the 32-dim audio latent.
|
||||
MINIMAX_H3_VIDEO_ROW_WIDTH = 96
|
||||
MINIMAX_H3_AUDIO_ROW_WIDTH = 32
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _minimax_h3_update_target_rows_(
|
||||
state: torch.Tensor,
|
||||
velocity: torch.Tensor,
|
||||
*,
|
||||
sigma_t: torch.Tensor,
|
||||
sigma_curr: float,
|
||||
sigma_ratio: torch.Tensor,
|
||||
one_minus_sigma_ratio: torch.Tensor,
|
||||
denoised_scratch: torch.Tensor,
|
||||
) -> None:
|
||||
torch.mul(sigma_t, velocity, out=denoised_scratch)
|
||||
torch.add(state, denoised_scratch, out=denoised_scratch)
|
||||
if sigma_curr == 0.0:
|
||||
return
|
||||
torch.mul(one_minus_sigma_ratio, denoised_scratch, out=velocity)
|
||||
torch.mul(sigma_ratio, state, out=state)
|
||||
torch.add(state, velocity, out=state)
|
||||
|
||||
|
||||
def _ulysses_ctx() -> tuple[int, int]:
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_ulysses_parallel_rank,
|
||||
get_ulysses_parallel_world_size,
|
||||
model_parallel_is_initialized,
|
||||
)
|
||||
|
||||
if not model_parallel_is_initialized():
|
||||
return 1, 0
|
||||
return get_ulysses_parallel_world_size(), get_ulysses_parallel_rank()
|
||||
|
||||
|
||||
def _build_local_embedding_layout(
|
||||
*,
|
||||
seq_len: int,
|
||||
text_pos: torch.Tensor,
|
||||
img_pos: torch.Tensor,
|
||||
audio_pos: torch.Tensor,
|
||||
world_size: int,
|
||||
rank: int,
|
||||
device: torch.device,
|
||||
) -> dict[str, torch.Tensor | int]:
|
||||
if seq_len % world_size:
|
||||
raise ValueError(
|
||||
f"packed seq_len {seq_len} not divisible by Ulysses world size "
|
||||
f"{world_size}"
|
||||
)
|
||||
local_seq_len = seq_len // world_size
|
||||
row_start = rank * local_seq_len
|
||||
row_stop = row_start + local_seq_len
|
||||
|
||||
def local_ids(pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
source_ids = torch.nonzero(
|
||||
(pos >= row_start) & (pos < row_stop),
|
||||
as_tuple=False,
|
||||
).view(-1)
|
||||
return source_ids.to(device), pos.index_select(0, source_ids).to(device)
|
||||
|
||||
text_source_start = min(row_start, int(text_pos.shape[0]))
|
||||
text_source_stop = min(row_stop, int(text_pos.shape[0]))
|
||||
_, img_global_ids = local_ids(img_pos)
|
||||
_, audio_global_ids = local_ids(audio_pos)
|
||||
return {
|
||||
"text_source_start": text_source_start,
|
||||
"text_source_stop": text_source_stop,
|
||||
"img_global_ids": img_global_ids,
|
||||
"img_row_ids": img_global_ids - row_start,
|
||||
"audio_global_ids": audio_global_ids,
|
||||
"audio_row_ids": audio_global_ids - row_start,
|
||||
}
|
||||
|
||||
|
||||
class MiniMaxH3DenoiseBranch:
|
||||
"""Static per-branch state: packed layout + fixed forward kwargs.
|
||||
|
||||
`packed` is a minimax_h3_packed_sequence(...) result (or equivalent layout
|
||||
dict); `text_embeddings` is the branch's [text_len, 5120] hidden states;
|
||||
`token_tags` must already carry any fl2va vision-span overrides.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
packed: dict[str, torch.Tensor],
|
||||
text_embeddings: torch.Tensor,
|
||||
token_tags: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
seq_len = int(packed["seq_len"])
|
||||
self.seq_len = seq_len
|
||||
self.img_pos = packed["img_pos"].view(-1).to(torch.long)
|
||||
self.audio_pos = packed["audio_pos"].view(-1).to(torch.long)
|
||||
self.update_mask = packed["update_mask"].view(-1).to(torch.bool)
|
||||
# ref2va: audio_pos may include reference-audio anchor rows
|
||||
# (audio_update_mask False); absent means all rows are targets.
|
||||
if "audio_update_mask" in packed:
|
||||
self.audio_update_mask = packed["audio_update_mask"].view(-1).to(torch.bool)
|
||||
else:
|
||||
self.audio_update_mask = torch.ones(
|
||||
self.audio_pos.shape[0], dtype=torch.bool
|
||||
)
|
||||
# Packed H3 layouts place reference rows before the generated suffix
|
||||
# for both modalities. Keep the suffix boundary once so the denoise
|
||||
# hot path can update contiguous views instead of gathering and
|
||||
# scattering the same static index tensors on every step.
|
||||
self.video_target_start = int((~self.update_mask).sum())
|
||||
self.audio_target_start = int((~self.audio_update_mask).sum())
|
||||
self.video_target_slice = slice(self.video_target_start, None)
|
||||
self.audio_target_slice = slice(self.audio_target_start, None)
|
||||
text_pos = packed["text_pos"].view(-1).to(torch.long)
|
||||
text_len = int(text_pos.shape[0])
|
||||
if list(text_embeddings.shape)[0] != text_len:
|
||||
raise ValueError(
|
||||
f"text_embeddings rows {list(text_embeddings.shape)} != "
|
||||
f"packed text_len {text_len}"
|
||||
)
|
||||
if int(token_tags.view(-1).shape[0]) != seq_len:
|
||||
raise ValueError(
|
||||
f"token_tags length {int(token_tags.view(-1).shape[0])} != "
|
||||
f"seq_len {seq_len}"
|
||||
)
|
||||
cu = packed["cu_seqlens"].to(torch.int32)
|
||||
self.img_pos_dev = self.img_pos.to(device)
|
||||
self.audio_pos_dev = self.audio_pos.to(device)
|
||||
self.update_mask_dev = self.update_mask.to(device)
|
||||
self.audio_update_mask_dev = self.audio_update_mask.to(device)
|
||||
# Resolve the remaining step-static packed-sequence and anchor row
|
||||
# sets once, keeping nonzero-driven work out of the hot loop.
|
||||
self.img_cond_seq_idx = self.img_pos_dev[~self.update_mask_dev]
|
||||
self.img_target_seq_idx = self.img_pos_dev[self.update_mask_dev]
|
||||
self.audio_target_seq_idx = self.audio_pos_dev[self.audio_update_mask_dev]
|
||||
self.audio_ref_seq_idx = self.audio_pos_dev[~self.audio_update_mask_dev]
|
||||
self.cond_row_idx = torch.nonzero(~self.update_mask_dev).view(-1)
|
||||
self.audio_ref_row_idx = torch.nonzero(~self.audio_update_mask_dev).view(-1)
|
||||
# rows that keep the video timestep each step: text, padding, and
|
||||
# video target rows — everything the three overwrite sets do not cover
|
||||
self.n_video_timestep_rows = (
|
||||
seq_len
|
||||
- int(self.img_cond_seq_idx.numel())
|
||||
- int(self.audio_target_seq_idx.numel())
|
||||
- int(self.audio_ref_seq_idx.numel())
|
||||
)
|
||||
# persistent packed-row buffers; every img/audio position is fully
|
||||
# rewritten by index_copy_ on the first forward_kwargs() call, then
|
||||
# only the target-row subset each step after (condition/reference
|
||||
# rows never change post-priming -- see forward_kwargs).
|
||||
self._x_buffer_primed = False
|
||||
self.x_buffer = torch.zeros(
|
||||
1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device
|
||||
)
|
||||
self.audio_x_buffer = torch.zeros(
|
||||
1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device
|
||||
)
|
||||
text_pos_dev = text_pos.to(device)
|
||||
ulysses_world_size, ulysses_rank = _ulysses_ctx()
|
||||
token_tags_host = token_tags.view(-1).to(dtype=torch.long)
|
||||
local_seq_len = seq_len // ulysses_world_size
|
||||
local_row_start = ulysses_rank * local_seq_len
|
||||
local_row_stop = local_row_start + local_seq_len
|
||||
self.local_row_slice = slice(local_row_start, local_row_stop)
|
||||
self.block_token_tags = (
|
||||
token_tags_host[local_row_start:local_row_stop].clamp(min=0).to(device)
|
||||
)
|
||||
self.static_kwargs: dict[str, Any] = {
|
||||
# Cast the fp64 position grid on the host. MiniMaxH3Rope casts it to
|
||||
# fp32 as its first op anyway, so the values are identical; doing the
|
||||
# cast on CPU also avoids platforms that cannot execute fp64 on
|
||||
# device (e.g. Iluvatar CoreX returns zeros for device fp64).
|
||||
"img_position_ids": packed["img_position_ids"][None]
|
||||
.to(torch.float32)
|
||||
.to(device),
|
||||
"update_mask": self.update_mask_dev,
|
||||
"block_token_tags": self.block_token_tags,
|
||||
"skip_mask_out_condition": True,
|
||||
"prompt_embeds": text_embeddings.to(device),
|
||||
"img_pos_info": {"position_ids": self.img_pos_dev},
|
||||
"audio_pos_info": {"position_ids": self.audio_pos_dev},
|
||||
"text_pos_info": {"position_ids": text_pos_dev},
|
||||
"img_pos_for_infer_output_info": {"position_ids": self.img_target_seq_idx},
|
||||
"local_embedding_layout": _build_local_embedding_layout(
|
||||
seq_len=seq_len,
|
||||
text_pos=text_pos,
|
||||
img_pos=self.img_pos,
|
||||
audio_pos=self.audio_pos,
|
||||
world_size=ulysses_world_size,
|
||||
rank=ulysses_rank,
|
||||
device=device,
|
||||
),
|
||||
"packed_seq_params": {
|
||||
"cu_seqlens_q": cu.to(device),
|
||||
"cu_seqlens_q_host": tuple(int(value) for value in cu.tolist()),
|
||||
"max_seqlen_q": int(cu[1]),
|
||||
},
|
||||
"refiner_packed_seq_params": {
|
||||
"cu_seqlens_q": torch.tensor(
|
||||
[0, text_len, text_len], dtype=torch.int32, device=device
|
||||
),
|
||||
"cu_seqlens_q_host": (0, text_len, text_len),
|
||||
"max_seqlen_q": text_len,
|
||||
},
|
||||
}
|
||||
|
||||
def forward_kwargs(
|
||||
self,
|
||||
*,
|
||||
video_rows: torch.Tensor,
|
||||
audio_rows: torch.Tensor,
|
||||
step_timesteps: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
|
||||
) -> dict[str, Any]:
|
||||
x = self.x_buffer
|
||||
audio_x = self.audio_x_buffer
|
||||
if not self._x_buffer_primed:
|
||||
# First step: condition/reference rows have just been pinned into
|
||||
# video_rows/audio_rows (see minimax_h3_denoise_loop) and never
|
||||
# change again, so this is the only step that needs the full
|
||||
# img/audio extent written into the persistent buffers.
|
||||
x[0].index_copy_(0, self.img_pos_dev, video_rows)
|
||||
audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows)
|
||||
self._x_buffer_primed = True
|
||||
else:
|
||||
# Later steps: only the target-row subset changed since the
|
||||
# buffers were primed; rewriting condition/reference rows again
|
||||
# would just copy the same bytes already sitting there.
|
||||
x[0].index_copy_(
|
||||
0, self.img_target_seq_idx, video_rows[self.video_target_slice]
|
||||
)
|
||||
audio_x[0].index_copy_(
|
||||
0, self.audio_target_seq_idx, audio_rows[self.audio_target_slice]
|
||||
)
|
||||
unique_timesteps, inverse_indices, block_combined_indices = step_timesteps
|
||||
return {
|
||||
**self.static_kwargs,
|
||||
"x": x,
|
||||
"audio_x": audio_x,
|
||||
"unique_timesteps": unique_timesteps,
|
||||
"inverse_indices": inverse_indices,
|
||||
"block_combined_indices": block_combined_indices,
|
||||
}
|
||||
|
||||
def _expand_step_timesteps(
|
||||
self,
|
||||
*,
|
||||
t_video: float,
|
||||
t_audio: float,
|
||||
imgvid_cond_timestep: float,
|
||||
audio_ref_cond_timestep: float,
|
||||
inverse_indices_by_pattern: dict[tuple[int, ...], torch.Tensor],
|
||||
block_combined_by_pattern: dict[tuple[int, ...], torch.Tensor],
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build step-local timestep and AdaLN index tensors.
|
||||
|
||||
Packed-sequence timestep semantics: non-media rows (text and padding)
|
||||
inherit the current video timestep, condition rows pin their noise-aug
|
||||
timesteps. The row->timestep layout is step-static, so instead of
|
||||
materializing the full timestep tensor and paying a device-syncing
|
||||
torch.unique per step, the at-most-four candidate values are deduped
|
||||
in fp32 on the host — torch.unique on the candidate tensor keeps exact
|
||||
fp32 collision semantics — and inverse indices are index_fill'ed from
|
||||
the static position sets.
|
||||
"""
|
||||
candidates: list[float] = []
|
||||
fill_groups: list[tuple[torch.Tensor, int]] = []
|
||||
base_slot = -1
|
||||
if self.n_video_timestep_rows > 0:
|
||||
base_slot = len(candidates)
|
||||
candidates.append(float(t_video))
|
||||
for seq_idx, value in (
|
||||
(self.img_cond_seq_idx, imgvid_cond_timestep),
|
||||
(self.audio_target_seq_idx, t_audio),
|
||||
(self.audio_ref_seq_idx, audio_ref_cond_timestep),
|
||||
):
|
||||
if seq_idx.numel() > 0:
|
||||
fill_groups.append((seq_idx, len(candidates)))
|
||||
candidates.append(float(value))
|
||||
unique_cpu, slot_to_unique = torch.unique(
|
||||
torch.tensor(candidates, dtype=torch.float32),
|
||||
sorted=True,
|
||||
return_inverse=True,
|
||||
)
|
||||
device = self.img_pos_dev.device
|
||||
base_index = int(slot_to_unique[base_slot]) if base_slot >= 0 else 0
|
||||
pattern = tuple(slot_to_unique.tolist())
|
||||
inverse_indices = inverse_indices_by_pattern.get(pattern)
|
||||
if inverse_indices is None:
|
||||
inverse_indices = torch.full(
|
||||
(self.seq_len,), base_index, dtype=torch.long, device=device
|
||||
)
|
||||
for seq_idx, slot in fill_groups:
|
||||
inverse_indices.index_fill_(0, seq_idx, int(slot_to_unique[slot]))
|
||||
inverse_indices_by_pattern[pattern] = inverse_indices
|
||||
block_combined = block_combined_by_pattern.get(pattern)
|
||||
if block_combined is None:
|
||||
block_combined = torch.add(
|
||||
self.block_token_tags,
|
||||
inverse_indices[self.local_row_slice],
|
||||
alpha=MINIMAX_H3_ADALN_MODALITY_NUM,
|
||||
)
|
||||
block_combined_by_pattern[pattern] = block_combined
|
||||
return unique_cpu.to(device), inverse_indices, block_combined
|
||||
|
||||
def prepare_timestep_plan(
|
||||
self,
|
||||
*,
|
||||
video_timesteps: list[float],
|
||||
audio_timesteps: list[float],
|
||||
imgvid_cond_noise_aug: float,
|
||||
audio_ref_cond_noise_aug: float,
|
||||
) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
|
||||
"""Stage every step's packed timestep state before denoising."""
|
||||
if len(video_timesteps) != len(audio_timesteps):
|
||||
raise ValueError("video/audio timestep plans must have equal length")
|
||||
inverse_indices_by_pattern: dict[tuple[int, ...], torch.Tensor] = {}
|
||||
block_combined_by_pattern: dict[tuple[int, ...], torch.Tensor] = {}
|
||||
return [
|
||||
self._expand_step_timesteps(
|
||||
t_video=t_video,
|
||||
t_audio=t_audio,
|
||||
imgvid_cond_timestep=max(t_video, imgvid_cond_noise_aug),
|
||||
audio_ref_cond_timestep=max(t_audio, audio_ref_cond_noise_aug),
|
||||
inverse_indices_by_pattern=inverse_indices_by_pattern,
|
||||
block_combined_by_pattern=block_combined_by_pattern,
|
||||
)
|
||||
for t_video, t_audio in zip(video_timesteps, audio_timesteps)
|
||||
]
|
||||
|
||||
|
||||
def minimax_h3_denoise_loop(
|
||||
*,
|
||||
model: Any,
|
||||
model_forward: (
|
||||
Callable[[Any, dict[str, Any], int], tuple[torch.Tensor, torch.Tensor]] | None
|
||||
) = None,
|
||||
positive: MiniMaxH3DenoiseBranch,
|
||||
initial_video_rows: torch.Tensor,
|
||||
initial_audio_rows: torch.Tensor,
|
||||
keyframe_cond_rows: torch.Tensor | None,
|
||||
audio_ref_rows: torch.Tensor | None = None,
|
||||
sigmas_video: list[float],
|
||||
sigmas_audio: list[float],
|
||||
device: torch.device,
|
||||
imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP,
|
||||
audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
|
||||
on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
|
||||
step_profiler: Callable[[int], AbstractContextManager] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run the full denoise loop; returns final (video_rows, audio_rows).
|
||||
|
||||
``initial_video_rows`` covers all image rows of the positive layout. For a
|
||||
conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to
|
||||
pin those rows across every step. The model's raw positive velocity is the
|
||||
update signal; MiniMax H3 only supports cfg-distilled checkpoints.
|
||||
``model_forward`` is the native-stage hook for residency/BCG runners and
|
||||
receives the zero-based loop step; the default keeps this helper
|
||||
independently testable with a plain callable.
|
||||
"""
|
||||
if len(sigmas_video) != len(sigmas_audio):
|
||||
raise ValueError("video/audio sigma schedules must have equal length")
|
||||
if len(sigmas_video) < 2:
|
||||
raise ValueError("sigma schedules need at least 2 entries")
|
||||
n_cond = positive.video_target_start
|
||||
if keyframe_cond_rows is None:
|
||||
if n_cond != 0:
|
||||
raise ValueError(
|
||||
f"layout has {n_cond} cond rows but keyframe_cond_rows is None"
|
||||
)
|
||||
else:
|
||||
if int(keyframe_cond_rows.shape[0]) != n_cond:
|
||||
raise ValueError(
|
||||
f"keyframe_cond_rows {int(keyframe_cond_rows.shape[0])} != "
|
||||
f"layout cond rows {n_cond}"
|
||||
)
|
||||
video_rows = initial_video_rows.to(device=device, dtype=torch.float32, copy=True)
|
||||
audio_rows = initial_audio_rows.to(device=device, dtype=torch.float32, copy=True)
|
||||
if int(video_rows.shape[0]) != int(positive.img_pos.shape[0]):
|
||||
raise ValueError(
|
||||
f"initial video rows {int(video_rows.shape[0])} != positive layout "
|
||||
f"rows {int(positive.img_pos.shape[0])}"
|
||||
)
|
||||
if int(audio_rows.shape[0]) != int(positive.audio_pos.shape[0]):
|
||||
raise ValueError(
|
||||
f"initial audio rows {int(audio_rows.shape[0])} != positive layout "
|
||||
f"rows {int(positive.audio_pos.shape[0])}"
|
||||
)
|
||||
n_audio_ref = positive.audio_target_start
|
||||
if audio_ref_rows is None:
|
||||
if n_audio_ref != 0:
|
||||
raise ValueError(
|
||||
f"layout has {n_audio_ref} audio ref rows but audio_ref_rows is None"
|
||||
)
|
||||
audio_anchor = None
|
||||
else:
|
||||
if int(audio_ref_rows.shape[0]) != n_audio_ref:
|
||||
raise ValueError(
|
||||
f"audio_ref_rows {int(audio_ref_rows.shape[0])} != layout "
|
||||
f"audio ref rows {n_audio_ref}"
|
||||
)
|
||||
audio_anchor = audio_ref_rows.to(device=device, dtype=torch.float32)
|
||||
cond_anchor = (
|
||||
keyframe_cond_rows.to(device=device, dtype=torch.float32)
|
||||
if keyframe_cond_rows is not None
|
||||
else None
|
||||
)
|
||||
if cond_anchor is not None:
|
||||
video_rows.index_copy_(0, positive.cond_row_idx, cond_anchor)
|
||||
if audio_anchor is not None:
|
||||
audio_rows.index_copy_(0, positive.audio_ref_row_idx, audio_anchor)
|
||||
|
||||
num_steps = len(sigmas_video) - 1
|
||||
video_target_slice = positive.video_target_slice
|
||||
audio_target_slice = positive.audio_target_slice
|
||||
video_timesteps = [1.0 - sigma for sigma in sigmas_video[:-1]]
|
||||
audio_timesteps = [1.0 - sigma for sigma in sigmas_audio[:-1]]
|
||||
# One H2D copy per schedule, preserving the previous Python-float
|
||||
# subtraction followed by fp32 conversion.
|
||||
video_step_t = torch.tensor(video_timesteps, dtype=torch.float32, device=device)
|
||||
audio_step_t = torch.tensor(audio_timesteps, dtype=torch.float32, device=device)
|
||||
timestep_plan = positive.prepare_timestep_plan(
|
||||
video_timesteps=video_timesteps,
|
||||
audio_timesteps=audio_timesteps,
|
||||
imgvid_cond_noise_aug=float(imgvid_cond_noise_aug_for_inference),
|
||||
audio_ref_cond_noise_aug=float(audio_cond_noise_aug_for_inference),
|
||||
)
|
||||
# match the scheduler's device-fp32 math once, then reuse one denoised
|
||||
# scratch per modality instead of allocating intermediates every step
|
||||
video_sigmas = torch.tensor(sigmas_video, dtype=torch.float32, device=device)
|
||||
audio_sigmas = torch.tensor(sigmas_audio, dtype=torch.float32, device=device)
|
||||
video_sigma_ratios = video_sigmas[1:] / video_sigmas[:-1]
|
||||
audio_sigma_ratios = audio_sigmas[1:] / audio_sigmas[:-1]
|
||||
video_sigma_t = 1.0 - video_step_t
|
||||
audio_sigma_t = 1.0 - audio_step_t
|
||||
video_one_minus_sigma_ratios = 1.0 - video_sigma_ratios
|
||||
audio_one_minus_sigma_ratios = 1.0 - audio_sigma_ratios
|
||||
video_denoised_scratch = torch.empty_like(video_rows[video_target_slice])
|
||||
audio_denoised_scratch = torch.empty_like(audio_rows[audio_target_slice])
|
||||
for step in range(num_steps):
|
||||
step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
|
||||
with step_cm:
|
||||
s_v = sigmas_video[step]
|
||||
s_a = sigmas_audio[step]
|
||||
|
||||
fk = positive.forward_kwargs(
|
||||
video_rows=video_rows,
|
||||
audio_rows=audio_rows,
|
||||
step_timesteps=timestep_plan[step],
|
||||
)
|
||||
with torch.inference_mode():
|
||||
if model_forward is None:
|
||||
v_video, v_audio = model(**fk)
|
||||
else:
|
||||
v_video, v_audio = model_forward(model, fk, step)
|
||||
# The model outputs are inference tensors. Keep their disposable
|
||||
# fp32 velocity updates in the same context so ``out=velocity``
|
||||
# can reuse the output storage without an extra clone.
|
||||
mv_video_t = v_video.float()
|
||||
mv_audio_t = v_audio[audio_target_slice].float()
|
||||
|
||||
video_target = video_rows[video_target_slice]
|
||||
_minimax_h3_update_target_rows_(
|
||||
video_target,
|
||||
mv_video_t,
|
||||
sigma_t=video_sigma_t[step],
|
||||
sigma_curr=s_v,
|
||||
sigma_ratio=video_sigma_ratios[step],
|
||||
one_minus_sigma_ratio=video_one_minus_sigma_ratios[step],
|
||||
denoised_scratch=video_denoised_scratch,
|
||||
)
|
||||
|
||||
audio_target = audio_rows[audio_target_slice]
|
||||
_minimax_h3_update_target_rows_(
|
||||
audio_target,
|
||||
mv_audio_t,
|
||||
sigma_t=audio_sigma_t[step],
|
||||
sigma_curr=s_a,
|
||||
sigma_ratio=audio_sigma_ratios[step],
|
||||
one_minus_sigma_ratio=audio_one_minus_sigma_ratios[step],
|
||||
denoised_scratch=audio_denoised_scratch,
|
||||
)
|
||||
if on_step is not None:
|
||||
on_step(step, video_rows, audio_rows)
|
||||
|
||||
return video_rows, audio_rows
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_AUDIO_REF_COND_TIMESTEP",
|
||||
"MINIMAX_H3_AUDIO_ROW_WIDTH",
|
||||
"MINIMAX_H3_IMGVID_COND_TIMESTEP",
|
||||
"MINIMAX_H3_VIDEO_ROW_WIDTH",
|
||||
"MiniMaxH3DenoiseBranch",
|
||||
"minimax_h3_denoise_loop",
|
||||
]
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 keyframe (imgvid) condition encoding.
|
||||
|
||||
Condition anchor row recipe:
|
||||
|
||||
- ``video_vae.encode_images(PIL, use_fp16_latent=True)`` under a
|
||||
scoped seed-42 RNG fork — the DiagonalGaussian is SAMPLED
|
||||
(use_mean=False) with seed 42, so the seed is part of
|
||||
the contract, not a convenience
|
||||
- normalize ``(z - latents_mean) / latents_std`` with the loader-injected
|
||||
``MiniMaxH3VideoVAEArchConfig`` values
|
||||
- patchify [1, 2, 2] into packed cond rows, fp32
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.models.vaes.minimax_h3_video import (
|
||||
MiniMaxH3VideoVAEArchConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_tokens import (
|
||||
minimax_h3_patchify_video_latent,
|
||||
)
|
||||
|
||||
MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
|
||||
MINIMAX_H3_KEYFRAME_PATCH_SIZE = (1, 2, 2)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def minimax_h3_scoped_encode_rng(seed: int, device: torch.device | None = None):
|
||||
"""Seed torch RNGs for a deterministic sampled VAE encode without leaking state.
|
||||
|
||||
The encode recipes seed the default torch generators right before a
|
||||
posterior-sampled VAE encode. Forking restores the process-global CPU and
|
||||
CUDA generators after the encode while preserving the exact sampled result.
|
||||
"""
|
||||
devices: list[torch.device] = []
|
||||
if device is not None and device.type == "cuda" and torch.cuda.is_available():
|
||||
devices = [device]
|
||||
with torch.random.fork_rng(devices=devices):
|
||||
torch.default_generator.manual_seed(int(seed))
|
||||
for forked_device in devices:
|
||||
with torch.cuda.device(forked_device):
|
||||
torch.cuda.manual_seed(int(seed))
|
||||
yield
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def minimax_h3_scoped_encode_fp32(video_vae: Any):
|
||||
"""Scope the video VAE to fp32 for one or more keyframe/reference encodes.
|
||||
|
||||
encode_keyframe_cond_rows and encode_reference_video_rows each also guard
|
||||
their own cast (skipping it when already fp32), so nesting this around a
|
||||
caller that does more than one encode -- FL2VA's two keyframes, ref2va's
|
||||
image reference plus video reference -- turns their per-call casts into
|
||||
no-ops instead of toggling the whole VAE's dtype once per encode.
|
||||
"""
|
||||
parameter = next(video_vae.parameters())
|
||||
prev_dtype = parameter.dtype
|
||||
if prev_dtype != torch.float32:
|
||||
video_vae.to(torch.float32)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if prev_dtype != torch.float32:
|
||||
video_vae.to(prev_dtype)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=None)
|
||||
def _cached_latent_mean_std(
|
||||
mean_values: tuple[float, ...],
|
||||
std_values: tuple[float, ...],
|
||||
view_shape: tuple[int, ...],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""CPU mean/std tensors for a fixed (values, shape) triple, built once.
|
||||
|
||||
arch_config.latents_mean/std are static config values fixed for the
|
||||
life of a loaded VAE, so every call with the same values reconstructs
|
||||
an identical tensor; cache it instead of rebuilding on every encode.
|
||||
"""
|
||||
mean = torch.tensor(mean_values).view(view_shape)
|
||||
std = torch.tensor(std_values).view(view_shape)
|
||||
return mean, std
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def minimax_h3_encode_keyframe_cond_rows(
|
||||
video_vae: Any,
|
||||
image: Any,
|
||||
arch_config: MiniMaxH3VideoVAEArchConfig,
|
||||
) -> torch.Tensor:
|
||||
"""Encode a target-canvas PIL image into packed imgvid cond rows.
|
||||
|
||||
Returns [n_rows, 24 * patch_h * patch_w] fp32 on CPU.
|
||||
"""
|
||||
seed = MINIMAX_H3_KEYFRAME_ENCODE_SEED
|
||||
# The encode recipe runs on fp32 weights. Normal H3 residency already keeps
|
||||
# the shared video VAE in fp32; retain the scoped cast for standalone use.
|
||||
parameter = next(video_vae.parameters())
|
||||
prev_dtype = parameter.dtype
|
||||
if prev_dtype != torch.float32:
|
||||
video_vae.to(torch.float32)
|
||||
try:
|
||||
with minimax_h3_scoped_encode_rng(seed, parameter.device):
|
||||
z = video_vae.encode_images(image, use_fp16_latent=True)[0]
|
||||
finally:
|
||||
if prev_dtype != torch.float32:
|
||||
video_vae.to(prev_dtype)
|
||||
z = z.cpu().float()
|
||||
if z.dim() == 4:
|
||||
z = z[None]
|
||||
latent_channels = arch_config.latent_channels
|
||||
if z.dim() != 5 or int(z.shape[1]) != latent_channels:
|
||||
raise ValueError(f"unexpected imgvid latent shape {list(z.shape)}")
|
||||
mean, std = _cached_latent_mean_std(
|
||||
tuple(arch_config.latents_mean),
|
||||
tuple(arch_config.latents_std),
|
||||
(1, latent_channels, 1, 1, 1),
|
||||
)
|
||||
z.sub_(mean).div_(std)
|
||||
rows = minimax_h3_patchify_video_latent(
|
||||
z, patch_size=list(MINIMAX_H3_KEYFRAME_PATCH_SIZE)
|
||||
)
|
||||
return rows.to(torch.float32)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_KEYFRAME_ENCODE_SEED",
|
||||
"MINIMAX_H3_KEYFRAME_PATCH_SIZE",
|
||||
"_cached_latent_mean_std",
|
||||
"minimax_h3_encode_keyframe_cond_rows",
|
||||
"minimax_h3_scoped_encode_rng",
|
||||
"minimax_h3_scoped_encode_fp32",
|
||||
]
|
||||
+913
@@ -0,0 +1,913 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Request-owned material URI localization for the MiniMax H3 pipeline.
|
||||
|
||||
The canonical MiniMax H3 contract intentionally carries semantic URIs rather
|
||||
than worker-local paths. Direct media consumers (Pillow, ffmpeg and
|
||||
torchaudio) cannot consume every URI scheme in that contract, so localization
|
||||
belongs at the model-specific material boundary. Materialized sources and
|
||||
derived work directories are registered on ``Req.extra`` and explicitly
|
||||
released by the encoder stages.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import math
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY = "minimax_h3_material_localization"
|
||||
MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY = "minimax_h3_material_probe_facts"
|
||||
MINIMAX_H3_TEMP_DIRS_EXTRA_KEY = "minimax_h3_request_temp_dirs"
|
||||
MINIMAX_H3_HTTP_READ_CHUNK_BYTES = 1024 * 1024
|
||||
MINIMAX_H3_BASE64_DECODE_CHUNK_CHARS = 1024 * 1024
|
||||
MINIMAX_H3_BASE64_HEADER_MAX_CHARS = 4 * 1024
|
||||
MINIMAX_H3_TAR_HEADER_MAX_ENCODED_CHARS = 64 * 1024
|
||||
|
||||
_BASE64_ALPHABET = frozenset(
|
||||
b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=_-"
|
||||
)
|
||||
|
||||
_DEFAULT_SUFFIX_BY_TYPE = {
|
||||
"image": ".png",
|
||||
"video": ".mp4",
|
||||
"video_audio": ".mp4",
|
||||
"audio": ".wav",
|
||||
}
|
||||
_SUFFIX_BY_MEDIA_TYPE = {
|
||||
"image/jpeg": ".jpg",
|
||||
"image/png": ".png",
|
||||
"image/webp": ".webp",
|
||||
"video/mp4": ".mp4",
|
||||
"video/quicktime": ".mov",
|
||||
"audio/mpeg": ".mp3",
|
||||
"audio/mp4": ".m4a",
|
||||
"audio/wav": ".wav",
|
||||
"audio/x-wav": ".wav",
|
||||
"audio/flac": ".flac",
|
||||
}
|
||||
|
||||
|
||||
def minimax_h3_register_temp_dir(batch: Any, path: str, *, owner: str) -> str:
|
||||
"""Register one request-owned directory and return *path* unchanged."""
|
||||
|
||||
registry = batch.extra.setdefault(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY, {})
|
||||
paths = registry.setdefault(str(owner), [])
|
||||
normalized = str(path)
|
||||
if normalized not in paths:
|
||||
paths.append(normalized)
|
||||
return normalized
|
||||
|
||||
|
||||
def minimax_h3_cleanup_temp_dirs(
|
||||
batch: Any, *, owners: Iterable[str] | None = None
|
||||
) -> None:
|
||||
"""Remove registered request directories, tolerating repeated cleanup."""
|
||||
|
||||
registry = batch.extra.get(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY)
|
||||
selected = (
|
||||
list(registry)
|
||||
if owners is None and isinstance(registry, dict)
|
||||
else [str(owner) for owner in (owners or ())]
|
||||
)
|
||||
if isinstance(registry, dict):
|
||||
for owner in selected:
|
||||
paths = registry.pop(owner, [])
|
||||
if isinstance(paths, (list, tuple)):
|
||||
for path in paths:
|
||||
shutil.rmtree(str(path), ignore_errors=True)
|
||||
if not registry:
|
||||
batch.extra.pop(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY, None)
|
||||
if owners is None or "material" in selected:
|
||||
batch.extra.pop(MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY, None)
|
||||
batch.extra.pop(MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY, None)
|
||||
|
||||
|
||||
def _base64_uri_payload_start(uri: str) -> tuple[int, str | None]:
|
||||
media_type = None
|
||||
if uri.startswith("data:"):
|
||||
separator = uri.find(",")
|
||||
if separator < 0:
|
||||
raise ValueError("data URI must contain a comma separator")
|
||||
if separator > MINIMAX_H3_BASE64_HEADER_MAX_CHARS:
|
||||
raise ValueError("data URI header is too large")
|
||||
header = uri[:separator]
|
||||
if ";base64" not in header:
|
||||
raise ValueError("data URI must use ;base64 encoding")
|
||||
media_type = header[5:].split(";", 1)[0].lower() or None
|
||||
payload_start = separator + 1
|
||||
elif uri.startswith("base64://"):
|
||||
payload_start = len("base64://")
|
||||
separator = uri.find(",", payload_start)
|
||||
if separator >= 0:
|
||||
if separator - payload_start > MINIMAX_H3_BASE64_HEADER_MAX_CHARS:
|
||||
raise ValueError("base64 URI header is too large")
|
||||
header = uri[payload_start:separator]
|
||||
media_type = header.split(";", 1)[0].lower() or None
|
||||
payload_start = separator + 1
|
||||
else: # pragma: no cover - guarded by the caller
|
||||
raise ValueError("not a base64 material URI")
|
||||
return payload_start, media_type
|
||||
|
||||
|
||||
def _iter_base64_payload_bytes(uri: str, payload_start: int):
|
||||
"""Yield validated, unquoted base64 bytes without copying the payload."""
|
||||
|
||||
index = payload_start
|
||||
while index < len(uri):
|
||||
character = uri[index]
|
||||
if character == "%":
|
||||
if index + 2 >= len(uri):
|
||||
raise ValueError("material URI has an invalid percent escape")
|
||||
try:
|
||||
value = int(uri[index + 1 : index + 3], 16)
|
||||
except ValueError as exc:
|
||||
raise ValueError("material URI has an invalid percent escape") from exc
|
||||
index += 3
|
||||
character = chr(value)
|
||||
else:
|
||||
index += 1
|
||||
|
||||
if character.isspace():
|
||||
continue
|
||||
if len(character) != 1 or ord(character) > 127:
|
||||
raise ValueError("material URI base64 payload must be ASCII")
|
||||
value = ord(character)
|
||||
if value not in _BASE64_ALPHABET:
|
||||
raise ValueError(
|
||||
f"material URI has an invalid base64 character {character!r}"
|
||||
)
|
||||
yield value
|
||||
|
||||
|
||||
def _parse_tar_member_uri(uri: str) -> tuple[Path, int, int, str | None]:
|
||||
if uri.startswith("tar+offset://"):
|
||||
prefix = "tar+offset://"
|
||||
elif uri.startswith("tar+b64header://"):
|
||||
prefix = "tar+b64header://"
|
||||
else:
|
||||
raise ValueError("unsupported tar material URI")
|
||||
try:
|
||||
tar_path, encoded_header = uri[len(prefix) :].rsplit(":", 1)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
"tar material URI must contain '<tar_path>:<encoded_header>'"
|
||||
) from exc
|
||||
if len(encoded_header) > MINIMAX_H3_TAR_HEADER_MAX_ENCODED_CHARS:
|
||||
raise ValueError("tar material URI encoded header is too large")
|
||||
padded = encoded_header + "=" * (-len(encoded_header) % 4)
|
||||
try:
|
||||
header = json.loads(
|
||||
base64.b64decode(
|
||||
padded.encode("ascii"), altchars=b"-_", validate=True
|
||||
).decode("utf-8")
|
||||
)
|
||||
except Exception as exc:
|
||||
raise ValueError("tar material URI has an invalid encoded header") from exc
|
||||
if not isinstance(header, dict):
|
||||
raise ValueError("tar material URI header must be a JSON object")
|
||||
if header.get("schema") != "sglang.tar_member_ref/v1":
|
||||
raise ValueError(
|
||||
f"unsupported tar material header schema: {header.get('schema')!r}"
|
||||
)
|
||||
try:
|
||||
offset = int(header["offset_data"])
|
||||
size = int(header["size"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
"tar material header requires integer offset_data and size"
|
||||
) from exc
|
||||
if offset < 0 or size < 0:
|
||||
raise ValueError("tar material offset_data and size must be non-negative")
|
||||
return (
|
||||
Path(tar_path).expanduser(),
|
||||
offset,
|
||||
size,
|
||||
str(header.get("member") or "") or None,
|
||||
)
|
||||
|
||||
|
||||
def _safe_suffix(value: str | None) -> str | None:
|
||||
if not value:
|
||||
return None
|
||||
suffix = Path(urllib.parse.urlsplit(value).path).suffix.lower()
|
||||
if suffix and len(suffix) <= 10 and suffix[1:].isalnum():
|
||||
return suffix
|
||||
return None
|
||||
|
||||
|
||||
def _checked_material_file(path: Path, *, label: str) -> str:
|
||||
"""Validate that a localized source exists and is non-empty."""
|
||||
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"{label} does not exist or is not a file: {path}")
|
||||
if path.stat().st_size <= 0:
|
||||
raise ValueError(f"{label} is empty: {path}")
|
||||
return str(path)
|
||||
|
||||
|
||||
def _parse_frame_rate(value: Any) -> float:
|
||||
if value in {None, "", "N/A", "0/0"}:
|
||||
return 0.0
|
||||
raw = str(value)
|
||||
try:
|
||||
if "/" in raw:
|
||||
numerator, denominator = raw.split("/", 1)
|
||||
denominator_value = float(denominator)
|
||||
parsed = float(numerator) / denominator_value if denominator_value else 0.0
|
||||
else:
|
||||
parsed = float(raw)
|
||||
return parsed if math.isfinite(parsed) else 0.0
|
||||
except (TypeError, ValueError, ZeroDivisionError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _parse_display_ratio(value: Any) -> float:
|
||||
if value in {None, "", "N/A", "0:1", "0/1"}:
|
||||
return 0.0
|
||||
raw = str(value).strip()
|
||||
separator = ":" if ":" in raw else "/" if "/" in raw else None
|
||||
try:
|
||||
if separator is None:
|
||||
ratio = float(raw)
|
||||
else:
|
||||
numerator, denominator = raw.split(separator, 1)
|
||||
ratio = float(numerator) / float(denominator)
|
||||
except (TypeError, ValueError, ZeroDivisionError):
|
||||
return 0.0
|
||||
return ratio if math.isfinite(ratio) and ratio > 0 else 0.0
|
||||
|
||||
|
||||
def _stream_rotation_degrees(stream: dict[str, Any]) -> float:
|
||||
values: list[Any] = []
|
||||
side_data = stream.get("side_data_list")
|
||||
if isinstance(side_data, list):
|
||||
values.extend(
|
||||
item.get("rotation") for item in side_data if isinstance(item, dict)
|
||||
)
|
||||
tags = stream.get("tags")
|
||||
if isinstance(tags, dict):
|
||||
values.append(tags.get("rotate"))
|
||||
for value in values:
|
||||
if value in {None, "", "N/A"}:
|
||||
continue
|
||||
try:
|
||||
rotation = float(value)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if math.isfinite(rotation):
|
||||
return rotation % 360.0
|
||||
return 0.0
|
||||
|
||||
|
||||
def _display_geometry(stream: dict[str, Any]) -> tuple[float, float, float, float]:
|
||||
"""Return square-pixel display width/height, SAR, and rotation."""
|
||||
|
||||
coded_width = int(stream.get("width") or 0)
|
||||
coded_height = int(stream.get("height") or 0)
|
||||
sar = _parse_display_ratio(stream.get("sample_aspect_ratio")) or 1.0
|
||||
dar = _parse_display_ratio(stream.get("display_aspect_ratio"))
|
||||
physical_height = float(coded_height)
|
||||
physical_width = dar * physical_height if dar > 0.0 else float(coded_width) * sar
|
||||
rotation = _stream_rotation_degrees(stream)
|
||||
quarter_turns = round(rotation / 90.0)
|
||||
if abs(rotation - quarter_turns * 90.0) <= 1e-6:
|
||||
if quarter_turns % 2:
|
||||
return physical_height, physical_width, sar, rotation
|
||||
return physical_width, physical_height, sar, rotation
|
||||
radians = math.radians(rotation)
|
||||
cosine = abs(math.cos(radians))
|
||||
sine = abs(math.sin(radians))
|
||||
display_width = physical_width * cosine + physical_height * sine
|
||||
display_height = physical_width * sine + physical_height * cosine
|
||||
return display_width, display_height, sar, rotation
|
||||
|
||||
|
||||
_FFPROBE_STREAM_ENTRIES = (
|
||||
"stream=codec_type,width,height,duration,sample_rate,channels,"
|
||||
"avg_frame_rate,r_frame_rate,nb_frames,sample_aspect_ratio,display_aspect_ratio"
|
||||
":stream_tags=rotate"
|
||||
)
|
||||
|
||||
# ffprobe gained the per-stream "stream_side_data" section in 6.0 and rejects
|
||||
# the whole -show_entries spec without it. Older builds (e.g. Ubuntu 22.04's
|
||||
# 4.4) report rotation through the "rotate" stream tag, which
|
||||
# _stream_rotation_degrees already reads, so drop the section on fallback.
|
||||
_FFPROBE_ENTRY_VARIANTS = (
|
||||
f"{_FFPROBE_STREAM_ENTRIES}:stream_side_data=rotation:format=format_name,duration",
|
||||
f"{_FFPROBE_STREAM_ENTRIES}:format=format_name,duration",
|
||||
)
|
||||
_ffprobe_entries: str | None = None
|
||||
|
||||
|
||||
def _ffprobe_media(path: str) -> dict[str, Any]:
|
||||
global _ffprobe_entries
|
||||
|
||||
variants = (
|
||||
(_ffprobe_entries,) if _ffprobe_entries is not None else _FFPROBE_ENTRY_VARIANTS
|
||||
)
|
||||
|
||||
last_error: Exception | None = None
|
||||
for entries in variants:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-protocol_whitelist",
|
||||
"file",
|
||||
"-format_whitelist",
|
||||
"mov,mp4,m4a,3gp,3g2,mj2,matroska,webm,wav,mp3,flac,ogg",
|
||||
"-show_entries",
|
||||
entries,
|
||||
"-of",
|
||||
"json",
|
||||
"-i",
|
||||
path,
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
# Only an unknown section name is worth retrying; a genuinely bad
|
||||
# input must fail on the first variant.
|
||||
if "No match for section" not in (exc.stderr or ""):
|
||||
raise
|
||||
last_error = exc
|
||||
continue
|
||||
_ffprobe_entries = entries
|
||||
return json.loads(result.stdout)
|
||||
|
||||
raise last_error # type: ignore[misc]
|
||||
|
||||
|
||||
def _validate_localized_media(
|
||||
path: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Probe one localized source and return facts used by MiniMax H3 admission.
|
||||
|
||||
This is deliberately a model-facing validity check: it verifies that the
|
||||
source is non-empty, parseable and contains the stream type requested by
|
||||
the condition. Generic transport/resource ceilings do not belong here;
|
||||
the model's shape and temporal contracts are resolved separately.
|
||||
"""
|
||||
|
||||
if condition_type == "image":
|
||||
try:
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
with Image.open(path) as image:
|
||||
coded_width, coded_height = image.size
|
||||
image_format = str(image.format or "").upper()
|
||||
if coded_width <= 0 or coded_height <= 0:
|
||||
raise ValueError("image has no positive dimensions")
|
||||
display_image = ImageOps.exif_transpose(image)
|
||||
width, height = display_image.size
|
||||
except Exception as exc:
|
||||
raise ValueError("MiniMax H3 image material is invalid") from exc
|
||||
if image_format not in {"JPEG", "PNG", "WEBP"}:
|
||||
raise ValueError("MiniMax H3 image material uses an unsupported format")
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError(
|
||||
"MiniMax H3 image material has no positive display geometry"
|
||||
)
|
||||
return {
|
||||
"condition_type": "image",
|
||||
"coded_width": int(coded_width),
|
||||
"coded_height": int(coded_height),
|
||||
"display_width": int(width),
|
||||
"display_height": int(height),
|
||||
"image_format": image_format,
|
||||
"exif_transposed": (coded_width, coded_height) != (width, height),
|
||||
}
|
||||
|
||||
if condition_type not in {"audio", "video", "video_audio"}:
|
||||
raise ValueError(f"unsupported MiniMax H3 condition type {condition_type!r}")
|
||||
try:
|
||||
payload = _ffprobe_media(path)
|
||||
except Exception as exc:
|
||||
raise ValueError("MiniMax H3 media material is invalid") from exc
|
||||
|
||||
streams = payload.get("streams") or []
|
||||
format_names = set(
|
||||
str((payload.get("format") or {}).get("format_name") or "").split(",")
|
||||
)
|
||||
allowed_formats = {
|
||||
"mov",
|
||||
"mp4",
|
||||
"m4a",
|
||||
"3gp",
|
||||
"3g2",
|
||||
"mj2",
|
||||
"matroska",
|
||||
"webm",
|
||||
"wav",
|
||||
"mp3",
|
||||
"flac",
|
||||
"ogg",
|
||||
}
|
||||
if not format_names or not format_names.issubset(allowed_formats):
|
||||
raise ValueError("MiniMax H3 media container format is not allowed")
|
||||
video_streams = [s for s in streams if s.get("codec_type") == "video"]
|
||||
audio_streams = [s for s in streams if s.get("codec_type") == "audio"]
|
||||
if condition_type in {"video", "video_audio"} and not video_streams:
|
||||
raise ValueError("MiniMax H3 video material has no video stream")
|
||||
if condition_type in {"audio", "video_audio"} and not audio_streams:
|
||||
raise ValueError("MiniMax H3 audio material has no audio stream")
|
||||
|
||||
primary_video: dict[str, Any] | None = None
|
||||
for stream in video_streams:
|
||||
try:
|
||||
width = int(stream.get("width") or 0)
|
||||
height = int(stream.get("height") or 0)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
"MiniMax H3 video material has invalid dimensions"
|
||||
) from exc
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError("MiniMax H3 video material has no positive dimensions")
|
||||
fps = _parse_frame_rate(stream.get("avg_frame_rate")) or _parse_frame_rate(
|
||||
stream.get("r_frame_rate")
|
||||
)
|
||||
if fps <= 0:
|
||||
raise ValueError("MiniMax H3 video material has no usable frame rate")
|
||||
if primary_video is None:
|
||||
primary_video = stream
|
||||
|
||||
for stream in audio_streams:
|
||||
try:
|
||||
sample_rate = int(stream.get("sample_rate") or 0)
|
||||
channels = int(stream.get("channels") or 0)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("MiniMax H3 audio material has invalid metadata") from exc
|
||||
if sample_rate <= 0:
|
||||
raise ValueError("MiniMax H3 audio material has no usable sample rate")
|
||||
if channels <= 0:
|
||||
raise ValueError("MiniMax H3 audio material has no usable channel count")
|
||||
|
||||
durations: list[float] = []
|
||||
for value in [
|
||||
(payload.get("format") or {}).get("duration"),
|
||||
*(stream.get("duration") for stream in streams),
|
||||
]:
|
||||
if value in {None, "", "N/A"}:
|
||||
continue
|
||||
try:
|
||||
duration = float(value)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if math.isfinite(duration) and duration > 0:
|
||||
durations.append(duration)
|
||||
if not durations:
|
||||
raise ValueError("MiniMax H3 media material has no positive duration")
|
||||
|
||||
duration_seconds = max(durations)
|
||||
facts: dict[str, Any] = {
|
||||
"condition_type": condition_type,
|
||||
"duration_seconds": duration_seconds,
|
||||
"has_audio": bool(audio_streams),
|
||||
}
|
||||
if primary_video is not None:
|
||||
coded_width = int(primary_video.get("width") or 0)
|
||||
coded_height = int(primary_video.get("height") or 0)
|
||||
display_width, display_height, sar, rotation = _display_geometry(primary_video)
|
||||
fps = _parse_frame_rate(primary_video.get("avg_frame_rate"))
|
||||
if fps <= 0:
|
||||
fps = _parse_frame_rate(primary_video.get("r_frame_rate"))
|
||||
raw_count = primary_video.get("nb_frames")
|
||||
try:
|
||||
frame_count = int(raw_count)
|
||||
except (TypeError, ValueError):
|
||||
frame_count = max(1, int(round(duration_seconds * fps)))
|
||||
if frame_count <= 0:
|
||||
frame_count = max(1, int(round(duration_seconds * fps)))
|
||||
try:
|
||||
video_duration_seconds = float(primary_video.get("duration"))
|
||||
except (TypeError, ValueError):
|
||||
video_duration_seconds = 0.0
|
||||
if not math.isfinite(video_duration_seconds) or video_duration_seconds <= 0:
|
||||
video_duration_seconds = frame_count / fps
|
||||
facts.update(
|
||||
{
|
||||
"coded_width": coded_width,
|
||||
"coded_height": coded_height,
|
||||
"display_width": display_width,
|
||||
"display_height": display_height,
|
||||
"sample_aspect_ratio": str(
|
||||
primary_video.get("sample_aspect_ratio") or "1:1"
|
||||
),
|
||||
"sample_aspect_ratio_value": sar,
|
||||
"display_aspect_ratio": display_width / display_height,
|
||||
"rotation_degrees": rotation,
|
||||
"fps": fps,
|
||||
"frame_count": frame_count,
|
||||
"video_duration_seconds": video_duration_seconds,
|
||||
}
|
||||
)
|
||||
if audio_streams:
|
||||
facts["audio_sample_rate"] = int(audio_streams[0].get("sample_rate") or 0)
|
||||
facts["audio_channels"] = int(audio_streams[0].get("channels") or 0)
|
||||
try:
|
||||
audio_duration_seconds = float(audio_streams[0].get("duration"))
|
||||
except (TypeError, ValueError):
|
||||
audio_duration_seconds = 0.0
|
||||
facts["audio_duration_seconds"] = (
|
||||
audio_duration_seconds
|
||||
if math.isfinite(audio_duration_seconds) and audio_duration_seconds > 0
|
||||
else duration_seconds
|
||||
)
|
||||
return facts
|
||||
|
||||
|
||||
def _validate_material_once(
|
||||
batch: Any,
|
||||
uri: str,
|
||||
path: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
) -> dict[str, Any]:
|
||||
probe_facts = batch.extra.setdefault(MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY, {})
|
||||
key = (uri, condition_type)
|
||||
cached = probe_facts.get(key)
|
||||
if isinstance(cached, dict):
|
||||
return cached
|
||||
facts = _validate_localized_media(path, condition_type=condition_type)
|
||||
if not isinstance(facts, dict):
|
||||
raise ValueError("MiniMax H3 material probe did not return facts")
|
||||
probe_facts[key] = facts
|
||||
return facts
|
||||
|
||||
|
||||
def _material_output_paths(
|
||||
batch: Any,
|
||||
*,
|
||||
condition_type: str,
|
||||
condition_index: int,
|
||||
source_name: str | None = None,
|
||||
media_type: str | None = None,
|
||||
) -> tuple[Path, Path]:
|
||||
suffix = (
|
||||
_safe_suffix(source_name)
|
||||
or _SUFFIX_BY_MEDIA_TYPE.get(str(media_type or "").lower())
|
||||
or _DEFAULT_SUFFIX_BY_TYPE.get(condition_type, ".bin")
|
||||
)
|
||||
output_path = (
|
||||
Path(_material_workdir(batch)) / f"condition_{int(condition_index):04d}{suffix}"
|
||||
)
|
||||
return output_path, output_path.with_name(output_path.name + ".partial")
|
||||
|
||||
|
||||
def _material_workdir(batch: Any) -> str:
|
||||
registry = batch.extra.setdefault(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY, {})
|
||||
material_dirs = registry.get("material")
|
||||
if isinstance(material_dirs, list) and material_dirs:
|
||||
return str(material_dirs[0])
|
||||
return minimax_h3_register_temp_dir(
|
||||
batch,
|
||||
tempfile.mkdtemp(prefix="minimax_h3_material_"),
|
||||
owner="material",
|
||||
)
|
||||
|
||||
|
||||
def _decode_base64_chunk(encoded: bytes | bytearray) -> bytes:
|
||||
try:
|
||||
return base64.b64decode(encoded, altchars=b"-_", validate=True)
|
||||
except Exception as exc:
|
||||
raise ValueError("material URI has an invalid base64 payload") from exc
|
||||
|
||||
|
||||
def _stream_base64_material(
|
||||
batch: Any,
|
||||
uri: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
condition_index: int,
|
||||
) -> str:
|
||||
payload_start, media_type = _base64_uri_payload_start(uri)
|
||||
|
||||
output_path, partial_path = _material_output_paths(
|
||||
batch,
|
||||
condition_type=condition_type,
|
||||
condition_index=condition_index,
|
||||
media_type=media_type,
|
||||
)
|
||||
total = 0
|
||||
encoded_size = 0
|
||||
padding = 0
|
||||
saw_padding = False
|
||||
encoded_chunk = bytearray()
|
||||
try:
|
||||
with partial_path.open("wb") as output:
|
||||
for value in _iter_base64_payload_bytes(uri, payload_start):
|
||||
encoded_size += 1
|
||||
if value == ord("="):
|
||||
saw_padding = True
|
||||
padding += 1
|
||||
if padding > 2:
|
||||
raise ValueError("material URI has invalid base64 padding")
|
||||
elif saw_padding:
|
||||
raise ValueError("material URI has data after base64 padding")
|
||||
encoded_chunk.append(value)
|
||||
if len(encoded_chunk) == MINIMAX_H3_BASE64_DECODE_CHUNK_CHARS:
|
||||
decoded_chunk = _decode_base64_chunk(encoded_chunk)
|
||||
total += len(decoded_chunk)
|
||||
output.write(decoded_chunk)
|
||||
encoded_chunk.clear()
|
||||
if encoded_size == 0:
|
||||
raise ValueError("material URI base64 payload is empty")
|
||||
if encoded_size % 4 == 1 or (padding and encoded_size % 4):
|
||||
raise ValueError("material URI has an invalid base64 payload length")
|
||||
if encoded_chunk:
|
||||
encoded_chunk.extend(b"=" * (-len(encoded_chunk) % 4))
|
||||
decoded_chunk = _decode_base64_chunk(encoded_chunk)
|
||||
total += len(decoded_chunk)
|
||||
output.write(decoded_chunk)
|
||||
decoded_size = (encoded_size * 3) // 4 - padding
|
||||
if decoded_size <= 0:
|
||||
raise ValueError("material URI decoded payload is empty")
|
||||
if total != decoded_size:
|
||||
raise ValueError(
|
||||
f"MiniMax H3 base64 decoded size {total} != expected {decoded_size}"
|
||||
)
|
||||
partial_path.replace(output_path)
|
||||
except Exception:
|
||||
partial_path.unlink(missing_ok=True)
|
||||
output_path.unlink(missing_ok=True)
|
||||
raise
|
||||
return str(output_path)
|
||||
|
||||
|
||||
def _stream_tar_member_material(
|
||||
batch: Any,
|
||||
uri: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
condition_index: int,
|
||||
) -> str:
|
||||
source_path, offset, size, member = _parse_tar_member_uri(uri)
|
||||
if size <= 0:
|
||||
raise ValueError("tar material payload is empty")
|
||||
if not source_path.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"tar material source does not exist or is not a file: {source_path}"
|
||||
)
|
||||
source_size = source_path.stat().st_size
|
||||
if offset > source_size or size > source_size - offset:
|
||||
available = max(0, source_size - offset)
|
||||
raise ValueError(
|
||||
f"tar material payload is truncated: expected {size} bytes, "
|
||||
f"only {available} available"
|
||||
)
|
||||
|
||||
output_path, partial_path = _material_output_paths(
|
||||
batch,
|
||||
condition_type=condition_type,
|
||||
condition_index=condition_index,
|
||||
source_name=member,
|
||||
)
|
||||
remaining = size
|
||||
try:
|
||||
with source_path.open("rb") as source, partial_path.open("wb") as output:
|
||||
source.seek(offset)
|
||||
while remaining:
|
||||
chunk = source.read(min(MINIMAX_H3_HTTP_READ_CHUNK_BYTES, remaining))
|
||||
if not chunk:
|
||||
raise ValueError(
|
||||
f"tar material payload is truncated with {remaining} bytes left"
|
||||
)
|
||||
if len(chunk) > remaining:
|
||||
raise ValueError("tar material reader returned too many bytes")
|
||||
output.write(chunk)
|
||||
remaining -= len(chunk)
|
||||
partial_path.replace(output_path)
|
||||
except Exception:
|
||||
partial_path.unlink(missing_ok=True)
|
||||
output_path.unlink(missing_ok=True)
|
||||
raise
|
||||
return str(output_path)
|
||||
|
||||
|
||||
def _http_media_type(response: Any) -> str | None:
|
||||
headers = getattr(response, "headers", None)
|
||||
if headers is None:
|
||||
return None
|
||||
get_content_type = getattr(headers, "get_content_type", None)
|
||||
if callable(get_content_type):
|
||||
value = get_content_type()
|
||||
else:
|
||||
value = headers.get("Content-Type") or headers.get("content-type")
|
||||
if isinstance(value, str):
|
||||
value = value.split(";", 1)[0].strip()
|
||||
return str(value).lower() if value else None
|
||||
|
||||
|
||||
def _stream_http_material(
|
||||
batch: Any,
|
||||
uri: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
condition_index: int,
|
||||
timeout_s: float,
|
||||
) -> str:
|
||||
# Use the repository's legacy urllib behavior. Model-specific material
|
||||
# localization intentionally does not perform the shared public SSRF
|
||||
# policy or a cumulative request deadline.
|
||||
with urllib.request.urlopen(uri, timeout=timeout_s) as response:
|
||||
media_type = _http_media_type(response)
|
||||
suffix = (
|
||||
_safe_suffix(uri)
|
||||
or _SUFFIX_BY_MEDIA_TYPE.get(str(media_type or "").lower())
|
||||
or _DEFAULT_SUFFIX_BY_TYPE.get(condition_type, ".bin")
|
||||
)
|
||||
output_path = (
|
||||
Path(_material_workdir(batch))
|
||||
/ f"condition_{int(condition_index):04d}{suffix}"
|
||||
)
|
||||
partial_path = output_path.with_name(output_path.name + ".partial")
|
||||
total = 0
|
||||
try:
|
||||
with partial_path.open("wb") as output:
|
||||
while True:
|
||||
chunk = response.read(MINIMAX_H3_HTTP_READ_CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
if not isinstance(chunk, bytes):
|
||||
raise TypeError(
|
||||
"HTTP material response.read() must return bytes, got "
|
||||
f"{type(chunk).__name__}"
|
||||
)
|
||||
total += len(chunk)
|
||||
output.write(chunk)
|
||||
if total == 0:
|
||||
raise ValueError(f"HTTP material body is empty: {uri}")
|
||||
partial_path.replace(output_path)
|
||||
except Exception:
|
||||
partial_path.unlink(missing_ok=True)
|
||||
output_path.unlink(missing_ok=True)
|
||||
raise
|
||||
return str(output_path)
|
||||
|
||||
|
||||
def minimax_h3_localize_material_uri(
|
||||
batch: Any,
|
||||
uri: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
condition_index: int,
|
||||
timeout_s: float = 120.0,
|
||||
) -> str:
|
||||
"""Return a local path for a canonical condition URI.
|
||||
|
||||
Local paths and local ``file://`` URIs are validated and returned without
|
||||
copying. HTTP(S), base64/data and direct tar-member URIs are materialized
|
||||
once per request and cached for all MiniMax H3 consumers.
|
||||
"""
|
||||
|
||||
if not isinstance(uri, str) or not uri:
|
||||
raise ValueError("condition URI must be a non-empty string")
|
||||
parsed = None
|
||||
for special_scheme in ("data", "base64", "tar+offset", "tar+b64header"):
|
||||
if uri.startswith(special_scheme + ":"):
|
||||
scheme = special_scheme
|
||||
break
|
||||
else:
|
||||
parsed = urllib.parse.urlsplit(uri)
|
||||
scheme = parsed.scheme
|
||||
|
||||
if scheme == "file":
|
||||
assert parsed is not None
|
||||
if parsed.netloc not in {"", "localhost"}:
|
||||
raise ValueError(f"file URI host must be local, got {parsed.netloc!r}")
|
||||
output_path = _checked_material_file(
|
||||
Path(urllib.parse.unquote(parsed.path)),
|
||||
label="MiniMax H3 material source",
|
||||
)
|
||||
_validate_material_once(batch, uri, output_path, condition_type=condition_type)
|
||||
return output_path
|
||||
if not scheme:
|
||||
output_path = _checked_material_file(
|
||||
Path(uri).expanduser(),
|
||||
label="MiniMax H3 material source",
|
||||
)
|
||||
_validate_material_once(batch, uri, output_path, condition_type=condition_type)
|
||||
return output_path
|
||||
if scheme == "s3":
|
||||
raise NotImplementedError(
|
||||
"MiniMax H3 s3:// material URIs require a configured artifact resolver"
|
||||
)
|
||||
|
||||
cache = batch.extra.setdefault(MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY, {})
|
||||
cached = cache.get(uri)
|
||||
if isinstance(cached, str):
|
||||
cached_path = Path(cached)
|
||||
if cached_path.exists():
|
||||
output_path = _checked_material_file(
|
||||
cached_path,
|
||||
label="cached MiniMax H3 material",
|
||||
)
|
||||
_validate_material_once(
|
||||
batch, uri, output_path, condition_type=condition_type
|
||||
)
|
||||
return output_path
|
||||
cache.pop(uri, None)
|
||||
|
||||
if scheme in {"http", "https"}:
|
||||
output_path = _stream_http_material(
|
||||
batch,
|
||||
uri,
|
||||
condition_type=condition_type,
|
||||
condition_index=condition_index,
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
try:
|
||||
_validate_material_once(
|
||||
batch, uri, output_path, condition_type=condition_type
|
||||
)
|
||||
except Exception:
|
||||
Path(output_path).unlink(missing_ok=True)
|
||||
raise
|
||||
cache[uri] = output_path
|
||||
return output_path
|
||||
|
||||
if scheme in {"data", "base64"}:
|
||||
output_path = _stream_base64_material(
|
||||
batch,
|
||||
uri,
|
||||
condition_type=condition_type,
|
||||
condition_index=condition_index,
|
||||
)
|
||||
try:
|
||||
_validate_material_once(
|
||||
batch, uri, output_path, condition_type=condition_type
|
||||
)
|
||||
except Exception:
|
||||
Path(output_path).unlink(missing_ok=True)
|
||||
raise
|
||||
cache[uri] = output_path
|
||||
return output_path
|
||||
|
||||
if scheme in {"tar+offset", "tar+b64header"}:
|
||||
output_path = _stream_tar_member_material(
|
||||
batch,
|
||||
uri,
|
||||
condition_type=condition_type,
|
||||
condition_index=condition_index,
|
||||
)
|
||||
try:
|
||||
_validate_material_once(
|
||||
batch, uri, output_path, condition_type=condition_type
|
||||
)
|
||||
except Exception:
|
||||
Path(output_path).unlink(missing_ok=True)
|
||||
raise
|
||||
cache[uri] = output_path
|
||||
return output_path
|
||||
|
||||
raise NotImplementedError(
|
||||
f"MiniMax H3 material localization does not support URI scheme {scheme!r}"
|
||||
)
|
||||
|
||||
|
||||
def minimax_h3_probe_material(
|
||||
batch: Any,
|
||||
uri: str,
|
||||
*,
|
||||
condition_type: str,
|
||||
condition_index: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Localize and return cached display-geometry facts for one condition."""
|
||||
|
||||
path = minimax_h3_localize_material_uri(
|
||||
batch,
|
||||
uri,
|
||||
condition_type=condition_type,
|
||||
condition_index=condition_index,
|
||||
)
|
||||
key = (uri, condition_type)
|
||||
facts = batch.extra.get(MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY, {}).get(key)
|
||||
if not isinstance(facts, dict) or not facts:
|
||||
raise RuntimeError(
|
||||
"MiniMax H3 material localization completed without cached probe facts"
|
||||
)
|
||||
return {"local_path": path, **facts}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_MATERIAL_CACHE_EXTRA_KEY",
|
||||
"MINIMAX_H3_MATERIAL_PROBE_EXTRA_KEY",
|
||||
"MINIMAX_H3_TEMP_DIRS_EXTRA_KEY",
|
||||
"minimax_h3_cleanup_temp_dirs",
|
||||
"minimax_h3_localize_material_uri",
|
||||
"minimax_h3_probe_material",
|
||||
"minimax_h3_register_temp_dir",
|
||||
]
|
||||
+502
@@ -0,0 +1,502 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 packed-sequence materialization from the validated workspace
|
||||
builder, covering fl2va and t2va layouts.
|
||||
|
||||
Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P].
|
||||
Builder rules:
|
||||
- block-derived position infos, update masks, token tags, and cu_seqlens
|
||||
- img_position_ids fp64 grid: text rows (row_idx,0,0); video/cond t counter
|
||||
continues text_len with temporal interp spans (frame_rescale 5/3 x
|
||||
frame_per_token (1,4,4,4,4)); each spatial sqrt_area axis uses evenly spaced
|
||||
coordinates excluding the right endpoint, then scales them by INTERP;
|
||||
audio channel-major blocks pinned to the w-grid extremes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
|
||||
MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
|
||||
MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES,
|
||||
)
|
||||
|
||||
_INTERP = 32
|
||||
_T_GROUP = 5
|
||||
_FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
|
||||
_FRAME_RESCALE = 5.0 / 3.0
|
||||
_PATCH_H = 2
|
||||
_PATCH_W = 2
|
||||
|
||||
|
||||
def _keyframe_cond_frame_indices(
|
||||
*,
|
||||
include_keyframe_cond: bool,
|
||||
keyframe_frame_indices: list[int] | tuple[int, ...] | None,
|
||||
) -> list[int]:
|
||||
if not include_keyframe_cond:
|
||||
if keyframe_frame_indices is not None:
|
||||
raise ValueError(
|
||||
"keyframe_frame_indices must be omitted when keyframe cond is not included"
|
||||
)
|
||||
return []
|
||||
if keyframe_frame_indices is None:
|
||||
raise ValueError("strict fl2va packed layout requires keyframe_frame_indices")
|
||||
if any(
|
||||
isinstance(value, bool) or not isinstance(value, int)
|
||||
for value in keyframe_frame_indices
|
||||
):
|
||||
raise ValueError(
|
||||
"strict fl2va packed layout requires integer keyframe_frame_indices"
|
||||
)
|
||||
out = list(keyframe_frame_indices)
|
||||
if tuple(out) not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES:
|
||||
raise ValueError(
|
||||
"strict fl2va packed layout requires keyframe_frame_indices in "
|
||||
f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {out!r}"
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _resolve_keyframe_frame_indices(
|
||||
frame_indices: Sequence[int],
|
||||
*,
|
||||
frame_count: int | None,
|
||||
) -> list[int]:
|
||||
if frame_indices and frame_count is None:
|
||||
raise ValueError(
|
||||
"frame_count is required when keyframe_frame_indices are provided"
|
||||
)
|
||||
if frame_count is None:
|
||||
return []
|
||||
if frame_count <= 0:
|
||||
raise ValueError("frame_count must be positive")
|
||||
seen: dict[int, int] = {}
|
||||
resolved: list[int] = []
|
||||
for block_index, semantic_index in enumerate(frame_indices):
|
||||
if semantic_index == -1:
|
||||
resolved_index = frame_count - 1
|
||||
elif 0 <= semantic_index < frame_count:
|
||||
resolved_index = semantic_index
|
||||
else:
|
||||
raise ValueError(
|
||||
f"keyframe frame index {semantic_index} must be -1 or in "
|
||||
f"[0, {frame_count})"
|
||||
)
|
||||
previous = seen.get(resolved_index)
|
||||
if previous is not None:
|
||||
raise ValueError(
|
||||
f"keyframe frame index at block {block_index} resolves to "
|
||||
f"{resolved_index}, already bound by block {previous}"
|
||||
)
|
||||
seen[resolved_index] = block_index
|
||||
resolved.append(resolved_index)
|
||||
return resolved
|
||||
|
||||
|
||||
def _temporal_position_span(temporal_length: int) -> float:
|
||||
"""Temporal position span for patch_t=1, in fp64.
|
||||
|
||||
NOTE: intentionally NOT merged with ``_video_t_span``. This variant sums
|
||||
via numpy (pairwise summation), matching the fl2va anchor computation,
|
||||
while ``_video_t_span`` sums sequentially, matching the ref2va
|
||||
t-origin accumulation. The two orders diverge in the last ulp
|
||||
from n=16 onward, so each path must keep its own summation order.
|
||||
"""
|
||||
spans = np.ones(int(temporal_length), dtype=np.float64) * _FRAME_RESCALE
|
||||
for token_index in range(_T_GROUP):
|
||||
spans[token_index::_T_GROUP] *= _FRAME_PER_TOKEN[token_index]
|
||||
return float(spans.sum())
|
||||
|
||||
|
||||
def minimax_h3_packed_sequence(
|
||||
*,
|
||||
text_len: int,
|
||||
latent_t: int,
|
||||
latent_h: int,
|
||||
latent_w: int,
|
||||
audio_t: int,
|
||||
audio_channel: int = 2,
|
||||
include_keyframe_cond: bool,
|
||||
keyframe_frame_indices: list[int] | tuple[int, ...] | None = None,
|
||||
frame_count: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build the packed-sequence structural fields for one CFG branch.
|
||||
|
||||
The used length is padded up to a multiple of 64.
|
||||
"""
|
||||
ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
|
||||
frame_rows = ph * pw
|
||||
cond_frame_indices = _keyframe_cond_frame_indices(
|
||||
include_keyframe_cond=include_keyframe_cond,
|
||||
keyframe_frame_indices=keyframe_frame_indices,
|
||||
)
|
||||
resolved_cond_frame_indices = _resolve_keyframe_frame_indices(
|
||||
cond_frame_indices,
|
||||
frame_count=frame_count,
|
||||
)
|
||||
cond_rows = len(cond_frame_indices) * frame_rows
|
||||
video_rows = latent_t * frame_rows
|
||||
audio_rows = audio_t * audio_channel
|
||||
used = text_len + cond_rows + audio_rows + video_rows
|
||||
seq_len = (
|
||||
(used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1)
|
||||
// MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
|
||||
* MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
|
||||
)
|
||||
|
||||
text_sl = slice(0, text_len)
|
||||
cond_sl = slice(text_len, text_len + cond_rows)
|
||||
audio_sl = slice(cond_sl.stop, cond_sl.stop + audio_rows)
|
||||
video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
|
||||
target_img_pos = torch.arange(video_sl.start, video_sl.stop)
|
||||
img_pos = (
|
||||
torch.cat([torch.arange(cond_sl.start, cond_sl.stop), target_img_pos])
|
||||
if cond_rows
|
||||
else target_img_pos
|
||||
)
|
||||
update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
|
||||
update_mask[cond_rows:] = True
|
||||
audio_pos = torch.arange(audio_sl.start, audio_sl.stop)
|
||||
text_pos = torch.arange(0, text_len)
|
||||
|
||||
g = torch.zeros(seq_len, 3, dtype=torch.float64)
|
||||
g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)
|
||||
|
||||
t_grid = _video_t_grid(latent_t, float(text_len))
|
||||
sqrt_area = np.sqrt(latent_h * latent_w)
|
||||
h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, sqrt_area)
|
||||
w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, sqrt_area)
|
||||
hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
|
||||
frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
|
||||
video_g = g[video_sl].view(latent_t, frame_rows, 3)
|
||||
video_g[:, :, 0] = t_grid[:, None]
|
||||
video_g[:, :, 1:] = frame[None]
|
||||
for block_index, pixel_index in enumerate(resolved_cond_frame_indices):
|
||||
sl = slice(
|
||||
cond_sl.start + block_index * frame_rows,
|
||||
cond_sl.start + (block_index + 1) * frame_rows,
|
||||
)
|
||||
if pixel_index == 0:
|
||||
cond_t = float(text_len)
|
||||
elif frame_count is not None and pixel_index == frame_count - 1:
|
||||
cond_t = (
|
||||
float(text_len) + _temporal_position_span(latent_t) - _FRAME_RESCALE
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"fl2va packed layout only supports first/last keyframe anchors, "
|
||||
f"got resolved frame index {pixel_index}"
|
||||
)
|
||||
g[sl, 0] = cond_t
|
||||
g[sl, 1:] = frame
|
||||
audio_t_grid = float(text_len) + torch.arange(audio_t, dtype=torch.float64)
|
||||
g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
|
||||
g[audio_sl.start : audio_sl.start + audio_t, 2] = float(w_grid[0])
|
||||
g[audio_sl.start + audio_t : audio_sl.stop, 2] = float(w_grid[-1])
|
||||
|
||||
token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING
|
||||
token_tags[text_sl] = 1 # TEXT (fl2va image-segment override happens upstream)
|
||||
token_tags[audio_sl] = 2 # AUDIO
|
||||
token_tags[img_pos] = 0 # VIDEO
|
||||
|
||||
cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
|
||||
return {
|
||||
"seq_len": seq_len,
|
||||
"img_pos": img_pos,
|
||||
"audio_pos": audio_pos,
|
||||
"text_pos": text_pos,
|
||||
"update_mask": update_mask,
|
||||
"img_position_ids": g,
|
||||
"token_tags": token_tags,
|
||||
"cu_seqlens": cu,
|
||||
}
|
||||
|
||||
|
||||
def _positive_int(
|
||||
block: Mapping[str, object],
|
||||
key: str,
|
||||
path: str,
|
||||
*,
|
||||
allow_zero: bool = False,
|
||||
) -> int:
|
||||
value = block.get(key)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"{path}.{key} must be an integer")
|
||||
if value < 0 or (value == 0 and not allow_zero):
|
||||
predicate = "non-negative" if allow_zero else "positive"
|
||||
raise ValueError(f"{path}.{key} must be {predicate}")
|
||||
return int(value)
|
||||
|
||||
|
||||
def _axis_from_sqrt_area(dim: int, patch: int, sqrt_area: float) -> torch.Tensor:
|
||||
ratio = dim / sqrt_area
|
||||
left = (1.0 - ratio) * 1.0 / 2.0
|
||||
right = left + ratio * 1.0
|
||||
grid = np.linspace(left, right, dim // patch, endpoint=False) * _INTERP
|
||||
return torch.from_numpy(grid).to(torch.float64)
|
||||
|
||||
|
||||
def _video_t_grid(n: int, origin: float) -> torch.Tensor:
|
||||
spans = torch.tensor(
|
||||
[_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)],
|
||||
dtype=torch.float64,
|
||||
)
|
||||
return origin + torch.cat(
|
||||
[torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)]
|
||||
)
|
||||
|
||||
|
||||
def _video_t_span(n: int) -> float:
|
||||
# Sequential fp64 summation on purpose — see _temporal_position_span for
|
||||
# why the two span implementations must not be unified.
|
||||
return sum(_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n))
|
||||
|
||||
|
||||
def _range_for_slice(sl: slice) -> torch.Tensor:
|
||||
return torch.arange(sl.start, sl.stop, dtype=torch.long)
|
||||
|
||||
|
||||
def _cat_ranges(parts: list[torch.Tensor]) -> torch.Tensor:
|
||||
if len(parts) == 1:
|
||||
return parts[0]
|
||||
if parts:
|
||||
return torch.cat(parts)
|
||||
return torch.empty(0, dtype=torch.long)
|
||||
|
||||
|
||||
def minimax_h3_packed_sequence_ref2va_blocks(
|
||||
*,
|
||||
text_len: int,
|
||||
latent_t: int,
|
||||
latent_h: int,
|
||||
latent_w: int,
|
||||
audio_t: int,
|
||||
ref_blocks: Sequence[Mapping[str, object]],
|
||||
audio_channel: int = 2,
|
||||
seq_len: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""General ref2va-family packed layout.
|
||||
|
||||
``ref_blocks`` are consumed in request/plan order:
|
||||
- ``{"kind": "image", "latent_h": H, "latent_w": W}``
|
||||
- ``{"kind": "audio", "ref_audio_t": T}``
|
||||
- ``{"kind": "video"|"video_audio", "ref_audio_t": T,
|
||||
"latent_t": RT, "latent_h": RH, "latent_w": RW}``
|
||||
|
||||
Video-bearing blocks pack their audio rows immediately before their video
|
||||
rows; both share the same temporal origin and advance by the longer of the
|
||||
audio and video spans. Standalone audio advances the target origin by its
|
||||
own T, and image blocks advance it by one integer slot.
|
||||
"""
|
||||
if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)):
|
||||
raise ValueError("ref_blocks must be a sequence")
|
||||
|
||||
parsed: list[dict[str, object]] = []
|
||||
ref_visual_rows = 0
|
||||
ref_audio_rows = 0
|
||||
for index, raw in enumerate(ref_blocks):
|
||||
path = f"ref_blocks[{index}]"
|
||||
if not isinstance(raw, Mapping):
|
||||
raise ValueError(f"{path} must be an object")
|
||||
kind = raw.get("kind", raw.get("type"))
|
||||
if not isinstance(kind, str) or not kind:
|
||||
raise ValueError(f"{path}.kind must be a non-empty string")
|
||||
if kind == "image":
|
||||
rh = _positive_int(raw, "latent_h", path)
|
||||
rw = _positive_int(raw, "latent_w", path)
|
||||
rows = (rh // _PATCH_H) * (rw // _PATCH_W)
|
||||
item = {"kind": kind, "latent_h": rh, "latent_w": rw, "rows": rows}
|
||||
ref_visual_rows += rows
|
||||
elif kind == "audio":
|
||||
rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
|
||||
rows = rt * audio_channel
|
||||
item = {"kind": kind, "ref_audio_t": rt, "audio_rows": rows}
|
||||
ref_audio_rows += rows
|
||||
elif kind in ("video", "video_audio"):
|
||||
rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True)
|
||||
vt = _positive_int(raw, "latent_t", path)
|
||||
vh = _positive_int(raw, "latent_h", path)
|
||||
vw = _positive_int(raw, "latent_w", path)
|
||||
frame_rows = (vh // _PATCH_H) * (vw // _PATCH_W)
|
||||
audio_rows = rt * audio_channel
|
||||
video_rows = vt * frame_rows
|
||||
item = {
|
||||
"kind": kind,
|
||||
"ref_audio_t": rt,
|
||||
"latent_t": vt,
|
||||
"latent_h": vh,
|
||||
"latent_w": vw,
|
||||
"frame_rows": frame_rows,
|
||||
"audio_rows": audio_rows,
|
||||
"video_rows": video_rows,
|
||||
}
|
||||
ref_audio_rows += audio_rows
|
||||
ref_visual_rows += video_rows
|
||||
else:
|
||||
raise ValueError(f"{path}.kind unsupported for ref2va: {kind!r}")
|
||||
parsed.append(item)
|
||||
|
||||
ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W
|
||||
frame_rows = ph * pw
|
||||
video_rows = latent_t * frame_rows
|
||||
audio_rows = audio_t * audio_channel
|
||||
ref_rows = ref_visual_rows + ref_audio_rows
|
||||
used = text_len + ref_rows + audio_rows + video_rows
|
||||
if seq_len is None:
|
||||
seq_len = (
|
||||
(used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1)
|
||||
// MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
|
||||
* MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT
|
||||
)
|
||||
if seq_len < used:
|
||||
raise ValueError(f"seq_len {seq_len} < used rows {used}")
|
||||
|
||||
text_sl = slice(0, text_len)
|
||||
cursor = text_len
|
||||
block_slices: list[dict[str, object]] = []
|
||||
for item in parsed:
|
||||
kind = str(item["kind"])
|
||||
if kind == "image":
|
||||
rows = int(item["rows"])
|
||||
visual_sl = slice(cursor, cursor + rows)
|
||||
cursor = visual_sl.stop
|
||||
block_slices.append({**item, "visual_sl": visual_sl})
|
||||
elif kind == "audio":
|
||||
rows = int(item["audio_rows"])
|
||||
audio_sl = slice(cursor, cursor + rows)
|
||||
cursor = audio_sl.stop
|
||||
block_slices.append({**item, "audio_sl": audio_sl})
|
||||
else:
|
||||
a_rows = int(item["audio_rows"])
|
||||
v_rows = int(item["video_rows"])
|
||||
audio_sl = slice(cursor, cursor + a_rows)
|
||||
visual_sl = slice(audio_sl.stop, audio_sl.stop + v_rows)
|
||||
cursor = visual_sl.stop
|
||||
block_slices.append({**item, "audio_sl": audio_sl, "visual_sl": visual_sl})
|
||||
|
||||
audio_sl = slice(cursor, cursor + audio_rows)
|
||||
video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows)
|
||||
ref_img_pos_parts: list[torch.Tensor] = []
|
||||
ref_audio_pos_parts: list[torch.Tensor] = []
|
||||
g = torch.zeros(seq_len, 3, dtype=torch.float64)
|
||||
g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64)
|
||||
|
||||
target_area = np.sqrt(latent_h * latent_w)
|
||||
h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, target_area)
|
||||
w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, target_area)
|
||||
hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij")
|
||||
target_frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1)
|
||||
|
||||
t_cursor = float(text_len)
|
||||
for item in block_slices:
|
||||
kind = str(item["kind"])
|
||||
if kind == "image":
|
||||
visual_sl = item["visual_sl"]
|
||||
assert isinstance(visual_sl, slice)
|
||||
ref_img_pos_parts.append(_range_for_slice(visual_sl))
|
||||
rh = int(item["latent_h"])
|
||||
rw = int(item["latent_w"])
|
||||
area = np.sqrt(rh * rw)
|
||||
ref_hh, ref_ww = torch.meshgrid(
|
||||
_axis_from_sqrt_area(rh, _PATCH_H, area),
|
||||
_axis_from_sqrt_area(rw, _PATCH_W, area),
|
||||
indexing="ij",
|
||||
)
|
||||
g[visual_sl, 0] = t_cursor
|
||||
g[visual_sl, 1] = ref_hh.reshape(-1)
|
||||
g[visual_sl, 2] = ref_ww.reshape(-1)
|
||||
t_cursor += 1.0
|
||||
elif kind == "audio":
|
||||
audio_ref_sl = item["audio_sl"]
|
||||
assert isinstance(audio_ref_sl, slice)
|
||||
ref_t = int(item["ref_audio_t"])
|
||||
ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
|
||||
ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
|
||||
g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
|
||||
if ref_t:
|
||||
g[audio_ref_sl.start : audio_ref_sl.start + ref_t, 2] = float(w_grid[0])
|
||||
g[audio_ref_sl.start + ref_t : audio_ref_sl.stop, 2] = float(w_grid[-1])
|
||||
t_cursor += float(ref_t)
|
||||
else:
|
||||
audio_ref_sl = item["audio_sl"]
|
||||
visual_sl = item["visual_sl"]
|
||||
assert isinstance(audio_ref_sl, slice)
|
||||
assert isinstance(visual_sl, slice)
|
||||
ref_t = int(item["ref_audio_t"])
|
||||
vt = int(item["latent_t"])
|
||||
vh = int(item["latent_h"])
|
||||
vw = int(item["latent_w"])
|
||||
ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl))
|
||||
ref_img_pos_parts.append(_range_for_slice(visual_sl))
|
||||
|
||||
ref_area = np.sqrt(vh * vw)
|
||||
rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area)
|
||||
rv_w_grid = _axis_from_sqrt_area(vw, _PATCH_W, ref_area)
|
||||
rv_hh, rv_ww = torch.meshgrid(rv_h_grid, rv_w_grid, indexing="ij")
|
||||
|
||||
ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64)
|
||||
g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel)
|
||||
if ref_t:
|
||||
g[audio_ref_sl.start : audio_ref_sl.start + ref_t, 2] = float(
|
||||
rv_w_grid[0]
|
||||
)
|
||||
g[audio_ref_sl.start + ref_t : audio_ref_sl.stop, 2] = float(
|
||||
rv_w_grid[-1]
|
||||
)
|
||||
|
||||
rv_frame = torch.stack([rv_hh.reshape(-1), rv_ww.reshape(-1)], dim=-1)
|
||||
rv_g = g[visual_sl].view(vt, int(item["frame_rows"]), 3)
|
||||
rv_g[:, :, 0] = _video_t_grid(vt, t_cursor)[:, None]
|
||||
rv_g[:, :, 1:] = rv_frame[None]
|
||||
t_cursor += max(float(ref_t), _video_t_span(vt))
|
||||
|
||||
audio_t_grid = t_cursor + torch.arange(audio_t, dtype=torch.float64)
|
||||
g[audio_sl, 0] = audio_t_grid.repeat(audio_channel)
|
||||
g[audio_sl.start : audio_sl.start + audio_t, 2] = float(w_grid[0])
|
||||
g[audio_sl.start + audio_t : audio_sl.stop, 2] = float(w_grid[-1])
|
||||
|
||||
video_g = g[video_sl].view(latent_t, frame_rows, 3)
|
||||
video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None]
|
||||
video_g[:, :, 1:] = target_frame[None]
|
||||
|
||||
target_img_pos = _range_for_slice(video_sl)
|
||||
target_audio_pos = _range_for_slice(audio_sl)
|
||||
img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos])
|
||||
audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos])
|
||||
|
||||
update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool)
|
||||
update_mask[ref_visual_rows:] = True
|
||||
audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool)
|
||||
audio_update_mask[ref_audio_rows:] = True
|
||||
text_pos = torch.arange(0, text_len)
|
||||
|
||||
token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING
|
||||
token_tags[text_sl] = 1 # TEXT
|
||||
token_tags[audio_pos] = 2 # AUDIO (refs + target)
|
||||
token_tags[img_pos] = 0 # VIDEO (refs + target)
|
||||
|
||||
cu = torch.tensor([0, used, seq_len], dtype=torch.int32)
|
||||
return {
|
||||
"seq_len": seq_len,
|
||||
"img_pos": img_pos,
|
||||
"audio_pos": audio_pos,
|
||||
"text_pos": text_pos,
|
||||
"update_mask": update_mask,
|
||||
"audio_update_mask": audio_update_mask,
|
||||
"img_position_ids": g,
|
||||
"token_tags": token_tags,
|
||||
"cu_seqlens": cu,
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"minimax_h3_packed_sequence",
|
||||
"minimax_h3_packed_sequence_ref2va_blocks",
|
||||
]
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _int_tuple(value: Sequence[int], name: str, length: int) -> tuple[int, ...]:
|
||||
if len(value) != length:
|
||||
raise ValueError(f"{name} must have length {length}, got {list(value)!r}")
|
||||
out = tuple(int(item) for item in value)
|
||||
if any(item <= 0 for item in out):
|
||||
raise ValueError(f"{name} values must be positive, got {list(value)!r}")
|
||||
return out
|
||||
|
||||
|
||||
def _rank(tensor: torch.Tensor, name: str, rank: int) -> None:
|
||||
if tensor.ndim != rank:
|
||||
raise ValueError(f"{name} must be rank {rank}, got shape={list(tensor.shape)}")
|
||||
|
||||
|
||||
def minimax_h3_patchify_video_latent(
|
||||
latent: torch.Tensor,
|
||||
*,
|
||||
patch_size: Sequence[int],
|
||||
) -> torch.Tensor:
|
||||
"""Pack SGLang video latent [B,C,T,H,W] into DiT token rows."""
|
||||
|
||||
_rank(latent, "video latent", 5)
|
||||
pt, ph, pw = _int_tuple(patch_size, "patch_size", 3)
|
||||
batch, channel, full_t, full_h, full_w = (int(dim) for dim in latent.shape)
|
||||
if full_t % pt or full_h % ph or full_w % pw:
|
||||
raise ValueError(
|
||||
"video latent spatial/time dims must be divisible by patch_size: "
|
||||
f"shape={list(latent.shape)}, patch_size={[pt, ph, pw]}"
|
||||
)
|
||||
t, h, w = full_t // pt, full_h // ph, full_w // pw
|
||||
packed = latent.reshape(batch, channel, t, pt, h, ph, w, pw)
|
||||
packed = torch.einsum("nctrhpwq->nthwcrpq", packed)
|
||||
return packed.reshape(batch * t * h * w, channel * pt * ph * pw).contiguous()
|
||||
|
||||
|
||||
def minimax_h3_unpatchify_video_tokens(
|
||||
rows: torch.Tensor,
|
||||
*,
|
||||
latent_shape: Sequence[int],
|
||||
patch_size: Sequence[int],
|
||||
) -> torch.Tensor:
|
||||
"""Unpack DiT video token rows into SGLang latent [B,C,T,H,W]."""
|
||||
|
||||
_rank(rows, "video token rows", 2)
|
||||
t, h, w, channel = _int_tuple(latent_shape, "latent_shape", 4)
|
||||
pt, ph, pw = _int_tuple(patch_size, "patch_size", 3)
|
||||
expected_dim = pt * ph * pw * channel
|
||||
if int(rows.shape[-1]) != expected_dim:
|
||||
raise ValueError(
|
||||
f"video token dim {int(rows.shape[-1])} != patch volume * channel "
|
||||
f"{expected_dim} for latent_shape={list(latent_shape)}, "
|
||||
f"patch_size={[pt, ph, pw]}"
|
||||
)
|
||||
rows_per_sample = t * h * w
|
||||
if int(rows.shape[0]) % rows_per_sample:
|
||||
raise ValueError(
|
||||
f"video rows {int(rows.shape[0])} must be divisible by t*h*w "
|
||||
f"{rows_per_sample} for latent_shape={list(latent_shape)}"
|
||||
)
|
||||
packed = rows.reshape(-1, t, h, w, channel, pt, ph, pw)
|
||||
latent = torch.einsum("nthwcrpq->nctrhpwq", packed)
|
||||
return latent.reshape(-1, channel, t * pt, h * ph, w * pw).contiguous()
|
||||
|
||||
|
||||
def minimax_h3_unpack_audio_tokens(
|
||||
rows: torch.Tensor,
|
||||
*,
|
||||
audio_t: int,
|
||||
audio_channel: int,
|
||||
) -> torch.Tensor:
|
||||
"""Unpack DiT audio token rows into SGLang audio VAE latent [C,latent_dim,T]."""
|
||||
|
||||
_rank(rows, "audio token rows", 2)
|
||||
audio_t = int(audio_t)
|
||||
audio_channel = int(audio_channel)
|
||||
if audio_t <= 0 or audio_channel <= 0:
|
||||
raise ValueError(
|
||||
f"audio_t and audio_channel must be positive, got {audio_t=} "
|
||||
f"{audio_channel=}"
|
||||
)
|
||||
if int(rows.shape[0]) != audio_t:
|
||||
raise ValueError(f"audio rows {int(rows.shape[0])} != audio_t {audio_t}")
|
||||
if audio_t % audio_channel:
|
||||
raise ValueError(
|
||||
f"audio_t must be divisible by audio_channel, got {audio_t=} "
|
||||
f"{audio_channel=}"
|
||||
)
|
||||
native = rows.reshape(audio_channel, audio_t // audio_channel, int(rows.shape[-1]))
|
||||
return native.permute(0, 2, 1).contiguous()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"minimax_h3_patchify_video_latent",
|
||||
"minimax_h3_unpack_audio_tokens",
|
||||
"minimax_h3_unpatchify_video_tokens",
|
||||
]
|
||||
+346
@@ -0,0 +1,346 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 probe -> resolve-once admission hook.
|
||||
|
||||
This module is intentionally data/CPU only. It localizes condition media,
|
||||
caches display-geometry facts, freezes every target/material canvas, and
|
||||
resolves the real aligned workload before a video job is published or sent to
|
||||
the scheduler.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.constants import (
|
||||
MINIMAX_H3_MAX_DURATION_SECONDS,
|
||||
MINIMAX_H3_MIN_DURATION_SECONDS,
|
||||
MINIMAX_H3_SUPPORTED_FPS,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
|
||||
MINIMAX_H3_TEMP_DIRS_EXTRA_KEY,
|
||||
minimax_h3_cleanup_temp_dirs,
|
||||
minimax_h3_probe_material,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
|
||||
MINIMAX_H3_BASE_SHORT_EDGE,
|
||||
MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY,
|
||||
MiniMaxH3ResolvedPlan,
|
||||
minimax_h3_plan_from_batch,
|
||||
minimax_h3_resolve_spatial_shape,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
|
||||
minimax_h3_align_frame_count,
|
||||
minimax_h3_audio_latent_t,
|
||||
minimax_h3_video_latent_t,
|
||||
)
|
||||
|
||||
MINIMAX_H3_PROBE_FACTS_EXTRA_KEY = "minimax_h3_probe_facts_by_condition"
|
||||
MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY = "minimax_h3_resolved_material_shapes"
|
||||
|
||||
|
||||
def _replace_plan_shape(
|
||||
plan: MiniMaxH3ResolvedPlan, shape: dict[str, Any]
|
||||
) -> MiniMaxH3ResolvedPlan:
|
||||
return MiniMaxH3ResolvedPlan(
|
||||
task=plan.task,
|
||||
prompt=plan.prompt,
|
||||
seed=plan.seed,
|
||||
materials=plan.materials,
|
||||
encoders=plan.encoders,
|
||||
branches=plan.branches,
|
||||
default_flow_shift=plan.default_flow_shift,
|
||||
default_audio_flow_shift=plan.default_audio_flow_shift,
|
||||
flow_shift=plan.flow_shift,
|
||||
audio_flow_shift=plan.audio_flow_shift,
|
||||
shape=shape,
|
||||
condition_mask=plan.condition_mask,
|
||||
)
|
||||
|
||||
|
||||
def _display_shape(facts: dict[str, Any], *, label: str) -> tuple[float, float]:
|
||||
try:
|
||||
width = float(facts["display_width"])
|
||||
height = float(facts["display_height"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ValueError(f"{label} has no usable display geometry") from exc
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError(f"{label} has no usable display geometry")
|
||||
return width, height
|
||||
|
||||
|
||||
def _resolve_deferred_spatial_shape(
|
||||
plan: MiniMaxH3ResolvedPlan,
|
||||
shape: dict[str, Any],
|
||||
probe_facts: dict[int, dict[str, Any]],
|
||||
) -> None:
|
||||
if str(shape.get("geometry")) != "deferred":
|
||||
return
|
||||
if plan.task == "fl2va":
|
||||
candidates = [
|
||||
material
|
||||
for material in plan.materials
|
||||
if material.material_chain == "image.target_canvas"
|
||||
]
|
||||
if len(candidates) not in {1, 2}:
|
||||
raise ValueError(
|
||||
f"fl2va requires one or two keyframe materials, got {len(candidates)}"
|
||||
)
|
||||
for material in candidates:
|
||||
if material.frame_index not in {0, -1}:
|
||||
raise ValueError(
|
||||
"fl2va deferred geometry requires semantic frame_index 0 or "
|
||||
f"-1, got {material.frame_index!r} for "
|
||||
f"conditions[{material.condition_index}]"
|
||||
)
|
||||
# Select by semantic time, not request/material iteration order. The
|
||||
# last-frame sentinel sorts after the first-frame anchor.
|
||||
source = min(
|
||||
candidates,
|
||||
key=lambda material: (
|
||||
material.frame_index == -1,
|
||||
int(material.frame_index),
|
||||
int(material.condition_index),
|
||||
),
|
||||
)
|
||||
geometry_source = (
|
||||
"first_keyframe" if source.frame_index == 0 else "last_keyframe"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"task {plan.task!r} has unsupported deferred target geometry")
|
||||
width, height = _display_shape(
|
||||
probe_facts[int(source.condition_index)],
|
||||
label=f"conditions[{source.condition_index}]",
|
||||
)
|
||||
shape.update(
|
||||
minimax_h3_resolve_spatial_shape(
|
||||
width=width,
|
||||
height=height,
|
||||
base_short_edge=int(shape["base_short_edge"]),
|
||||
)
|
||||
)
|
||||
shape["geometry_source"] = geometry_source
|
||||
shape["geometry_source_condition_index"] = int(source.condition_index)
|
||||
if source.frame_index is not None:
|
||||
shape["geometry_source_frame_index"] = int(source.frame_index)
|
||||
|
||||
|
||||
def _resolve_deferred_temporal_shape(
|
||||
plan: MiniMaxH3ResolvedPlan,
|
||||
shape: dict[str, Any],
|
||||
probe_facts: dict[int, dict[str, Any]],
|
||||
) -> None:
|
||||
if str(shape.get("temporal")) != "deferred_from_audio_reference":
|
||||
return
|
||||
sources = [
|
||||
material
|
||||
for material in plan.materials
|
||||
if material.condition_type in {"audio", "video", "video_audio"}
|
||||
and bool(probe_facts[int(material.condition_index)].get("has_audio"))
|
||||
]
|
||||
if len(sources) != 1:
|
||||
raise ValueError(
|
||||
"audio-derived target duration requires exactly one probed "
|
||||
f"condition with an audio stream, got {len(sources)}"
|
||||
)
|
||||
source = sources[0]
|
||||
facts = probe_facts[int(source.condition_index)]
|
||||
try:
|
||||
duration_seconds = float(facts["audio_duration_seconds"]) - float(
|
||||
source.start_time_seconds
|
||||
)
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ValueError("audio reference has no positive probed duration") from exc
|
||||
if duration_seconds <= 0:
|
||||
raise ValueError("audio reference has no positive probed duration")
|
||||
if not (
|
||||
MINIMAX_H3_MIN_DURATION_SECONDS
|
||||
<= duration_seconds
|
||||
<= MINIMAX_H3_MAX_DURATION_SECONDS
|
||||
):
|
||||
raise ValueError(
|
||||
"audio reference duration must be in "
|
||||
f"[{MINIMAX_H3_MIN_DURATION_SECONDS:g}, "
|
||||
f"{MINIMAX_H3_MAX_DURATION_SECONDS:g}] seconds, got {duration_seconds:g}"
|
||||
)
|
||||
fps = MINIMAX_H3_SUPPORTED_FPS
|
||||
frame_count = minimax_h3_align_frame_count(int(round(duration_seconds * fps)))
|
||||
aligned_duration = frame_count / fps
|
||||
shape.update(
|
||||
{
|
||||
"temporal": "resolved_from_audio_reference",
|
||||
"duration_seconds": aligned_duration,
|
||||
"frame_count": frame_count,
|
||||
"video_latent_t": minimax_h3_video_latent_t(frame_count),
|
||||
"audio_latent_t": minimax_h3_audio_latent_t(aligned_duration),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _validate_reference_start_times(
|
||||
plan: MiniMaxH3ResolvedPlan,
|
||||
probe_facts: dict[int, dict[str, Any]],
|
||||
) -> None:
|
||||
for material in plan.materials:
|
||||
start_time_seconds = float(material.start_time_seconds)
|
||||
if start_time_seconds == 0:
|
||||
continue
|
||||
condition_index = int(material.condition_index)
|
||||
facts = probe_facts[condition_index]
|
||||
try:
|
||||
video_duration_seconds = float(facts["video_duration_seconds"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
f"conditions[{condition_index}].start_time_seconds requires "
|
||||
"a video with a positive probed duration"
|
||||
) from exc
|
||||
if start_time_seconds >= video_duration_seconds:
|
||||
raise ValueError(
|
||||
f"conditions[{condition_index}].start_time_seconds must be less "
|
||||
f"than the video duration {video_duration_seconds:g}, got "
|
||||
f"{start_time_seconds:g}"
|
||||
)
|
||||
if bool(facts.get("has_audio")):
|
||||
try:
|
||||
audio_duration_seconds = float(facts["audio_duration_seconds"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
f"conditions[{condition_index}] has no usable audio duration"
|
||||
) from exc
|
||||
if start_time_seconds >= audio_duration_seconds:
|
||||
raise ValueError(
|
||||
f"conditions[{condition_index}].start_time_seconds must be "
|
||||
f"less than the soundtrack duration {audio_duration_seconds:g}, "
|
||||
f"got {start_time_seconds:g}"
|
||||
)
|
||||
|
||||
|
||||
def _resolved_work_frame_count(
|
||||
plan: MiniMaxH3ResolvedPlan,
|
||||
shape: dict[str, Any],
|
||||
probe_facts: dict[int, dict[str, Any]],
|
||||
) -> int:
|
||||
if shape.get("frame_count") is not None:
|
||||
return int(shape["frame_count"])
|
||||
raise ValueError("MiniMax H3 target frame count remained unresolved before queue")
|
||||
|
||||
|
||||
def _preserve_prequeue_material_dirs(batch: Any) -> None:
|
||||
temp_registry = batch.extra.get(MINIMAX_H3_TEMP_DIRS_EXTRA_KEY)
|
||||
if isinstance(temp_registry, dict) and "material" in temp_registry:
|
||||
# Multi-output dispatch shallow-copies request extras. Keep the
|
||||
# single localized source closure owned by the API request until
|
||||
# every output finishes; encoder-stage "material" cleanup must not
|
||||
# delete it after the first expanded output.
|
||||
paths = temp_registry.pop("material")
|
||||
prequeue_paths = temp_registry.setdefault("prequeue_material", [])
|
||||
for path in paths if isinstance(paths, list) else ():
|
||||
if path not in prequeue_paths:
|
||||
prequeue_paths.append(path)
|
||||
|
||||
|
||||
def minimax_h3_prepare_for_queue(batch: Any) -> MiniMaxH3ResolvedPlan:
|
||||
"""Freeze MiniMax H3 media/shape facts before queue admission."""
|
||||
|
||||
try:
|
||||
plan = minimax_h3_plan_from_batch(batch)
|
||||
if plan is None:
|
||||
raise ValueError(
|
||||
"MiniMax H3 pre-queue validation requires a canonical request"
|
||||
)
|
||||
|
||||
probe_facts: dict[int, dict[str, Any]] = {}
|
||||
for material in plan.materials:
|
||||
probe_facts[int(material.condition_index)] = minimax_h3_probe_material(
|
||||
batch,
|
||||
material.uri,
|
||||
condition_type=material.condition_type,
|
||||
condition_index=int(material.condition_index),
|
||||
)
|
||||
batch.extra[MINIMAX_H3_PROBE_FACTS_EXTRA_KEY] = probe_facts
|
||||
|
||||
shape = dict(plan.shape)
|
||||
_validate_reference_start_times(plan, probe_facts)
|
||||
_resolve_deferred_spatial_shape(plan, shape, probe_facts)
|
||||
_resolve_deferred_temporal_shape(plan, shape, probe_facts)
|
||||
if str(shape.get("geometry")) != "resolved_v2":
|
||||
raise ValueError(
|
||||
"MiniMax H3 target geometry remained unresolved before queue"
|
||||
)
|
||||
|
||||
material_shapes: dict[int, dict[str, Any]] = {}
|
||||
for material in plan.materials:
|
||||
condition_index = int(material.condition_index)
|
||||
if material.material_chain == "image.target_canvas":
|
||||
resolved = {
|
||||
key: shape[key]
|
||||
for key in (
|
||||
"geometry",
|
||||
"shape_policy_version",
|
||||
"base_short_edge",
|
||||
"effective_short_edge",
|
||||
"size_mode",
|
||||
"max_pixels",
|
||||
"multiple",
|
||||
"rounding",
|
||||
"width",
|
||||
"height",
|
||||
)
|
||||
if key in shape
|
||||
}
|
||||
elif material.material_chain in {
|
||||
"video.reference_preserve",
|
||||
"video_audio.reference_preserve",
|
||||
}:
|
||||
width, height = _display_shape(
|
||||
probe_facts[condition_index],
|
||||
label=f"conditions[{condition_index}]",
|
||||
)
|
||||
resolved = minimax_h3_resolve_spatial_shape(
|
||||
width=width,
|
||||
height=height,
|
||||
base_short_edge=MINIMAX_H3_BASE_SHORT_EDGE,
|
||||
)
|
||||
elif material.material_chain == "image.reference_preserve":
|
||||
width, height = _display_shape(
|
||||
probe_facts[condition_index],
|
||||
label=f"conditions[{condition_index}]",
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.reference_encoding import (
|
||||
minimax_h3_resolve_reference_image_shape,
|
||||
)
|
||||
|
||||
resolved = minimax_h3_resolve_reference_image_shape(
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
else:
|
||||
continue
|
||||
resolved = dict(resolved)
|
||||
resolved["condition_index"] = condition_index
|
||||
material_shapes[condition_index] = resolved
|
||||
batch.extra[MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY] = material_shapes
|
||||
|
||||
work_frames = _resolved_work_frame_count(plan, shape, probe_facts)
|
||||
resolved_plan = _replace_plan_shape(plan, shape)
|
||||
batch.extra[MINIMAX_H3_RESOLVED_PLAN_EXTRA_KEY] = resolved_plan
|
||||
# Req delegates these fields to SamplingParams. Freezing them here makes
|
||||
# queue metadata and dynamic-batch signatures use the same final shape
|
||||
# that the MiniMax H3 stages consume.
|
||||
batch.width = int(shape["width"])
|
||||
batch.height = int(shape["height"])
|
||||
batch.fps = MINIMAX_H3_SUPPORTED_FPS
|
||||
batch.num_frames = int(work_frames)
|
||||
_preserve_prequeue_material_dirs(batch)
|
||||
return resolved_plan
|
||||
except Exception:
|
||||
minimax_h3_cleanup_temp_dirs(batch)
|
||||
batch.extra.pop(MINIMAX_H3_PROBE_FACTS_EXTRA_KEY, None)
|
||||
batch.extra.pop(MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY, None)
|
||||
raise
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MINIMAX_H3_PROBE_FACTS_EXTRA_KEY",
|
||||
"MINIMAX_H3_RESOLVED_MATERIAL_SHAPES_EXTRA_KEY",
|
||||
"minimax_h3_prepare_for_queue",
|
||||
]
|
||||
+278
@@ -0,0 +1,278 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax H3 Qwen presentation building.
|
||||
|
||||
Builds the positive presentation token stream:
|
||||
- fl2va: '<Picture 1>: ' label + vision block (<|vision_start|> +
|
||||
N*<|image_pad|> + <|vision_end|>) + prompt text.
|
||||
- t2va: prompt text only (no vision block).
|
||||
Prompt text passes through verbatim (no stripping or rewriting).
|
||||
|
||||
All presentation variants are emitted through the shared ``_Presentation``
|
||||
accumulator so ids and AdaLN token tags cannot drift apart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
VISION_START = "<|vision_start|>"
|
||||
VISION_END = "<|vision_end|>"
|
||||
IMAGE_PAD = "<|image_pad|>"
|
||||
VIDEO_PAD = "<|video_pad|>"
|
||||
|
||||
_TEXT_TAG = 1
|
||||
_VIDEO_TAG = 0
|
||||
|
||||
|
||||
def _text_ids(tokenizer: Any, text: str) -> list[int]:
|
||||
return list(tokenizer(text, add_special_tokens=False)["input_ids"])
|
||||
|
||||
|
||||
def _vision_block_ids(tokenizer: Any, pad_token: str, count: int) -> list[int]:
|
||||
return (
|
||||
[tokenizer.convert_tokens_to_ids(VISION_START)]
|
||||
+ [tokenizer.convert_tokens_to_ids(pad_token)] * int(count)
|
||||
+ [tokenizer.convert_tokens_to_ids(VISION_END)]
|
||||
)
|
||||
|
||||
|
||||
class _Presentation:
|
||||
"""Accumulates aligned (ids, token_tags) presentation segments."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.ids: list[int] = []
|
||||
self.tags: list[int] = []
|
||||
|
||||
def text(self, token_ids: list[int]) -> None:
|
||||
self.ids += token_ids
|
||||
self.tags += [_TEXT_TAG] * len(token_ids)
|
||||
|
||||
def vision(self, token_ids: list[int]) -> None:
|
||||
self.ids += token_ids
|
||||
self.tags += [_VIDEO_TAG] * len(token_ids)
|
||||
|
||||
def build(self) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return (
|
||||
torch.tensor(self.ids, dtype=torch.long),
|
||||
torch.tensor(self.tags, dtype=torch.long),
|
||||
)
|
||||
|
||||
|
||||
def _timestamped_video_blocks(
|
||||
presentation: _Presentation,
|
||||
tokenizer: Any,
|
||||
*,
|
||||
counts: Sequence[int],
|
||||
timestamps: Sequence[float],
|
||||
context: str,
|
||||
) -> None:
|
||||
"""Emit per-temporal-block ``<{t:.1f} seconds>`` text + VIDEO vision."""
|
||||
|
||||
counts = [int(value) for value in counts]
|
||||
timestamps = [float(value) for value in timestamps]
|
||||
if not counts or len(counts) != len(timestamps):
|
||||
raise ValueError(f"{context}video block token counts and timestamps must align")
|
||||
for count, timestamp in zip(counts, timestamps):
|
||||
if count <= 0:
|
||||
raise ValueError(f"{context}video block token count must be positive")
|
||||
presentation.text(_text_ids(tokenizer, f"<{timestamp:.1f} seconds>"))
|
||||
presentation.vision(_vision_block_ids(tokenizer, VIDEO_PAD, count))
|
||||
|
||||
|
||||
def minimax_h3_text_only_ids(tokenizer: Any, prompt: str) -> torch.Tensor:
|
||||
"""t2va presentation: verbatim prompt, no special tokens."""
|
||||
if not prompt:
|
||||
raise ValueError("prompt must be non-empty")
|
||||
return torch.tensor(_text_ids(tokenizer, prompt), dtype=torch.long)
|
||||
|
||||
|
||||
def minimax_h3_multi_image_presentation(
|
||||
tokenizer: Any,
|
||||
*,
|
||||
prompt: str,
|
||||
image_token_counts: list[int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if not image_token_counts:
|
||||
raise ValueError("image_token_counts must be non-empty")
|
||||
presentation = _Presentation()
|
||||
for index, count in enumerate(image_token_counts, start=1):
|
||||
if int(count) <= 0:
|
||||
raise ValueError("image_token_count must be positive")
|
||||
presentation.text(_text_ids(tokenizer, f"<Picture {index}>: "))
|
||||
presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
|
||||
presentation.text(_text_ids(tokenizer, prompt))
|
||||
return presentation.build()
|
||||
|
||||
|
||||
def minimax_h3_ref2va_presentation(
|
||||
tokenizer: Any,
|
||||
*,
|
||||
prompt: str,
|
||||
condition_labels: list[tuple[str, int]],
|
||||
image_token_count: int | list[int] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""ref2va positive presentation:
|
||||
|
||||
per condition in request order — image i: ``<Picture i>: `` label followed
|
||||
by the vision block; audio j: ``<Audio j>: `` label only (audio content
|
||||
never enters Qwen) — then the verbatim prompt. Returns (ids, token_tags)
|
||||
with the vision block tagged VIDEO(0) and everything else TEXT(1).
|
||||
|
||||
condition_labels: [("image", 1), ("audio", 1), ...] with 1-based ordinals
|
||||
per type.
|
||||
"""
|
||||
return minimax_h3_ref2va_video_presentation(
|
||||
tokenizer,
|
||||
prompt=prompt,
|
||||
condition_labels=condition_labels,
|
||||
image_token_count=image_token_count,
|
||||
video_block_token_counts=None,
|
||||
video_block_timestamps=None,
|
||||
)
|
||||
|
||||
|
||||
def _as_int_list(value: int | Sequence[int] | None, *, name: str) -> list[int]:
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, int):
|
||||
return [int(value)]
|
||||
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
|
||||
raise ValueError(f"{name} must be an int or a sequence of ints")
|
||||
return [int(item) for item in value]
|
||||
|
||||
|
||||
def _as_nested_int_list(
|
||||
value: Sequence[int] | Sequence[Sequence[int]] | None,
|
||||
*,
|
||||
name: str,
|
||||
) -> list[list[int]]:
|
||||
if value is None:
|
||||
return []
|
||||
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
|
||||
raise ValueError(f"{name} must be a sequence")
|
||||
if len(value) == 0:
|
||||
return []
|
||||
first = value[0]
|
||||
if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
|
||||
out: list[list[int]] = []
|
||||
for group in value:
|
||||
if not isinstance(group, Sequence) or isinstance(group, (str, bytes)):
|
||||
raise ValueError(f"{name} must not mix nested and flat entries")
|
||||
out.append([int(item) for item in group])
|
||||
return out
|
||||
return [[int(item) for item in value]]
|
||||
|
||||
|
||||
def _as_nested_float_list(
|
||||
value: Sequence[float] | Sequence[Sequence[float]] | None,
|
||||
*,
|
||||
name: str,
|
||||
) -> list[list[float]]:
|
||||
if value is None:
|
||||
return []
|
||||
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
|
||||
raise ValueError(f"{name} must be a sequence")
|
||||
if len(value) == 0:
|
||||
return []
|
||||
first = value[0]
|
||||
if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
|
||||
out: list[list[float]] = []
|
||||
for group in value:
|
||||
if not isinstance(group, Sequence) or isinstance(group, (str, bytes)):
|
||||
raise ValueError(f"{name} must not mix nested and flat entries")
|
||||
out.append([float(item) for item in group])
|
||||
return out
|
||||
return [[float(item) for item in value]]
|
||||
|
||||
|
||||
def minimax_h3_ref2va_video_presentation(
|
||||
tokenizer: Any,
|
||||
*,
|
||||
prompt: str,
|
||||
condition_labels: list[tuple[str, int]],
|
||||
image_token_count: int | list[int] | None,
|
||||
video_block_token_counts: list[int] | list[list[int]] | None,
|
||||
video_block_timestamps: list[float] | list[list[float]] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""ref2va (optionally with video refs) positive presentation:
|
||||
|
||||
per condition in request order —
|
||||
- image i: ``<Picture i>: `` label + one image vision block;
|
||||
- audio j: ``<Audio j>: `` label only (audio content never enters Qwen);
|
||||
- video k: ``<Video k>: `` label, then per temporal block a timestamp
|
||||
text ``<{t:.1f} seconds>`` followed by a VIDEO vision block
|
||||
(<|vision_start|> + <|video_pad|> x n + <|vision_end|>). Timestamps are
|
||||
the mean of each merged frame pair (Qwen3VL temporal merge 2; odd frame
|
||||
counts repeat the last frame), emitting the
|
||||
``<0.2 seconds>`` ..
|
||||
``<4.0 seconds>`` sequence — note Python bankers-rounding at .1f.
|
||||
then the verbatim prompt. Vision blocks are tagged VIDEO(0), everything
|
||||
else TEXT(1).
|
||||
"""
|
||||
if not prompt:
|
||||
raise ValueError("prompt must be non-empty")
|
||||
presentation = _Presentation()
|
||||
image_token_counts = _as_int_list(image_token_count, name="image_token_count")
|
||||
video_counts_by_ref = _as_nested_int_list(
|
||||
video_block_token_counts,
|
||||
name="video_block_token_counts",
|
||||
)
|
||||
video_timestamps_by_ref = _as_nested_float_list(
|
||||
video_block_timestamps,
|
||||
name="video_block_timestamps",
|
||||
)
|
||||
if len(video_counts_by_ref) != len(video_timestamps_by_ref):
|
||||
raise ValueError("video block token counts and timestamps must align")
|
||||
image_seen = 0
|
||||
video_seen = 0
|
||||
for cond_type, ordinal in condition_labels:
|
||||
if cond_type == "image":
|
||||
image_seen += 1
|
||||
if image_seen > len(image_token_counts):
|
||||
raise ValueError("image_token_count required for an image reference")
|
||||
count = int(image_token_counts[image_seen - 1])
|
||||
if count <= 0:
|
||||
raise ValueError("image_token_count required for an image reference")
|
||||
presentation.text(_text_ids(tokenizer, f"<Picture {ordinal}>: "))
|
||||
presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count))
|
||||
elif cond_type == "audio":
|
||||
presentation.text(_text_ids(tokenizer, f"<Audio {ordinal}>: "))
|
||||
elif cond_type == "video":
|
||||
video_seen += 1
|
||||
if video_seen > len(video_counts_by_ref):
|
||||
raise ValueError(
|
||||
"video reference requires block token counts and timestamps"
|
||||
)
|
||||
counts = video_counts_by_ref[video_seen - 1]
|
||||
timestamps = video_timestamps_by_ref[video_seen - 1]
|
||||
if not counts or not timestamps:
|
||||
raise ValueError(
|
||||
"video reference requires block token counts and timestamps"
|
||||
)
|
||||
presentation.text(_text_ids(tokenizer, f"<Video {ordinal}>: "))
|
||||
_timestamped_video_blocks(
|
||||
presentation,
|
||||
tokenizer,
|
||||
counts=counts,
|
||||
timestamps=timestamps,
|
||||
context="",
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported ref2va condition type {cond_type!r}")
|
||||
if image_seen != len(image_token_counts):
|
||||
raise ValueError("unused image_token_count entries")
|
||||
if video_seen != len(video_counts_by_ref):
|
||||
raise ValueError("unused video block token count entries")
|
||||
presentation.text(_text_ids(tokenizer, prompt))
|
||||
return presentation.build()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"minimax_h3_multi_image_presentation",
|
||||
"minimax_h3_ref2va_presentation",
|
||||
"minimax_h3_ref2va_video_presentation",
|
||||
"minimax_h3_text_only_ids",
|
||||
]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user