[Feature][MiniMax-H3] Support diffusion continuous batching (#5810)
Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
15 files changed
+1834
-140
No files matched your search
@@ -101,8 +101,16 @@ set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
|
||||
`--max-num-seqs >1`. See the
|
||||
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||
for its validated configuration. Helios supports single-request step execution only: use
|
||||
`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's
|
||||
documentation and source for the latest support status.
|
||||
`--step-execution --max-num-seqs 1` for Helios. MiniMax H3 supports step-wise
|
||||
continuous batching by packing co-batched requests into one sequence that keeps
|
||||
a separate attention document per request; that layout needs a backend which
|
||||
honors the packed `cu_seqlens` metadata, so run it with
|
||||
`--diffusion-attention-backend FLASH_ATTN` (other backends stay correct but
|
||||
fall back to one transformer forward per request). Batching H3 does not improve
|
||||
its throughput — see the measured numbers in the
|
||||
[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
|
||||
— so keep `--max-num-seqs 1` unless you need step-level scheduling. Consult the
|
||||
selected pipeline's documentation and source for the latest support status.
|
||||
|
||||
## Streaming Output
|
||||
|
||||
|
||||
@@ -380,6 +380,63 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
|
||||
`task=ref2va` routes to
|
||||
`Ref2VA/transformer`. T2VA uses the FL2VA DiT.
|
||||
|
||||
### Step execution and continuous batching
|
||||
|
||||
H3 implements the step-wise execution contract, so the scheduler can admit and
|
||||
retire requests between denoise steps instead of running one request end to
|
||||
end. Add the feature gate, then raise `--max-num-seqs` to co-batch:
|
||||
|
||||
```bash
|
||||
--step-execution --max-num-seqs 4
|
||||
```
|
||||
|
||||
Co-batched requests are packed into a single sequence that keeps one attention
|
||||
document per request, so a batch costs one DiT forward. That packing requires
|
||||
`--diffusion-attention-backend FLASH_ATTN`; other backends fall back to one
|
||||
forward per request. `--max-num-seqs 1` keeps the conservative single-request
|
||||
step path. Cache acceleration (`--cache-backend`) is not available in step mode.
|
||||
|
||||
!!! warning "Co-batching does not improve H3 throughput for large simultaneous requests"
|
||||
Keep `--max-num-seqs 1` unless you specifically need scheduler-level control
|
||||
(admitting and retiring requests between denoise steps). Measured on two
|
||||
H100s (TP2, BF16, 672x384, 209 frames, 30 steps, 4 requests at concurrency
|
||||
4 with `--request-rate inf`, i.e. all four submitted at time 0; one packed
|
||||
request is 16384 rows):
|
||||
|
||||
| Configuration | Wall time | Mean latency | Peak memory |
|
||||
|---------------|-----------|--------------|-------------|
|
||||
| request mode | 174.8 s | 111.5 s | 72.4 GB |
|
||||
| `--step-execution --max-num-seqs 1` | 179.0 s | 113.8 s | 72.4 GB |
|
||||
| `--step-execution --max-num-seqs 4` | 182.1 s | 175.7 s | 78.3 GB |
|
||||
|
||||
A single H3 denoise step is a compute-bound dense GEMM over an already long
|
||||
packed sequence, so fusing N requests costs N times the FLOPs and buys almost
|
||||
no amortization — unlike LLM decoding, which is memory-bandwidth bound.
|
||||
Going from one request per step to four cuts the per-request denoise cost
|
||||
only from 1.323 s to 1.291 s (2.4%), which the step-mode bookkeeping then
|
||||
spends. Mean latency degrades further because co-batched requests finish
|
||||
together instead of staggered. Quantization moves the absolute numbers
|
||||
without changing this: with online `int8` the same workload runs in 153.3 s
|
||||
at 56.9 GB (request mode), and `--max-num-seqs 4` is still 5.0% slower than
|
||||
request mode.
|
||||
|
||||
Two workloads outside this table are unmeasured and may behave differently:
|
||||
|
||||
- **Staggered arrivals (admission latency).** A single 16384-row H3 request
|
||||
already saturates the GPUs, so submitting all requests at time 0 is the
|
||||
one arrival pattern where co-batching cannot win: it can only bunch
|
||||
completions. In request mode a new request queues behind the whole
|
||||
in-flight generation (~45 s at this size); step mode admits at the next
|
||||
denoise-step boundary (~1.3 s). Whether that shows up as a wall-clock
|
||||
benefit under Poisson arrivals is not yet measured for H3. To reproduce,
|
||||
run request mode vs `--step-execution --max-num-seqs 4` with
|
||||
`diffusion_benchmark_serving.py --request-rate 0.05` (roughly one request
|
||||
every 20 s) and report mean / p95 latency, which then includes queueing.
|
||||
- **Small requests.** The numbers above are drawn from 672x384 / 209-frame
|
||||
requests. A short clip at lower resolution packs a few thousand rows and
|
||||
may not saturate the hardware; co-batching may amortize better there.
|
||||
This is also unmeasured.
|
||||
|
||||
### Online FP8 quantization
|
||||
|
||||
MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
|
||||
@@ -924,7 +981,12 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
|
||||
- TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached.
|
||||
- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
|
||||
both task-specific DiTs, and loads shared components once from `FL2VA`.
|
||||
- H3 currently executes one generation request per diffusion batch.
|
||||
- Request mode executes one generation request per diffusion batch. Use
|
||||
`--step-execution` with `--max-num-seqs N` to admit several requests at once
|
||||
(see
|
||||
[Execution modes](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/execution_modes.md));
|
||||
step
|
||||
mode does not support `cache_backend`.
|
||||
- The first regional-compile request is a warmup and should not be included in
|
||||
steady-state performance measurements.
|
||||
- The serving path accepts fewer references than the model supports. H3 documents up
|
||||
|
||||
@@ -821,6 +821,82 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
|
||||
)
|
||||
|
||||
|
||||
def _fake_packed_attention(
|
||||
backend_name: str,
|
||||
*,
|
||||
prefix_kv_slicing: bool = False,
|
||||
supports_multi_doc: bool = False,
|
||||
):
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
)
|
||||
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = prefix_kv_slicing
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return backend_name
|
||||
|
||||
@classmethod
|
||||
def supports_multi_doc_packed_varlen(cls) -> bool:
|
||||
return supports_multi_doc
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
attn_backend = FakeBackend
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.metadata = None
|
||||
|
||||
def forward(self, query, key, value, metadata):
|
||||
self.metadata = metadata
|
||||
return query
|
||||
|
||||
attention = object.__new__(MiniMaxH3Attention)
|
||||
torch.nn.Module.__init__(attention)
|
||||
attention.attention = FakeAttention()
|
||||
return attention
|
||||
|
||||
|
||||
def test_packed_attention_drops_prefix_semantics_for_a_co_batched_layout():
|
||||
"""A batch's valid rows are block-diagonal, so no prefix can describe them."""
|
||||
attention = _fake_packed_attention("FLASH_ATTN", supports_multi_doc=True)
|
||||
q = torch.randn(12, 2, 4)
|
||||
|
||||
attention._run_packed_attention(
|
||||
q,
|
||||
q,
|
||||
q,
|
||||
# Two requests: rows + padding tail each.
|
||||
cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32),
|
||||
max_seqlen=5,
|
||||
packed_total=12,
|
||||
num_requests=2,
|
||||
)
|
||||
|
||||
assert attention.attention.metadata.attn_mask is None
|
||||
assert attention.attention.metadata.extra["valid_kv_length"] == 12
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend_name", ["TORCH_SDPA", "FLASH_ATTN"])
|
||||
def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_name):
|
||||
"""Backend names alone must not allow cross-request attention."""
|
||||
attention = _fake_packed_attention(backend_name, supports_multi_doc=False)
|
||||
q = torch.randn(12, 2, 4)
|
||||
|
||||
with pytest.raises(ValueError, match=backend_name):
|
||||
attention._run_packed_attention(
|
||||
q,
|
||||
q,
|
||||
q,
|
||||
cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32),
|
||||
max_seqlen=5,
|
||||
packed_total=12,
|
||||
num_requests=2,
|
||||
)
|
||||
|
||||
|
||||
def test_reference_image_resize_contract():
|
||||
from PIL import Image
|
||||
|
||||
|
||||
@@ -0,0 +1,471 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""MiniMax H3 step-wise execution (continuous batching) contract tests.
|
||||
|
||||
These run on CPU against a stand-in DiT, so they cover the packing, the
|
||||
per-request scheduler math, and the runner-facing state wiring without needing
|
||||
checkpoint weights.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
_HIDDEN = 8
|
||||
|
||||
|
||||
class _SegmentMeanModel:
|
||||
"""Stand-in DiT whose output depends on packed-document boundaries.
|
||||
|
||||
Every row is mixed with the mean of its own ``cu_seqlens`` document and with
|
||||
its own timestep, so a wrong document boundary, row offset, or timestep
|
||||
assignment shows up as a numeric difference instead of passing silently.
|
||||
"""
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
x = kwargs["x"][0]
|
||||
audio_x = kwargs["audio_x"][0]
|
||||
bounds = kwargs["packed_seq_params"]["cu_seqlens_q"].tolist()
|
||||
img_pos = kwargs["img_pos_info"]["position_ids"]
|
||||
audio_pos = kwargs["audio_pos_info"]["position_ids"]
|
||||
|
||||
pooled_video = torch.zeros_like(x)
|
||||
pooled_audio = torch.zeros_like(audio_x)
|
||||
for start, stop in zip(bounds[:-1], bounds[1:]):
|
||||
if stop <= start:
|
||||
continue
|
||||
pooled_video[start:stop] = x[start:stop].mean(dim=0, keepdim=True)
|
||||
pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True)
|
||||
|
||||
row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1)
|
||||
video = (pooled_video + 0.5 * x + row_timesteps)[img_pos]
|
||||
audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos]
|
||||
if not kwargs.get("skip_mask_out_condition", False):
|
||||
video = video * kwargs["update_mask"].view(-1).unsqueeze(-1)
|
||||
return video, audio
|
||||
|
||||
|
||||
def _make_branch(*, text_len: int, latent_t: int, latent_h: int, latent_w: int, audio_t: int, seed: int):
|
||||
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import MiniMaxH3DenoiseBranch
|
||||
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence
|
||||
|
||||
packed = minimax_h3_packed_sequence(
|
||||
text_len=text_len,
|
||||
latent_t=latent_t,
|
||||
latent_h=latent_h,
|
||||
latent_w=latent_w,
|
||||
audio_t=audio_t,
|
||||
include_keyframe_cond=False,
|
||||
)
|
||||
generator = torch.Generator().manual_seed(seed)
|
||||
text_embeddings = torch.randn(text_len, _HIDDEN, generator=generator, dtype=torch.float32)
|
||||
branch = MiniMaxH3DenoiseBranch(
|
||||
packed=packed,
|
||||
text_embeddings=text_embeddings,
|
||||
token_tags=packed["token_tags"],
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
video_rows = torch.randn(int(branch.img_pos.shape[0]), 96, generator=generator, dtype=torch.float32)
|
||||
audio_rows = torch.randn(int(branch.audio_pos.shape[0]), 32, generator=generator, dtype=torch.float32)
|
||||
return branch, video_rows, audio_rows
|
||||
|
||||
|
||||
def _sigmas(num_points: int, shift: float) -> list[float]:
|
||||
from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
|
||||
|
||||
return minimax_h3_time_shift_sigmas(num_steps=num_points, shift_scale=shift)
|
||||
|
||||
|
||||
def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
|
||||
state = StepRequestState(request_id=request_id, sampling=SimpleNamespace())
|
||||
state.latents = video_rows.clone()
|
||||
state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32)
|
||||
state.step_index = 0
|
||||
state.extra = {
|
||||
mod._STEP_BRANCH: branch,
|
||||
# Co-batched requests must share one DiT instance, or denoise_step()
|
||||
# treats the batch as mixed-task and falls back to one forward each.
|
||||
mod._STEP_TRANSFORMER: model,
|
||||
mod._STEP_AUDIO_ROWS: audio_rows.clone(),
|
||||
mod._STEP_COND_ANCHOR: None,
|
||||
mod._STEP_AUDIO_ANCHOR: None,
|
||||
mod._STEP_SIGMAS_VIDEO: sigmas_video,
|
||||
mod._STEP_SIGMAS_AUDIO: sigmas_audio,
|
||||
}
|
||||
return state
|
||||
|
||||
|
||||
def _step_pipeline(model, *, packed_batch_supported: bool = True):
|
||||
"""A pipeline instance carrying only what the step methods touch."""
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
pipeline.transformer = model
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline._transformer_for_task = lambda task: model
|
||||
pipeline._packed_batch_supported = lambda transformer: packed_batch_supported
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_step_execution_matches_request_mode_denoise_loop():
|
||||
"""Stepping through the contract must reproduce the request-mode loop."""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
|
||||
|
||||
model = _SegmentMeanModel()
|
||||
branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5)
|
||||
sigmas_video = _sigmas(6, 12.0)
|
||||
sigmas_audio = _sigmas(6, 3.0)
|
||||
|
||||
reference_video, reference_audio = minimax_h3_denoise_loop(
|
||||
model=model,
|
||||
positive=branch,
|
||||
initial_video_rows=video_rows,
|
||||
initial_audio_rows=audio_rows,
|
||||
keyframe_cond_rows=None,
|
||||
sigmas_video=sigmas_video,
|
||||
sigmas_audio=sigmas_audio,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
pipeline = _step_pipeline(model)
|
||||
state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
|
||||
input_batch = SimpleNamespace(states=(state,))
|
||||
|
||||
steps = 0
|
||||
while not state.denoise_completed:
|
||||
noise_pred = pipeline.denoise_step(input_batch, states=[state])
|
||||
pipeline.step_scheduler(state, noise_pred)
|
||||
steps += 1
|
||||
|
||||
assert steps == len(sigmas_video) - 1
|
||||
assert state.total_steps == steps
|
||||
torch.testing.assert_close(state.latents, reference_video)
|
||||
torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
|
||||
|
||||
|
||||
def test_batched_step_execution_matches_independent_requests():
|
||||
"""Two co-batched requests must land where they would have landed alone."""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
|
||||
model = _SegmentMeanModel()
|
||||
specs = [
|
||||
dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6),
|
||||
# Exactly 64 packed rows exercises the no-padding-tail boundary case.
|
||||
dict(text_len=46, latent_t=7, latent_h=2, latent_w=4, audio_t=2, seed=7),
|
||||
]
|
||||
# Different step counts, so the batch composition changes mid-flight.
|
||||
schedules = [(_sigmas(6, 12.0), _sigmas(6, 3.0)), (_sigmas(4, 12.0), _sigmas(4, 3.0))]
|
||||
|
||||
pipeline = _step_pipeline(model)
|
||||
|
||||
alone: list[tuple[torch.Tensor, torch.Tensor]] = []
|
||||
for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules):
|
||||
branch, video_rows, audio_rows = _make_branch(**spec)
|
||||
state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
|
||||
while not state.denoise_completed:
|
||||
pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
|
||||
alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS]))
|
||||
|
||||
states = []
|
||||
for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)):
|
||||
branch, video_rows, audio_rows = _make_branch(**spec)
|
||||
if index == 1:
|
||||
assert branch.used_len == branch.seq_len
|
||||
states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio))
|
||||
|
||||
active = list(states)
|
||||
while active:
|
||||
noise_pred = pipeline.denoise_step(SimpleNamespace(states=tuple(active)), states=active)
|
||||
offset = 0
|
||||
for state in active:
|
||||
rows = state.latents.shape[0]
|
||||
pipeline.step_scheduler(state, noise_pred[offset : offset + rows])
|
||||
offset += rows
|
||||
assert offset == noise_pred.shape[0]
|
||||
# Finished requests leave the batch, exactly like the runner drops them.
|
||||
active = [state for state in active if not state.denoise_completed]
|
||||
|
||||
for state, (expected_video, expected_audio) in zip(states, alone):
|
||||
torch.testing.assert_close(state.latents, expected_video)
|
||||
torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], expected_audio)
|
||||
|
||||
|
||||
def test_both_modes_publish_denoise_progress_for_gated_attention():
|
||||
"""TRTLLM's skip gate and RAINFUSION's warmup stay dense without these."""
|
||||
from vllm_omni.diffusion import forward_context as fc
|
||||
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
|
||||
|
||||
model = _SegmentMeanModel()
|
||||
branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=14)
|
||||
sigmas_video = _sigmas(4, 12.0)
|
||||
sigmas_audio = _sigmas(4, 3.0)
|
||||
|
||||
published: list[tuple[int | None, float | None]] = []
|
||||
|
||||
class _Recorder:
|
||||
def __init__(self):
|
||||
self.denoise_step_idx = None
|
||||
self.denoise_timestep = None
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
object.__setattr__(self, name, value)
|
||||
if name == "denoise_timestep":
|
||||
published.append((self.denoise_step_idx, value))
|
||||
|
||||
recorder = _Recorder()
|
||||
published.clear() # drop the pair __init__ recorded
|
||||
original = fc._forward_context
|
||||
fc._forward_context = recorder
|
||||
try:
|
||||
minimax_h3_denoise_loop(
|
||||
model=model,
|
||||
positive=branch,
|
||||
initial_video_rows=video_rows,
|
||||
initial_audio_rows=audio_rows,
|
||||
keyframe_cond_rows=None,
|
||||
sigmas_video=sigmas_video,
|
||||
sigmas_audio=sigmas_audio,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
request_mode = list(published)
|
||||
|
||||
published.clear()
|
||||
pipeline = _step_pipeline(model)
|
||||
state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
|
||||
while not state.denoise_completed:
|
||||
pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
|
||||
step_mode = list(published)
|
||||
finally:
|
||||
fc._forward_context = original
|
||||
|
||||
num_steps = len(sigmas_video) - 1
|
||||
expected = [(step, sigmas_video[step]) for step in range(num_steps)]
|
||||
assert request_mode == expected + [(None, None)]
|
||||
# Step mode has no loop to close, so it publishes only the per-step pairs.
|
||||
assert step_mode == expected
|
||||
|
||||
|
||||
def test_mixed_step_batch_leaves_gated_attention_dense():
|
||||
"""A batch spanning different steps has no single progress point to publish."""
|
||||
from vllm_omni.diffusion import forward_context as fc
|
||||
|
||||
model = _SegmentMeanModel()
|
||||
first, first_video, first_audio = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=15)
|
||||
second, second_video, second_audio = _make_branch(
|
||||
text_len=5, latent_t=3, latent_h=6, latent_w=4, audio_t=2, seed=16
|
||||
)
|
||||
sigmas_video, sigmas_audio = _sigmas(6, 12.0), _sigmas(6, 3.0)
|
||||
states = [
|
||||
_make_state("req-0", model, first, first_video, first_audio, sigmas_video, sigmas_audio),
|
||||
_make_state("req-1", model, second, second_video, second_audio, sigmas_video, sigmas_audio),
|
||||
]
|
||||
states[1].step_index = 2 # admitted later, so it trails the batch
|
||||
|
||||
recorder = SimpleNamespace(denoise_step_idx="unset", denoise_timestep="unset")
|
||||
original = fc._forward_context
|
||||
fc._forward_context = recorder
|
||||
try:
|
||||
_step_pipeline(model).denoise_step(SimpleNamespace(states=tuple(states)), states=states)
|
||||
finally:
|
||||
fc._forward_context = original
|
||||
|
||||
assert recorder.denoise_step_idx is None
|
||||
assert recorder.denoise_timestep is None
|
||||
|
||||
|
||||
def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
|
||||
branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8)
|
||||
sigmas_video = _sigmas(6, 12.0)
|
||||
sigmas_audio = _sigmas(6, 3.0)
|
||||
context = {
|
||||
"height": 96,
|
||||
"width": 64,
|
||||
"latent_t": 2,
|
||||
"latent_h": 4,
|
||||
"latent_w": 6,
|
||||
"audio_t": 3,
|
||||
**{key: None for key in mod._MINIMAX_H3_DENOISE_INPUT_KEYS},
|
||||
}
|
||||
|
||||
pipeline = _step_pipeline(_SegmentMeanModel())
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
|
||||
monkeypatch.setattr(
|
||||
mod.MiniMaxH3Pipeline,
|
||||
"_build_denoise_inputs",
|
||||
lambda self, **_: {
|
||||
"branch": branch,
|
||||
"video_rows": video_rows,
|
||||
"audio_rows": audio_rows,
|
||||
"cond_anchor": None,
|
||||
"audio_anchor": None,
|
||||
"sigmas_video": sigmas_video,
|
||||
"sigmas_audio": sigmas_audio,
|
||||
},
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
|
||||
state = StepRequestState(
|
||||
request_id="req-0",
|
||||
sampling=SimpleNamespace(num_outputs_per_prompt=1),
|
||||
prompt="a prompt",
|
||||
)
|
||||
pipeline.prepare_encode(state)
|
||||
|
||||
# The runner slices the batched velocity by this row count.
|
||||
assert state.latents.shape == (int(branch.img_pos.shape[0]), 96)
|
||||
assert state.total_steps == len(sigmas_video) - 1
|
||||
assert state.step_index == 0
|
||||
assert state.do_true_cfg is False
|
||||
torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0]))
|
||||
assert state.extra[mod._STEP_BRANCH] is branch
|
||||
assert state.extra[mod._STEP_SHAPE]["height"] == 96
|
||||
|
||||
|
||||
def test_prepare_encode_rejects_request_mode_only_features():
|
||||
"""Multi-output and DLO have no representation in the step contract."""
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
# A request state carries exactly one latent tensor.
|
||||
multi_output = StepRequestState(
|
||||
request_id="req-0",
|
||||
sampling=SimpleNamespace(num_outputs_per_prompt=2),
|
||||
prompt="a prompt",
|
||||
)
|
||||
with pytest.raises(OmniClientError, match="one output per request"):
|
||||
_step_pipeline(_SegmentMeanModel()).prepare_encode(multi_output)
|
||||
|
||||
# Distributed layerwise offload streams the DiT around a whole denoise loop.
|
||||
single_output = StepRequestState(
|
||||
request_id="req-1",
|
||||
sampling=SimpleNamespace(num_outputs_per_prompt=1),
|
||||
prompt="a prompt",
|
||||
)
|
||||
dlo_pipeline = _step_pipeline(_SegmentMeanModel())
|
||||
dlo_pipeline._dlo_residency_controller = object()
|
||||
with pytest.raises(ValueError, match="distributed layerwise offload"):
|
||||
dlo_pipeline.prepare_encode(single_output)
|
||||
|
||||
|
||||
def test_prepare_encode_rejects_high_quality_cache_dit():
|
||||
"""quality=high installs a transformer-scoped Cache-DiT profile that
|
||||
would leak across interleaved or co-batched step-mode requests."""
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
state = StepRequestState(
|
||||
request_id="req-hi",
|
||||
sampling=SimpleNamespace(num_outputs_per_prompt=1, quality="high"),
|
||||
prompt="a prompt",
|
||||
)
|
||||
with pytest.raises(OmniClientError, match="quality=high"):
|
||||
_step_pipeline(_SegmentMeanModel()).prepare_encode(state)
|
||||
|
||||
|
||||
def _fake_attention_module(
|
||||
*,
|
||||
use_ring: bool,
|
||||
backend: str = "FLASH_ATTN",
|
||||
supports_multi_doc: bool = True,
|
||||
):
|
||||
"""Build a bare MiniMaxH3Attention with the attributes ``_packed_batch_supported`` reads.
|
||||
|
||||
``supports_multi_doc`` models the platform-dependent capability probe on
|
||||
``AttentionBackend`` (e.g. FLASH_ATTN returns True on CUDA/ROCm/MUSA and
|
||||
False on NPU/XPU); ``backend`` remains only for logging/back-compat.
|
||||
"""
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3Attention
|
||||
|
||||
attn = object.__new__(MiniMaxH3Attention)
|
||||
attn.attention = SimpleNamespace(
|
||||
attn_backend=SimpleNamespace(
|
||||
get_name=lambda: backend,
|
||||
supports_multi_doc_packed_varlen=lambda: supports_multi_doc,
|
||||
),
|
||||
use_ring=use_ring,
|
||||
)
|
||||
return attn
|
||||
|
||||
|
||||
class _FakeTransformer:
|
||||
"""The minimal ``modules()`` protocol ``_packed_batch_supported`` walks."""
|
||||
|
||||
def __init__(self, attentions):
|
||||
self._attentions = list(attentions)
|
||||
|
||||
def modules(self):
|
||||
return iter([self, *self._attentions])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"attention",
|
||||
[
|
||||
_fake_attention_module(use_ring=True),
|
||||
_fake_attention_module(use_ring=False, backend="XFORMERS", supports_multi_doc=False),
|
||||
_fake_attention_module(use_ring=False, backend="FLASH_ATTN", supports_multi_doc=False),
|
||||
],
|
||||
)
|
||||
def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
|
||||
|
||||
|
||||
def test_broadcast_rank0_exception_single_rank_reraises():
|
||||
"""Single-rank execution has no group; the helper just reraises."""
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_broadcast_rank0_exception,
|
||||
)
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
_broadcast_rank0_exception(None)
|
||||
with pytest.raises(OmniClientError, match="bad ref"):
|
||||
_broadcast_rank0_exception(OmniClientError("bad ref"))
|
||||
|
||||
|
||||
def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
|
||||
"""A rank-0 error becomes a matching client error on every other DiT rank."""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
def fake_rank_world(rank):
|
||||
return lambda: (object(), rank, 4)
|
||||
|
||||
def make_broadcast(rank0_payload):
|
||||
def fake_broadcast(payload_list, *, src, group):
|
||||
payload_list[0] = rank0_payload
|
||||
|
||||
return fake_broadcast
|
||||
|
||||
monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
|
||||
err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
|
||||
rank0_payload = {
|
||||
"type": type(err).__name__,
|
||||
"message": str(err),
|
||||
"status_code": err.status_code,
|
||||
"error_type": err.error_type,
|
||||
}
|
||||
monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
|
||||
with pytest.raises(OmniClientError) as rank0_info:
|
||||
mod._broadcast_rank0_exception(err)
|
||||
assert rank0_info.value is err
|
||||
|
||||
monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
|
||||
with pytest.raises(OmniClientError) as rank2_info:
|
||||
mod._broadcast_rank0_exception(None)
|
||||
assert rank2_info.value.status_code == 422
|
||||
assert rank2_info.value.error_type == "UnprocessableEntityError"
|
||||
assert "invalid reference-video file" in str(rank2_info.value)
|
||||
@@ -25,7 +25,7 @@ from vllm_omni.diffusion.sched import (
|
||||
StepScheduler,
|
||||
)
|
||||
from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
|
||||
from vllm_omni.diffusion.worker.utils import RunnerOutput
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
@@ -1559,6 +1559,39 @@ class TestStepScheduler:
|
||||
assert sched_output.num_running_reqs == 2
|
||||
assert sched_output.num_waiting_reqs == 0
|
||||
|
||||
def test_mixed_batch_error_finishes_only_failed_request(self) -> None:
|
||||
scheduler = StepScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
|
||||
|
||||
req_bad = scheduler.add_request(_make_step_request("bad", num_inference_steps=2))
|
||||
req_good = scheduler.add_request(_make_step_request("good", num_inference_steps=2))
|
||||
first = scheduler.schedule()
|
||||
|
||||
finished = scheduler.update_from_output(
|
||||
first,
|
||||
BatchRunnerOutput.from_list(
|
||||
[
|
||||
_make_step_output(
|
||||
req_bad,
|
||||
step_index=0,
|
||||
finished=True,
|
||||
error="invalid request input",
|
||||
),
|
||||
_make_step_output(req_good, step_index=1),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
assert finished == {req_bad}
|
||||
assert scheduler.get_request_state(req_bad).status == DiffusionRequestStatus.FINISHED_ERROR
|
||||
assert scheduler.get_request_state(req_good).status == DiffusionRequestStatus.RUNNING
|
||||
assert scheduler.get_request_state(req_good).req.sampling_params.step_index == 1
|
||||
|
||||
second = scheduler.schedule()
|
||||
assert second.finished_req_ids == {req_bad}
|
||||
assert _new_ids(second) == []
|
||||
assert _cached_ids(second) == [req_good]
|
||||
|
||||
def test_step_batch_rejects_different_sampling_key(self) -> None:
|
||||
scheduler = StepScheduler()
|
||||
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
|
||||
|
||||
@@ -47,6 +47,7 @@ from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
||||
from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -102,6 +103,22 @@ class _StepPipeline:
|
||||
return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
|
||||
|
||||
|
||||
class _PerRequestErrorStepPipeline(_StepPipeline):
|
||||
def prepare_encode(self, state, **kwargs):
|
||||
if state.prompt == "fail-prepare":
|
||||
raise OmniClientError(
|
||||
"invalid request input",
|
||||
status_code=422,
|
||||
error_type="UnprocessableEntityError",
|
||||
)
|
||||
return super().prepare_encode(state, **kwargs)
|
||||
|
||||
def denoise_step(self, input_batch, **kwargs):
|
||||
del kwargs
|
||||
self.denoise_calls += 1
|
||||
return torch.ones_like(input_batch.latents)
|
||||
|
||||
|
||||
class _ProfilingStepPipeline(_StepPipeline):
|
||||
enable_diffusion_pipeline_profiler = True
|
||||
|
||||
@@ -697,6 +714,72 @@ class TestRunner:
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
assert len(result) == 2
|
||||
|
||||
def test_prepare_error_isolated_from_healthy_request(self, monkeypatch):
|
||||
runner = _make_runner()
|
||||
runner.pipeline = _PerRequestErrorStepPipeline()
|
||||
bad_req = _make_step_request()
|
||||
bad_req.request_id = "req-bad"
|
||||
bad_req.prompt = "fail-prepare"
|
||||
good_req = _make_step_request()
|
||||
good_req.request_id = "req-good"
|
||||
good_req.prompt = "valid"
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(
|
||||
runner,
|
||||
_make_batch_scheduler_output([bad_req, good_req]),
|
||||
)
|
||||
|
||||
bad_output = result.get_request_output("req-bad")
|
||||
assert bad_output.finished is True
|
||||
assert bad_output.result is not None
|
||||
assert bad_output.result.error == "invalid request input"
|
||||
assert bad_output.result.error_status_code == 422
|
||||
assert bad_output.result.error_type == "UnprocessableEntityError"
|
||||
good_output = result.get_request_output("req-good")
|
||||
assert good_output.finished is False
|
||||
assert good_output.step_index == 1
|
||||
assert good_output.result is None
|
||||
assert set(runner.state_cache) == {"req-good"}
|
||||
assert runner.pipeline.denoise_calls == 1
|
||||
assert runner.pipeline.scheduler_calls == 1
|
||||
|
||||
def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch):
|
||||
"""P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep)
|
||||
can raise on rank 0 while non-zero ranks return cleanly and then hang
|
||||
on a downstream ``dist.broadcast``. Beyond the pipeline-side guard,
|
||||
the runner all-reduces a per-request failure flag across the DiT
|
||||
group so a rank that ``prepare_encode`` returned cleanly on still
|
||||
skips the request and does not enter the denoise step batch.
|
||||
"""
|
||||
runner = _make_runner()
|
||||
runner.pipeline = _StepPipeline() # locally always succeeds
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
# Simulate a peer DiT rank having reported failure. We patch the
|
||||
# helper directly (rather than ``torch.distributed`` broadly) because
|
||||
# ``execute_stepwise`` also queries ``torch.distributed.get_rank``.
|
||||
monkeypatch.setattr(
|
||||
model_runner_module,
|
||||
"_dit_any_rank_failed",
|
||||
lambda local_failed: True,
|
||||
)
|
||||
|
||||
req = _make_step_request(num_inference_steps=1)
|
||||
result = DiffusionModelRunner.execute_stepwise(
|
||||
runner,
|
||||
_make_scheduler_output(req),
|
||||
)
|
||||
|
||||
output = result.get_request_output("req-1")
|
||||
assert output.finished is True
|
||||
assert output.result is not None
|
||||
assert output.result.error is not None
|
||||
assert "another DiT rank" in output.result.error
|
||||
# A rank that skipped the request must not have driven a denoise step.
|
||||
assert runner.pipeline.denoise_calls == 0
|
||||
assert runner.state_cache == {}
|
||||
|
||||
def test_receives_kv_payload_before_prepare_encode(self, monkeypatch):
|
||||
runner = _make_runner()
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
@@ -36,6 +36,23 @@ class AttentionBackend(ABC):
|
||||
"""
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def supports_multi_doc_packed_varlen(cls) -> bool:
|
||||
"""Whether this backend keeps N-document packed boundaries isolated.
|
||||
|
||||
When True, the backend consumes ``AttentionMetadata.extra`` cu_seqlens
|
||||
as a genuine block-diagonal attention plan (a dedicated varlen kernel,
|
||||
not a padding-mask rebuild), so a caller may pack multiple real
|
||||
requests into one forward without attention crossing document
|
||||
boundaries. When False, callers packing more than one real document
|
||||
must run one forward per document; otherwise a backend that only
|
||||
supports a ``[real, pad]`` two-document contract, or that ignores
|
||||
cu_seqlens outright, will silently attend across request boundaries.
|
||||
Backends whose kernel selection is platform-dependent must consult
|
||||
``current_omni_platform``.
|
||||
"""
|
||||
return False
|
||||
|
||||
# ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
|
||||
# Platform resolution rejects an explicit selection outside this set, so a
|
||||
# hardware-specific backend fails before the model is built.
|
||||
|
||||
@@ -37,6 +37,20 @@ class FlashAttentionBackend(AttentionBackend):
|
||||
# constructing it there.
|
||||
return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
|
||||
|
||||
@classmethod
|
||||
def supports_multi_doc_packed_varlen(cls) -> bool:
|
||||
# CUDA / ROCm / MUSA all route through ``forward_cuda``, which
|
||||
# dispatches ``_forward_varlen_packed`` -> ``flash_attn_varlen_func``
|
||||
# over the caller's cu_seqlens without a mask, so an arbitrary
|
||||
# N-document packing keeps its boundaries. NPU's forward path
|
||||
# (``_resolve_packed_seq_npu``) only accepts a ``[real, pad]``
|
||||
# two-document layout (cu_seqlens shape <= 3) and otherwise falls
|
||||
# back to a padding-mask rebuild that spans the whole packed row,
|
||||
# so N>=2 real documents would silently attend across request
|
||||
# boundaries. XPU never consumes cu_seqlens from ``extra``. Both
|
||||
# therefore stay False here.
|
||||
return current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa()
|
||||
|
||||
@classmethod
|
||||
def supports_attention_mask(cls) -> bool:
|
||||
return True
|
||||
|
||||
@@ -1304,6 +1304,9 @@ def extract_minimax_h3_context(
|
||||
psp = _required_kwarg(kwargs, "packed_seq_params")
|
||||
cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
|
||||
max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
|
||||
# Mirror forward()'s host-side scalar so the refiner sees the number of
|
||||
# packed requests without reading cu_seqlens off-device.
|
||||
num_requests = int(module._psp_optional(psp, "num_requests", 1))
|
||||
refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
|
||||
refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
|
||||
refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
|
||||
@@ -1330,6 +1333,7 @@ def extract_minimax_h3_context(
|
||||
text_pos=text_pos.to(device),
|
||||
refiner_cu_seqlens=refiner_cu.to(device),
|
||||
refiner_max_seqlen=refiner_max,
|
||||
num_requests=num_requests,
|
||||
seq_len=seq_len,
|
||||
device=device,
|
||||
local_span=(0, seq_len),
|
||||
|
||||
@@ -244,6 +244,17 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
|
||||
_forward_context.denoise_step_idx = step_idx
|
||||
|
||||
|
||||
def set_forward_context_denoise_timestep(timestep: float | None) -> None:
|
||||
"""Set the normalized (descending, 1 -> 0) denoise timestep.
|
||||
|
||||
Timestep-gated attention features read this; pipelines that drive their own
|
||||
denoise loop can publish it directly instead of going through
|
||||
:meth:`DenoiseProgressMixin.record_denoise_step`.
|
||||
"""
|
||||
if _forward_context is not None:
|
||||
_forward_context.denoise_timestep = None if timestep is None else float(timestep)
|
||||
|
||||
|
||||
class DenoiseProgressMixin:
|
||||
def record_denoise_step(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Batched multi-request packing for MiniMax H3 step-wise execution.
|
||||
|
||||
Request mode forwards exactly one packed sequence per denoise step. Step mode
|
||||
(continuous batching) may hold several requests at once, so this module
|
||||
concatenates their packed layouts into a single sequence and rebuilds
|
||||
``cu_seqlens`` with one document per request plus that request's
|
||||
alignment-padding tail. Attention therefore never crosses a request boundary,
|
||||
while every request keeps its own RoPE coordinates, token tags, and timesteps.
|
||||
A batch of one is not rebuilt at all -- it delegates to the request-mode layout.
|
||||
|
||||
Row order is the request order, so the concatenated video velocity returned by
|
||||
the DiT can be split back per request by row count alone -- which is exactly
|
||||
what the runner does with ``StepRequestState.latents``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .denoise_loop import (
|
||||
MINIMAX_H3_AUDIO_ROW_WIDTH,
|
||||
MINIMAX_H3_VIDEO_ROW_WIDTH,
|
||||
MiniMaxH3DenoiseBranch,
|
||||
)
|
||||
|
||||
|
||||
def minimax_h3_batched_forward_kwargs(
|
||||
*,
|
||||
branches: Sequence[MiniMaxH3DenoiseBranch],
|
||||
video_rows: Sequence[torch.Tensor],
|
||||
audio_rows: Sequence[torch.Tensor],
|
||||
t_video: Sequence[float],
|
||||
t_audio: Sequence[float],
|
||||
imgvid_cond_timesteps: Sequence[float],
|
||||
audio_ref_cond_timesteps: Sequence[float],
|
||||
) -> dict[str, Any]:
|
||||
"""Build one DiT forward kwargs dict covering every request in the batch.
|
||||
|
||||
A single request delegates to :meth:`MiniMaxH3DenoiseBranch.forward_kwargs`,
|
||||
so step mode and request mode are the same call rather than two layouts that
|
||||
happen to agree -- which also keeps ``video_token_layout`` (the sparse-
|
||||
attention hint) on the single-request step path.
|
||||
"""
|
||||
if len(branches) == 1:
|
||||
return branches[0].forward_kwargs(
|
||||
video_rows=video_rows[0],
|
||||
audio_rows=audio_rows[0],
|
||||
t_video=t_video[0],
|
||||
t_audio=t_audio[0],
|
||||
imgvid_cond_timestep=imgvid_cond_timesteps[0],
|
||||
audio_ref_cond_timestep=audio_ref_cond_timesteps[0],
|
||||
)
|
||||
|
||||
device = branches[0].device
|
||||
total_seq = sum(branch.seq_len for branch in branches)
|
||||
|
||||
img_pos_parts: list[torch.Tensor] = []
|
||||
audio_pos_parts: list[torch.Tensor] = []
|
||||
text_pos_parts: list[torch.Tensor] = []
|
||||
update_parts: list[torch.Tensor] = []
|
||||
tag_parts: list[torch.Tensor] = []
|
||||
position_id_parts: list[torch.Tensor] = []
|
||||
text_embed_parts: list[torch.Tensor] = []
|
||||
cu_bounds: list[int] = [0]
|
||||
refiner_bounds: list[int] = [0]
|
||||
max_seqlen = 0
|
||||
refiner_max_seqlen = 0
|
||||
|
||||
# Non-media rows (text and alignment padding) inherit their request's video
|
||||
# timestep, matching the single-request packed-sequence semantics.
|
||||
timesteps = torch.empty(total_seq, dtype=torch.float32, device=device)
|
||||
|
||||
seq_offset = 0
|
||||
text_offset = 0
|
||||
for index, branch in enumerate(branches):
|
||||
img_pos = branch.img_pos_dev + seq_offset
|
||||
audio_pos = branch.audio_pos_dev + seq_offset
|
||||
img_pos_parts.append(img_pos)
|
||||
audio_pos_parts.append(audio_pos)
|
||||
text_pos_parts.append(branch.text_pos_dev + seq_offset)
|
||||
update_parts.append(branch.update_mask_dev)
|
||||
tag_parts.append(branch.token_tags_dev)
|
||||
position_id_parts.append(branch.img_position_ids_dev)
|
||||
text_embed_parts.append(branch.text_embeddings_dev)
|
||||
|
||||
timesteps[seq_offset : seq_offset + branch.seq_len] = float(t_video[index])
|
||||
timesteps[img_pos[branch.update_mask_dev]] = float(t_video[index])
|
||||
timesteps[img_pos[~branch.update_mask_dev]] = float(imgvid_cond_timesteps[index])
|
||||
timesteps[audio_pos[branch.audio_update_mask_dev]] = float(t_audio[index])
|
||||
timesteps[audio_pos[~branch.audio_update_mask_dev]] = float(audio_ref_cond_timesteps[index])
|
||||
|
||||
# One document for the request's real rows and one for its padding tail,
|
||||
# matching the single-request layout. The tail bound is emitted even when
|
||||
# the 64-row alignment left it empty, so that a request always
|
||||
# contributes exactly two bounds regardless of its length.
|
||||
cu_bounds.append(seq_offset + branch.used_len)
|
||||
cu_bounds.append(seq_offset + branch.seq_len)
|
||||
max_seqlen = max(max_seqlen, branch.used_len, branch.seq_len - branch.used_len)
|
||||
|
||||
text_offset += branch.text_len
|
||||
refiner_bounds.append(text_offset)
|
||||
refiner_max_seqlen = max(refiner_max_seqlen, branch.text_len)
|
||||
seq_offset += branch.seq_len
|
||||
|
||||
img_pos_all = torch.cat(img_pos_parts)
|
||||
audio_pos_all = torch.cat(audio_pos_parts)
|
||||
|
||||
x = torch.zeros(1, total_seq, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
|
||||
x[0].index_copy_(0, img_pos_all, torch.cat([rows.to(device=device, dtype=torch.float32) for rows in video_rows]))
|
||||
audio_x = torch.zeros(1, total_seq, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
|
||||
audio_x[0].index_copy_(
|
||||
0,
|
||||
audio_pos_all,
|
||||
torch.cat([rows.to(device=device, dtype=torch.float32) for rows in audio_rows]),
|
||||
)
|
||||
|
||||
unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True)
|
||||
return {
|
||||
"x": x,
|
||||
"audio_x": audio_x,
|
||||
"img_position_ids": torch.cat(position_id_parts)[None],
|
||||
"update_mask": torch.cat(update_parts),
|
||||
"token_tags": torch.cat(tag_parts),
|
||||
"skip_mask_out_condition": False,
|
||||
"prompt_embeds": torch.cat(text_embed_parts),
|
||||
"img_pos_info": {"position_ids": img_pos_all},
|
||||
"audio_pos_info": {"position_ids": audio_pos_all},
|
||||
"text_pos_info": {"position_ids": torch.cat(text_pos_parts)},
|
||||
"img_pos_for_infer_output_info": {"position_ids": img_pos_all},
|
||||
"packed_seq_params": {
|
||||
"cu_seqlens_q": torch.tensor(cu_bounds, dtype=torch.int32, device=device),
|
||||
"max_seqlen_q": max_seqlen,
|
||||
# Attention needs the request count on the host to know its valid
|
||||
# rows are block-diagonal rather than a prefix. It describes the
|
||||
# whole forward, so the DiT reads it once here and applies it to the
|
||||
# token refiner too -- the refiner's text-only bounds cannot be told
|
||||
# apart from a single padded request on their own.
|
||||
"num_requests": len(branches),
|
||||
},
|
||||
"refiner_packed_seq_params": {
|
||||
"cu_seqlens_q": torch.tensor(refiner_bounds, dtype=torch.int32, device=device),
|
||||
"max_seqlen_q": refiner_max_seqlen,
|
||||
},
|
||||
"unique_timesteps": unique_timesteps,
|
||||
"inverse_indices": inverse_indices,
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"minimax_h3_batched_forward_kwargs",
|
||||
]
|
||||
@@ -15,7 +15,10 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
set_forward_context_denoise_step_idx,
|
||||
set_forward_context_denoise_timestep,
|
||||
)
|
||||
|
||||
from .scheduling_minimax_h3_euler_ancestral import (
|
||||
minimax_h3_euler_eta0_step,
|
||||
@@ -31,6 +34,17 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96
|
||||
MINIMAX_H3_AUDIO_ROW_WIDTH = 32
|
||||
|
||||
|
||||
def minimax_h3_publish_denoise_progress(step: int | None, sigma_video: float | None) -> None:
|
||||
"""Publish denoise progress for step-gated attention features.
|
||||
|
||||
Both execution modes must publish the same pair: the step index drives the
|
||||
dense warmup of RAINFUSION_ATTN, and the normalized descending timestep
|
||||
drives the TRTLLM_ATTN skip gate, which stays dense while it is unset.
|
||||
"""
|
||||
set_forward_context_denoise_step_idx(step)
|
||||
set_forward_context_denoise_timestep(sigma_video)
|
||||
|
||||
|
||||
class MiniMaxH3DenoiseBranch:
|
||||
"""Static per-branch state: packed layout + fixed forward kwargs.
|
||||
|
||||
@@ -64,26 +78,39 @@ class MiniMaxH3DenoiseBranch:
|
||||
if int(token_tags.view(-1).shape[0]) != seq_len:
|
||||
raise ValueError(f"token_tags length {int(token_tags.view(-1).shape[0])} != seq_len {seq_len}")
|
||||
cu = packed["cu_seqlens"].to(torch.int32)
|
||||
self.device = device
|
||||
# ``used_len`` is the real (non-padding) document length; step-mode
|
||||
# batching concatenates layouts and needs both bounds as Python ints so
|
||||
# it can rebuild ``cu_seqlens`` without a device sync per step.
|
||||
self.used_len = int(cu[1])
|
||||
self.text_len = text_len
|
||||
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)
|
||||
self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
|
||||
self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
|
||||
text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
|
||||
self.text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
|
||||
self.token_tags_dev = token_tags.view(-1).to(torch.long).to(device)
|
||||
self.img_position_ids_dev = packed["img_position_ids"].to(device)
|
||||
self.text_embeddings_dev = text_embeddings.to(device)
|
||||
self.static_kwargs: dict[str, Any] = {
|
||||
"img_position_ids": packed["img_position_ids"][None].to(device),
|
||||
"img_position_ids": self.img_position_ids_dev[None],
|
||||
"update_mask": self.update_mask_dev,
|
||||
"token_tags": token_tags.view(-1).to(torch.long).to(device),
|
||||
"token_tags": self.token_tags_dev,
|
||||
"skip_mask_out_condition": False,
|
||||
"prompt_embeds": text_embeddings.to(device),
|
||||
"prompt_embeds": self.text_embeddings_dev,
|
||||
"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},
|
||||
"text_pos_info": {"position_ids": self.text_pos_dev},
|
||||
"img_pos_for_infer_output_info": {"position_ids": self.img_pos_dev},
|
||||
"packed_seq_params": {
|
||||
"cu_seqlens_q": cu.to(device),
|
||||
"max_seqlen_q": int(cu[1]),
|
||||
"max_seqlen_q": self.used_len,
|
||||
# One request: valid rows are a prefix, so attention may use a
|
||||
# KV prefix length or a 1-D pad mask. See batched_packing for the
|
||||
# co-batched layout, where neither can describe the valid rows.
|
||||
"num_requests": 1,
|
||||
},
|
||||
"refiner_packed_seq_params": {
|
||||
"cu_seqlens_q": torch.tensor([0, text_len, text_len], dtype=torch.int32, device=device),
|
||||
@@ -136,34 +163,21 @@ class MiniMaxH3DenoiseBranch:
|
||||
}
|
||||
|
||||
|
||||
def minimax_h3_denoise_loop(
|
||||
def minimax_h3_prepare_denoise_rows(
|
||||
*,
|
||||
model: Any,
|
||||
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],
|
||||
audio_ref_rows: torch.Tensor | None,
|
||||
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_end: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
|
||||
on_step_start: Callable[[int, float, float], 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).
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
|
||||
"""Validate the initial rows against the layout and pin the condition 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.
|
||||
Returns ``(video_rows, audio_rows, cond_anchor, audio_anchor)`` with the
|
||||
anchors already written into their rows, which is the state both the
|
||||
request-mode loop and step-mode ``prepare_encode()`` start from.
|
||||
"""
|
||||
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 = int((~positive.update_mask).sum())
|
||||
if keyframe_cond_rows is None:
|
||||
if n_cond != 0:
|
||||
@@ -197,19 +211,58 @@ def minimax_h3_denoise_loop(
|
||||
video_rows[~update] = cond_anchor
|
||||
if audio_anchor is not None:
|
||||
audio_rows[~audio_update] = audio_anchor
|
||||
return video_rows, audio_rows, cond_anchor, audio_anchor
|
||||
|
||||
|
||||
def minimax_h3_denoise_loop(
|
||||
*,
|
||||
model: Any,
|
||||
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.
|
||||
"""
|
||||
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")
|
||||
video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows(
|
||||
positive=positive,
|
||||
initial_video_rows=initial_video_rows,
|
||||
initial_audio_rows=initial_audio_rows,
|
||||
keyframe_cond_rows=keyframe_cond_rows,
|
||||
audio_ref_rows=audio_ref_rows,
|
||||
device=device,
|
||||
)
|
||||
update = positive.update_mask_dev
|
||||
audio_update = positive.audio_update_mask_dev
|
||||
|
||||
num_steps = len(sigmas_video) - 1
|
||||
for step in range(num_steps):
|
||||
step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
|
||||
with step_cm:
|
||||
# Publish the step index so step-gated attention features (e.g. the
|
||||
# dense warmup of RAINFUSION_ATTN) can see where we are.
|
||||
set_forward_context_denoise_step_idx(step)
|
||||
s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
|
||||
s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
|
||||
if on_step_start is not None:
|
||||
# Attention gates use the scheduler-style descending timestep.
|
||||
on_step_start(step, s_v, s_a)
|
||||
# Publish where we are so step-gated attention features (the dense
|
||||
# warmup of RAINFUSION_ATTN, the timestep gate of TRTLLM_ATTN) can
|
||||
# see it. Gates use the scheduler-style descending timestep, which
|
||||
# for a rectified-flow schedule is the video sigma.
|
||||
minimax_h3_publish_denoise_progress(step, s_v)
|
||||
t_v, t_a = 1.0 - s_v, 1.0 - s_a
|
||||
imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference))
|
||||
audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference))
|
||||
@@ -250,10 +303,10 @@ def minimax_h3_denoise_loop(
|
||||
audio_rows[audio_update] = new_audio
|
||||
if audio_anchor is not None:
|
||||
audio_rows[~audio_update] = audio_anchor # per-step audio ref reset
|
||||
if on_step_end is not None:
|
||||
on_step_end(step, video_rows, audio_rows)
|
||||
if on_step is not None:
|
||||
on_step(step, video_rows, audio_rows)
|
||||
|
||||
set_forward_context_denoise_step_idx(None)
|
||||
minimax_h3_publish_denoise_progress(None, None)
|
||||
return video_rows, audio_rows
|
||||
|
||||
|
||||
@@ -264,4 +317,6 @@ __all__ = [
|
||||
"MINIMAX_H3_VIDEO_ROW_WIDTH",
|
||||
"MiniMaxH3DenoiseBranch",
|
||||
"minimax_h3_denoise_loop",
|
||||
"minimax_h3_prepare_denoise_rows",
|
||||
"minimax_h3_publish_denoise_progress",
|
||||
]
|
||||
@@ -56,6 +56,26 @@ if TYPE_CHECKING:
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
# Packed multi-request forwards require the attention backend to actually
|
||||
# consume cu_seqlens as a block-diagonal plan (not a padding-mask rebuild that
|
||||
# spans the full packed row). The pipeline gates on this capability before
|
||||
# packing, and ``_run_packed_attention`` re-checks it per forward; a name-only
|
||||
# gate would let FLASH_ATTN's NPU/XPU code paths through even though those
|
||||
# variants would silently attend across request boundaries.
|
||||
def _attention_isolates_packed_requests(attention_layer: Any) -> bool:
|
||||
"""True if this attention layer keeps N-document packed boundaries.
|
||||
|
||||
Requires a backend advertising ``supports_multi_doc_packed_varlen`` *and*
|
||||
that the layer is not running under ring sequence parallelism (the ring
|
||||
kernel dispatches through its own attention that ignores the packed
|
||||
cu_seqlens regardless of the configured backend).
|
||||
"""
|
||||
backend = getattr(attention_layer, "attn_backend", None)
|
||||
if backend is None or not backend.supports_multi_doc_packed_varlen():
|
||||
return False
|
||||
return not getattr(attention_layer, "use_ring", False)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiniMaxH3DiTArchConfig:
|
||||
num_layers: int = 50
|
||||
@@ -433,6 +453,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
packed_total: int,
|
||||
num_requests: int = 1,
|
||||
video_layout: VideoTokenLayout | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run packed attention as a small eager island.
|
||||
@@ -442,27 +463,47 @@ class MiniMaxH3Attention(nn.Module):
|
||||
narrow lets regional compile fuse projections, norms, RoPE, and the
|
||||
surrounding DiT block without repeated graph breaks.
|
||||
"""
|
||||
# max_seqlen is already the first (real) packed document length. Do
|
||||
# not read the CUDA cu_seqlens scalars here: this function runs once
|
||||
# per layer and .item() would serialize every attention launch.
|
||||
# max_seqlen is already the longest packed document length. Do not read
|
||||
# the CUDA cu_seqlens scalars here: this function runs once per layer
|
||||
# and .item() would serialize every attention launch. ``num_requests``
|
||||
# is carried as a Python int for the same reason.
|
||||
if not 0 < max_seqlen <= packed_total:
|
||||
raise ValueError(
|
||||
f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}"
|
||||
)
|
||||
used = min(max_seqlen, packed_total)
|
||||
attn_mask = None
|
||||
# Ring attention can dispatch to a different implementation from the
|
||||
# configured backend, so the no-mask fast paths are local-only.
|
||||
# supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
|
||||
# supports_packed_mask_free: backend consumes the packed metadata
|
||||
# without ever reading attn_mask (CUDA packed varlen, NPU
|
||||
# npu_attn_varlen opt-in with its own fallback rebuild).
|
||||
no_mask = not getattr(self.attention, "use_ring", False) and (
|
||||
self.attention.attn_backend.supports_prefix_kv_slicing
|
||||
or self.attention.attn_backend.supports_packed_mask_free()
|
||||
)
|
||||
if used < packed_total and not no_mask:
|
||||
attn_mask = torch.arange(packed_total, device=q.device)[None] < used
|
||||
if num_requests > 1:
|
||||
# A step-mode batch packs one document per request, so its valid
|
||||
# rows are block-diagonal rather than a prefix: neither a KV prefix
|
||||
# length nor a 1-D key mask can describe them. Such a layout is
|
||||
# only correct on a backend that actually attends by cu_seqlens as
|
||||
# a block-diagonal plan. Check the capability (not the backend
|
||||
# name): FLASH_ATTN's NPU/XPU variants would otherwise silently
|
||||
# fall back to a padding-mask rebuild that spans the whole packed
|
||||
# row and attend across request boundaries.
|
||||
if not _attention_isolates_packed_requests(self.attention):
|
||||
backend_name = self.attention.attn_backend.get_name()
|
||||
raise ValueError(
|
||||
f"MiniMax H3 packed a {num_requests}-request batch, but the resolved "
|
||||
f"attention ({backend_name}, use_ring={getattr(self.attention, 'use_ring', False)}) "
|
||||
"does not isolate multi-document packed cu_seqlens. Run one request "
|
||||
"per forward on this backend."
|
||||
)
|
||||
used = packed_total
|
||||
else:
|
||||
used = min(max_seqlen, packed_total)
|
||||
# Ring attention can dispatch to a different implementation from the
|
||||
# configured backend, so the no-mask fast paths are local-only.
|
||||
# supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
|
||||
# supports_packed_mask_free: backend consumes the packed metadata
|
||||
# without ever reading attn_mask (CUDA packed varlen, NPU
|
||||
# npu_attn_varlen opt-in with its own fallback rebuild).
|
||||
no_mask = not getattr(self.attention, "use_ring", False) and (
|
||||
self.attention.attn_backend.supports_prefix_kv_slicing
|
||||
or self.attention.attn_backend.supports_packed_mask_free()
|
||||
)
|
||||
if used < packed_total and not no_mask:
|
||||
attn_mask = torch.arange(packed_total, device=q.device)[None] < used
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=attn_mask,
|
||||
extra={
|
||||
@@ -498,6 +539,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
packed_total: int | None = None,
|
||||
num_requests: int = 1,
|
||||
sp_seq_lens: list[int] | None = None,
|
||||
video_layout: VideoTokenLayout | None = None,
|
||||
) -> torch.Tensor:
|
||||
@@ -533,9 +575,9 @@ class MiniMaxH3Attention(nn.Module):
|
||||
self.q_norm.variance_epsilon,
|
||||
)
|
||||
|
||||
# The packed layout uses a second document for alignment padding.
|
||||
# Local/Ulysses backends unpad it, while Ring keeps aligned rows for
|
||||
# fixed-size P2P buffers.
|
||||
# Each request contributes a document for its rows plus one for its
|
||||
# alignment padding. Local/Ulysses backends unpad it, while Ring keeps
|
||||
# aligned rows for fixed-size P2P buffers.
|
||||
out = self._run_packed_attention(
|
||||
q,
|
||||
k,
|
||||
@@ -546,6 +588,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
# backend receives the global sequence after all-to-all, so carry
|
||||
# its Python length explicitly instead of inferring it from q.
|
||||
packed_total=packed_total if packed_total is not None else q.shape[0],
|
||||
num_requests=num_requests,
|
||||
video_layout=video_layout,
|
||||
)
|
||||
out = out.reshape(total, self.num_heads * self.head_dim)
|
||||
@@ -673,12 +716,14 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
num_requests: int = 1,
|
||||
) -> torch.Tensor:
|
||||
x = x + self.attn(
|
||||
self.norm1(x),
|
||||
rope_table=None,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
num_requests=num_requests,
|
||||
)
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x
|
||||
@@ -711,9 +756,10 @@ class MiniMaxH3TokenRefiner(nn.Module):
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
num_requests: int = 1,
|
||||
) -> torch.Tensor:
|
||||
for block in self.blocks:
|
||||
x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
|
||||
x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, num_requests=num_requests)
|
||||
return self.final_norm(x)
|
||||
|
||||
|
||||
@@ -759,6 +805,7 @@ class MiniMaxH3DiTBlock(nn.Module):
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
packed_total: int,
|
||||
num_requests: int = 1,
|
||||
sp_seq_lens: list[int] | None = None,
|
||||
video_layout: VideoTokenLayout | None = None,
|
||||
) -> torch.Tensor:
|
||||
@@ -793,6 +840,7 @@ class MiniMaxH3DiTBlock(nn.Module):
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
packed_total=packed_total,
|
||||
num_requests=num_requests,
|
||||
sp_seq_lens=sp_seq_lens,
|
||||
video_layout=video_layout,
|
||||
)
|
||||
@@ -1145,6 +1193,11 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
raise ValueError(f"{key}.{field} is required")
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _psp_optional(psp: Any, field: str, default: Any) -> Any:
|
||||
value = psp.get(field) if isinstance(psp, dict) else getattr(psp, field, None)
|
||||
return default if value is None else value
|
||||
|
||||
def _embed(
|
||||
self,
|
||||
*,
|
||||
@@ -1160,10 +1213,16 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
seq_len: int,
|
||||
device: torch.device,
|
||||
local_span: tuple[int, int],
|
||||
num_requests: int = 1,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Build this rank's packed multimodal embedding rows.
|
||||
|
||||
Returns (decoder_input [S_local, H] bf16, t_emb [M, t_dim] fp32).
|
||||
|
||||
``num_requests`` defaults to a single packed request so callers that
|
||||
pre-date the continuous-batching change (e.g. TeaCache's extractor
|
||||
contract) do not silently miss the kwarg. ``forward()`` reads the real
|
||||
value from ``packed_seq_params["num_requests"]``.
|
||||
"""
|
||||
local_start, local_len = local_span
|
||||
local_end = local_start + local_len
|
||||
@@ -1200,6 +1259,7 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
text_embed,
|
||||
cu_seqlens=refiner_cu_seqlens,
|
||||
max_seqlen=refiner_max_seqlen,
|
||||
num_requests=num_requests,
|
||||
)
|
||||
if text_local_indices is not None:
|
||||
text_embed = text_embed.index_select(0, text_local_indices)
|
||||
@@ -1269,6 +1329,10 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
psp = _required_kwarg(kwargs, "packed_seq_params")
|
||||
cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
|
||||
max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
|
||||
# How many requests share this packed sequence. Carried as a host int so
|
||||
# attention never reads cu_seqlens scalars off the device; a producer
|
||||
# that omits it is packing a single request.
|
||||
num_requests = int(self._psp_optional(psp, "num_requests", 1))
|
||||
refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
|
||||
refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
|
||||
refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
|
||||
@@ -1305,6 +1369,7 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
text_pos=text_pos.to(device),
|
||||
refiner_cu_seqlens=refiner_cu.to(device),
|
||||
refiner_max_seqlen=refiner_max,
|
||||
num_requests=num_requests,
|
||||
seq_len=seq_len,
|
||||
device=device,
|
||||
local_span=local_span,
|
||||
@@ -1339,6 +1404,7 @@ class MiniMaxH3DiTModel(nn.Module):
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
packed_total=seq_len,
|
||||
num_requests=num_requests,
|
||||
video_layout=video_layout,
|
||||
)
|
||||
if local_len == seq_len:
|
||||
|
||||
@@ -57,7 +57,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
|
||||
)
|
||||
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.errors import OmniClientError, client_error_from_metadata
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf_specific,
|
||||
)
|
||||
@@ -69,14 +69,24 @@ from vllm_omni.quantization.component_config import (
|
||||
resolve_encoder_quant_config as _resolve_encoder_quant_config,
|
||||
)
|
||||
|
||||
from .batched_packing import minimax_h3_batched_forward_kwargs
|
||||
from .condition_noise import (
|
||||
minimax_h3_audio_cond_noise_aug_rows,
|
||||
minimax_h3_imgvid_cond_noise_aug_rows,
|
||||
)
|
||||
from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop
|
||||
from .denoise_loop import (
|
||||
MiniMaxH3DenoiseBranch,
|
||||
minimax_h3_denoise_loop,
|
||||
minimax_h3_prepare_denoise_rows,
|
||||
minimax_h3_publish_denoise_progress,
|
||||
)
|
||||
from .encoder import MiniMaxH3Qwen3VLEncoder
|
||||
from .lora import load_minimax_h3_turbo_lora
|
||||
from .minimax_h3_transformer import MiniMaxH3DiTModel
|
||||
from .minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
MiniMaxH3DiTModel,
|
||||
_attention_isolates_packed_requests,
|
||||
)
|
||||
from .packed_sequence import (
|
||||
minimax_h3_packed_sequence,
|
||||
minimax_h3_packed_sequence_ref2va_blocks,
|
||||
@@ -103,6 +113,10 @@ from .reference_video import (
|
||||
validate_reference_audio_files,
|
||||
validate_reference_audio_waveforms,
|
||||
)
|
||||
from .scheduling_minimax_h3_euler_ancestral import (
|
||||
minimax_h3_euler_eta0_step,
|
||||
minimax_h3_rf_v_to_x0,
|
||||
)
|
||||
from .time_request import (
|
||||
MINIMAX_H3_SHAPE_PLANNER,
|
||||
minimax_h3_align_frame_count,
|
||||
@@ -110,6 +124,10 @@ from .time_request import (
|
||||
)
|
||||
from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -201,6 +219,68 @@ def _resolve_minimax_h3_model_root(
|
||||
)
|
||||
|
||||
|
||||
# Keys of ``_prepare_request_inputs`` that feed ``diffuse`` / ``_build_denoise_inputs``.
|
||||
_MINIMAX_H3_DENOISE_INPUT_KEYS = (
|
||||
"task",
|
||||
"text_embeddings",
|
||||
"text_tags",
|
||||
"seed",
|
||||
"latent_t",
|
||||
"latent_h",
|
||||
"latent_w",
|
||||
"audio_t",
|
||||
"num_frames",
|
||||
"num_steps",
|
||||
"video_shift",
|
||||
"audio_shift",
|
||||
"base_schedule",
|
||||
"visual_condition",
|
||||
"visual_condition_shape",
|
||||
"audio_condition",
|
||||
"ref_audio_t",
|
||||
"ref_blocks",
|
||||
"visual_condition_shapes",
|
||||
"audio_condition_lengths",
|
||||
"keyframe_frame_indices",
|
||||
)
|
||||
|
||||
# ``StepRequestState.extra`` keys owned by the step-execution path.
|
||||
_STEP_BRANCH = "minimax_h3_branch"
|
||||
_STEP_AUDIO_ROWS = "minimax_h3_audio_rows"
|
||||
_STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred"
|
||||
_STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video"
|
||||
_STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio"
|
||||
_STEP_COND_ANCHOR = "minimax_h3_cond_anchor"
|
||||
_STEP_AUDIO_ANCHOR = "minimax_h3_audio_anchor"
|
||||
_STEP_SHAPE = "minimax_h3_shape"
|
||||
_STEP_TRANSFORMER = "minimax_h3_transformer"
|
||||
|
||||
|
||||
def _minimax_h3_step_schedule(state: StepRequestState) -> dict[str, float]:
|
||||
"""Return the sigma/timestep values this request needs for its current step.
|
||||
|
||||
Mirrors the per-iteration arithmetic of ``minimax_h3_denoise_loop`` so step
|
||||
mode and request mode advance identically.
|
||||
"""
|
||||
step = int(state.step_index)
|
||||
sigmas_video = state.extra[_STEP_SIGMAS_VIDEO]
|
||||
sigmas_audio = state.extra[_STEP_SIGMAS_AUDIO]
|
||||
sigma_video = float(sigmas_video[step])
|
||||
sigma_audio = float(sigmas_audio[step])
|
||||
t_video = 1.0 - sigma_video
|
||||
t_audio = 1.0 - sigma_audio
|
||||
return {
|
||||
"sigma_video": sigma_video,
|
||||
"sigma_video_next": float(sigmas_video[step + 1]),
|
||||
"sigma_audio": sigma_audio,
|
||||
"sigma_audio_next": float(sigmas_audio[step + 1]),
|
||||
"t_video": t_video,
|
||||
"t_audio": t_audio,
|
||||
"imgvid_cond_timestep": max(t_video, MINIMAX_H3_IMGVID_COND_TIMESTEP),
|
||||
"audio_ref_cond_timestep": max(t_audio, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
|
||||
}
|
||||
|
||||
|
||||
def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None:
|
||||
"""Read a partition's distilled schedule. An absent key means legacy uniform."""
|
||||
return DMD2SigmaSchedule.from_metadata(release)
|
||||
@@ -444,6 +524,58 @@ def _dit_rank_world() -> tuple[Any, int, int]:
|
||||
return group, dist.get_rank(group), dist.get_world_size(group)
|
||||
|
||||
|
||||
def _broadcast_rank0_exception(exc: Exception | None) -> None:
|
||||
"""Synchronize a rank-0-only exception across every DiT rank.
|
||||
|
||||
H3 reference-video preparation runs only on rank 0; the other DiT ranks
|
||||
return ``None`` without touching disk. When rank 0 raises inside that
|
||||
path it exits :meth:`prepare_encode` before reaching the downstream
|
||||
``dist.broadcast`` calls, and non-zero ranks then hang on those
|
||||
collectives forever. Every rank calls this helper right after the
|
||||
rank-0-only work, before any subsequent collective, so all ranks either
|
||||
raise the same error together or all continue.
|
||||
"""
|
||||
group, rank, world_size = _dit_rank_world()
|
||||
if world_size == 1:
|
||||
if exc is not None:
|
||||
raise exc
|
||||
return
|
||||
if rank == 0:
|
||||
if exc is None:
|
||||
payload: list[Any] = [None]
|
||||
else:
|
||||
payload = [
|
||||
{
|
||||
"type": type(exc).__name__,
|
||||
"message": str(exc),
|
||||
"status_code": getattr(exc, "status_code", None),
|
||||
"error_type": getattr(exc, "error_type", None),
|
||||
}
|
||||
]
|
||||
else:
|
||||
payload = [None]
|
||||
dist.broadcast_object_list(payload, src=0, group=group)
|
||||
info = payload[0]
|
||||
if info is None:
|
||||
return
|
||||
if rank == 0:
|
||||
assert exc is not None
|
||||
raise exc
|
||||
# Rebuild a matching client-facing error on non-zero ranks so the runner's
|
||||
# per-request try/except records the same 4xx status as rank 0. The exact
|
||||
# subclass need not survive the wire; the message and status suffice.
|
||||
status_code = info.get("status_code")
|
||||
error_type = info.get("error_type")
|
||||
message = f"[rank 0] {info['type']}: {info['message']}"
|
||||
if status_code is not None:
|
||||
raise client_error_from_metadata(
|
||||
message,
|
||||
status_code=int(status_code),
|
||||
error_type=error_type,
|
||||
)
|
||||
raise RuntimeError(message)
|
||||
|
||||
|
||||
def _broadcast_tensor(
|
||||
tensor: torch.Tensor | None,
|
||||
*,
|
||||
@@ -548,6 +680,8 @@ class MiniMaxH3Pipeline(
|
||||
):
|
||||
"""CFG-distilled joint video/audio generation for MiniMax H3."""
|
||||
|
||||
supports_step_execution: ClassVar[bool] = True
|
||||
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
|
||||
_vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"]
|
||||
@@ -564,6 +698,9 @@ class MiniMaxH3Pipeline(
|
||||
"_encode_reference_audio_conditions",
|
||||
"diffuse",
|
||||
"decode",
|
||||
"prepare_encode",
|
||||
"denoise_step",
|
||||
"post_decode",
|
||||
]
|
||||
dummy_run_num_frames: ClassVar[int] = 0
|
||||
# Only distilled releases pin a schedule, so the default keeps the legacy
|
||||
@@ -1538,7 +1675,7 @@ class MiniMaxH3Pipeline(
|
||||
if controller is not None and enabled:
|
||||
controller.offload_resident_layers()
|
||||
|
||||
def diffuse(
|
||||
def _build_denoise_inputs(
|
||||
self,
|
||||
*,
|
||||
task: str,
|
||||
@@ -1562,7 +1699,12 @@ class MiniMaxH3Pipeline(
|
||||
visual_condition_shapes: list[tuple[int, int, int]] | None = None,
|
||||
audio_condition_lengths: list[int] | None = None,
|
||||
keyframe_frame_indices: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> dict[str, Any]:
|
||||
"""Build the packed layout, initial rows, anchors, and sigma schedules.
|
||||
|
||||
Shared by request-mode :meth:`diffuse` and step-mode
|
||||
:meth:`prepare_encode` so both paths start from identical state.
|
||||
"""
|
||||
initial_video, initial_audio = self._initial_noise(
|
||||
seed=seed,
|
||||
latent_t=latent_t,
|
||||
@@ -1662,32 +1804,34 @@ class MiniMaxH3Pipeline(
|
||||
shift_scale=audio_shift,
|
||||
base_schedule=base_schedule,
|
||||
)
|
||||
transformer = self._transformer_for_task(task)
|
||||
# The static DLO plan keeps leading blocks resident only for the
|
||||
# primary ``transformer``. In combined mode ``transformers_ref`` is
|
||||
# fully streamed, so a Ref2VA request must not stage the inactive
|
||||
# FL2VA transformer resident blocks.
|
||||
with self._resident_dit_layers_on_device(enabled=transformer is self.transformer):
|
||||
with self.progress_bar(total=len(video_sigmas) - 1) as progress:
|
||||
video_rows, audio_rows = minimax_h3_denoise_loop(
|
||||
model=transformer,
|
||||
positive=branch,
|
||||
initial_video_rows=initial_video,
|
||||
initial_audio_rows=initial_audio,
|
||||
keyframe_cond_rows=visual_anchor,
|
||||
audio_ref_rows=audio_anchor,
|
||||
sigmas_video=video_sigmas,
|
||||
sigmas_audio=audio_sigmas,
|
||||
device=self.device,
|
||||
imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
|
||||
audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
|
||||
on_step_start=lambda step, video_sigma, audio_sigma: self.record_denoise_step(
|
||||
step,
|
||||
normalized_timestep=video_sigma,
|
||||
),
|
||||
on_step_end=lambda step, video, audio: progress.update(),
|
||||
)
|
||||
return {
|
||||
"branch": branch,
|
||||
# The request-mode loop moves these onto the device itself; step mode
|
||||
# keeps them resident across steps, so normalize once for both.
|
||||
"video_rows": initial_video.to(device=self.device, dtype=torch.float32),
|
||||
"audio_rows": initial_audio.to(device=self.device, dtype=torch.float32),
|
||||
"cond_anchor": (
|
||||
None if visual_anchor is None else visual_anchor.to(device=self.device, dtype=torch.float32)
|
||||
),
|
||||
"audio_anchor": (
|
||||
None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32)
|
||||
),
|
||||
"sigmas_video": video_sigmas,
|
||||
"sigmas_audio": audio_sigmas,
|
||||
}
|
||||
|
||||
def _unpack_denoised_rows(
|
||||
self,
|
||||
branch: MiniMaxH3DenoiseBranch,
|
||||
video_rows: torch.Tensor,
|
||||
audio_rows: torch.Tensor,
|
||||
*,
|
||||
latent_t: int,
|
||||
latent_h: int,
|
||||
latent_w: int,
|
||||
audio_t: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Select the target rows and unpack them back into VAE latents."""
|
||||
target_video = video_rows[branch.update_mask_dev]
|
||||
video_latent = minimax_h3_unpatchify_video_tokens(
|
||||
target_video,
|
||||
@@ -1707,6 +1851,83 @@ class MiniMaxH3Pipeline(
|
||||
)
|
||||
return video_latent, audio_latent
|
||||
|
||||
def diffuse(
|
||||
self,
|
||||
*,
|
||||
task: str,
|
||||
text_embeddings: torch.Tensor,
|
||||
text_tags: torch.Tensor,
|
||||
seed: int,
|
||||
latent_t: int,
|
||||
latent_h: int,
|
||||
latent_w: int,
|
||||
audio_t: int,
|
||||
num_frames: int,
|
||||
num_steps: int,
|
||||
video_shift: float,
|
||||
audio_shift: float,
|
||||
base_schedule: Sequence[float] | None,
|
||||
visual_condition: torch.Tensor | None,
|
||||
visual_condition_shape: tuple[int, int, int] | None,
|
||||
audio_condition: torch.Tensor | None,
|
||||
ref_audio_t: int | None,
|
||||
ref_blocks: list[dict[str, Any]] | None = None,
|
||||
visual_condition_shapes: list[tuple[int, int, int]] | None = None,
|
||||
audio_condition_lengths: list[int] | None = None,
|
||||
keyframe_frame_indices: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
inputs = self._build_denoise_inputs(
|
||||
task=task,
|
||||
text_embeddings=text_embeddings,
|
||||
text_tags=text_tags,
|
||||
seed=seed,
|
||||
latent_t=latent_t,
|
||||
latent_h=latent_h,
|
||||
latent_w=latent_w,
|
||||
audio_t=audio_t,
|
||||
num_frames=num_frames,
|
||||
num_steps=num_steps,
|
||||
video_shift=video_shift,
|
||||
audio_shift=audio_shift,
|
||||
base_schedule=base_schedule,
|
||||
visual_condition=visual_condition,
|
||||
visual_condition_shape=visual_condition_shape,
|
||||
audio_condition=audio_condition,
|
||||
ref_audio_t=ref_audio_t,
|
||||
ref_blocks=ref_blocks,
|
||||
visual_condition_shapes=visual_condition_shapes,
|
||||
audio_condition_lengths=audio_condition_lengths,
|
||||
keyframe_frame_indices=keyframe_frame_indices,
|
||||
)
|
||||
branch = inputs["branch"]
|
||||
transformer = self._transformer_for_task(task)
|
||||
with self._resident_dit_layers_on_device(enabled=transformer is self.transformer):
|
||||
with self.progress_bar(total=len(inputs["sigmas_video"]) - 1) as progress:
|
||||
video_rows, audio_rows = minimax_h3_denoise_loop(
|
||||
model=transformer,
|
||||
positive=branch,
|
||||
initial_video_rows=inputs["video_rows"],
|
||||
initial_audio_rows=inputs["audio_rows"],
|
||||
keyframe_cond_rows=inputs["cond_anchor"],
|
||||
audio_ref_rows=inputs["audio_anchor"],
|
||||
sigmas_video=inputs["sigmas_video"],
|
||||
sigmas_audio=inputs["sigmas_audio"],
|
||||
device=self.device,
|
||||
imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
|
||||
audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
|
||||
on_step=lambda step, video, audio: progress.update(),
|
||||
)
|
||||
|
||||
return self._unpack_denoised_rows(
|
||||
branch,
|
||||
video_rows,
|
||||
audio_rows,
|
||||
latent_t=latent_t,
|
||||
latent_h=latent_h,
|
||||
latent_w=latent_w,
|
||||
audio_t=audio_t,
|
||||
)
|
||||
|
||||
def decode(
|
||||
self,
|
||||
video_latent: torch.Tensor,
|
||||
@@ -1727,11 +1948,9 @@ class MiniMaxH3Pipeline(
|
||||
audio = self.audio_vae.decode_latent(audio_latent)
|
||||
return video, audio
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
|
||||
if len(request.prompts) != 1:
|
||||
raise OmniClientError("MiniMax H3 supports one request at a time")
|
||||
raw_prompt = request.prompts[0]
|
||||
@staticmethod
|
||||
def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]:
|
||||
"""Split a request prompt into its text and multimodal parts."""
|
||||
if isinstance(raw_prompt, str):
|
||||
prompt = raw_prompt
|
||||
multi_modal_data: dict[str, Any] = {}
|
||||
@@ -1740,8 +1959,21 @@ class MiniMaxH3Pipeline(
|
||||
multi_modal_data = raw_prompt.get("multi_modal_data") or {}
|
||||
if not prompt:
|
||||
raise OmniClientError("MiniMax H3 requires a non-empty prompt")
|
||||
return prompt, multi_modal_data
|
||||
|
||||
sampling = request.sampling_params
|
||||
def _prepare_request_inputs(
|
||||
self,
|
||||
*,
|
||||
prompt: str,
|
||||
multi_modal_data: dict[str, Any],
|
||||
sampling: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Resolve the task and output shape, then run every request-level encode.
|
||||
|
||||
Shared by request-mode :meth:`forward` and step-mode
|
||||
:meth:`prepare_encode`; the returned mapping feeds :meth:`diffuse` and
|
||||
:meth:`_build_denoise_inputs` unchanged.
|
||||
"""
|
||||
quality = sampling.quality
|
||||
logger.debug("MiniMax H3 request quality=%s", quality)
|
||||
extra = sampling.extra_args or {}
|
||||
@@ -1818,12 +2050,24 @@ class MiniMaxH3Pipeline(
|
||||
video_count = 0
|
||||
if raw_videos is not None:
|
||||
video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1
|
||||
prepared_videos = self._prepare_reference_videos(
|
||||
raw_videos,
|
||||
target_frame_count=num_frames,
|
||||
workdir=workdir,
|
||||
start_time_seconds=extra.get("start_time_seconds"),
|
||||
)
|
||||
# File-based reference-video prep runs only on rank 0; other
|
||||
# ranks return None without touching disk. If rank 0 raises
|
||||
# (e.g. invalid file, unsupported codec) it must not exit
|
||||
# ``prepare_encode`` before the downstream broadcasts below --
|
||||
# non-zero ranks would then deadlock on them forever. Capture
|
||||
# the exception here and let every rank agree on the outcome
|
||||
# before starting any subsequent collective.
|
||||
prep_error: Exception | None = None
|
||||
try:
|
||||
prepared_videos = self._prepare_reference_videos(
|
||||
raw_videos,
|
||||
target_frame_count=num_frames,
|
||||
workdir=workdir,
|
||||
start_time_seconds=extra.get("start_time_seconds"),
|
||||
)
|
||||
except Exception as exc:
|
||||
prep_error = exc
|
||||
_broadcast_rank0_exception(prep_error)
|
||||
has_audio_tensor = torch.zeros(
|
||||
video_count,
|
||||
dtype=torch.long,
|
||||
@@ -1960,33 +2204,60 @@ class MiniMaxH3Pipeline(
|
||||
)
|
||||
self._cache_dit_runtime.prepare(quality_plan.cache_dit)
|
||||
num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt)
|
||||
return {
|
||||
"task": task,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"num_frames": num_frames,
|
||||
"latent_t": latent_t,
|
||||
"latent_h": height // 16,
|
||||
"latent_w": width // 16,
|
||||
"audio_t": audio_t,
|
||||
"text_embeddings": text_embeddings,
|
||||
"text_tags": text_tags,
|
||||
"visual_condition": visual_condition,
|
||||
"visual_condition_shape": visual_shape,
|
||||
"audio_condition": audio_condition,
|
||||
"ref_audio_t": ref_audio_t,
|
||||
"ref_blocks": ref_blocks,
|
||||
"visual_condition_shapes": visual_shapes,
|
||||
"audio_condition_lengths": audio_lengths,
|
||||
"keyframe_frame_indices": keyframe_frame_indices,
|
||||
"seed": seed,
|
||||
"num_steps": num_steps,
|
||||
"video_shift": video_shift,
|
||||
"audio_shift": audio_shift,
|
||||
"base_schedule": base_schedule,
|
||||
"num_outputs": num_outputs,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _denoise_kwargs(context: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Select the denoise-input arguments from a prepared request context."""
|
||||
return {key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS}
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
|
||||
if len(request.prompts) != 1:
|
||||
raise OmniClientError("MiniMax H3 supports one request at a time")
|
||||
prompt, multi_modal_data = self._extract_prompt(request.prompts[0])
|
||||
context = self._prepare_request_inputs(
|
||||
prompt=prompt,
|
||||
multi_modal_data=multi_modal_data,
|
||||
sampling=request.sampling_params,
|
||||
)
|
||||
denoise_kwargs = self._denoise_kwargs(context)
|
||||
num_outputs = context["num_outputs"]
|
||||
videos = []
|
||||
audios = []
|
||||
for output_seed in _minimax_h3_output_seeds(seed, num_outputs):
|
||||
video_latent, audio_latent = self.diffuse(
|
||||
task=task,
|
||||
text_embeddings=text_embeddings,
|
||||
text_tags=text_tags,
|
||||
seed=output_seed,
|
||||
latent_t=latent_t,
|
||||
latent_h=height // 16,
|
||||
latent_w=width // 16,
|
||||
audio_t=audio_t,
|
||||
num_frames=num_frames,
|
||||
num_steps=num_steps,
|
||||
video_shift=video_shift,
|
||||
audio_shift=audio_shift,
|
||||
base_schedule=base_schedule,
|
||||
visual_condition=visual_condition,
|
||||
visual_condition_shape=visual_shape,
|
||||
audio_condition=audio_condition,
|
||||
ref_audio_t=ref_audio_t,
|
||||
ref_blocks=ref_blocks,
|
||||
visual_condition_shapes=visual_shapes,
|
||||
audio_condition_lengths=audio_lengths,
|
||||
keyframe_frame_indices=keyframe_frame_indices,
|
||||
for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs):
|
||||
video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed})
|
||||
video, audio = self.decode(
|
||||
video_latent,
|
||||
audio_latent,
|
||||
height=context["height"],
|
||||
width=context["width"],
|
||||
)
|
||||
video, audio = self.decode(video_latent, audio_latent, height=height, width=width)
|
||||
videos.append(video)
|
||||
audios.append(audio)
|
||||
video = torch.cat(videos, dim=0)
|
||||
@@ -1997,6 +2268,302 @@ class MiniMaxH3Pipeline(
|
||||
stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Step-wise execution (continuous batching)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _packed_batch_supported(transformer: MiniMaxH3DiTModel) -> bool:
|
||||
"""Whether every attention in this DiT honors multi-document cu_seqlens.
|
||||
|
||||
A packed batch is only isolated if *all* of them do: the token refiner
|
||||
runs under its own attention role and can resolve to a different backend
|
||||
from the DiT blocks. Ring sequence parallelism dispatches through
|
||||
``RingParallelAttention``, whose kernels ignore the packed
|
||||
``cu_seqlens`` metadata regardless of the configured backend; packing
|
||||
multiple requests under ring would let attention cross document
|
||||
boundaries, so any layer running ring disqualifies the batch.
|
||||
|
||||
The gate probes a per-backend capability rather than a fixed backend
|
||||
name: FLASH_ATTN, for example, only isolates arbitrary N-document
|
||||
packed cu_seqlens on CUDA/ROCm/MUSA. Its NPU path only accepts a
|
||||
``[real, pad]`` two-document layout and its XPU path ignores
|
||||
cu_seqlens outright — either would silently attend across request
|
||||
boundaries.
|
||||
"""
|
||||
attentions = [module for module in transformer.modules() if isinstance(module, MiniMaxH3Attention)]
|
||||
if not attentions:
|
||||
return False
|
||||
return all(_attention_isolates_packed_requests(module.attention) for module in attentions)
|
||||
|
||||
def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState:
|
||||
"""Run every request-level stage once and seed the per-request step state."""
|
||||
del kwargs
|
||||
# Two request-mode features have no place in the shared step contract:
|
||||
# a request state carries exactly one latent tensor, and distributed
|
||||
# layerwise offload streams the DiT around one whole denoise loop rather
|
||||
# than around a single scheduler-driven step.
|
||||
num_outputs = _resolve_minimax_h3_num_outputs(state.sampling.num_outputs_per_prompt)
|
||||
if num_outputs != 1:
|
||||
raise OmniClientError(
|
||||
f"MiniMax H3 step execution produces one output per request, got num_outputs_per_prompt={num_outputs}"
|
||||
)
|
||||
if getattr(self, "_dlo_residency_controller", None) is not None:
|
||||
raise ValueError(
|
||||
"MiniMax H3 step execution is not compatible with distributed layerwise offload; "
|
||||
"the resident-layer window spans a whole denoise loop, so per-step streaming would "
|
||||
"reload the DiT every step. Drop --step-execution or --enable-distributed-layerwise-offload."
|
||||
)
|
||||
# Request-scoped Cache-DiT (quality=high) mutates hook state on the
|
||||
# shared transformer rather than on ``StepRequestState``. In step mode
|
||||
# two requests can interleave denoise steps, or be co-batched into a
|
||||
# single forward, and the second one would then re-enter the DiT with
|
||||
# cache buffers shaped for the first. Reject the profile here rather
|
||||
# than let it corrupt outputs at runtime; startup-configured Cache-DiT
|
||||
# is already blocked in ``DiffusionModelRunner.execute_stepwise``.
|
||||
if getattr(state.sampling, "quality", None) == "high":
|
||||
raise OmniClientError(
|
||||
"MiniMax H3 step execution does not support the high-quality Cache-DiT profile "
|
||||
"(quality=high); its hooks live on the shared transformer, so interleaved or "
|
||||
"co-batched requests would reuse incompatible cache state. Drop --step-execution "
|
||||
"or omit quality=high."
|
||||
)
|
||||
prompt, multi_modal_data = self._extract_prompt(state.prompt)
|
||||
context = self._prepare_request_inputs(
|
||||
prompt=prompt,
|
||||
multi_modal_data=multi_modal_data,
|
||||
sampling=state.sampling,
|
||||
)
|
||||
inputs = self._build_denoise_inputs(**self._denoise_kwargs(context))
|
||||
|
||||
sigmas_video = inputs["sigmas_video"]
|
||||
sigmas_audio = inputs["sigmas_audio"]
|
||||
if len(sigmas_video) < 2:
|
||||
raise OmniClientError(
|
||||
f"MiniMax H3 step execution needs at least one denoise step, got num_inference_steps="
|
||||
f"{len(sigmas_video) - 1}"
|
||||
)
|
||||
|
||||
branch = inputs["branch"]
|
||||
video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows(
|
||||
positive=branch,
|
||||
initial_video_rows=inputs["video_rows"],
|
||||
initial_audio_rows=inputs["audio_rows"],
|
||||
keyframe_cond_rows=inputs["cond_anchor"],
|
||||
audio_ref_rows=inputs["audio_anchor"],
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
# The denoise loop consumes sigma pairs, so the schedule carries one more
|
||||
# point than there are steps. ``timesteps`` holds the video branch because
|
||||
# the shared contract gives a request exactly one timestep sequence; the
|
||||
# audio schedule rides along in ``extra``.
|
||||
state.timesteps = torch.tensor(
|
||||
[1.0 - sigma for sigma in sigmas_video[:-1]],
|
||||
dtype=torch.float32,
|
||||
device=self.device,
|
||||
)
|
||||
state.step_index = 0
|
||||
# Video rows are the batched tensor the runner slices per request; audio
|
||||
# rows have a different width, so they stay request-private.
|
||||
state.latents = video_rows
|
||||
state.do_true_cfg = False # H3 checkpoints are CFG-distilled.
|
||||
state.extra.update(
|
||||
{
|
||||
_STEP_BRANCH: branch,
|
||||
_STEP_TRANSFORMER: self._transformer_for_task(context["task"]),
|
||||
_STEP_AUDIO_ROWS: audio_rows,
|
||||
_STEP_COND_ANCHOR: cond_anchor,
|
||||
_STEP_AUDIO_ANCHOR: audio_anchor,
|
||||
_STEP_SIGMAS_VIDEO: sigmas_video,
|
||||
_STEP_SIGMAS_AUDIO: sigmas_audio,
|
||||
_STEP_SHAPE: {
|
||||
"height": context["height"],
|
||||
"width": context["width"],
|
||||
"latent_t": context["latent_t"],
|
||||
"latent_h": context["latent_h"],
|
||||
"latent_w": context["latent_w"],
|
||||
"audio_t": context["audio_t"],
|
||||
},
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def denoise_step(
|
||||
self,
|
||||
input_batch: InputBatch,
|
||||
*,
|
||||
states: Sequence[StepRequestState] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor | None:
|
||||
"""Run one denoise forward covering every request in the batch.
|
||||
|
||||
Requests are concatenated into a single packed sequence that keeps one
|
||||
attention document each, so the whole batch costs one DiT forward.
|
||||
Backends that ignore ``cu_seqlens`` cannot express that isolation, so
|
||||
they fall back to one forward per request.
|
||||
"""
|
||||
del kwargs
|
||||
batch_states = list(states if states is not None else input_batch.states)
|
||||
|
||||
branches = [state.extra[_STEP_BRANCH] for state in batch_states]
|
||||
video_rows = [state.latents for state in batch_states]
|
||||
audio_rows = [state.extra[_STEP_AUDIO_ROWS] for state in batch_states]
|
||||
schedules = [_minimax_h3_step_schedule(state) for state in batch_states]
|
||||
transformers = [state.extra[_STEP_TRANSFORMER] for state in batch_states]
|
||||
mixed_transformers = len({id(transformer) for transformer in transformers}) > 1
|
||||
|
||||
# Both execution modes must publish denoise progress for step-gated
|
||||
# attention features. Requests can differ in both step index and sigma
|
||||
# schedule, so a batch that is not at one single point has nothing to
|
||||
# publish and those gates stay dense -- which is their safe default.
|
||||
progress = {(state.step_index, schedule["sigma_video"]) for state, schedule in zip(batch_states, schedules)}
|
||||
minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None)))
|
||||
|
||||
if len(batch_states) > 1 and (mixed_transformers or not self._packed_batch_supported(transformers[0])):
|
||||
if mixed_transformers:
|
||||
logger.warning_once(
|
||||
"MiniMax H3 step batch contains requests for different task-specific DiTs; "
|
||||
"running %d requests one forward at a time.",
|
||||
len(batch_states),
|
||||
)
|
||||
elif any(
|
||||
getattr(getattr(module, "attention", None), "use_ring", False)
|
||||
for module in transformers[0].modules()
|
||||
if isinstance(module, MiniMaxH3Attention)
|
||||
):
|
||||
logger.warning_once(
|
||||
"MiniMax H3 step batching is disabled when ring attention is active: "
|
||||
"the ring kernels ignore packed cu_seqlens and would attend across request "
|
||||
"boundaries. Running %d requests one forward at a time.",
|
||||
len(batch_states),
|
||||
)
|
||||
else:
|
||||
logger.warning_once(
|
||||
"MiniMax H3 step batching needs every attention on a backend that isolates "
|
||||
"packed multi-document cu_seqlens (see AttentionBackend."
|
||||
"supports_multi_doc_packed_varlen); running %d requests one forward at a time.",
|
||||
len(batch_states),
|
||||
)
|
||||
video_parts: list[torch.Tensor] = []
|
||||
audio_parts: list[torch.Tensor] = []
|
||||
for index, branch in enumerate(branches):
|
||||
forward_kwargs = branch.forward_kwargs(
|
||||
video_rows=video_rows[index],
|
||||
audio_rows=audio_rows[index],
|
||||
t_video=schedules[index]["t_video"],
|
||||
t_audio=schedules[index]["t_audio"],
|
||||
imgvid_cond_timestep=schedules[index]["imgvid_cond_timestep"],
|
||||
audio_ref_cond_timestep=schedules[index]["audio_ref_cond_timestep"],
|
||||
)
|
||||
request_video, request_audio = transformers[index](**forward_kwargs)
|
||||
video_parts.append(request_video)
|
||||
audio_parts.append(request_audio)
|
||||
video_velocity = torch.cat(video_parts)
|
||||
audio_velocity = torch.cat(audio_parts)
|
||||
else:
|
||||
forward_kwargs = minimax_h3_batched_forward_kwargs(
|
||||
branches=branches,
|
||||
video_rows=video_rows,
|
||||
audio_rows=audio_rows,
|
||||
t_video=[schedule["t_video"] for schedule in schedules],
|
||||
t_audio=[schedule["t_audio"] for schedule in schedules],
|
||||
imgvid_cond_timesteps=[schedule["imgvid_cond_timestep"] for schedule in schedules],
|
||||
audio_ref_cond_timesteps=[schedule["audio_ref_cond_timestep"] for schedule in schedules],
|
||||
)
|
||||
logger.debug(
|
||||
"MiniMax H3 denoise step: %d request(s) packed into %d rows",
|
||||
len(batch_states),
|
||||
int(forward_kwargs["x"].shape[1]),
|
||||
)
|
||||
video_velocity, audio_velocity = transformers[0](**forward_kwargs)
|
||||
|
||||
# The shared contract carries one velocity tensor per step, and audio rows
|
||||
# are a different width than video rows, so hand the audio branch to
|
||||
# step_scheduler() through request-private state.
|
||||
audio_parts_by_request = torch.split(audio_velocity, [int(branch.audio_pos.shape[0]) for branch in branches])
|
||||
for state, request_audio in zip(batch_states, audio_parts_by_request, strict=True):
|
||||
state.extra[_STEP_AUDIO_NOISE_PRED] = request_audio
|
||||
return video_velocity
|
||||
|
||||
def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
|
||||
"""Apply one Euler-eta0 update to this request's video and audio rows."""
|
||||
del kwargs
|
||||
# denoise_step() stages the audio half of this step's velocity; popping
|
||||
# it keeps a second step_scheduler() call from reusing a stale one.
|
||||
audio_noise_pred = state.extra.pop(_STEP_AUDIO_NOISE_PRED)
|
||||
|
||||
branch = state.extra[_STEP_BRANCH]
|
||||
schedule = _minimax_h3_step_schedule(state)
|
||||
update = branch.update_mask_dev
|
||||
audio_update = branch.audio_update_mask_dev
|
||||
video_rows = state.latents
|
||||
audio_rows = state.extra[_STEP_AUDIO_ROWS]
|
||||
cond_anchor = state.extra[_STEP_COND_ANCHOR]
|
||||
audio_anchor = state.extra[_STEP_AUDIO_ANCHOR]
|
||||
device = video_rows.device
|
||||
|
||||
x0_video = minimax_h3_rf_v_to_x0(
|
||||
video_rows[update],
|
||||
noise_pred.float()[update],
|
||||
torch.tensor(schedule["t_video"], dtype=torch.float32, device=device),
|
||||
)
|
||||
new_video = minimax_h3_euler_eta0_step(
|
||||
video_rows[update],
|
||||
x0_video,
|
||||
sigma_curr=schedule["sigma_video"],
|
||||
sigma_next=schedule["sigma_video_next"],
|
||||
)
|
||||
video_rows = video_rows.clone()
|
||||
video_rows[update] = new_video
|
||||
if cond_anchor is not None:
|
||||
video_rows[~update] = cond_anchor # per-step imgvid cond reset
|
||||
|
||||
x0_audio = minimax_h3_rf_v_to_x0(
|
||||
audio_rows[audio_update],
|
||||
audio_noise_pred.float()[audio_update],
|
||||
torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device),
|
||||
)
|
||||
new_audio = minimax_h3_euler_eta0_step(
|
||||
audio_rows[audio_update],
|
||||
x0_audio,
|
||||
sigma_curr=schedule["sigma_audio"],
|
||||
sigma_next=schedule["sigma_audio_next"],
|
||||
)
|
||||
audio_rows = audio_rows.clone()
|
||||
audio_rows[audio_update] = new_audio
|
||||
if audio_anchor is not None:
|
||||
audio_rows[~audio_update] = audio_anchor # per-step audio ref reset
|
||||
|
||||
state.latents = video_rows
|
||||
state.extra[_STEP_AUDIO_ROWS] = audio_rows
|
||||
state.step_index += 1
|
||||
|
||||
def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput:
|
||||
"""Unpack the denoised rows and run the joint video/audio VAE decode."""
|
||||
del kwargs
|
||||
shape = state.extra[_STEP_SHAPE]
|
||||
video_latent, audio_latent = self._unpack_denoised_rows(
|
||||
state.extra[_STEP_BRANCH],
|
||||
state.latents,
|
||||
state.extra[_STEP_AUDIO_ROWS],
|
||||
latent_t=shape["latent_t"],
|
||||
latent_h=shape["latent_h"],
|
||||
latent_w=shape["latent_w"],
|
||||
audio_t=shape["audio_t"],
|
||||
)
|
||||
video, audio = self.decode(
|
||||
video_latent,
|
||||
audio_latent,
|
||||
height=shape["height"],
|
||||
width=shape["width"],
|
||||
)
|
||||
return DiffusionOutput(
|
||||
output=(video, audio),
|
||||
post_process_func=get_minimax_h3_post_process_func(self.od_config),
|
||||
stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MiniMaxH3Pipeline",
|
||||
|
||||
@@ -81,6 +81,31 @@ if TYPE_CHECKING:
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _dit_any_rank_failed(local_failed: bool) -> bool:
|
||||
"""All-reduce a per-request failure flag across the DiT process group.
|
||||
|
||||
Every DiT rank must reach this point in lockstep; the caller wraps
|
||||
``pipeline.prepare_encode`` so both the success and the exception paths
|
||||
call the helper. In single-rank execution this collapses to the local
|
||||
flag with no collectives issued.
|
||||
"""
|
||||
if not torch.distributed.is_initialized():
|
||||
return local_failed
|
||||
try:
|
||||
from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
|
||||
|
||||
group = get_dit_group()
|
||||
except (AssertionError, ImportError):
|
||||
group = None
|
||||
if group is None:
|
||||
return local_failed
|
||||
signal = torch.tensor(1 if local_failed else 0, dtype=torch.int32)
|
||||
if current_omni_platform.is_available():
|
||||
signal = signal.to(device=current_omni_platform.device_type)
|
||||
torch.distributed.all_reduce(signal, op=torch.distributed.ReduceOp.MAX, group=group)
|
||||
return bool(signal.item())
|
||||
|
||||
|
||||
def _normalize_pipeline_outputs(
|
||||
outputs: object,
|
||||
*,
|
||||
@@ -853,25 +878,70 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
|
||||
return resolved, new_request_ids
|
||||
|
||||
def _prepare_batch_inputs(self, states: list[StepRequestState], new_request_ids: list[str]) -> InputBatch:
|
||||
def _prepare_batch_inputs(
|
||||
self,
|
||||
states: list[StepRequestState],
|
||||
new_request_ids: list[str],
|
||||
) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]:
|
||||
# process new reqs
|
||||
prepared_states: list[StepRequestState] = []
|
||||
error_outputs: list[RunnerOutput] = []
|
||||
for state in states:
|
||||
if state.request_id in new_request_ids:
|
||||
self._initialize_generator(state.sampling)
|
||||
clear_pipeline_stage_durations(self.pipeline)
|
||||
# encode
|
||||
self.pipeline.prepare_encode(state)
|
||||
merge_stage_durations(
|
||||
state,
|
||||
consume_pipeline_stage_durations(self.pipeline),
|
||||
)
|
||||
# Everything that runs before ``_dit_any_rank_failed`` must be
|
||||
# inside the try: an exception in ``_initialize_generator`` or
|
||||
# ``clear_pipeline_stage_durations`` on one rank would skip the
|
||||
# all-reduce here while every peer proceeds into it, and the
|
||||
# peers then hang on the NCCL collective until timeout.
|
||||
per_req_exc: BaseException | None = None
|
||||
try:
|
||||
self._initialize_generator(state.sampling)
|
||||
clear_pipeline_stage_durations(self.pipeline)
|
||||
# encode
|
||||
self.pipeline.prepare_encode(state)
|
||||
merge_stage_durations(
|
||||
state,
|
||||
consume_pipeline_stage_durations(self.pipeline),
|
||||
)
|
||||
except Exception as exc:
|
||||
per_req_exc = exc
|
||||
# Pipelines that do rank-0-only work (e.g. MiniMax H3
|
||||
# reference-video prep) must broadcast per-request failures
|
||||
# internally so downstream collectives stay in step; even so,
|
||||
# cross-check that every DiT rank agrees so a rank-local error
|
||||
# (or a future pipeline that omits the guard) does not leave
|
||||
# the process group half-way through a new request.
|
||||
if _dit_any_rank_failed(per_req_exc is not None):
|
||||
self.state_cache.pop(state.request_id, None)
|
||||
if per_req_exc is None:
|
||||
per_req_exc = RuntimeError(
|
||||
f"Stepwise preparation failed on another DiT rank for {state.request_id}"
|
||||
)
|
||||
logger.error(
|
||||
"Stepwise request preparation failed for %s: %s",
|
||||
state.request_id,
|
||||
per_req_exc,
|
||||
exc_info=isinstance(per_req_exc, Exception),
|
||||
)
|
||||
error_outputs.append(
|
||||
RunnerOutput(
|
||||
request_id=state.request_id,
|
||||
step_index=state.step_index,
|
||||
finished=True,
|
||||
result=DiffusionOutput.from_exception(per_req_exc),
|
||||
)
|
||||
)
|
||||
continue
|
||||
prepared_states.append(state)
|
||||
|
||||
if not prepared_states:
|
||||
return prepared_states, None, error_outputs
|
||||
input_batch = InputBatch.make_batch(
|
||||
states,
|
||||
prepared_states,
|
||||
cached_batch=getattr(self, "input_batch", None),
|
||||
)
|
||||
self.input_batch = input_batch
|
||||
return input_batch
|
||||
return prepared_states, input_batch, error_outputs
|
||||
|
||||
def _update_states_after(
|
||||
self,
|
||||
@@ -978,7 +1048,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
and current_omni_platform.is_available()
|
||||
):
|
||||
current_omni_platform.reset_peak_memory_stats()
|
||||
input_batch = self._prepare_batch_inputs(states, new_request_ids)
|
||||
states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids)
|
||||
if input_batch is None:
|
||||
return BatchRunnerOutput.from_list(runner_output_list)
|
||||
attn_metadata = {}
|
||||
|
||||
with set_forward_context(
|
||||
@@ -995,7 +1067,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
denoise_stage_durations,
|
||||
)
|
||||
|
||||
runner_output_list = []
|
||||
pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
|
||||
if noise_pred is None and pipeline_interrupted:
|
||||
for state in states:
|
||||
@@ -1048,6 +1119,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
offset = offset + row_num
|
||||
except Exception as per_req_exc:
|
||||
offset = offset + row_num
|
||||
self.state_cache.pop(req.request_id, None)
|
||||
logger.error(
|
||||
"Stepwise per-request error for %s: %s",
|
||||
req.request_id,
|
||||
@@ -1059,7 +1131,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
request_id=req.request_id,
|
||||
step_index=req.step_index,
|
||||
finished=True,
|
||||
result=DiffusionOutput(error=str(per_req_exc)),
|
||||
result=DiffusionOutput.from_exception(per_req_exc),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
Reference in new issue
Block a user