Skip to content

Commit 15568f2

Browse files
SolitaryThinkerSolitaryThinker
authored
[perf]: H3 torch.compile + CUDA graphs (1.2-1.3x) with denoising step marking (#1689)
1 parent 126a75a commit 15568f2

4 files changed

Lines changed: 81 additions & 44 deletions

File tree

‎examples/inference/basic/basic_minimax_h3_t2v.py‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
from fastvideo import VideoGenerator
1010
from fastvideo.api import (
11+
CompileConfig,
1112
EngineConfig,
1213
GenerationRequest,
1314
GeneratorConfig,
@@ -29,6 +30,13 @@ def parse_args() -> argparse.Namespace:
2930
parser.add_argument("--steps", type=int, default=50)
3031
parser.add_argument("--seed", type=int, default=0)
3132
parser.add_argument("--num-gpus", type=int, default=4)
33+
parser.add_argument("--torch-compile", action="store_true",
34+
help="torch.compile the DiT transformer path")
35+
parser.add_argument("--compile-mode", default=None,
36+
help='torch.compile mode, e.g. "reduce-overhead" for CUDA graphs')
37+
parser.add_argument("--repeats", type=int, default=1,
38+
help="generate N times; with --torch-compile the first run pays "
39+
"compilation, so steady-state is the last repeat")
3240
return parser.parse_args()
3341

3442

@@ -51,11 +59,14 @@ def main() -> None:
5159
vae=True,
5260
pin_cpu_memory=False,
5361
),
62+
compile=CompileConfig(
63+
enabled=args.torch_compile,
64+
mode=args.compile_mode,
65+
),
5466
),
5567
))
5668
try:
57-
result = generator.generate(
58-
GenerationRequest(
69+
request = GenerationRequest(
5970
prompt=args.prompt,
6071
negative_prompt="",
6172
sampling=SamplingConfig(
@@ -73,8 +84,17 @@ def main() -> None:
7384
save_video=True,
7485
return_frames=False,
7586
),
76-
))
87+
)
88+
result = generator.generate(request)
7789
print(f"Output written to: {result.video_path}")
90+
if result.generation_time is not None:
91+
# machine-readable: benchmark harnesses parse this line to separate
92+
# generation from model-load time (last occurrence = steady state)
93+
print(f"Generation time: {result.generation_time:.2f}s")
94+
for _ in range(args.repeats - 1):
95+
result = generator.generate(request)
96+
if result.generation_time is not None:
97+
print(f"Generation time: {result.generation_time:.2f}s")
7898
finally:
7999
generator.shutdown()
80100

‎fastvideo/configs/models/dits/minimax_h3.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ class MiniMaxH3ArchConfig(DiTArchConfig):
2525
_supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
2626
AttentionBackendEnum.TORCH_SDPA,
2727
AttentionBackendEnum.FLASH_ATTN,
28+
# FP4-quantized QK attention (fa4_fp4 on sm_100/sm_103, cutlass on
29+
# sm_12x). Enabled for speed experiments; output quality against the
30+
# SSIM references is not yet validated.
31+
AttentionBackendEnum.ATTN_QAT_INFER,
2832
)
2933

3034
param_names_mapping: dict = field(

‎fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py‎

Lines changed: 50 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,16 @@
33

44
from __future__ import annotations
55

6+
import contextlib
67
from typing import Any
78

9+
810
import torch
911

1012
from fastvideo.distributed import get_local_torch_device
1113
from fastvideo.fastvideo_args import FastVideoArgs
1214
from fastvideo.forward_context import set_forward_context
15+
from fastvideo.profiler import get_global_controller
1316
from fastvideo.hooks.activation_trace import trace_step
1417
from fastvideo.pipelines.basic.minimax_h3.packing import (
1518
MINIMAX_H3_KEYFRAME_NOISE_AUG,
@@ -98,44 +101,54 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
98101
text_indices = layout.text_indices.to(device)
99102
prompt_embeds = batch.prompt_embeds[0].to(device)
100103

104+
controller = get_global_controller()
105+
denoise_region = (controller.region("profiler_region_inference_denoising")
106+
if controller is not None else contextlib.nullcontext())
101107
try:
102-
for index, (video_timestep, audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps,
103-
strict=True)):
104-
unique_timesteps, timestep_indices = row_timestep_plan[index]
105-
with trace_step(index), set_forward_context(
106-
current_timestep=index,
107-
attn_metadata=None,
108-
forward_batch=batch,
109-
):
110-
video_velocity, audio_velocity = self.transformer(
111-
hidden_states=batch.latents[None],
112-
audio_hidden_states=batch.audio_latents[None],
113-
encoder_hidden_states=prompt_embeds,
114-
timestep=unique_timesteps,
115-
timestep_indices=timestep_indices,
116-
token_tags=token_tags,
117-
position_ids=position_ids,
118-
video_indices=video_indices,
119-
audio_indices=audio_indices,
120-
text_indices=text_indices,
121-
)
122-
123-
video_start = layout.num_condition_video_rows
124-
audio_start = layout.num_condition_audio_rows
125-
batch.latents[video_start:] = self.scheduler.step(
126-
video_velocity[0, video_start:].float(),
127-
video_timestep,
128-
batch.latents[video_start:],
129-
return_dict=False,
130-
)[0]
131-
batch.audio_latents[audio_start:] = self.audio_scheduler.step(
132-
audio_velocity[0, audio_start:].float(),
133-
audio_timestep,
134-
batch.audio_latents[audio_start:],
135-
return_dict=False,
136-
)[0]
137-
batch.step_index = index
138-
batch.timestep = video_timestep
108+
with denoise_region:
109+
for index, (video_timestep, audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps,
110+
strict=True)):
111+
unique_timesteps, timestep_indices = row_timestep_plan[index]
112+
# Under torch.compile(mode="reduce-overhead") each denoising
113+
# step must be marked, or cudagraph trees flag cross-step
114+
# reuse of pooled outputs as "accessing tensor output of
115+
# CUDAGraphs that has been overwritten" (surfaces at sp=1;
116+
# sp>1 is masked by collective-induced graph breaks).
117+
torch.compiler.cudagraph_mark_step_begin()
118+
with trace_step(index), set_forward_context(
119+
current_timestep=index,
120+
attn_metadata=None,
121+
forward_batch=batch,
122+
):
123+
video_velocity, audio_velocity = self.transformer(
124+
hidden_states=batch.latents[None],
125+
audio_hidden_states=batch.audio_latents[None],
126+
encoder_hidden_states=prompt_embeds,
127+
timestep=unique_timesteps,
128+
timestep_indices=timestep_indices,
129+
token_tags=token_tags,
130+
position_ids=position_ids,
131+
video_indices=video_indices,
132+
audio_indices=audio_indices,
133+
text_indices=text_indices,
134+
)
135+
136+
video_start = layout.num_condition_video_rows
137+
audio_start = layout.num_condition_audio_rows
138+
batch.latents[video_start:] = self.scheduler.step(
139+
video_velocity[0, video_start:].float(),
140+
video_timestep,
141+
batch.latents[video_start:],
142+
return_dict=False,
143+
)[0]
144+
batch.audio_latents[audio_start:] = self.audio_scheduler.step(
145+
audio_velocity[0, audio_start:].float(),
146+
audio_timestep,
147+
batch.audio_latents[audio_start:],
148+
return_dict=False,
149+
)[0]
150+
batch.step_index = index
151+
batch.timestep = video_timestep
139152
finally:
140153
if bool(getattr(fastvideo_args, "dit_layerwise_offload", False)):
141154
manager = getattr(self.transformer, "_layerwise_offload_manager", None)

‎fastvideo/profiler.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -128,10 +128,10 @@ def set_global_controller(controller: TorchProfilerController | None) -> None:
128128
# name="profiler_region_inference_pre_denoising",
129129
# description="Pre-denoising inference steps (conditioning, preprocessing).",
130130
# )
131-
# register_profiler_region(
132-
# name="profiler_region_inference_denoising",
133-
# description="The main inference denoising loop.",
134-
# )
131+
register_profiler_region(
132+
name="profiler_region_inference_denoising",
133+
description="The main inference denoising loop.",
134+
)
135135
# register_profiler_region(
136136
# name="profiler_region_inference_post_denoising",
137137
# description=

0 commit comments

Comments
 (0)