Compare commits
2 commits
837c16871d
...
bd92baeb46
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bd92baeb46 | ||
|
|
1d5faa8f16 |
56 changed files with 4339 additions and 88 deletions
109
AUDIO_BOUNDARY_INVESTIGATION.md
Normal file
109
AUDIO_BOUNDARY_INVESTIGATION.md
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
# H3 Startup Audio Investigation
|
||||
|
||||
Investigation date: 2026-08-21
|
||||
|
||||
## Reproduction Cases
|
||||
|
||||
- Affected dialogue: base 12-step, Sage2, seed `440420`, tagged dialogue.
|
||||
- Clean control: base 12-step, Sage2, seed `440421`, immediate nightclub music.
|
||||
- Prompt-format control: affected dialogue prompt and seed with only the two
|
||||
`<d>[English]...</d>` spans replaced by quoted speech.
|
||||
|
||||
All diagnostic assets are under `/home/daniel/StoryStudioAssets/H3-output/h3-baselines`.
|
||||
|
||||
## Findings
|
||||
|
||||
1. The affected transient is already present in the retained lossless WAV. AAC
|
||||
encoding and MP4 muxing are not the source.
|
||||
2. The AudioVAE is not the primary source. Moving the affected first four audio
|
||||
latent frames to frame 40 in an otherwise near-silent latent carrier creates
|
||||
a similarly strong event at `1.0s`: `-18.44 dBFS` peak versus `-19.12 dBFS`
|
||||
when the same frames are placed at the start.
|
||||
3. AudioVAE boundary context changes the exact waveform, but the event remains.
|
||||
This makes the decoder a secondary shaper rather than the origin.
|
||||
4. Zero normalized latents are not silence. The official AudioVAE decodes them
|
||||
to approximately `-26.17 dBFS` RMS in the first 100ms, so zero replacement is
|
||||
not a valid repair.
|
||||
5. Repeating affected latent frame 4 produces near-silence (`-56.02 dBFS` RMS),
|
||||
but unconditional replacement is unsafe because valid music begins in the
|
||||
same first four frames in the clean control.
|
||||
6. There is no evidence for end-to-start wraparound in this sample. First/last
|
||||
four-frame latent cosine is `-0.122`; first/last 100ms PCM correlation is
|
||||
`0.006`.
|
||||
7. In the tagged-dialogue denoising trace, the unwanted onset is near silence
|
||||
through step 6 and begins growing materially at step 7 (`audio sigma 0.751`).
|
||||
It reaches `-18.94 dBFS` peak and `-34.23 dBFS` RMS in the final first 100ms.
|
||||
8. Clean immediate music is strongly predicted from step 1. This distinguishes
|
||||
legitimate onset generation from the late-forming dialogue artifact.
|
||||
9. Replacing only tagged dialogue with quoted speech suppresses the final first
|
||||
100ms by about `21.3 dB` peak and `15.9 dB` RMS. The quoted result remains
|
||||
below `-40.28 dBFS` peak and retains two later non-silent speech regions.
|
||||
|
||||
## Current Diagnosis
|
||||
|
||||
The startup artifact is encoded into the first four generated audio latent
|
||||
frames during late denoising. Dialogue markup is a strong trigger for the
|
||||
reproduced seed. The AudioVAE renders and contextually shapes the event but does
|
||||
not create it independently.
|
||||
|
||||
This does not prove that markup is the only trigger or that quoted dialogue is
|
||||
universally clean. It does establish markup as a repeatable trigger for this
|
||||
prompt family. Automatic prompt transformation still requires subjective speech,
|
||||
word-accuracy, and lip-sync review.
|
||||
|
||||
## Ten-Seed Prompt-Format Sweep
|
||||
|
||||
A matched sweep used seeds `440420` through `440429`, Sage2, base 12-step
|
||||
beta/RES sampling, and identical dialogue semantics. Each seed generated one
|
||||
tagged and one quoted audio latent and lossless waveform.
|
||||
|
||||
- Quoted speech reduced first-100ms peak level for all 10 seeds.
|
||||
- Median peak reduction was `21.55 dB`; mean was `20.11 dB`.
|
||||
- Median RMS reduction was `17.50 dB`; mean was `17.93 dB`.
|
||||
- Tagged speech exceeded `-40 dBFS` peak in the first 100ms for 9/10 seeds.
|
||||
- Quoted speech exceeded that threshold for 0/10 seeds.
|
||||
- Tagged speech became active within 100ms for 9/10 seeds; quoted speech did so
|
||||
for 0/10 seeds.
|
||||
- A greater-than-10dB boundary decay occurred for 8/10 tagged cases and 0/10
|
||||
quoted cases.
|
||||
- Simple first-frame and first-four-frame latent magnitude/delta features overlap
|
||||
between groups and cannot safely detect the artifact by themselves.
|
||||
|
||||
The complete report and paired WAV/latent files are in
|
||||
`h3-baselines/audio-dialogue-format-sweep`. A subjective listening pass on
|
||||
2026-08-21 judged all ten quoted WAVs good. Quoted dialogue is therefore the
|
||||
project default; full-video lip-sync validation remains pending.
|
||||
|
||||
## Diagnostic Assets
|
||||
|
||||
- `audio-diagnostic-affected-dialogue-864x480-141f-base12-sage2-seed440420.latent.pt`
|
||||
- `audio-diagnostic-affected-dialogue-864x480-141f-base12-sage2-seed440420.wav`
|
||||
- `audio-diagnostic-clean-nightclub-864x480-141f-base12-sage2-seed440421.latent.pt`
|
||||
- `audio-diagnostic-clean-nightclub-864x480-141f-base12-sage2-seed440421.wav`
|
||||
- `audio-diagnostic-dialogue-quoted-base12-sage2-seed440420.wav`
|
||||
- `audio-diagnostic-affected-dialogue-sage2-denoise-trace.pt`
|
||||
- `audio-diagnostic-affected-dialogue-sage2-denoise-trace.json`
|
||||
- `audio-diagnostic-clean-nightclub-sage2-denoise-trace.pt`
|
||||
- `audio-diagnostic-clean-nightclub-sage2-denoise-trace.json`
|
||||
- `audio-diagnostic-dialogue-quoted-sage2-denoise-trace.pt`
|
||||
- `audio-diagnostic-dialogue-quoted-sage2-denoise-trace.json`
|
||||
- `audio-vae-boundary-probes/report.json`
|
||||
|
||||
## Next Experiments
|
||||
|
||||
1. Generate selected full videos to compare lip-sync and prompt adherence with
|
||||
the new quoted-dialogue default.
|
||||
2. Add one ambience-only prompt and one immediate-impact sound prompt to prevent
|
||||
a detector from equating quiet starts with correctness.
|
||||
3. Compare tagged-versus-quoted Qwen conditioning and per-step first-four-frame
|
||||
denoiser outputs to localize the conditioning pathway.
|
||||
4. Prototype a selective late-step boundary re-denoise only after a reliable
|
||||
latent classifier exists.
|
||||
5. Reject unconditional trimming, fading, zeroing, or fixed-frame replacement.
|
||||
|
||||
## Measurement Correction
|
||||
|
||||
FFmpeg's `apsnr` results previously recorded for attention-backend audio were
|
||||
inconsistent with direct decoded-PCM array comparisons and must not be used.
|
||||
Future audio comparisons must decode each stream to aligned float PCM and compute
|
||||
error metrics directly.
|
||||
114
CURRENT_STATE.md
Normal file
114
CURRENT_STATE.md
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
# H3 Runtime Current State
|
||||
|
||||
Status date: 2026-08-22
|
||||
|
||||
This document is the canonical snapshot of implemented scope and remaining work.
|
||||
Historical handoffs in `PLAN.md` and `PARITY.md` may describe older states.
|
||||
|
||||
## Implemented And Validated
|
||||
|
||||
- Single-GPU prompt-only T2VA with joint video/audio generation.
|
||||
- First-frame I2VA, last-frame L2VA, and first/last FL2VA through the shared
|
||||
keyframe-conditioning path.
|
||||
- Qwen text and vision conditioning, token refinement, video VAE encoding, H3
|
||||
packed denoising, beta/RES sampling, video/audio decoding, and final MP4 mux.
|
||||
- Resident HTTP runtime with warmup, readiness reporting, request-level backend
|
||||
selection, timing stages, optional latent saving, and diagnostic intermediates.
|
||||
- SageAttention2 as the default quality backend.
|
||||
- SDPA, forced cuDNN SDPA, FlashAttention-4, Sage3, Comfy Kitchen INT8, KJ Sage,
|
||||
head-sliced, and Sol-Attn experimental backends.
|
||||
- Official FL2VA Turbo 4-step and 8-step adapters.
|
||||
- Optional resident H3-native latent upscaling.
|
||||
- Experimental EasyCache and H3-Cache delta-reuse modes.
|
||||
- Quoted dialogue as the project prompt default. In a matched 10-seed test,
|
||||
quoted dialogue eliminated immediate first-100ms activity in all ten cases and
|
||||
all ten quoted WAVs passed subjective review.
|
||||
- Ragged Ulysses sequence parallelism with 2/4/6/8-rank transport tests.
|
||||
- Sequence-sharded 50-block execution and distributed final projection.
|
||||
- True H3 NVFP4 tensor parallelism for attention QKV/output and MLP FC1/FC2.
|
||||
- Automatic visible-GPU launchers and 1/2/4/6/8 benchmark matrix tooling.
|
||||
- Real-checkpoint one-rank Ulysses-versus-TP identity at 864x480, 141 frames,
|
||||
and 12 steps, including exact video and audio latent equality.
|
||||
|
||||
## Primary Missing Scope
|
||||
|
||||
### Full Ref2VA
|
||||
|
||||
- Arbitrary reference image, video, and audio inputs.
|
||||
- Reference-audio encoder and reference soundtrack conditioning.
|
||||
- Reference identity/voice blocks in the standalone packer.
|
||||
- Ref2VA position, modality, and scheduling contracts.
|
||||
- Direct-versus-Comfy full Ref2VA per-step and final-output parity benchmark.
|
||||
|
||||
### Explicit Task API
|
||||
|
||||
- Named `task` selection for T2VA, I2VA, L2VA, FL2VA, and Ref2VA.
|
||||
- Mode-specific request schemas and incompatible-input validation.
|
||||
- Intermediate keyframe anchors beyond the current first/last restriction.
|
||||
|
||||
### Distributed Execution
|
||||
|
||||
- Real NCCL transport and output parity above one rank.
|
||||
- 2/4/6/8-GPU topology and performance sweeps on one Blackwell machine.
|
||||
- Distributed resident-service orchestration; the current launcher is batch
|
||||
generation through `torchrun`.
|
||||
- x86 SageAttention2 packaging; RunPod validation initially uses SDPA.
|
||||
|
||||
### Owned Performance Kernels
|
||||
|
||||
- H3-specific attention backend optimized for real GB10 tensor shapes.
|
||||
- Blackwell-native CUTLASS/CuTe or cuBLASLt NVFP4 GEMMs.
|
||||
- CUDA graph capture and shape buckets.
|
||||
- Fused Q/K RMSNorm, RoPE, and layout work on the Sage2 quality path.
|
||||
|
||||
## Quality Work Remaining
|
||||
|
||||
- Generate full quoted-dialogue videos and validate wording, voice consistency,
|
||||
speech timing, and lip-sync before closing the startup-audio work.
|
||||
- Complete strict per-step LightX2V parity for Turbo adapters.
|
||||
- Add real-adapter Turbo end-to-end fixtures.
|
||||
- Add the optional target-resolution refinement stage after latent upscaling.
|
||||
- Resolve or formally bound upscaler ringing, texture, chromatic-edge, and
|
||||
identity changes.
|
||||
- Run full-size cache threshold and quality sweeps before enabling caches for
|
||||
production output.
|
||||
- Keep Sage3, Sol-Attn, INT8, and other approximate backends quality-gated.
|
||||
- Fix the inactive fused Sol QKV-layout path, which currently references an
|
||||
undefined `qkv` value. The deployed native Sol layout does not use this path.
|
||||
|
||||
## Production Work Remaining
|
||||
|
||||
- Asynchronous jobs, queueing, progress, cancellation, and timeouts.
|
||||
- Strict request validation, including Boolean fields and mode combinations.
|
||||
- Input/output path sandboxing, request-size limits, authentication, and TLS.
|
||||
- Configurable FPS, video codec, audio codec, sample rate, and media policy.
|
||||
- Container healthcheck, restart policy, resource limits, durable structured
|
||||
request logs, and runtime metrics.
|
||||
- Batch generation and an intentional worker/concurrency model.
|
||||
|
||||
## Validation And Packaging Gaps
|
||||
|
||||
- GPU end-to-end fixtures for T2VA, I2VA, L2VA, and FL2VA.
|
||||
- Full Ref2VA, AudioVAE waveform, cache, HTTP API, real Turbo, real upscaler,
|
||||
attention-quality, CUDA-graph, and distributed tests.
|
||||
- Reproducible local fixtures for parity evidence currently stored on Spark/SMB.
|
||||
- Explicit package declarations/checks for NumPy, SciPy, Pillow, and FFmpeg.
|
||||
- A standalone base image if removing the Comfy-derived image becomes a product
|
||||
requirement; the current denoising path still intentionally uses Comfy Kitchen
|
||||
kernels.
|
||||
- Align Docker `H3_MODEL_PATH` and `RuntimeConfig`; the environment variable is
|
||||
currently not consumed by the runtime default.
|
||||
|
||||
## Recommended Execution Order
|
||||
|
||||
1. Validate full quoted-dialogue video lip-sync and close the audio prompt change.
|
||||
2. Correct the inactive fused Sol path.
|
||||
3. Add explicit task schemas and automated single-GPU mode tests.
|
||||
4. Implement full Ref2VA, including reference-audio encoding.
|
||||
5. Build the H3-specific attention backend and CUDA graph buckets.
|
||||
6. Harden the service API and operational deployment.
|
||||
7. Complete RunPod NCCL validation and distributed scaling benchmarks.
|
||||
|
||||
The current single-GPU T2VA/FL2VA runtime is mature. Distributed execution is
|
||||
implemented and CPU/one-GPU validated, with real multi-GPU NCCL results still
|
||||
blocked on an eight-GPU host. The other largest gap is standalone Ref2VA.
|
||||
134
DISTRIBUTED.md
Normal file
134
DISTRIBUTED.md
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
# Distributed H3 Execution
|
||||
|
||||
The runtime supports two single-node distributed denoiser modes:
|
||||
|
||||
- `ulysses`: token-sharded blocks with QKV sequence-to-head all-to-all and the
|
||||
inverse head-to-sequence all-to-all around attention.
|
||||
- `tensor`: token-sharded blocks plus true NVFP4 tensor parallelism across QKV,
|
||||
attention output, MLP FC1, and MLP FC2.
|
||||
|
||||
Both modes preserve the native packed H3 sequence and reject empty token or head
|
||||
partitions. They do not add semantic padding. Temporary collective padding used
|
||||
by ragged row gathering is removed before model operations.
|
||||
|
||||
## Partition Contract
|
||||
|
||||
H3 has 56 attention heads of width 128. Head ownership is balanced and can be
|
||||
ragged:
|
||||
|
||||
| GPUs | Heads per rank |
|
||||
| --- | --- |
|
||||
| 2 | 28, 28 |
|
||||
| 4 | 14, 14, 14, 14 |
|
||||
| 6 | 10, 10, 9, 9, 9, 9 |
|
||||
| 8 | 7, 7, 7, 7, 7, 7, 7, 7 |
|
||||
|
||||
Token ranges use the same quotient/remainder partitioning. Segment boundaries
|
||||
are clipped to each rank's token interval and rebased before AdaLN modulation
|
||||
and residual gating.
|
||||
|
||||
In tensor mode, QKV output rows are selected by local head ownership. Attention
|
||||
output and MLP FC2 input columns are sliced on NVFP4 alignment boundaries; each
|
||||
rank computes a partial output and a ragged reduce-scatter sums and assigns token
|
||||
rows. MLP FC1 selects corresponding local ranges from both the gate and value
|
||||
halves. Projection bias is omitted from rank partials and added exactly once
|
||||
after reduction.
|
||||
|
||||
## Launch
|
||||
|
||||
The launcher uses every visible GPU when the world size is omitted:
|
||||
|
||||
```bash
|
||||
tools/run_distributed_t2va.sh ulysses
|
||||
tools/run_distributed_t2va.sh tensor
|
||||
```
|
||||
|
||||
Pass an explicit count and attention backend when needed:
|
||||
|
||||
```bash
|
||||
tools/run_distributed_t2va.sh tensor 8 sdpa
|
||||
```
|
||||
|
||||
Relevant environment variables:
|
||||
|
||||
- `H3_WORLD_SIZE`: fallback world size when no positional count is supplied.
|
||||
- `H3_DISTRIBUTED_BENCHMARK`: benchmark JSON path.
|
||||
- `H3_DISTRIBUTED_OUTPUT`: report and optional latent output directory.
|
||||
- `H3_MODEL_PATH` and `H3_TEXT_ENCODER_PATH`: checkpoint paths.
|
||||
- `H3_SAVE_LATENTS=0`: write reports without retaining large latent files.
|
||||
|
||||
Run every feasible target count on the current machine with:
|
||||
|
||||
```bash
|
||||
H3_SAVE_LATENTS=0 tools/run_distributed_matrix.sh
|
||||
```
|
||||
|
||||
The default matrix runs Ulysses and tensor modes at 1, 2, 4, 6, and 8 GPUs and
|
||||
skips counts larger than the visible device count. Override the lists with
|
||||
`H3_GPU_COUNTS` and `H3_DISTRIBUTED_MODES`.
|
||||
|
||||
The lower-level transport benchmark does not load H3 weights:
|
||||
|
||||
```bash
|
||||
tools/run_ulysses_benchmark.sh
|
||||
tools/run_ulysses_benchmark.sh 6 sdpa
|
||||
```
|
||||
|
||||
## Validation
|
||||
|
||||
Automated Gloo tests cover:
|
||||
|
||||
- Ulysses transport identity at 2, 4, 6, and 8 ranks.
|
||||
- Two-rank distributed SDPA parity.
|
||||
- Ragged gather and reduce-scatter behavior.
|
||||
- Distributed final-projection parity.
|
||||
- TP attention and MLP parity at 2 and 6 ranks.
|
||||
- Packed NVFP4 column- and row-shard layout preservation.
|
||||
|
||||
On GB10, the real 50-block NVFP4 checkpoint completed both one-rank distributed
|
||||
paths at 864x480, 141 frames, 12 steps, seed 440420. Ulysses and tensor modes
|
||||
produced identical video and audio tensors with zero maximum absolute error.
|
||||
This validates integration and the world-size-one identity path, but it does not
|
||||
replace multi-GPU NCCL parity testing.
|
||||
|
||||
## RunPod
|
||||
|
||||
The target is one eight-GPU RTX PRO 6000 Blackwell machine. Query current stock
|
||||
using the guarded API v2 client:
|
||||
|
||||
```powershell
|
||||
$env:RUNPOD_API_KEY = "..."
|
||||
python .\tools\runpod_api.py catalog --count 8
|
||||
```
|
||||
|
||||
After choosing an available data center, create the pod explicitly:
|
||||
|
||||
```powershell
|
||||
python .\tools\runpod_api.py create --count 8 --datacenter US-XX-N --yes
|
||||
```
|
||||
|
||||
The client defaults to the server-edition RTX PRO 6000 Blackwell and RunPod's
|
||||
x86_64 CUDA 13.0, Torch 2.9.1 image. It creates persistent workspace storage and
|
||||
enables SSH. Creation and termination require `--yes` to avoid accidental spend
|
||||
or data loss. Use `get` to poll status and `terminate POD_ID --yes` when finished.
|
||||
|
||||
Transfer this checkout plus these two existing checkpoints to the pod workspace:
|
||||
|
||||
- `minimax_h3_fl2va_pruned_nvfp4.safetensors` (12,528,636,800 bytes)
|
||||
- `qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` (15,687,142,551 bytes)
|
||||
|
||||
Install the project dependencies in the RunPod image, set the checkpoint path
|
||||
variables, run `test_distributed.py`, then run the transport and generation
|
||||
matrices. `Dockerfile.runpod` provides the equivalent reproducible custom image
|
||||
once an x86_64 image builder and registry are available.
|
||||
|
||||
SageAttention2 is not currently packaged in the generic x86 image, so cloud
|
||||
correctness and scaling start with SDPA. Sage2 can be measured after an x86
|
||||
wheel is added without changing the distributed layout.
|
||||
|
||||
## Remaining GPU Gates
|
||||
|
||||
- Real NCCL identity at 2, 4, 6, and 8 GPUs.
|
||||
- Distributed-versus-single latent parity above one rank.
|
||||
- Full 1/2/4/6/8 timing, transport, and memory reports.
|
||||
- Quality comparison after selecting an x86 attention backend.
|
||||
24
Dockerfile.runpod
Normal file
24
Dockerfile.runpod
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
# x86_64 CUDA 13 image for RunPod Blackwell distributed validation.
|
||||
FROM pytorch/pytorch:2.9.1-cuda13.0-cudnn9-devel
|
||||
|
||||
ARG COMFY_KITCHEN_VERSION=0.2.31
|
||||
|
||||
RUN apt-get update \
|
||||
&& DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ffmpeg git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /opt/h3-blackwell-runtime
|
||||
COPY . .
|
||||
|
||||
RUN python -m pip install --no-cache-dir \
|
||||
"comfy-kitchen==${COMFY_KITCHEN_VERSION}" \
|
||||
"fastsafetensors>=0.1.10" \
|
||||
"numpy>=2.0" \
|
||||
"pillow>=11" \
|
||||
"safetensors>=0.5.0" \
|
||||
"scipy>=1.14" \
|
||||
"transformers>=4.51,<5"
|
||||
|
||||
ENV PYTHONPATH=/opt/h3-blackwell-runtime/src
|
||||
ENTRYPOINT []
|
||||
CMD ["bash"]
|
||||
|
|
@ -15,6 +15,12 @@ RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.31
|
|||
|
||||
RUN python -m pip install --no-cache-dir "fastsafetensors>=0.1.10"
|
||||
|
||||
# Official CuTeDSL FlashAttention-4 beta with CUDA 13 Blackwell support.
|
||||
RUN python -m pip install --no-cache-dir --pre \
|
||||
"flash-attn-4[cu13]==4.0.0b27" \
|
||||
"nvidia-cutlass-dsl[cu13]==4.6.2" \
|
||||
"quack-kernels==0.6.4"
|
||||
|
||||
RUN python -m pip uninstall -y pynvml \
|
||||
&& python -m pip install --no-cache-dir nvidia-ml-py
|
||||
|
||||
|
|
|
|||
71
FLASH4.md
Normal file
71
FLASH4.md
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
# FlashAttention-4
|
||||
|
||||
The hot runtime exposes the official FlashAttention-4 CuTeDSL Blackwell kernel
|
||||
as the strict `flash4` attention backend. It does not silently fall back to
|
||||
SageAttention or SDPA.
|
||||
|
||||
## Versions
|
||||
|
||||
- `flash-attn-4==4.0.0b27`
|
||||
- `nvidia-cutlass-dsl==4.6.2` with the CUDA 13 extra
|
||||
- `quack-kernels==0.6.4`
|
||||
- Upstream: <https://github.com/Dao-AILab/flash-attention>
|
||||
- License: BSD-3-Clause
|
||||
|
||||
The package is a beta release. The Spark image pins these versions rather than
|
||||
tracking the latest prerelease.
|
||||
|
||||
## API
|
||||
|
||||
Select Flash4 per request:
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": "A simple cinematic scene.",
|
||||
"output": "/output/flash4-example.mp4",
|
||||
"width": 960,
|
||||
"height": 544,
|
||||
"frames": 124,
|
||||
"steps": 12,
|
||||
"attention": "flash4"
|
||||
}
|
||||
```
|
||||
|
||||
H3's main attention blocks call Flash4 in native
|
||||
`[batch, sequence, heads, head_dim]` layout. The adapter requires CUDA, matching
|
||||
FP16/BF16 Q/K/V tensors, matching shapes, and H3's head dimension of 128. Any
|
||||
violation raises an error instead of changing kernels.
|
||||
|
||||
## GB10 Validation
|
||||
|
||||
Flash4 dispatches its SM120 forward implementation on the GB10, which PyTorch
|
||||
reports as compute capability 12.1.
|
||||
|
||||
| Validation | Result |
|
||||
| --- | --- |
|
||||
| BF16 `[1,257,8,128]` vs SDPA | max abs `0.001953125`, mean abs `5.04e-05` |
|
||||
| BF16 `[1,20480,56,128]` vs SDPA | max abs `0.000244141`, mean abs `1.96e-08` |
|
||||
| Tiny joint AV H3 smoke | passed, no fallback |
|
||||
| 28 runtime contracts | passed |
|
||||
|
||||
Matched base 12-step talking benchmark at `960x544x124`, seed `440410`:
|
||||
|
||||
| Backend | Sampling | Relative to Flash4 |
|
||||
| --- | ---: | ---: |
|
||||
| Sage2 | `114.67s` | `-19.4%` |
|
||||
| Flash4 | `142.23s` | baseline |
|
||||
| PyTorch SDPA | `144.40s` | `+1.5%` |
|
||||
|
||||
Flash4 is therefore close to SDPA and materially slower than Sage2 for this H3
|
||||
workload on GB10. Its value is as another exact-style attention and audio-quality
|
||||
comparison, not as the current speed default.
|
||||
|
||||
The matched Flash4 proof contains H.264 video and stereo 32 kHz AAC audio:
|
||||
|
||||
```text
|
||||
/home/daniel/StoryStudioAssets/H3-output/h3-native-upscale-t2v/talking-woman-base12-flash4-960x544.mp4
|
||||
```
|
||||
|
||||
The current CuTeDSL package emits a one-time `AuxData` JIT argument warning.
|
||||
Compilation and inference still complete successfully, but the warning should
|
||||
be rechecked when upgrading FlashAttention-4 or CUTLASS DSL.
|
||||
149
H3_LATENT_UPSCALER.md
Normal file
149
H3_LATENT_UPSCALER.md
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
# H3 Native Latent Upscaler
|
||||
|
||||
The direct runtime can apply LBH-123-AI's learned 3D upscaler to a saved,
|
||||
normalized MiniMax H3 video latent without a ComfyUI dependency. This is an
|
||||
experimental spatial upscale. It preserves latent time and does not interpolate
|
||||
frames.
|
||||
|
||||
Upstream implementation and checkpoint:
|
||||
|
||||
- <https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler>
|
||||
- <https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler>
|
||||
- Revision: `13ccf95d85d120bdbc92c05b1247a6e147bf54bf`
|
||||
- FP16 SHA-256: `043e5a48e161610ef6c3ea974645220354d06fa618abca15f76d084812eb55c2`
|
||||
- License: Apache-2.0
|
||||
|
||||
## Usage
|
||||
|
||||
The hot service supports the complete generation, latent upscale, video/audio
|
||||
decode, and mux process in one request:
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": "A simple cinematic scene.",
|
||||
"output": "/output/h3-native-upscale/example.mp4",
|
||||
"width": 960,
|
||||
"height": 544,
|
||||
"frames": 124,
|
||||
"turbo": "4step",
|
||||
"upscale": 2.0,
|
||||
"mux_audio": true,
|
||||
"keep_intermediates": false
|
||||
}
|
||||
```
|
||||
|
||||
The response reports `source_width`, `source_height`, final `width` and
|
||||
`height`, and a `latent_upscaled` timing stage. H3 audio bypasses the spatial
|
||||
upscaler and is decoded and muxed normally.
|
||||
|
||||
Only the final muxed MP4 is retained by default. Set `keep_intermediates` to
|
||||
`true` only when a separate WAV and video-only MP4 are needed for diagnostics;
|
||||
`save_latent` remains a separate opt-in field.
|
||||
|
||||
For offline diagnostics, save the normalized latent during generation by
|
||||
setting `save_latent`, then run:
|
||||
|
||||
```bash
|
||||
python tools/upscale_video_latent.py \
|
||||
--latent /output/source.pt \
|
||||
--model /upscaler/minimax_h3_latent_upscaler_3d_fp16.safetensors \
|
||||
--output /output/source-h3latent2x.pt \
|
||||
--scale 2 \
|
||||
--precision float16
|
||||
|
||||
python tools/decode_video_latent.py \
|
||||
--latent /output/source-h3latent2x.pt \
|
||||
--output /output/source-h3latent2x.mp4 \
|
||||
--vae-dtype float16 \
|
||||
--vae-tile-size 256
|
||||
```
|
||||
|
||||
The upscaled state retains other source metadata and any audio latent. Its
|
||||
`width` and `height` fields are updated from the resulting H3 latent dimensions.
|
||||
|
||||
## GB10 Validation
|
||||
|
||||
Validated on 2026-08-20 with the Posly hero image, Turbo 4-step v1.1, Sage2,
|
||||
seed `440408`, and 124 frames:
|
||||
|
||||
| Stage | Result |
|
||||
| --- | --- |
|
||||
| Source | `960x544`, latent `[1,24,37,34,60]` |
|
||||
| Learned 3D upscale | `10.94s`, latent `[1,24,37,68,120]` |
|
||||
| Checkpoint load | `1.36s` |
|
||||
| Tiled FP16 decode command | `137.67s`, `1920x1088` H.264 |
|
||||
| Source vs downscaled result | SSIM `0.924986` |
|
||||
|
||||
The learned stage adds visible edge, face, and food detail. It also produces a
|
||||
harsher illustrated texture, ringing/chromatic edges in some high-contrast
|
||||
areas, and small facial changes. It is not fidelity-safe for protected UI,
|
||||
text, QR codes, or identity-critical material without visual review.
|
||||
|
||||
The upstream example follows latent upscaling with a second target-resolution
|
||||
H3 sampling pass using sigmas `0.9035, 0.6316, 0.3158, 0.0`. That refinement
|
||||
was deliberately excluded from this first test so the upscaler itself could be
|
||||
evaluated. It adds generation cost and another opportunity to alter content.
|
||||
|
||||
## Proofs
|
||||
|
||||
Spark output directory:
|
||||
|
||||
```text
|
||||
/home/daniel/StoryStudioAssets/H3-output/posly/upscale-h3-native/
|
||||
```
|
||||
|
||||
Key files:
|
||||
|
||||
- `01-posly-hero-turbo4-source-960x544.mp4`
|
||||
- `01-posly-hero-turbo4-h3latent2x-1920x1088.mp4`
|
||||
- `source-frame60-right-crop-lanczos.png`
|
||||
- `h3latent2x-frame60-right-crop.png`
|
||||
|
||||
This result is promising for cinematic footage but does not replace the proven
|
||||
LTX 2.3 chain yet. Run a matched LTX 2.3 comparison and a carefully bounded H3
|
||||
refinement test before selecting a default production path.
|
||||
|
||||
## Integrated T2V Validation
|
||||
|
||||
The single-request path was validated with a matched 5.17-second talking-head
|
||||
T2V pair at seed `440410`. Both requests sampled at `960x544`, applied the
|
||||
resident 2x upscaler, decoded at `1920x1088`, and muxed native H3 audio into the
|
||||
only retained artifact.
|
||||
|
||||
| Variant | Sampling | Latent upscale | VAE decode | Wall time |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| Turbo 4-step v1.1 | `56.54s` | `7.52s` | `82.74s` | `163.02s` |
|
||||
| Turbo 8-step v1.0 | `113.41s` | `7.28s` | `83.69s` | `220.06s` |
|
||||
| Base 12-step beta/RES | `114.67s` | `7.26s` | `83.34s` | `221.66s` |
|
||||
|
||||
Both outputs contain H.264 video at 24 fps and stereo AAC audio at 32 kHz. The
|
||||
4-step audio measured `-31.4 dB` mean / `-10.9 dB` peak; the 8-step audio
|
||||
measured `-23.0 dB` mean / `-5.0 dB` peak; and the base 12-step audio measured
|
||||
`-22.4 dB` mean / `-5.2 dB` peak.
|
||||
|
||||
```text
|
||||
/home/daniel/StoryStudioAssets/H3-output/h3-native-upscale-t2v/talking-woman-h3native2x-1920x1088.mp4
|
||||
/home/daniel/StoryStudioAssets/H3-output/h3-native-upscale-t2v/talking-woman-turbo8-h3native2x-1920x1088.mp4
|
||||
/home/daniel/StoryStudioAssets/H3-output/h3-native-upscale-t2v/talking-woman-base12-h3native2x-1920x1088.mp4
|
||||
```
|
||||
|
||||
### Attention Audio Isolation
|
||||
|
||||
A matched base 12-step run used direct PyTorch SDPA instead of Sage2. Upscaling
|
||||
was disabled, but that does not affect the audio comparison because the learned
|
||||
upscaler runs after joint AV sampling and only receives the separated video
|
||||
latent.
|
||||
|
||||
| Attention | Sampling | Final size | Audio level | Wall time |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| Sage2 | `114.67s` | `1920x1088` after post-sample 2x | `-22.4 dB` mean / `-5.2 dB` peak | `221.66s` |
|
||||
| PyTorch SDPA | `144.40s` | `960x544`, no upscale | `-22.5 dB` mean / `-2.3 dB` peak | `171.21s` |
|
||||
|
||||
The SDPA comparison file is:
|
||||
|
||||
```text
|
||||
/home/daniel/StoryStudioAssets/H3-output/h3-native-upscale-t2v/talking-woman-base12-sdpa-960x544.mp4
|
||||
```
|
||||
|
||||
Signal levels alone cannot establish dialogue quality or lip sync; the matched
|
||||
clips require listening review to determine whether Sage2 is the cause.
|
||||
167
H3_PROMPT_GUIDE.md
Normal file
167
H3_PROMPT_GUIDE.md
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
# H3 Prompt Guide
|
||||
|
||||
Project guidance for MiniMax H3 audiovisual prompts. This document combines
|
||||
MiniMax's official prompt format with observations from our direct-runtime and
|
||||
ComfyUI tests.
|
||||
|
||||
## Core Principle
|
||||
|
||||
Treat an H3 prompt as an audiovisual timeline, not as a loose prose image
|
||||
description. For T2VA and keyframe modes, use these exact fields in this order:
|
||||
|
||||
```text
|
||||
integrated_multimodal_description: ...
|
||||
|
||||
overall_soundscape: ...
|
||||
|
||||
non_diegetic_music: ...
|
||||
```
|
||||
|
||||
- `integrated_multimodal_description` contains shots, visible actions,
|
||||
speakers, exact dialogue, singing, and synchronized diegetic events.
|
||||
- `overall_soundscape` contains ambience, Foley, and non-verbal human sounds
|
||||
such as breathing, laughter, and panting.
|
||||
- `non_diegetic_music` contains only music heard by the audience rather than
|
||||
the characters. Use `N/A` when no score is wanted.
|
||||
|
||||
T2VA begins directly with these fields. I2VA, FL2VA, and L2VA add the official
|
||||
picture-alignment instruction before them. Ref2VA uses a different six-section
|
||||
format and should follow the official reference guide.
|
||||
|
||||
## Dialogue Syntax
|
||||
|
||||
Give every speaking character a stable speaker ID such as `(S1)` or `(S2)`.
|
||||
Describe the speaker, voice, action, and delivery outside the dialogue block.
|
||||
Use quoted speech as the project default and state the language outside it:
|
||||
|
||||
```text
|
||||
The woman with a low, clear voice and measured pace (S1) says in English in a
|
||||
flat, matter-of-fact delivery: "The meeting starts at three."
|
||||
```
|
||||
|
||||
MiniMax's official guide recommends `<d>[Language]...</d>`. Do not use that
|
||||
markup by default in this project. In a matched 10-seed Sage2 sweep, tagged
|
||||
dialogue produced immediate first-100ms activity in 9/10 cases, while quoted
|
||||
dialogue produced it in 0/10. All ten quoted WAVs passed subjective listening.
|
||||
See [`AUDIO_BOUNDARY_INVESTIGATION.md`](AUDIO_BOUNDARY_INVESTIGATION.md).
|
||||
|
||||
Rules:
|
||||
|
||||
- Preserve the dialogue wording and punctuation exactly.
|
||||
- Specify the language explicitly in prose, such as `says in English`.
|
||||
- Describe pitch, timbre, pace, volume, accent, and emotional restraint only
|
||||
when useful.
|
||||
- Keep the same speaker ID across shots.
|
||||
- Use a compound ID such as `(S1,S2)` only when speakers vocalize together.
|
||||
- Reserve `<d>[Language]...</d>` for controlled compatibility experiments until
|
||||
the upstream startup-audio defect is resolved.
|
||||
- Describe visible text explicitly without reusing dialogue syntax, for example
|
||||
`a sign visibly reads MEETING ROOM`.
|
||||
- For voiceover, use `says in an off-screen voiceover` and state that the
|
||||
visible character's lips remain completely closed.
|
||||
|
||||
## Lip Sync And End State
|
||||
|
||||
Describe speech as a physical action and specify what happens when it ends:
|
||||
|
||||
```text
|
||||
Her mouth movements naturally synchronize with each spoken word. Immediately
|
||||
after the final word, her lips meet and her jaw ceases speaking motion. She
|
||||
remains silent through the final frame.
|
||||
```
|
||||
|
||||
This is especially important in short clips. Without an explicit end state, H3
|
||||
may continue mouth motion or add a non-verbal vocal reaction after the line.
|
||||
|
||||
For multiple speakers, state whose lips move and whose remain still during each
|
||||
line and pause.
|
||||
|
||||
## Avoiding Unwanted Laughter
|
||||
|
||||
Laughter is a non-verbal human sound, so control it in `overall_soundscape`, not
|
||||
inside the dialogue block:
|
||||
|
||||
```text
|
||||
overall_soundscape: Quiet, dry indoor room tone with a faint ventilation hum.
|
||||
No laughter, chuckling, giggling, sighing, gasping, audible breathing, filler
|
||||
sounds, audience reaction, or other voices.
|
||||
```
|
||||
|
||||
Also reinforce the visible end state in the shot description: closed lips,
|
||||
stopped jaw motion, neutral expression, and silence through the final frame.
|
||||
|
||||
These exclusions are project guidance, not a guaranteed negative-prompt
|
||||
mechanism. H3 jointly generates video and audio, so unwanted sounds remain
|
||||
probabilistic.
|
||||
|
||||
Avoid positive cues that can conflict with the exclusion. In our first dialogue
|
||||
test, `She smiles slightly` repeatedly became laughter across Sage2, PyTorch
|
||||
SDPA, and FlashAttention-4. Since the behavior survived attention changes, the
|
||||
prompt was the stronger common cause. Use observable neutral behavior instead:
|
||||
|
||||
```text
|
||||
She maintains a neutral, composed, closed-mouth expression.
|
||||
```
|
||||
|
||||
## Short Dialogue Tests
|
||||
|
||||
For a five-second diagnostic clip:
|
||||
|
||||
- Use one visible speaker and one short sentence.
|
||||
- Use one static shot with no cuts.
|
||||
- Choose emotionally neutral wording.
|
||||
- Request a flat or matter-of-fact delivery.
|
||||
- Specify the post-dialogue closed-mouth state.
|
||||
- Permit only a simple, positive acoustic bed such as dry room tone.
|
||||
- Set `non_diegetic_music: N/A`.
|
||||
|
||||
Do not use `overall_soundscape: N/A` for a speaking test. The official guide
|
||||
reserves it for complete silence throughout the video.
|
||||
|
||||
## Reusable Neutral T2VA Prompt
|
||||
|
||||
```text
|
||||
integrated_multimodal_description: [Shot 1] Live-action, cinematic, a static
|
||||
medium close-up frames exactly one adult woman seated alone in an otherwise
|
||||
empty, quiet cafe. Warm natural window light falls evenly across her face. She
|
||||
maintains a neutral, composed, closed-mouth expression and looks steadily
|
||||
toward the camera. The woman with a low, clear voice and slow, even speaking
|
||||
pace (S1) physically speaks once in a flat, matter-of-fact delivery:
|
||||
"The meeting starts at three." Her mouth movements naturally
|
||||
synchronize with each spoken word. Immediately after the final word, her lips
|
||||
meet and her jaw ceases speaking motion. She remains silent and maintains the
|
||||
same neutral expression through the final frame. The camera remains completely
|
||||
static with no cuts.
|
||||
|
||||
overall_soundscape: Quiet, dry indoor room tone with a faint ventilation hum.
|
||||
No laughter, chuckling, giggling, smiling vocalization, sighing, gasping,
|
||||
audible breathing, filler sounds, audience reaction, or other voices.
|
||||
|
||||
non_diegetic_music: N/A
|
||||
```
|
||||
|
||||
## Prompt Checklist
|
||||
|
||||
- Correct task mode selected.
|
||||
- Exact field names and order preserved.
|
||||
- Every shot has observable visual and audible events.
|
||||
- Every speaker has a stable ID.
|
||||
- Dialogue uses quoted speech.
|
||||
- Spoken language, delivery, and voice are stated outside the quotation.
|
||||
- `<d>[Language]...</d>` is avoided unless explicitly testing official syntax.
|
||||
- Lip motion and post-speech closure are explicit.
|
||||
- Ambience and non-verbal sounds are in `overall_soundscape`.
|
||||
- Music is isolated in `non_diegetic_music`.
|
||||
- Script length fits the requested duration.
|
||||
- Positive emotional cues do not conflict with audio exclusions.
|
||||
|
||||
## Official Sources
|
||||
|
||||
- Base prompt guide:
|
||||
<https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/docs/VIDEO_PROMPT_WRITING_GUIDE_base_en.md>
|
||||
- Reference prompt guide:
|
||||
<https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/docs/VIDEO_PROMPT_WRITING_GUIDE_ref_en.md>
|
||||
- Official prompt-writing skill:
|
||||
<https://github.com/MiniMax-AI/MiniMax-H3/tree/main/.agents/skills/h3-prompt-writing>
|
||||
- Official model repository:
|
||||
<https://huggingface.co/MiniMaxAI/MiniMax-H3>
|
||||
15
PARITY.md
15
PARITY.md
|
|
@ -180,9 +180,9 @@ gate and is exact.
|
|||
|
||||
| Component | Implemented | Known limitation |
|
||||
| --- | --- | --- |
|
||||
| Text-only Qwen | Yes | No vision encoder, MRoPE, image/video expansion, reference labels, or modality tags |
|
||||
| Qwen text and vision conditioning | Yes | Text-only prompt parity is established; first/last keyframe vision conditioning is implemented, while arbitrary reference video/audio remains missing |
|
||||
| Token refiner | Yes | Bit-exact from captured 5376-wide refiner input through both blocks and final RMSNorm; Qwen-to-refiner projection boundary is still not separately captured |
|
||||
| Prompt-only FL2VA packer | Yes | Bit-exact for the coherent captured text-only FL2VA DiT input; no keyframe/reference condition rows |
|
||||
| T2VA/keyframe packer | Yes | Bit-exact for the coherent captured prompt-only DiT input; first-frame, last-frame, and first/last keyframe condition rows are implemented |
|
||||
| H3 DiT backbone | Yes | Bit-exact through all 50 blocks from the coherent assembled FL2VA input; requires the standalone Comfy Kitchen fused Q/K RMSNorm + split-half RoPE operator |
|
||||
| H3 final layer | Yes | Bit-exact final AdaLN, target-row modulation, and video/audio patch rows; Comfy materializes the AdaLN and output-head biases through BF16 |
|
||||
| H3 DiT | Yes | Strict all-block numeric parity not achieved |
|
||||
|
|
@ -190,18 +190,17 @@ gate and is exact.
|
|||
| Video VAE decoder | Yes | Direct VAE temporal assembly matches upstream after overlap fix. FP16 is the default Comfy-equivalent runtime path; cat benchmark VAE decode is `25.085s`. Use FP32 only for exact direct diagnostics. |
|
||||
| Audio VAE/decode/mux | Yes | Direct decoder-only MiniMax H3 audio VAE returns stereo `32000 Hz` waveform and muxes with generated video. Native audio latent scaling is fixed. |
|
||||
| End-to-end prompt-only FL2VA preview | Yes | Apples-to-apples warm cat benchmark is at Comfy parity: Comfy warm `150.26s`; direct warm after text conditioning `149.304s`; direct warm including text conditioning `151.465s`. |
|
||||
| Full Ref2VA | No | References, vision conditioning, VAE encode, audio, and muxing are unimplemented |
|
||||
| Full Ref2VA | No | Image keyframes, vision conditioning, VAE encode, generated audio, and muxing are implemented; arbitrary reference video/audio and identity/voice conditioning remain missing |
|
||||
|
||||
## Remaining Gates, In Dependency Order
|
||||
|
||||
Only these are outstanding. Do not recapture or revisit rows marked complete
|
||||
unless the checkpoint, Comfy version, prompt, or backend changes.
|
||||
|
||||
1. **Feature/performance work.** Prompt-only FL2VA is now closed against the
|
||||
warm Comfy baseline. Optimize load/caching/sampling, then evaluate Sage3,
|
||||
CUDA graphs, and multi-GPU execution.
|
||||
2. **Full Ref2VA support.** Add references, vision conditioning, VAE encode,
|
||||
and reference-path validation gates.
|
||||
1. **Distributed validation.** Run real NCCL parity and performance sweeps at
|
||||
2/4/6/8 GPUs; CPU transport and one-GPU real-checkpoint identity are complete.
|
||||
2. **Full Ref2VA support.** Add arbitrary reference video/audio, identity and
|
||||
voice conditioning, and reference-path validation gates.
|
||||
|
||||
## Existing Tools And Their Intended Gate
|
||||
|
||||
|
|
|
|||
37
PLAN.md
37
PLAN.md
|
|
@ -1,5 +1,8 @@
|
|||
# H3 Blackwell Runtime Plan
|
||||
|
||||
Current implementation status is tracked in [`CURRENT_STATE.md`](CURRENT_STATE.md).
|
||||
The dated handoffs below are retained as historical investigation records.
|
||||
|
||||
## Goal
|
||||
|
||||
Build a direct MiniMax H3 Ref2VA runtime for Blackwell and Grace Blackwell that consumes the current Comfy safetensors checkpoints while removing ComfyUI and Raylight from the denoising critical path.
|
||||
|
|
@ -58,7 +61,9 @@ Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Fe
|
|||
3. Add exact memory/lifetime optimizations next: `kj_head_sliced` and `kj_chunked_ffn`. These must preserve the validated direct outputs before being kept.
|
||||
4. Evaluate prior H3-tested attention candidates as standalone adapters: `sol_attn` and `kj_sage`.
|
||||
- `kj_sage` is implemented as explicit SageAttention mode backends: `kj_sage_cuda`, `kj_sage_triton`, `kj_sage_fp8`, and `kj_sage_fp8pp`; all passed hot-runtime smoke tests.
|
||||
- `sol_attn` is still blocked on locating/adding the standalone Sol-Attn source or package. It is not installed in the Spark image and is not present in this repository.
|
||||
- `sol_attn` is implemented through the pinned ComfyUI Triton source vendored
|
||||
into the Spark image. The native QKV layout is deployed; the inactive fused
|
||||
layout remains a known repair item.
|
||||
5. Evaluate approximate denoiser caches only after exact baselines are recorded: `easycache` and `h3_cache`.
|
||||
- Initial direct cache modes are implemented as opt-in approximate sampler modes. They reuse cached denoised deltas and report skipped-step stats; full-size quality/threshold sweeps are still required before using them for production output.
|
||||
6. Keep every backend explicit per run, with separate quality and timing records for sampling, VAE, audio, and end-to-end output.
|
||||
|
|
@ -109,6 +114,36 @@ DGX Spark caveat:
|
|||
- Record attention, GEMM, communication, VAE, and end-to-end timings separately.
|
||||
- Treat SageAttention3 as an experimental quality-gated kernel for H3.
|
||||
|
||||
## 2026-08-15 Handoff
|
||||
|
||||
Current quality baseline:
|
||||
|
||||
- `sage2` is the correct visual path for the 960x544, 124-frame, 12-step cat prompt at seed `440407`.
|
||||
- Fresh Spark hot-runtime Sage2/Vortex A/B asset: `\\192.168.1.162\StoryStudioAssets\H3-output\h3-blackwell-runtime\cat-sage2-vortex-5s-440407-ab.mp4`.
|
||||
- `sol_attn` is faster but visually wrong for the same prompt/seed. Keep it experimental until localized and quality-gated.
|
||||
- Sol-native BSHD layout is exact versus the old Sol path, so the visual issue is likely Sol's approximation/configuration, not the layout removal.
|
||||
|
||||
Tomorrow priority:
|
||||
|
||||
1. Diagnose `sol_attn` quality failure.
|
||||
- Save per-step latents for `sage2` and `sol_attn` with identical prompt, seed, dimensions, and sampler settings.
|
||||
- Measure per-step max/mean/RMSE/cosine divergence.
|
||||
- Try hybrid schedules: early steps `sage2`, later steps `sol_attn`; also try Sage2 every Nth step.
|
||||
- Sweep Sol params only after the divergence point is known: `H3_SOL_TAU`, `H3_SOL_THRESH_TYPE`, `H3_SOL_INT8_QK`, `H3_SOL_INT8_PV`.
|
||||
2. Optimize the quality path (`sage2`).
|
||||
- Use `sage2` as visual/numerical reference.
|
||||
- Attack Q/K/V HND layout and Q/K RMS+RoPE first; a fused Q/K RMS+RoPE + Sage HND layout kernel is the most plausible exact win.
|
||||
- Keep Sage3 experimental; previous measurements were slower and visually/numerically different.
|
||||
3. Add durable hot-runtime request logging.
|
||||
- Log `generate_start`, `generate_complete`, and `generate_failed` to stdout with output path, seed, attention, dimensions, and timings.
|
||||
- This prevents losing timing data when a client disconnects before reading the synchronous response.
|
||||
4. Check LTX 2.5 downloads.
|
||||
- Detached Spark container: `ltx25-nvfp4-download`.
|
||||
- Expected files under `/home/daniel/aeon-spark-test/h3/comfy-models`:
|
||||
- `diffusion_models/ltx-2.5-22b-distilled-transformer-nvfp4.safetensors`
|
||||
- `text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors`
|
||||
- Official `Lightricks/LTX-2.5` does not list a matching NVFP4 text encoder; the int8 ConvRot Gemma4 encoder is the matching low-VRAM official text encoder.
|
||||
|
||||
## 2026-08-13 VAE Debug Handoff
|
||||
|
||||
Resolved on `vae-decode-optimization`: direct VAE temporal overlap constants now match upstream, audio decode/mux is implemented, and Comfy-equivalent FP16 video VAE is the default runtime path. The 960x544x124 cat benchmark now matches warm Comfy performance: Comfy `150.26s`, direct `149.304s` after text conditioning, direct VAE decode `25.085s`.
|
||||
|
|
|
|||
202
README.md
202
README.md
|
|
@ -1,79 +1,165 @@
|
|||
# H3 Blackwell Runtime
|
||||
|
||||
Direct MiniMax H3 Ref2VA runtime research project. ComfyUI is the checkpoint and correctness oracle, not the target runtime.
|
||||
Direct MiniMax H3 audiovisual inference for NVIDIA Blackwell and Grace Blackwell.
|
||||
The runtime consumes the current Comfy-format NVFP4 checkpoints while keeping
|
||||
ComfyUI out of the inference critical path. ComfyUI remains the checkpoint and
|
||||
correctness oracle.
|
||||
|
||||
## First Gate
|
||||
## Current Status
|
||||
|
||||
Inspect the mounted H3 NVFP4 safetensors headers before designing an importer:
|
||||
Implemented and validated:
|
||||
|
||||
```powershell
|
||||
python .\tools\inspect_safetensors.py /runpod-volume/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors
|
||||
- Prompt-only T2VA with jointly generated video and stereo audio.
|
||||
- First-frame I2VA, last-frame L2VA, and first/last-frame FL2VA conditioning.
|
||||
- Qwen3-VL text and image conditioning, token refinement, H3 packed denoising,
|
||||
beta/RES sampling, video and audio VAE decode, H.264/AAC encoding, and muxing.
|
||||
- Resident HTTP runtime with startup warmup, model reuse, request timing, and
|
||||
request-level attention selection.
|
||||
- SageAttention2 correctness default plus SDPA, forced cuDNN SDPA,
|
||||
FlashAttention-4, Sage3, KJ Sage, Comfy Kitchen INT8, head-sliced, and Sol-Attn
|
||||
experimental backends.
|
||||
- Official FL2VA Turbo 4-step and 8-step adapters.
|
||||
- Optional H3-native 3D latent upscaling.
|
||||
- Opt-in EasyCache and H3-Cache experiments.
|
||||
- Ragged Ulysses sequence parallelism and true NVFP4 tensor parallelism with
|
||||
automatic 1/2/4/6/8-GPU launch tooling.
|
||||
|
||||
The main remaining feature gap is full arbitrary Ref2VA, especially reference
|
||||
video/audio, identity, and voice conditioning. Multi-GPU code is CPU- and
|
||||
one-GPU-validated; real 2/4/6/8-GPU NCCL scaling measurements are still pending.
|
||||
See [`CURRENT_STATE.md`](CURRENT_STATE.md) for the canonical detailed status.
|
||||
|
||||
## Checkpoints
|
||||
|
||||
The default runtime paths are:
|
||||
|
||||
```text
|
||||
/models/minimax_h3_fl2va_pruned_nvfp4.safetensors
|
||||
/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors
|
||||
/vae/ae.safetensors
|
||||
/vae/mini_vae.safetensors
|
||||
```
|
||||
|
||||
Write the output to `artifacts/checkpoints/` on the mounted volume. The result must identify packed weights, scales, and tensor naming before any kernel conversion work begins.
|
||||
|
||||
## Benchmark Contract
|
||||
|
||||
`benchmarks/ref2va-960x544-124f.json` is the single-GPU performance contract. Record direct-runner results as JSON and compare them with:
|
||||
|
||||
```powershell
|
||||
python .\tools\compare_benchmark.py --result direct-result.json
|
||||
```
|
||||
The denoiser and Qwen checkpoints use Comfy Kitchen NVFP4 layouts. Do not
|
||||
convert or dequantize them during loading.
|
||||
|
||||
## DGX Spark
|
||||
|
||||
`Dockerfile.spark` and `compose.spark.yml` prepare an ARM64 GB10 development image using the existing AEON CUDA 13/SageAttention3 base. The compose target opens a shell only; it does not start inference.
|
||||
|
||||
### Forgejo Pulls From Spark
|
||||
|
||||
The Spark checkout uses Forgejo through the host's published local SSH port and a dedicated key:
|
||||
|
||||
```bash
|
||||
cd /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime
|
||||
git config core.sshCommand 'ssh -i ~/.ssh/id_ed25519_forgejo_h3 -o IdentitiesOnly=yes'
|
||||
git remote set-url origin ssh://git@127.0.0.1:2222/daniel/h3-blackwell-runtime.git
|
||||
git pull --ff-only origin master
|
||||
```
|
||||
|
||||
The private key remains on Spark at `~/.ssh/id_ed25519_forgejo_h3`; only its public key is registered in Forgejo.
|
||||
|
||||
## Runtime Output
|
||||
|
||||
Generation and latent-decode tools are quiet by default: they suppress ffmpeg banners and only print compact JSON summaries. Use these flags when debugging:
|
||||
|
||||
- `--progress`: print per-step sampler timing in `tools/direct_t2v_preview.py`.
|
||||
- `--profile-memory`: print memory checkpoints in `tools/direct_t2v_preview.py`.
|
||||
- `--ffmpeg-loglevel info`: show ffmpeg details instead of the default `error` level.
|
||||
- `--quiet`: suppress JSON summary lines.
|
||||
- `--vae-dtype float16`: use Comfy-style FP16 video VAE decode in `tools/direct_t2v_preview.py` or `tools/decode_video_latent.py`; this is the default runtime path. Use `--vae-dtype float32` only for exact direct-path diagnostics. `tools/direct_t2v_preview.py` also accepts `H3_VAE_DTYPE`.
|
||||
- `--vae-tile-size 256`: set the direct video VAE spatial tile size. `tools/direct_t2v_preview.py` also accepts `H3_VAE_TILE_SIZE`.
|
||||
|
||||
Standalone `tools/compare_*`, `tools/trace_*`, `tools/inspect_*`, and `tools/patch_comfy_*` scripts are debugging utilities and remain opt-in by being separate commands.
|
||||
|
||||
## Hot Runtime Service
|
||||
|
||||
`tools/serve_hot_runtime.py` keeps Qwen, H3, video VAE, and audio VAE resident in one process. Start the optional Spark service with:
|
||||
`Dockerfile.spark` and `compose.spark.yml` provide the ARM64 CUDA 13 runtime used
|
||||
on GB10. Build and start the resident service from the Spark checkout:
|
||||
|
||||
```bash
|
||||
docker compose -f compose.spark.yml build h3-hot-runtime
|
||||
docker compose -f compose.spark.yml up -d h3-hot-runtime
|
||||
curl http://127.0.0.1:8001/ready
|
||||
```
|
||||
|
||||
Use `GET /ready` to confirm resident model readiness. Use `POST /generate` with JSON fields like `prompt`, `output`, `width`, `height`, `frames`, `steps`, `seed`, and optional `attention`. Supported request-level attention values are reported by `/ready`; switching attention does not reload model weights. The hot image includes Sage2, forced cuDNN SDPA, and Comfy Kitchen INT8 attention. Sage2 is the default based on the 960x544x124 GB10 benchmark and the existing parity baseline.
|
||||
The service listens on container port 8000 and Spark host port 8001. It keeps
|
||||
Qwen, H3, both VAEs, Turbo adapters, and the optional latent upscaler resident.
|
||||
SageAttention2 is the default because it matches the established Comfy quality
|
||||
baseline.
|
||||
|
||||
The Spark hot service also keeps the official FL2VA Turbo adapters resident. Set `turbo` to `"4step"` for v1.1 768p (shift 6/3) or `"8step"` for v1.0 (shift 12/3). The matching step count is selected by default and enforced when `steps` is supplied. Set `turbo` to `null` or `"none"` for the base beta/RES path. Turbo uses its separate uniform training-Euler schedule and cannot be combined with denoiser caching.
|
||||
## HTTP API
|
||||
|
||||
See [`TURBO.md`](TURBO.md) for artifact hashes, implementation details, API examples, validation evidence, and matched GB10 performance results.
|
||||
`GET /health` and `GET /ready` report readiness, loaded options, warmup results,
|
||||
and attention backend status. `POST /generate` performs one serialized request:
|
||||
|
||||
Exact memory/lifetime options:
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:8001/generate \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"prompt": "A quiet medium shot of a woman by a rain-streaked window. She says, \"We should leave before dawn.\"",
|
||||
"output": "/output/h3-blackwell-runtime/example.mp4",
|
||||
"width": 864,
|
||||
"height": 480,
|
||||
"frames": 141,
|
||||
"steps": 12,
|
||||
"seed": 440420,
|
||||
"attention": "sage2"
|
||||
}'
|
||||
```
|
||||
|
||||
- `attention: "kj_head_sliced"` slices attention heads and runs the slice backend from `H3_HEAD_SLICE_BACKEND` (`sage2` by default) with `H3_HEAD_SLICE_SIZE` heads per slice (`8` by default).
|
||||
- `attention: "cudnn_sdpa"` forces cuDNN SDPA with no fallback to another PyTorch kernel.
|
||||
- `attention: "ck_int8"` uses Comfy Kitchen's approximate INT8 Q/K/V attention kernel.
|
||||
- `attention: "sol_attn"` routes eligible H3 attention calls through the pinned ComfyUI Sol-Attn Triton kernel vendored into the Spark image. Configure with `H3_SOL_TAU` (`1.3`), `H3_SOL_MIN_TOKENS` (`4096`), `H3_SOL_THRESH_TYPE` (`diag`), `H3_SOL_INT8_QK`, `H3_SOL_INT8_PV`, `H3_SOL_FALLBACK` (`sage2`), and `H3_SOL_STRICT`.
|
||||
- `--mlp-chunks N` on `tools/serve_hot_runtime.py` or `tools/direct_t2v_preview.py` chunks H3 SwiGLU rows exactly to reduce peak activation memory. Default is `1` (disabled).
|
||||
Optional request fields include `first_frame`, `last_frame`, `turbo`, `upscale`,
|
||||
`mux_audio`, `save_latent`, `keep_intermediates`, `cache_mode`, and cache tuning
|
||||
parameters. Keyframes accept an on-disk image path, base64 payload, or data URL.
|
||||
The presence of first and last frames selects I2VA, L2VA, or FL2VA behavior; a
|
||||
named task field is not yet exposed.
|
||||
|
||||
Approximate cache options are opt-in and must be quality-gated per prompt:
|
||||
Set `turbo` to `"4step"` or `"8step"`; the service enforces the corresponding
|
||||
step count and Turbo schedule. Set `upscale` to `2.0` to run the resident learned
|
||||
latent upscaler before video decode. Turbo and denoiser caching cannot be
|
||||
combined.
|
||||
|
||||
- `cache_mode: "easycache"` reuses cached denoised deltas while cumulative latent input change stays below `cache_threshold`.
|
||||
- `cache_mode: "h3_cache"` reuses cached denoised deltas when the current per-step latent input change is below `cache_threshold`.
|
||||
- Both modes accept `cache_start_percent`, `cache_end_percent`, and `cache_subsample_factor` in `POST /generate`; the CLI exposes equivalent `--cache-*` flags.
|
||||
Quoted dialogue is the project prompt default. Tagged `<d>[English]...` dialogue
|
||||
is a repeatable startup-audio trigger and should not be used as the default.
|
||||
Never apply unconditional audio trimming or fading because valid sound can begin
|
||||
at the first sample. See [`H3_PROMPT_GUIDE.md`](H3_PROMPT_GUIDE.md) and
|
||||
[`AUDIO_BOUNDARY_INVESTIGATION.md`](AUDIO_BOUNDARY_INVESTIGATION.md).
|
||||
|
||||
## Distributed Execution
|
||||
|
||||
Two batch-generation modes are available through `torchrun`:
|
||||
|
||||
- `ulysses`: token-sharded blocks with ragged sequence-to-head all-to-all around
|
||||
attention.
|
||||
- `tensor`: sequence-sharded residuals plus NVFP4-sharded QKV, attention output,
|
||||
MLP FC1, and MLP FC2 projections.
|
||||
|
||||
Use all visible GPUs or pass an explicit world size:
|
||||
|
||||
```bash
|
||||
tools/run_distributed_t2va.sh ulysses
|
||||
tools/run_distributed_t2va.sh tensor 8 sdpa
|
||||
H3_SAVE_LATENTS=0 tools/run_distributed_matrix.sh
|
||||
```
|
||||
|
||||
The matrix runs feasible 1/2/4/6/8-GPU configurations and skips counts larger
|
||||
than the visible device count. Six ranks use ragged head ownership
|
||||
`[10, 10, 9, 9, 9, 9]`; no semantic token padding is introduced.
|
||||
|
||||
See [`DISTRIBUTED.md`](DISTRIBUTED.md) for the collective contracts, validation
|
||||
evidence, environment variables, RunPod provisioning client, and x86 packaging
|
||||
status.
|
||||
|
||||
## Validation
|
||||
|
||||
Run the contract suite with:
|
||||
|
||||
```bash
|
||||
python -m unittest discover -s tests -v
|
||||
```
|
||||
|
||||
The current suite has 38 passing tests. Distributed tests cover 2/4/6/8-rank
|
||||
transport identity, ragged collectives, SDPA parity, final projection parity,
|
||||
NVFP4 shard layout, and 2/6-rank TP attention/MLP math. On GB10, real-checkpoint
|
||||
one-rank Ulysses and tensor paths produced exactly equal video and audio latents
|
||||
at 864x480, 141 frames, 12 steps, seed 440420.
|
||||
|
||||
Matched GB10 backend results and parity evidence are recorded in [`PLAN.md`](PLAN.md)
|
||||
and [`PARITY.md`](PARITY.md). Standalone `tools/compare_*`, `tools/trace_*`,
|
||||
`tools/inspect_*`, and `tools/patch_comfy_*` commands are diagnostic utilities,
|
||||
not runtime startup requirements.
|
||||
|
||||
## Documentation
|
||||
|
||||
- [`CURRENT_STATE.md`](CURRENT_STATE.md): canonical implemented and missing scope.
|
||||
- [`DISTRIBUTED.md`](DISTRIBUTED.md): Ulysses, TP, launchers, and RunPod workflow.
|
||||
- [`H3_PROMPT_GUIDE.md`](H3_PROMPT_GUIDE.md): audiovisual prompting contract.
|
||||
- [`AUDIO_BOUNDARY_INVESTIGATION.md`](AUDIO_BOUNDARY_INVESTIGATION.md): startup
|
||||
audio localization and prompt-format evidence.
|
||||
- [`TURBO.md`](TURBO.md): official Turbo adapters and schedules.
|
||||
- [`H3_LATENT_UPSCALER.md`](H3_LATENT_UPSCALER.md): learned latent upscaler.
|
||||
- [`FLASH4.md`](FLASH4.md): FlashAttention-4 integration and benchmark.
|
||||
- [`PARITY.md`](PARITY.md): direct-versus-Comfy evidence ledger.
|
||||
- [`PLAN.md`](PLAN.md): historical investigation and future kernel plan.
|
||||
|
||||
## Known Gaps
|
||||
|
||||
- Full arbitrary Ref2VA reference video/audio and identity/voice conditioning.
|
||||
- Explicit task schemas and stricter production request validation.
|
||||
- Real multi-GPU NCCL parity and 1/2/4/6/8 scaling results.
|
||||
- x86 SageAttention2 packaging for the generic RunPod image; SDPA is the initial
|
||||
cloud validation backend.
|
||||
- Production queueing, cancellation, authentication, TLS, metrics, and durable
|
||||
job state.
|
||||
- Full quality sweeps for approximate attention, cache, Turbo, and upscaler paths.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,50 @@
|
|||
{
|
||||
"name": "audio-dialogue-format-sweep-sage2-seeds440420-440429",
|
||||
"measured_at": "2026-08-21",
|
||||
"hardware": "NVIDIA GB10",
|
||||
"model": "minimax_h3_fl2va_pruned_nvfp4.safetensors",
|
||||
"attention": "sage2",
|
||||
"scheduler": "beta",
|
||||
"sampler": "res_multistep",
|
||||
"steps": 12,
|
||||
"resolution": [864, 480],
|
||||
"frames": 141,
|
||||
"seeds": [440420, 440421, 440422, 440423, 440424, 440425, 440426, 440427, 440428, 440429],
|
||||
"comparison": "Identical dialogue semantics with <d>[English]...</d> markup versus quoted speech",
|
||||
"cases": 20,
|
||||
"aggregate": {
|
||||
"quoted_first_100ms_peak_reduction_db": {
|
||||
"mean": 20.112007323193758,
|
||||
"median": 21.550717420448272,
|
||||
"minimum": 3.800458970207643,
|
||||
"maximum": 27.93839234881301
|
||||
},
|
||||
"quoted_first_100ms_rms_reduction_db": {
|
||||
"mean": 17.932329146978255,
|
||||
"median": 17.504237701162523,
|
||||
"minimum": 3.147146068931228,
|
||||
"maximum": 27.10315527948587
|
||||
},
|
||||
"first_100ms_peak_above_minus_40_dbfs": {
|
||||
"tagged": 9,
|
||||
"quoted": 0
|
||||
},
|
||||
"first_activity_under_100ms": {
|
||||
"tagged": 9,
|
||||
"quoted": 0
|
||||
},
|
||||
"boundary_decay_above_10db": {
|
||||
"tagged": 8,
|
||||
"quoted": 0
|
||||
}
|
||||
},
|
||||
"latent_feature_result": "Simple first-frame and first-block RMS/delta ranges overlap between tagged and quoted cases; no scalar latent threshold is justified.",
|
||||
"subjective_audio_review": {
|
||||
"reviewed_at": "2026-08-21",
|
||||
"quoted_cases_reviewed": 10,
|
||||
"result": "All ten quoted WAVs judged good"
|
||||
},
|
||||
"full_video_lip_sync_review_required": true,
|
||||
"output_directory": "/home/daniel/StoryStudioAssets/H3-output/h3-baselines/audio-dialogue-format-sweep",
|
||||
"full_report": "/home/daniel/StoryStudioAssets/H3-output/h3-baselines/audio-dialogue-format-sweep/report.json"
|
||||
}
|
||||
|
|
@ -0,0 +1,349 @@
|
|||
{
|
||||
"name": "t2va-dialogue-864x480-141f-base12-sdpa-seed440420",
|
||||
"measured_at": "2026-08-20",
|
||||
"measurement": "Warm resident direct-runtime request including conditioning, sampling, video/audio decode, encoding, and mux",
|
||||
"hardware": "NVIDIA GB10",
|
||||
"runtime_image": "sha256:d842225c9c5036647ee5cab1fae2e093357cc8041bd8617e33a82c81d0bcc874",
|
||||
"torch": "2.9.1+cu130",
|
||||
"model": "minimax_h3_fl2va_pruned_nvfp4.safetensors",
|
||||
"task": "t2va",
|
||||
"resolution": [864, 480],
|
||||
"frames": 141,
|
||||
"fps": 24,
|
||||
"duration_seconds": 5.875,
|
||||
"steps": 12,
|
||||
"scheduler": "beta",
|
||||
"sampler": "res_multistep",
|
||||
"seed": 440420,
|
||||
"attention": "sdpa",
|
||||
"turbo": null,
|
||||
"upscale": null,
|
||||
"mux_audio": true,
|
||||
"keep_intermediates": false,
|
||||
"vae_dtype": "float16",
|
||||
"vae_tile_size": 256,
|
||||
"mlp_chunks": 1,
|
||||
"prompt": "integrated_multimodal_description: [Shot 1] Live-action, cinematic, a static medium two-shot frames exactly two adults seated across from each other at a small table in a quiet, otherwise empty meeting room. A composed adult man sits on the left and a composed adult woman sits on the right. Both maintain neutral, attentive expressions. The man with a low, clear baritone voice and measured speaking pace (S1) looks toward the woman and says in a calm, matter-of-fact delivery: <d>[English] The north entrance closes at six.</d> During S1's line, only his lips and jaw move; the woman's lips remain completely closed. Immediately after his final word, his lips meet and his jaw ceases speaking motion. After a brief silent pause, the woman with a clear alto voice and measured speaking pace (S2) looks toward the man and replies in a calm, matter-of-fact delivery: <d>[English] Then we should leave by five thirty.</d> During S2's line, only her lips and jaw move; the man's lips remain completely closed. Immediately after her final word, her lips meet and her jaw ceases speaking motion. Both remain silent with neutral, closed-mouth expressions through the final frame. There is no overlapping speech. The camera remains completely static with no cuts.\n\noverall_soundscape: Quiet, dry indoor room tone with a faint ventilation hum. Only S1 and S2 are audible, one at a time. No laughter, chuckling, giggling, smiling vocalization, sighing, gasping, audible breathing, filler sounds, audience reaction, narration, or other voices.\n\nnon_diegetic_music: N/A",
|
||||
"timings_seconds": {
|
||||
"latents_initialized": 0.018280818010680377,
|
||||
"text_conditioned": 3.531974215002265,
|
||||
"sampled": 125.30079188900709,
|
||||
"vae_decoded": 23.84665890400356,
|
||||
"pixels_cpu": 0.06607734999852255,
|
||||
"raw_write": 2.75871228199685,
|
||||
"video_encode": 0.39305945999512915,
|
||||
"audio_decoded": 0.3400787889986532,
|
||||
"audio_raw_write": 0.0008342489891219884,
|
||||
"audio_encode": 0.034553833000245504,
|
||||
"mux": 0.19705491400964092,
|
||||
"request": 156.48807670301176,
|
||||
"wall": 156.56338787100685
|
||||
},
|
||||
"throughput": {
|
||||
"sampling_frames_per_second": 1.1252921699401506,
|
||||
"wall_frames_per_second": 0.9005936950992043,
|
||||
"sampling_realtime_factor": 21.327794364086314,
|
||||
"wall_realtime_factor": 26.649087297192655
|
||||
},
|
||||
"cache": {
|
||||
"mode": null,
|
||||
"threshold": 0.0,
|
||||
"skipped_steps": 0,
|
||||
"rates": []
|
||||
},
|
||||
"output": {
|
||||
"linux_path": "/home/daniel/StoryStudioAssets/H3-output/h3-baselines/dialogue-two-character-864x480-141f-base12-sdpa-seed440420.mp4",
|
||||
"unc_path": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-baselines\\dialogue-two-character-864x480-141f-base12-sdpa-seed440420.mp4",
|
||||
"sha256": "2eb61334978b42bcec00ec68449f3f39139d9e24ef5ba70ef2240dc8bf7116e0",
|
||||
"size_bytes": 429062,
|
||||
"container_bit_rate": 584254,
|
||||
"duration_seconds": 5.875,
|
||||
"video": {
|
||||
"codec": "h264",
|
||||
"profile": "High",
|
||||
"pixel_format": "yuv420p",
|
||||
"width": 864,
|
||||
"height": 480,
|
||||
"fps": 24,
|
||||
"frames": 141,
|
||||
"bit_rate": 439001
|
||||
},
|
||||
"audio": {
|
||||
"codec": "aac",
|
||||
"profile": "LC",
|
||||
"sample_rate": 32000,
|
||||
"channels": 2,
|
||||
"channel_layout": "stereo",
|
||||
"bit_rate": 136335,
|
||||
"mean_volume_db": -25.8,
|
||||
"max_volume_db": -4.3,
|
||||
"integrated_loudness_lufs": -21.18,
|
||||
"true_peak_dbtp": -4.34,
|
||||
"loudness_range_lu": 4.10,
|
||||
"loudness_threshold_lufs": -33.94
|
||||
}
|
||||
},
|
||||
"attention_backend_matrix": {
|
||||
"reference_backend": "sdpa",
|
||||
"comparison_note": "Video PSNR and decoded-PCM audio SNR measure output drift from SDPA, not subjective quality. Audio SNR is computed directly from aligned float PCM arrays; earlier FFmpeg apsnr results above 170 dB were invalid and have been replaced.",
|
||||
"runs": {
|
||||
"sdpa": {
|
||||
"sampled_seconds": 125.30079188900709,
|
||||
"request_seconds": 156.48807670301176,
|
||||
"wall_seconds": 156.56338787100685,
|
||||
"sampling_speedup_percent_vs_sdpa": 0.0,
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-sdpa-seed440420.mp4",
|
||||
"sha256": "2eb61334978b42bcec00ec68449f3f39139d9e24ef5ba70ef2240dc8bf7116e0",
|
||||
"size_bytes": 429062
|
||||
},
|
||||
"sage2": {
|
||||
"sampled_seconds": 101.26120180900034,
|
||||
"request_seconds": 132.59474182801205,
|
||||
"wall_seconds": 132.66185540499282,
|
||||
"sampling_speedup_percent_vs_sdpa": 19.185505,
|
||||
"video_psnr_average_db_vs_sdpa": 22.4225,
|
||||
"audio_pcm_snr_db_vs_sdpa": [2.515815, 2.492644],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-sage2-seed440420.mp4",
|
||||
"sha256": "c3e0691586c4101987f240c35a79637965af75088731108037cd084be945ba9d",
|
||||
"size_bytes": 433450
|
||||
},
|
||||
"cudnn_sdpa": {
|
||||
"sampled_seconds": 125.72722270998929,
|
||||
"request_seconds": 157.14237468996726,
|
||||
"wall_seconds": 157.2081264879962,
|
||||
"sampling_speedup_percent_vs_sdpa": -0.340326,
|
||||
"video_psnr_average_db_vs_sdpa": 22.093143,
|
||||
"audio_pcm_snr_db_vs_sdpa": [2.056909, 2.162249],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-cudnn_sdpa-seed440420.mp4",
|
||||
"sha256": "cd17833c12a99cd71d1c0a9137a1dd9feba03943980d6eb8fff296cd9f9529a6",
|
||||
"size_bytes": 429636
|
||||
},
|
||||
"ck_int8": {
|
||||
"sampled_seconds": 104.05008868798905,
|
||||
"request_seconds": 135.57199566195777,
|
||||
"wall_seconds": 135.62453711099806,
|
||||
"sampling_speedup_percent_vs_sdpa": 16.959752,
|
||||
"video_psnr_average_db_vs_sdpa": 22.989999,
|
||||
"audio_pcm_snr_db_vs_sdpa": [2.308146, 2.398058],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-ck_int8-seed440420.mp4",
|
||||
"sha256": "a54f9f3f2ba351d31b0b807646aea2740aae7c13c03d4eb7d44182aab3cf5618",
|
||||
"size_bytes": 423190
|
||||
},
|
||||
"sage3": {
|
||||
"sampled_seconds": 109.8837421490025,
|
||||
"request_seconds": 142.34848491402227,
|
||||
"wall_seconds": 142.41395058500348,
|
||||
"sampling_speedup_percent_vs_sdpa": 12.304032,
|
||||
"video_psnr_average_db_vs_sdpa": 19.63725,
|
||||
"audio_pcm_snr_db_vs_sdpa": [-3.047541, -2.395828],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-sage3-seed440420.mp4",
|
||||
"sha256": "3763ca4b79399037d997d5a37c795fdaf233aca88fdb90e122bf25450c75dddd",
|
||||
"size_bytes": 596651
|
||||
},
|
||||
"sage3_mean": {
|
||||
"sampled_seconds": 109.88214365398744,
|
||||
"request_seconds": 141.45796999098093,
|
||||
"wall_seconds": 141.51303354099218,
|
||||
"sampling_speedup_percent_vs_sdpa": 12.305308,
|
||||
"video_psnr_average_db_vs_sdpa": 19.63725,
|
||||
"audio_pcm_snr_db_vs_sdpa": [-3.047541, -2.395828],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-sage3_mean-seed440420.mp4",
|
||||
"sha256": "3763ca4b79399037d997d5a37c795fdaf233aca88fdb90e122bf25450c75dddd",
|
||||
"size_bytes": 596651,
|
||||
"identical_output_to": "sage3"
|
||||
},
|
||||
"kj_sage_cuda": {
|
||||
"sampled_seconds": 113.68463072601298,
|
||||
"request_seconds": 145.39383504101716,
|
||||
"wall_seconds": 145.45757150900317,
|
||||
"sampling_speedup_percent_vs_sdpa": 9.270621,
|
||||
"video_psnr_average_db_vs_sdpa": 22.83642,
|
||||
"audio_pcm_snr_db_vs_sdpa": [2.039297, 2.098551],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-kj_sage_cuda-seed440420.mp4",
|
||||
"sha256": "d85aaca677d2c678604b0aa2aba9bbf4921934bfc2970c20d0498b1762373ab5",
|
||||
"size_bytes": 427568
|
||||
},
|
||||
"kj_sage_triton": {
|
||||
"sampled_seconds": 118.28390433100867,
|
||||
"request_seconds": 150.37281119299587,
|
||||
"wall_seconds": 150.43365036998875,
|
||||
"sampling_speedup_percent_vs_sdpa": 5.600034,
|
||||
"video_psnr_average_db_vs_sdpa": 22.086796,
|
||||
"audio_pcm_snr_db_vs_sdpa": [5.920034, 6.184134],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-kj_sage_triton-seed440420.mp4",
|
||||
"sha256": "9fe63aafa51aae8ecc0940ba75fbc13a978a50091fd8f0c336c2bb04e3658a54",
|
||||
"size_bytes": 430038
|
||||
},
|
||||
"kj_sage_fp8": {
|
||||
"sampled_seconds": 102.6086639950081,
|
||||
"request_seconds": 134.1615118059999,
|
||||
"wall_seconds": 134.2236167689989,
|
||||
"sampling_speedup_percent_vs_sdpa": 18.110123,
|
||||
"video_psnr_average_db_vs_sdpa": 22.814833,
|
||||
"audio_pcm_snr_db_vs_sdpa": [3.442309, 3.542901],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-kj_sage_fp8-seed440420.mp4",
|
||||
"sha256": "8c5bc2a8bf7e496c2108194479fe60e67883dea0ab50e95f3ab4bac916274816",
|
||||
"size_bytes": 420189
|
||||
},
|
||||
"kj_sage_fp8pp": {
|
||||
"sampled_seconds": 103.24360375599645,
|
||||
"request_seconds": 134.77945204998832,
|
||||
"wall_seconds": 134.84802630099875,
|
||||
"sampling_speedup_percent_vs_sdpa": 17.603391,
|
||||
"video_psnr_average_db_vs_sdpa": 22.296461,
|
||||
"audio_pcm_snr_db_vs_sdpa": [3.893655, 4.120454],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-kj_sage_fp8pp-seed440420.mp4",
|
||||
"sha256": "71643c7163e5dc16d64ac0c5d567bdd846f7944471eb0a97fd54f9a1990c24e9",
|
||||
"size_bytes": 433672
|
||||
},
|
||||
"kj_head_sliced": {
|
||||
"sampled_seconds": 103.97128305501246,
|
||||
"request_seconds": 135.67128342803335,
|
||||
"wall_seconds": 135.71258915099315,
|
||||
"sampling_speedup_percent_vs_sdpa": 17.022645,
|
||||
"video_psnr_average_db_vs_sdpa": 21.679726,
|
||||
"audio_pcm_snr_db_vs_sdpa": [1.20378, 1.17592],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-kj_head_sliced-seed440420.mp4",
|
||||
"sha256": "48ee4da6c0a91f18714973ca23e26f1844b302024c8111946772cbc9a4a4dd83",
|
||||
"size_bytes": 421728
|
||||
},
|
||||
"sol_attn": {
|
||||
"sampled_seconds": 96.45773334198748,
|
||||
"request_seconds": 128.10418028896675,
|
||||
"wall_seconds": 128.17140150099294,
|
||||
"sampling_speedup_percent_vs_sdpa": 23.019055,
|
||||
"video_psnr_average_db_vs_sdpa": 16.653003,
|
||||
"audio_pcm_snr_db_vs_sdpa": [1.186243, 1.042278],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-sol_attn-seed440420.mp4",
|
||||
"sha256": "8b999630210692301549d11d20d9319a6af5873a60d4cd735035e2ed79c469d6",
|
||||
"size_bytes": 570667,
|
||||
"note": "Experimental sparse Triton dispatch; eligible calls use sol_attn and shorter calls fall back."
|
||||
},
|
||||
"flash4": {
|
||||
"sampled_seconds": 128.5334454880067,
|
||||
"request_seconds": 160.91555801399227,
|
||||
"wall_seconds": 160.96679471699463,
|
||||
"sampling_speedup_percent_vs_sdpa": -2.579915,
|
||||
"video_psnr_average_db_vs_sdpa": 21.238297,
|
||||
"audio_pcm_snr_db_vs_sdpa": [3.230108, 3.821894],
|
||||
"output_file": "dialogue-two-character-864x480-141f-base12-flash4-seed440420.mp4",
|
||||
"sha256": "cfb8a046cb79f8448988e7392ca748c0a08a2a034a43f6ab55a9058faa1a72ef",
|
||||
"size_bytes": 449687
|
||||
}
|
||||
}
|
||||
},
|
||||
"turbo_sage2_matrix": {
|
||||
"reference": "Base 12-step Sage2 from attention_backend_matrix",
|
||||
"runs": {
|
||||
"4step": {
|
||||
"adapter": "Official FL2VA Turbo 4-step v1.1 768p",
|
||||
"steps": 4,
|
||||
"attention": "sage2",
|
||||
"sampled_seconds": 51.661242722009774,
|
||||
"request_seconds": 83.01122188399313,
|
||||
"wall_seconds": 83.06598260099418,
|
||||
"sampling_speedup_percent_vs_base_sage2": 48.982195,
|
||||
"wall_speedup_percent_vs_base_sage2": 37.38518,
|
||||
"output_file": "dialogue-two-character-864x480-141f-turbo4step-sage2-seed440420.mp4",
|
||||
"sha256": "b02c29a8732d88c7dcf8e85462afcbe3aa83328541ce1647995e7d0d4188744a",
|
||||
"size_bytes": 455325,
|
||||
"container_bit_rate": 620017,
|
||||
"audio": {
|
||||
"integrated_loudness_lufs": -31.78,
|
||||
"true_peak_dbtp": -14.55,
|
||||
"first_0_25_seconds_peak_dbfs": -39.394894,
|
||||
"first_0_25_seconds_rms_dbfs": -57.733654,
|
||||
"first_sample_dbfs": [-54.110104, -50.553467]
|
||||
}
|
||||
},
|
||||
"8step": {
|
||||
"adapter": "Official FL2VA Turbo 8-step v1.0",
|
||||
"steps": 8,
|
||||
"attention": "sage2",
|
||||
"sampled_seconds": 102.70967868898879,
|
||||
"request_seconds": 134.31820170898573,
|
||||
"wall_seconds": 134.36813215899747,
|
||||
"sampling_speedup_percent_vs_base_sage2": -1.430436,
|
||||
"wall_speedup_percent_vs_base_sage2": -1.286185,
|
||||
"output_file": "dialogue-two-character-864x480-141f-turbo8step-sage2-seed440420.mp4",
|
||||
"sha256": "21f450848c2ded8c02054cc241c8080d55a17979022eadfd892f3ce5205833e6",
|
||||
"size_bytes": 585250,
|
||||
"container_bit_rate": 796936,
|
||||
"audio": {
|
||||
"integrated_loudness_lufs": -24.39,
|
||||
"true_peak_dbtp": -7.18,
|
||||
"first_0_25_seconds_peak_dbfs": -21.305752,
|
||||
"first_0_25_seconds_rms_dbfs": -41.479687,
|
||||
"first_sample_dbfs": [-35.859436, -39.793118]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"comfy_sdpa_comparison": {
|
||||
"comfyui_version": "0.31.0",
|
||||
"attention": "PyTorch attention without --use-sage-attention",
|
||||
"prompt_id": "29c45a77-efdf-4cc8-8ae9-9b8a40823b41",
|
||||
"execution_seconds": 152.767,
|
||||
"wall_seconds": 153.03360149999207,
|
||||
"wall_difference_from_direct_percent": -2.25454138353405,
|
||||
"output": {
|
||||
"linux_path": "/home/daniel/StoryStudioAssets/H3-output/h3-baselines/comfy-sdpa-dialogue-two-character-864x480-141f-base12-seed440420_00001_.mp4",
|
||||
"unc_path": "\\\\192.168.1.162\\StoryStudioAssets\\H3-output\\h3-baselines\\comfy-sdpa-dialogue-two-character-864x480-141f-base12-seed440420_00001_.mp4",
|
||||
"sha256": "f4510d6b8ac06b8f2cb22172454c6977a214e061dcc5539ace68e1cb70abc2ee",
|
||||
"size_bytes": 443723,
|
||||
"container_bit_rate": 604218,
|
||||
"duration_seconds": 5.875,
|
||||
"video": {
|
||||
"codec": "h264",
|
||||
"profile": "High",
|
||||
"pixel_format": "yuv420p",
|
||||
"width": 864,
|
||||
"height": 480,
|
||||
"fps": 24,
|
||||
"frames": 141,
|
||||
"bit_rate": 457452
|
||||
},
|
||||
"audio": {
|
||||
"codec": "aac",
|
||||
"profile": "LC",
|
||||
"sample_rate": 32000,
|
||||
"channels": 2,
|
||||
"channel_layout": "stereo",
|
||||
"bit_rate": 133407,
|
||||
"mean_volume_db": -25.7,
|
||||
"max_volume_db": -4.3,
|
||||
"integrated_loudness_lufs": -21.12,
|
||||
"true_peak_dbtp": -4.25,
|
||||
"loudness_range_lu": 4.10,
|
||||
"loudness_threshold_lufs": -33.88
|
||||
}
|
||||
},
|
||||
"first_0_25_seconds": {
|
||||
"direct": {
|
||||
"peak_level_db": -19.429245,
|
||||
"rms_level_db": -33.601728,
|
||||
"max_sample_difference": 0.022571,
|
||||
"peak_count": 2,
|
||||
"nan_count": 0,
|
||||
"inf_count": 0
|
||||
},
|
||||
"comfy": {
|
||||
"peak_level_db": -19.261140,
|
||||
"rms_level_db": -33.545443,
|
||||
"max_sample_difference": 0.023026,
|
||||
"peak_count": 2,
|
||||
"nan_count": 0,
|
||||
"inf_count": 0
|
||||
},
|
||||
"decoded_pcm_snr_db": {
|
||||
"channel_1": 31.498066,
|
||||
"channel_2": 29.897251
|
||||
},
|
||||
"conclusion": "No amplitude clipping. Direct and Comfy closely match during the first 250ms and both exhibit the startup artifact, but they are not numerically identical. The shared defect originates before runtime-specific audio encoding and mux."
|
||||
},
|
||||
"full_decoded_pcm_snr_db": {
|
||||
"channel_1": 21.414368,
|
||||
"channel_2": 26.592979
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
{
|
||||
"name": "t2va-dialogue-quoted-864x480-141f-base12-sage2-seed440420",
|
||||
"measured_at": "2026-08-21",
|
||||
"hardware": "NVIDIA GB10",
|
||||
"model": "minimax_h3_fl2va_pruned_nvfp4.safetensors",
|
||||
"task": "t2va",
|
||||
"resolution": [864, 480],
|
||||
"frames": 141,
|
||||
"fps": 24,
|
||||
"duration_seconds": 5.875,
|
||||
"steps": 12,
|
||||
"scheduler": "beta",
|
||||
"sampler": "res_multistep",
|
||||
"seed": 440420,
|
||||
"attention": "sage2",
|
||||
"turbo": null,
|
||||
"prompt_change": "Only the two <d>[English]...</d> spans were replaced with quoted dialogue.",
|
||||
"prompt": "integrated_multimodal_description: [Shot 1] Live-action, cinematic, a static medium two-shot frames exactly two adults seated across from each other at a small table in a quiet, otherwise empty meeting room. A composed adult man sits on the left and a composed adult woman sits on the right. Both maintain neutral, attentive expressions. The man with a low, clear baritone voice and measured speaking pace (S1) looks toward the woman and says in a calm, matter-of-fact delivery: \"The north entrance closes at six.\" During S1's line, only his lips and jaw move; the woman's lips remain completely closed. Immediately after his final word, his lips meet and his jaw ceases speaking motion. After a brief silent pause, the woman with a clear alto voice and measured speaking pace (S2) looks toward the man and replies in a calm, matter-of-fact delivery: \"Then we should leave by five thirty.\" During S2's line, only her lips and jaw move; the man's lips remain completely closed. Immediately after her final word, her lips meet and her jaw ceases speaking motion. Both remain silent with neutral, closed-mouth expressions through the final frame. There is no overlapping speech. The camera remains completely static with no cuts.\n\noverall_soundscape: Quiet, dry indoor room tone with a faint ventilation hum. Only S1 and S2 are audible, one at a time. No laughter, chuckling, giggling, smiling vocalization, sighing, gasping, audible breathing, filler sounds, audience reaction, narration, or other voices.\n\nnon_diegetic_music: N/A"
|
||||
}
|
||||
12
benchmarks/t2va-distributed-smoke-256x256-9f-1step.json
Normal file
12
benchmarks/t2va-distributed-smoke-256x256-9f-1step.json
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
{
|
||||
"name": "t2va-distributed-smoke-256x256-9f-1step",
|
||||
"mode": "t2va",
|
||||
"prompt": "A paper windmill turns steadily on a plain table. Quiet room tone.",
|
||||
"resolution": [
|
||||
256,
|
||||
256
|
||||
],
|
||||
"frames": 9,
|
||||
"steps": 1,
|
||||
"seed": 440420
|
||||
}
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
{
|
||||
"name": "t2va-nightclub-music-onset-864x480-141f-base12-sage2-seed440421",
|
||||
"measured_at": "2026-08-21",
|
||||
"hardware": "NVIDIA GB10",
|
||||
"model": "minimax_h3_fl2va_pruned_nvfp4.safetensors",
|
||||
"task": "t2va",
|
||||
"resolution": [864, 480],
|
||||
"frames": 141,
|
||||
"fps": 24,
|
||||
"duration_seconds": 5.875,
|
||||
"steps": 12,
|
||||
"scheduler": "beta",
|
||||
"sampler": "res_multistep",
|
||||
"seed": 440421,
|
||||
"attention": "sage2",
|
||||
"turbo": null,
|
||||
"prompt": "integrated_multimodal_description: 0.0-2.0s: Inside a packed underground nightclub, the camera glides low across a crowded dance floor toward a raised DJ booth. A clean four-on-the-floor kick and deep bassline begin immediately at 0.0s. Cyan and magenta strobes strike precisely on the beat while dancers move in synchronized rhythm. 2.0-4.0s: The camera sweeps around the DJ as she works the mixer, one hand adjusting a filter while the other raises toward the crowd. The house groove remains continuous and coherent; crisp hi-hats enter over the kick and bass. The crowd cheers naturally beneath the music. 4.0-5.875s: The camera pushes close to the mixer and then tilts up as the room erupts under a bright white strobe hit. The beat continues without interruption, ending on an energetic club moment. No dialogue. overall_soundscape: Loud but clean diegetic nightclub house music coming from the venue sound system, beginning exactly at the first frame, with a steady kick, deep controlled bass, crisp hi-hats, room reflections, dancing footsteps, and a lively crowd. No clipping, crackling, popping, startup noise, gibberish, or speech. non_diegetic_music: N/A",
|
||||
"diagnostic_capture": {
|
||||
"output_file": "audio-diagnostic-clean-nightclub-864x480-141f-base12-sage2-seed440421.mp4",
|
||||
"latent_file": "audio-diagnostic-clean-nightclub-864x480-141f-base12-sage2-seed440421.latent.pt",
|
||||
"wav_file": "audio-diagnostic-clean-nightclub-864x480-141f-base12-sage2-seed440421.wav",
|
||||
"sampled_seconds": 104.5049012459931,
|
||||
"wall_seconds": 135.91074104901054
|
||||
}
|
||||
}
|
||||
46
compose.qwen38-vllm.yml
Normal file
46
compose.qwen38-vllm.yml
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
services:
|
||||
qwen38-vllm:
|
||||
image: ghcr.io/aeon-7/aeon-vllm-ultimate:latest
|
||||
container_name: qwen38-vllm
|
||||
restart: unless-stopped
|
||||
gpus: all
|
||||
ipc: host
|
||||
shm_size: 16g
|
||||
network_mode: host
|
||||
volumes:
|
||||
- /home/daniel/aeon-spark-test/h3/models/Qwen3.8-27B-NVFP4:/model:ro
|
||||
- /home/daniel/aeon-spark-test/h3/vllm-cache:/root/.cache/vllm
|
||||
entrypoint: ["vllm"]
|
||||
command:
|
||||
- serve
|
||||
- /model
|
||||
- --served-model-name
|
||||
- qwen38
|
||||
- --host
|
||||
- 0.0.0.0
|
||||
- --port
|
||||
- "8002"
|
||||
- --tensor-parallel-size
|
||||
- "1"
|
||||
- --gpu-memory-utilization
|
||||
- "0.45"
|
||||
- --max-model-len
|
||||
- "262144"
|
||||
- --max-num-seqs
|
||||
- "4"
|
||||
- --max-num-batched-tokens
|
||||
- "8192"
|
||||
- --enable-chunked-prefill
|
||||
- --enable-prefix-caching
|
||||
- --kv-cache-dtype
|
||||
- fp8_e4m3
|
||||
- --reasoning-parser
|
||||
- qwen3
|
||||
- --tool-call-parser
|
||||
- qwen3_xml
|
||||
- --enable-auto-tool-choice
|
||||
- --distributed-executor-backend
|
||||
- mp
|
||||
- --speculative-config
|
||||
- '{"method":"mtp","num_speculative_tokens":5}'
|
||||
- --trust-remote-code
|
||||
84
compose.spark-stack.yml
Normal file
84
compose.spark-stack.yml
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
services:
|
||||
h3-hot-runtime:
|
||||
profiles: ["h3"]
|
||||
image: h3-blackwell-runtime:dev
|
||||
container_name: h3-hot-runtime
|
||||
restart: unless-stopped
|
||||
gpus: all
|
||||
volumes:
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/diffusion_models:/models:ro
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/text_encoders:/text-encoders:ro
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
|
||||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
ports:
|
||||
- "8001:8000"
|
||||
environment:
|
||||
H3_DISABLE_MMAP: "1"
|
||||
H3_NVFP4_SCALE_BACKEND: "vortex"
|
||||
H3_NVFP4_SCALE_VERSION: "1"
|
||||
H3_SOL_QKV_LAYOUT: "native"
|
||||
command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--warmup"]
|
||||
|
||||
qwen38-vllm:
|
||||
profiles: ["qwen"]
|
||||
image: ghcr.io/aeon-7/aeon-vllm-ultimate:latest
|
||||
container_name: qwen38-vllm
|
||||
restart: unless-stopped
|
||||
gpus: all
|
||||
ipc: host
|
||||
shm_size: 16g
|
||||
network_mode: host
|
||||
volumes:
|
||||
- /home/daniel/aeon-spark-test/h3/models/Qwen3.8-27B-NVFP4:/model:ro
|
||||
- /home/daniel/aeon-spark-test/h3/vllm-cache:/root/.cache/vllm
|
||||
entrypoint: ["vllm"]
|
||||
command:
|
||||
- serve
|
||||
- /model
|
||||
- --served-model-name
|
||||
- qwen38
|
||||
- --host
|
||||
- 0.0.0.0
|
||||
- --port
|
||||
- "8002"
|
||||
- --tensor-parallel-size
|
||||
- "1"
|
||||
- --gpu-memory-utilization
|
||||
- "0.45"
|
||||
- --max-model-len
|
||||
- "262144"
|
||||
- --max-num-seqs
|
||||
- "4"
|
||||
- --max-num-batched-tokens
|
||||
- "8192"
|
||||
- --enable-chunked-prefill
|
||||
- --enable-prefix-caching
|
||||
- --kv-cache-dtype
|
||||
- fp8_e4m3
|
||||
- --reasoning-parser
|
||||
- qwen3
|
||||
- --tool-call-parser
|
||||
- qwen3_xml
|
||||
- --enable-auto-tool-choice
|
||||
- --distributed-executor-backend
|
||||
- mp
|
||||
- --speculative-config
|
||||
- '{"method":"mtp","num_speculative_tokens":5}'
|
||||
- --trust-remote-code
|
||||
|
||||
flux2-klein-comfy:
|
||||
profiles: ["flux"]
|
||||
image: ghcr.io/aeon-7/comfyui-aeon-spark:slim
|
||||
container_name: flux2-klein-comfy
|
||||
restart: unless-stopped
|
||||
gpus: all
|
||||
ipc: host
|
||||
shm_size: 16g
|
||||
volumes:
|
||||
- /home/daniel/aeon-spark-test/workspace:/workspace/ComfyUI:rw
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/workspace/ComfyUI/output:rw
|
||||
- /home/daniel/aeon-spark-test/extra_model_paths.yaml:/opt/ComfyUI/extra_model_paths.yaml:ro
|
||||
- /home/daniel/comfy-data/models:/shared-models:ro
|
||||
ports:
|
||||
- "8192:8188"
|
||||
|
|
@ -11,6 +11,7 @@ services:
|
|||
- /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
|
||||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/aeon-spark-test/h3/models/MiniMax-H3-Turbo-FL2VA/latest:/turbo:ro
|
||||
- /home/daniel/aeon-spark-test/h3/models/Minimax-H3-Latent-Upscaler:/latent-upscaler:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
command: ["sleep", "infinity"]
|
||||
h3-hot-runtime:
|
||||
|
|
@ -25,6 +26,7 @@ services:
|
|||
- /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
|
||||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/aeon-spark-test/h3/models/MiniMax-H3-Turbo-FL2VA/latest:/turbo:ro
|
||||
- /home/daniel/aeon-spark-test/h3/models/Minimax-H3-Latent-Upscaler:/latent-upscaler:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
ports:
|
||||
- "8001:8000"
|
||||
|
|
@ -33,4 +35,4 @@ services:
|
|||
H3_NVFP4_SCALE_BACKEND: "vortex"
|
||||
H3_NVFP4_SCALE_VERSION: "1"
|
||||
H3_SOL_QKV_LAYOUT: "native"
|
||||
command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--turbo-lora", "4step=/turbo/minimax_h3_fl2v_turbo_4step_v1.1_768p_comfyui_bf16.safetensors", "--turbo-lora", "8step=/turbo/minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors", "--warmup"]
|
||||
command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--turbo-lora", "4step=/turbo/minimax_h3_fl2v_turbo_4step_v1.1_768p_comfyui_bf16.safetensors", "--turbo-lora", "8step=/turbo/minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors", "--latent-upscaler", "/latent-upscaler/minimax_h3_latent_upscaler_3d_fp16.safetensors", "--warmup"]
|
||||
|
|
|
|||
|
|
@ -11,5 +11,8 @@ dependencies = [
|
|||
"transformers>=4.51,<5"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
flash4 = ["flash-attn-4[cu13]==4.0.0b27"]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""Direct H3 self-attention using packed NVFP4 linears and SageAttention3."""
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
from torch import nn
|
||||
|
|
@ -8,9 +10,12 @@ from torch import nn
|
|||
from .checkpoint import H3Checkpoint
|
||||
from .nvfp4 import Nvfp4Linear
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .distributed import SequenceParallelContext
|
||||
|
||||
AVAILABLE_BACKENDS = ("sage2", "cudnn_sdpa", "ck_int8", "sdpa", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn")
|
||||
PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "kj_chunked_ffn")
|
||||
|
||||
AVAILABLE_BACKENDS = ("sage2", "cudnn_sdpa", "ck_int8", "sdpa", "flash4", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn")
|
||||
PLANNED_BACKENDS = ("easycache", "h3_cache", "kj_chunked_ffn")
|
||||
DEFAULT_ATTENTION_BACKEND = os.getenv("H3_DEFAULT_ATTENTION", "sage2")
|
||||
|
||||
|
||||
|
|
@ -20,7 +25,7 @@ def attention_backend_status() -> dict[str, str]:
|
|||
status.update({"cudnn_sdpa": "available: forced cuDNN SDPA with no backend fallback"})
|
||||
status.update({"ck_int8": "available: approximate Comfy Kitchen INT8 Q/K/V attention"})
|
||||
status.update({"sol_attn": "experimental: sparse Triton attention for eligible non-causal H3 attention calls; falls back below H3_SOL_MIN_TOKENS unless H3_SOL_STRICT=1"})
|
||||
status.update({"flash4": "planned: exact Blackwell kernel adapter"})
|
||||
status.update({"flash4": "available: official FlashAttention-4 CuTeDSL Blackwell kernel (strict, no fallback)"})
|
||||
status.update({"easycache": "planned: approximate denoiser cache"})
|
||||
status.update({"h3_cache": "planned: approximate H3-specific cache"})
|
||||
status.update({"kj_chunked_ffn": "available: exact H3 MLP row chunking via H3_MLP_CHUNKS or runtime args"})
|
||||
|
|
@ -85,6 +90,23 @@ def qkv_to_bshd(qkv: torch.Tensor, heads: int, head_dim: int) -> tuple[torch.Ten
|
|||
)
|
||||
|
||||
|
||||
def run_flash4_attention_bshd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor:
|
||||
"""Run official FlashAttention-4 on `[batch, sequence, heads, dim]` tensors."""
|
||||
if not q.is_cuda or not k.is_cuda or not v.is_cuda:
|
||||
raise ValueError("FlashAttention-4 requires CUDA tensors")
|
||||
if q.dtype not in {torch.float16, torch.bfloat16} or k.dtype != q.dtype or v.dtype != q.dtype:
|
||||
raise ValueError("FlashAttention-4 requires matching FP16 or BF16 Q/K/V tensors")
|
||||
if q.shape != k.shape or q.shape != v.shape:
|
||||
raise ValueError("FlashAttention-4 requires matching Q/K/V shapes")
|
||||
if q.shape[-1] != 128:
|
||||
raise ValueError(f"H3 FlashAttention-4 requires head dim 128, got {q.shape[-1]}")
|
||||
|
||||
from flash_attn.cute import flash_attn_func
|
||||
|
||||
result = flash_attn_func(q.contiguous(), k.contiguous(), v.contiguous(), causal=is_causal)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend: str, is_causal: bool) -> torch.Tensor:
|
||||
"""Run one `[batch, heads, sequence, dim]` attention operation."""
|
||||
if backend == "sol_attn":
|
||||
|
|
@ -103,6 +125,10 @@ def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend:
|
|||
from sageattention import sageattn
|
||||
|
||||
return sageattn(q, k, v, is_causal=is_causal, tensor_layout="HND", smooth_k=False)
|
||||
if backend == "flash4":
|
||||
return run_flash4_attention_bshd(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=is_causal,
|
||||
).transpose(1, 2)
|
||||
if backend == "cudnn_sdpa":
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
|
|
@ -212,17 +238,44 @@ class H3SageAttention(nn.Module):
|
|||
backend=backend,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, rope_rotation: torch.Tensor) -> torch.Tensor:
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
rope_rotation: torch.Tensor,
|
||||
sequence_parallel: "SequenceParallelContext | None" = None,
|
||||
tensor_parallel: "SequenceParallelContext | None" = None,
|
||||
) -> torch.Tensor:
|
||||
if x.ndim != 2:
|
||||
raise ValueError("H3 attention expects `[sequence, hidden]` input.")
|
||||
if sequence_parallel is not None and tensor_parallel is not None:
|
||||
raise ValueError("choose Ulysses sequence parallelism or tensor parallelism, not both")
|
||||
if tensor_parallel is not None:
|
||||
return self._forward_tensor_parallel(x, rope_rotation, tensor_parallel)
|
||||
sequence = x.shape[0]
|
||||
inner = self.heads * self.head_dim
|
||||
q, k, v = self.qkv_proj(x).split(inner, dim=-1)
|
||||
qkv = self.qkv_proj(x)
|
||||
q, k, v = qkv.split(inner, dim=-1)
|
||||
q = q.view(1, sequence, self.heads, self.head_dim)
|
||||
k = k.view(1, sequence, self.heads, self.head_dim)
|
||||
v = v.view(1, sequence, self.heads, self.head_dim)
|
||||
|
||||
q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps)
|
||||
if sequence_parallel is not None:
|
||||
q, k, v = sequence_parallel.seq_to_heads(q, k, v)
|
||||
if self.backend == "sol_attn":
|
||||
out = run_sol_attention_bshd(q, k, v, is_causal=False)
|
||||
elif self.backend == "flash4":
|
||||
out = run_flash4_attention_bshd(q, k, v, is_causal=False)
|
||||
else:
|
||||
out = run_attention(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
backend=self.backend,
|
||||
is_causal=False,
|
||||
).transpose(1, 2)
|
||||
local_out = sequence_parallel.heads_to_seq(out)
|
||||
return self.out_proj(local_out.reshape(sequence, inner))
|
||||
if self.backend == "sol_attn":
|
||||
if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() == "fused":
|
||||
q, k, v = qkv_to_bshd(qkv, self.heads, self.head_dim)
|
||||
|
|
@ -230,6 +283,9 @@ class H3SageAttention(nn.Module):
|
|||
q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
|
||||
out = run_sol_attention_bshd(q, k, v, is_causal=False)
|
||||
return self.out_proj(out.reshape(sequence, inner).contiguous())
|
||||
if self.backend == "flash4":
|
||||
out = run_flash4_attention_bshd(q, k, v, is_causal=False)
|
||||
return self.out_proj(out.reshape(sequence, inner).contiguous())
|
||||
|
||||
q = q.transpose(1, 2).contiguous()
|
||||
k = k.transpose(1, 2).contiguous()
|
||||
|
|
@ -237,3 +293,42 @@ class H3SageAttention(nn.Module):
|
|||
|
||||
out = run_attention(q, k, v, backend=self.backend, is_causal=False)
|
||||
return self.out_proj(out.transpose(1, 2).reshape(sequence, inner).contiguous())
|
||||
|
||||
def _forward_tensor_parallel(
|
||||
self,
|
||||
local_x: torch.Tensor,
|
||||
local_rotation: torch.Tensor,
|
||||
context: "SequenceParallelContext",
|
||||
) -> torch.Tensor:
|
||||
"""Run local-head attention with column/row-parallel NVFP4 projections."""
|
||||
local_sequence = local_x.shape[0]
|
||||
full_x = context.all_gather_rows(local_x)
|
||||
full_rotation = context.all_gather_rows(local_rotation[0]).unsqueeze(0)
|
||||
inner = self.heads * self.head_dim
|
||||
q, k, v = self.qkv_proj(full_x).split(inner, dim=-1)
|
||||
q = q.view(1, context.sequence_length, self.heads, self.head_dim)
|
||||
k = k.view(1, context.sequence_length, self.heads, self.head_dim)
|
||||
v = v.view(1, context.sequence_length, self.heads, self.head_dim)
|
||||
q, k = rms_rope_split_half_(
|
||||
q, k, full_rotation, self.q_norm_weight, self.k_norm_weight, self.eps,
|
||||
)
|
||||
if self.backend == "sol_attn":
|
||||
out = run_sol_attention_bshd(q, k, v, is_causal=False)
|
||||
elif self.backend == "flash4":
|
||||
out = run_flash4_attention_bshd(q, k, v, is_causal=False)
|
||||
else:
|
||||
out = run_attention(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
backend=self.backend,
|
||||
is_causal=False,
|
||||
).transpose(1, 2)
|
||||
partial = self.out_proj(out.reshape(context.sequence_length, inner).contiguous())
|
||||
local_output = context.reduce_scatter_rows(partial)
|
||||
bias = getattr(self, "tensor_parallel_output_bias", None)
|
||||
if bias is not None:
|
||||
local_output = local_output + bias.to(local_output)
|
||||
if local_output.shape[0] != local_sequence:
|
||||
raise RuntimeError("tensor-parallel attention returned the wrong local token count")
|
||||
return local_output
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
import torch
|
||||
from torch import nn
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .adaln import H3CurveAdaLN
|
||||
from .attention import DEFAULT_ATTENTION_BACKEND
|
||||
|
|
@ -9,6 +10,9 @@ from .block import H3DiTBlock
|
|||
from .checkpoint import H3Checkpoint
|
||||
from .rope import h3_rope_rotation
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .distributed import SequenceParallelContext
|
||||
|
||||
|
||||
class H3DenoiserBackbone(nn.Module):
|
||||
"""Execute H3 transformer blocks over an already packed Ref2VA hidden sequence."""
|
||||
|
|
@ -35,8 +39,29 @@ class H3DenoiserBackbone(nn.Module):
|
|||
timesteps: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
segments: list[tuple[int, int, int]],
|
||||
sequence_parallel: "SequenceParallelContext | None" = None,
|
||||
tensor_parallel: "SequenceParallelContext | None" = None,
|
||||
) -> torch.Tensor:
|
||||
if sequence_parallel is not None and tensor_parallel is not None:
|
||||
raise ValueError("choose Ulysses sequence parallelism or tensor parallelism, not both")
|
||||
parallel = sequence_parallel or tensor_parallel
|
||||
if parallel is not None:
|
||||
if hidden.shape[0] != parallel.local_token_length:
|
||||
raise ValueError(
|
||||
f"local hidden length {hidden.shape[0]} does not match sequence-parallel "
|
||||
f"partition {parallel.local_token_length}"
|
||||
)
|
||||
if position_ids.shape[0] != hidden.shape[0]:
|
||||
raise ValueError("local position IDs must match local hidden rows")
|
||||
segments = parallel.localize_segments(segments)
|
||||
rotation = h3_rope_rotation(position_ids.to(hidden.device), self.inv_freq, hidden.dtype)
|
||||
for block, adaln in zip(self.blocks, self.adaln, strict=True):
|
||||
hidden = block(hidden, rotation, *adaln(timesteps), segments)
|
||||
hidden = block(
|
||||
hidden,
|
||||
rotation,
|
||||
*adaln(timesteps),
|
||||
segments,
|
||||
sequence_parallel,
|
||||
tensor_parallel,
|
||||
)
|
||||
return hidden
|
||||
|
|
|
|||
|
|
@ -2,11 +2,15 @@
|
|||
|
||||
import torch
|
||||
from torch import nn
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .attention import DEFAULT_ATTENTION_BACKEND, H3SageAttention, rms_norm
|
||||
from .checkpoint import H3Checkpoint
|
||||
from .nvfp4 import Nvfp4Linear
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .distributed import SequenceParallelContext
|
||||
|
||||
|
||||
def modulate_segments(
|
||||
x: torch.Tensor,
|
||||
|
|
@ -49,7 +53,19 @@ class H3SwiGLU(nn.Module):
|
|||
checkpoint.nvfp4_linear(f"{prefix}.fc2", output_dtype=output_dtype),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
tensor_parallel: "SequenceParallelContext | None" = None,
|
||||
) -> torch.Tensor:
|
||||
if tensor_parallel is not None:
|
||||
full_x = tensor_parallel.all_gather_rows(x)
|
||||
partial = self._forward_chunk(full_x)
|
||||
local_output = tensor_parallel.reduce_scatter_rows(partial)
|
||||
bias = getattr(self, "tensor_parallel_output_bias", None)
|
||||
if bias is not None:
|
||||
local_output = local_output + bias.to(local_output)
|
||||
return local_output
|
||||
if self.chunks > 1 and x.shape[0] >= self.chunk_threshold:
|
||||
return torch.cat([self._forward_chunk(chunk) for chunk in x.chunk(self.chunks, dim=0)], dim=0)
|
||||
return self._forward_chunk(x)
|
||||
|
|
@ -109,8 +125,15 @@ class H3DiTBlock(nn.Module):
|
|||
scale_mlp: torch.Tensor,
|
||||
gate_mlp: torch.Tensor,
|
||||
segments: list[tuple[int, int, int]],
|
||||
sequence_parallel: "SequenceParallelContext | None" = None,
|
||||
tensor_parallel: "SequenceParallelContext | None" = None,
|
||||
) -> torch.Tensor:
|
||||
h = modulate_segments(rms_norm(x, self.norm1_weight, self.norm_eps), shift_msa, scale_msa, segments)
|
||||
x = gate_segments(x, self.attention(h, rope_rotation), gate_msa, segments)
|
||||
x = gate_segments(
|
||||
x,
|
||||
self.attention(h, rope_rotation, sequence_parallel, tensor_parallel),
|
||||
gate_msa,
|
||||
segments,
|
||||
)
|
||||
h = modulate_segments(rms_norm(x, self.norm2_weight, self.norm_eps), shift_mlp, scale_mlp, segments)
|
||||
return gate_segments(x, self.mlp(h), gate_mlp, segments)
|
||||
return gate_segments(x, self.mlp(h, tensor_parallel), gate_mlp, segments)
|
||||
|
|
|
|||
|
|
@ -59,6 +59,10 @@ class H3Checkpoint:
|
|||
value = checkpoint.get_tensor(name)
|
||||
return value.to(dtype=dtype) if dtype is not None else value
|
||||
|
||||
def release_cache(self) -> None:
|
||||
"""Release tensors retained by whole-file loading after modules are built."""
|
||||
self._no_mmap_tensors = None
|
||||
|
||||
def nvfp4_linear(self, prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
|
||||
names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias", "pre_quant_scale")
|
||||
tensors = {}
|
||||
|
|
|
|||
|
|
@ -2,12 +2,16 @@
|
|||
|
||||
import torch
|
||||
from torch import nn
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .attention import DEFAULT_ATTENTION_BACKEND
|
||||
from .backbone import H3DenoiserBackbone
|
||||
from .checkpoint import H3Checkpoint
|
||||
from .final import H3FinalLayer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .distributed import SequenceParallelContext
|
||||
|
||||
|
||||
class H3PackedDenoiser(nn.Module):
|
||||
"""Run the H3 transformer once its Ref2VA payload has been packed into hidden rows."""
|
||||
|
|
@ -35,3 +39,67 @@ class H3PackedDenoiser(nn.Module):
|
|||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
hidden = self.backbone(hidden, timesteps, position_ids, segments)
|
||||
return self.final_layer(hidden, timesteps, video_segment, audio_segment)
|
||||
|
||||
def forward_sequence_parallel(
|
||||
self,
|
||||
full_hidden: torch.Tensor,
|
||||
timesteps: torch.Tensor,
|
||||
full_position_ids: torch.Tensor,
|
||||
segments: list[tuple[int, int, int]],
|
||||
video_segment: tuple[int, int, int],
|
||||
audio_segment: tuple[int, int, int],
|
||||
context: "SequenceParallelContext",
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run all 50 blocks with token-sharded activations and Ulysses attention."""
|
||||
if full_hidden.shape[0] != context.sequence_length:
|
||||
raise ValueError("packed hidden length does not match sequence-parallel context")
|
||||
start, stop = context.local_token_range
|
||||
local_hidden = full_hidden[start:stop].contiguous()
|
||||
local_positions = full_position_ids[start:stop].contiguous()
|
||||
del full_hidden, full_position_ids
|
||||
local_hidden = self.backbone(
|
||||
local_hidden,
|
||||
timesteps,
|
||||
local_positions,
|
||||
segments,
|
||||
sequence_parallel=context,
|
||||
)
|
||||
return self.final_layer.forward_sequence_parallel(
|
||||
local_hidden,
|
||||
timesteps,
|
||||
video_segment,
|
||||
audio_segment,
|
||||
context,
|
||||
)
|
||||
|
||||
def forward_tensor_parallel(
|
||||
self,
|
||||
full_hidden: torch.Tensor,
|
||||
timesteps: torch.Tensor,
|
||||
full_position_ids: torch.Tensor,
|
||||
segments: list[tuple[int, int, int]],
|
||||
video_segment: tuple[int, int, int],
|
||||
audio_segment: tuple[int, int, int],
|
||||
context: "SequenceParallelContext",
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run TP-sharded NVFP4 linears with ragged sequence-sharded residuals."""
|
||||
if full_hidden.shape[0] != context.sequence_length:
|
||||
raise ValueError("packed hidden length does not match tensor-parallel context")
|
||||
start, stop = context.local_token_range
|
||||
local_hidden = full_hidden[start:stop].contiguous()
|
||||
local_positions = full_position_ids[start:stop].contiguous()
|
||||
del full_hidden, full_position_ids
|
||||
local_hidden = self.backbone(
|
||||
local_hidden,
|
||||
timesteps,
|
||||
local_positions,
|
||||
segments,
|
||||
tensor_parallel=context,
|
||||
)
|
||||
return self.final_layer.forward_sequence_parallel(
|
||||
local_hidden,
|
||||
timesteps,
|
||||
video_segment,
|
||||
audio_segment,
|
||||
context,
|
||||
)
|
||||
|
|
|
|||
324
src/h3_blackwell_runtime/distributed.py
Normal file
324
src/h3_blackwell_runtime/distributed.py
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
"""Ragged Ulysses sequence-parallel transport for H3 inference."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
def balanced_ranges(total: int, parts: int) -> tuple[tuple[int, int], ...]:
|
||||
"""Split ``total`` ordered items into balanced contiguous non-empty ranges."""
|
||||
if parts < 1:
|
||||
raise ValueError("parts must be positive")
|
||||
if total < parts:
|
||||
raise ValueError(f"cannot split {total} items into {parts} non-empty ranges")
|
||||
base, extra = divmod(total, parts)
|
||||
lengths = [base + (rank < extra) for rank in range(parts)]
|
||||
ranges = []
|
||||
start = 0
|
||||
for length in lengths:
|
||||
stop = start + int(length)
|
||||
ranges.append((start, stop))
|
||||
start = stop
|
||||
return tuple(ranges)
|
||||
|
||||
|
||||
def range_lengths(ranges: tuple[tuple[int, int], ...]) -> tuple[int, ...]:
|
||||
return tuple(stop - start for start, stop in ranges)
|
||||
|
||||
|
||||
def localize_segments(
|
||||
segments: list[tuple[int, int, int]],
|
||||
shard_start: int,
|
||||
shard_stop: int,
|
||||
) -> list[tuple[int, int, int]]:
|
||||
"""Clip global H3 AdaLN segments to one contiguous token shard."""
|
||||
localized = []
|
||||
for start, stop, row in segments:
|
||||
local_start = max(start, shard_start)
|
||||
local_stop = min(stop, shard_stop)
|
||||
if local_start < local_stop:
|
||||
localized.append((local_start - shard_start, local_stop - shard_start, row))
|
||||
return localized
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SequenceParallelContext:
|
||||
"""One rank's ragged token and attention-head ownership."""
|
||||
|
||||
group: dist.ProcessGroup | None
|
||||
rank: int
|
||||
world_size: int
|
||||
token_ranges: tuple[tuple[int, int], ...]
|
||||
head_ranges: tuple[tuple[int, int], ...]
|
||||
head_dim: int
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
sequence_length: int,
|
||||
heads: int,
|
||||
head_dim: int,
|
||||
*,
|
||||
group: dist.ProcessGroup | None = None,
|
||||
) -> "SequenceParallelContext":
|
||||
if not dist.is_initialized():
|
||||
raise RuntimeError("torch.distributed process group is not initialized")
|
||||
world_size = dist.get_world_size(group)
|
||||
rank = dist.get_rank(group)
|
||||
return cls(
|
||||
group=group,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
token_ranges=balanced_ranges(sequence_length, world_size),
|
||||
head_ranges=balanced_ranges(heads, world_size),
|
||||
head_dim=head_dim,
|
||||
)
|
||||
|
||||
@property
|
||||
def sequence_length(self) -> int:
|
||||
return self.token_ranges[-1][1]
|
||||
|
||||
@property
|
||||
def heads(self) -> int:
|
||||
return self.head_ranges[-1][1]
|
||||
|
||||
@property
|
||||
def token_lengths(self) -> tuple[int, ...]:
|
||||
return range_lengths(self.token_ranges)
|
||||
|
||||
@property
|
||||
def head_lengths(self) -> tuple[int, ...]:
|
||||
return range_lengths(self.head_ranges)
|
||||
|
||||
@property
|
||||
def local_token_range(self) -> tuple[int, int]:
|
||||
return self.token_ranges[self.rank]
|
||||
|
||||
@property
|
||||
def local_head_range(self) -> tuple[int, int]:
|
||||
return self.head_ranges[self.rank]
|
||||
|
||||
@property
|
||||
def local_token_length(self) -> int:
|
||||
start, stop = self.local_token_range
|
||||
return stop - start
|
||||
|
||||
@property
|
||||
def local_head_count(self) -> int:
|
||||
start, stop = self.local_head_range
|
||||
return stop - start
|
||||
|
||||
def localize_segments(self, segments: list[tuple[int, int, int]]) -> list[tuple[int, int, int]]:
|
||||
return localize_segments(segments, *self.local_token_range)
|
||||
|
||||
def seq_to_heads(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Exchange local tokens for full-sequence Q/K/V over locally owned heads.
|
||||
|
||||
Inputs use BSHD layout ``[1, local_tokens, all_heads, head_dim]``. Outputs
|
||||
use ``[1, all_tokens, local_heads, head_dim]``.
|
||||
"""
|
||||
expected = (1, self.local_token_length, self.heads, self.head_dim)
|
||||
if tuple(q.shape) != expected or tuple(k.shape) != expected or tuple(v.shape) != expected:
|
||||
raise ValueError(
|
||||
f"sequence-parallel Q/K/V must each have shape {expected}; "
|
||||
f"got {tuple(q.shape)}, {tuple(k.shape)}, {tuple(v.shape)}"
|
||||
)
|
||||
if q.dtype != k.dtype or q.dtype != v.dtype or q.device != k.device or q.device != v.device:
|
||||
raise ValueError("sequence-parallel Q/K/V must share dtype and device")
|
||||
if self.world_size == 1:
|
||||
return q, k, v
|
||||
|
||||
local_qkv = torch.stack((q[0], k[0], v[0]), dim=1)
|
||||
send_chunks = [
|
||||
local_qkv[:, :, start:stop, :].contiguous().view(-1)
|
||||
for start, stop in self.head_ranges
|
||||
]
|
||||
input_splits = [chunk.numel() for chunk in send_chunks]
|
||||
send = torch.cat(send_chunks)
|
||||
|
||||
output_splits = [
|
||||
token_length * 3 * self.local_head_count * self.head_dim
|
||||
for token_length in self.token_lengths
|
||||
]
|
||||
receive = torch.empty(sum(output_splits), dtype=q.dtype, device=q.device)
|
||||
dist.all_to_all_single(
|
||||
receive,
|
||||
send,
|
||||
output_split_sizes=output_splits,
|
||||
input_split_sizes=input_splits,
|
||||
group=self.group,
|
||||
)
|
||||
|
||||
source_chunks = []
|
||||
offset = 0
|
||||
for token_length, count in zip(self.token_lengths, output_splits, strict=True):
|
||||
source_chunks.append(
|
||||
receive[offset : offset + count].view(
|
||||
token_length, 3, self.local_head_count, self.head_dim,
|
||||
)
|
||||
)
|
||||
offset += count
|
||||
full_qkv = torch.cat(source_chunks, dim=0)
|
||||
full_q, full_k, full_v = full_qkv.unbind(dim=1)
|
||||
return full_q.unsqueeze(0), full_k.unsqueeze(0), full_v.unsqueeze(0)
|
||||
|
||||
def heads_to_seq(self, output: torch.Tensor) -> torch.Tensor:
|
||||
"""Exchange full-sequence local-head output back to local-token all-head output.
|
||||
|
||||
Input is BSHD ``[1, all_tokens, local_heads, head_dim]``. The return value
|
||||
is ``[local_tokens, all_heads, head_dim]``.
|
||||
"""
|
||||
expected = (1, self.sequence_length, self.local_head_count, self.head_dim)
|
||||
if tuple(output.shape) != expected:
|
||||
raise ValueError(f"sequence-parallel output must have shape {expected}, got {tuple(output.shape)}")
|
||||
if self.world_size == 1:
|
||||
return output[0]
|
||||
|
||||
output = output[0]
|
||||
send_chunks = []
|
||||
input_splits = []
|
||||
token_offset = 0
|
||||
for token_length in self.token_lengths:
|
||||
chunk = output[token_offset : token_offset + token_length].contiguous().view(-1)
|
||||
send_chunks.append(chunk)
|
||||
input_splits.append(chunk.numel())
|
||||
token_offset += token_length
|
||||
send = torch.cat(send_chunks)
|
||||
|
||||
output_splits = [
|
||||
self.local_token_length * head_length * self.head_dim
|
||||
for head_length in self.head_lengths
|
||||
]
|
||||
receive = torch.empty(sum(output_splits), dtype=output.dtype, device=output.device)
|
||||
dist.all_to_all_single(
|
||||
receive,
|
||||
send,
|
||||
output_split_sizes=output_splits,
|
||||
input_split_sizes=input_splits,
|
||||
group=self.group,
|
||||
)
|
||||
|
||||
head_chunks = []
|
||||
offset = 0
|
||||
for head_length, count in zip(self.head_lengths, output_splits, strict=True):
|
||||
head_chunks.append(
|
||||
receive[offset : offset + count].view(
|
||||
self.local_token_length, head_length, self.head_dim,
|
||||
)
|
||||
)
|
||||
offset += count
|
||||
return torch.cat(head_chunks, dim=1).contiguous()
|
||||
|
||||
def target_intersection(self, target_start: int, target_stop: int) -> tuple[int, int]:
|
||||
"""Return one global target span's bounds relative to this token shard."""
|
||||
shard_start, shard_stop = self.local_token_range
|
||||
start = max(target_start, shard_start)
|
||||
stop = min(target_stop, shard_stop)
|
||||
if start >= stop:
|
||||
return (0, 0)
|
||||
return (start - shard_start, stop - shard_start)
|
||||
|
||||
def target_counts(self, target_start: int, target_stop: int) -> tuple[int, ...]:
|
||||
"""Return ordered target-row counts contributed by every token rank."""
|
||||
counts = []
|
||||
for shard_start, shard_stop in self.token_ranges:
|
||||
counts.append(max(0, min(target_stop, shard_stop) - max(target_start, shard_start)))
|
||||
return tuple(counts)
|
||||
|
||||
def all_gather_target_rows(
|
||||
self,
|
||||
local_rows: torch.Tensor,
|
||||
target_start: int,
|
||||
target_stop: int,
|
||||
) -> torch.Tensor:
|
||||
"""Gather a global target span's projected rows onto every rank.
|
||||
|
||||
Padding is transport-only and is removed before concatenation; it is never
|
||||
exposed to attention or model semantics.
|
||||
"""
|
||||
counts = self.target_counts(target_start, target_stop)
|
||||
if local_rows.ndim != 2:
|
||||
raise ValueError("target rows must be rank-2 [rows, features]")
|
||||
if local_rows.shape[0] != counts[self.rank]:
|
||||
raise ValueError(
|
||||
f"rank {self.rank} must contribute {counts[self.rank]} target rows, "
|
||||
f"got {local_rows.shape[0]}"
|
||||
)
|
||||
if self.world_size == 1:
|
||||
return local_rows
|
||||
max_rows = max(counts)
|
||||
padded = torch.zeros(
|
||||
max_rows, local_rows.shape[1], dtype=local_rows.dtype, device=local_rows.device,
|
||||
)
|
||||
if local_rows.shape[0]:
|
||||
padded[: local_rows.shape[0]].copy_(local_rows)
|
||||
gathered = [torch.empty_like(padded) for _ in range(self.world_size)]
|
||||
dist.all_gather(gathered, padded, group=self.group)
|
||||
return torch.cat(
|
||||
[rows[:count] for rows, count in zip(gathered, counts, strict=True) if count],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
def all_gather_rows(self, local_rows: torch.Tensor) -> torch.Tensor:
|
||||
"""Gather ragged token rows on every rank without exposing padding to the model."""
|
||||
if local_rows.shape[0] != self.local_token_length:
|
||||
raise ValueError(
|
||||
f"rank {self.rank} must contribute {self.local_token_length} rows, "
|
||||
f"got {local_rows.shape[0]}"
|
||||
)
|
||||
if self.world_size == 1:
|
||||
return local_rows
|
||||
max_rows = max(self.token_lengths)
|
||||
padded = torch.zeros(
|
||||
(max_rows, *local_rows.shape[1:]),
|
||||
dtype=local_rows.dtype,
|
||||
device=local_rows.device,
|
||||
)
|
||||
padded[: local_rows.shape[0]].copy_(local_rows)
|
||||
gathered = [torch.empty_like(padded) for _ in range(self.world_size)]
|
||||
dist.all_gather(gathered, padded, group=self.group)
|
||||
return torch.cat(
|
||||
[rows[:count] for rows, count in zip(gathered, self.token_lengths, strict=True)],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
def reduce_scatter_rows(self, partial_full_rows: torch.Tensor) -> torch.Tensor:
|
||||
"""Sum tensor-parallel partials and return this rank's ragged token rows."""
|
||||
if partial_full_rows.shape[0] != self.sequence_length:
|
||||
raise ValueError(
|
||||
f"partial rows must cover sequence length {self.sequence_length}, "
|
||||
f"got {partial_full_rows.shape[0]}"
|
||||
)
|
||||
if self.world_size == 1:
|
||||
return partial_full_rows
|
||||
trailing_shape = partial_full_rows.shape[1:]
|
||||
row_width = partial_full_rows[0].numel()
|
||||
send_chunks = []
|
||||
input_splits = []
|
||||
offset = 0
|
||||
for token_length in self.token_lengths:
|
||||
chunk = partial_full_rows[offset : offset + token_length].contiguous().view(-1)
|
||||
send_chunks.append(chunk)
|
||||
input_splits.append(chunk.numel())
|
||||
offset += token_length
|
||||
send = torch.cat(send_chunks)
|
||||
|
||||
output_splits = [self.local_token_length * row_width] * self.world_size
|
||||
receive = torch.empty(sum(output_splits), dtype=send.dtype, device=send.device)
|
||||
dist.all_to_all_single(
|
||||
receive,
|
||||
send,
|
||||
output_split_sizes=output_splits,
|
||||
input_split_sizes=input_splits,
|
||||
group=self.group,
|
||||
)
|
||||
contributions = receive.view(self.world_size, self.local_token_length, *trailing_shape)
|
||||
return contributions.sum(dim=0)
|
||||
|
|
@ -3,10 +3,14 @@
|
|||
import torch
|
||||
import torch.nn.functional as functional
|
||||
from torch import nn
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .attention import rms_norm
|
||||
from .checkpoint import H3Checkpoint
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .distributed import SequenceParallelContext
|
||||
|
||||
|
||||
class H3FinalLayer(nn.Module):
|
||||
def __init__(
|
||||
|
|
@ -70,3 +74,43 @@ class H3FinalLayer(nn.Module):
|
|||
functional.linear(video_hidden, self.video_weight, self.video_bias),
|
||||
functional.linear(audio_hidden, self.audio_weight, self.audio_bias),
|
||||
)
|
||||
|
||||
def forward_sequence_parallel(
|
||||
self,
|
||||
local_hidden: torch.Tensor,
|
||||
timesteps: torch.Tensor,
|
||||
video_segment: tuple[int, int, int],
|
||||
audio_segment: tuple[int, int, int],
|
||||
context: "SequenceParallelContext",
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Project local target intersections and gather compact AV rows on every rank."""
|
||||
position = timesteps.float().clamp(0, 1) * (self.curve_table.shape[0] - 1)
|
||||
lower = position.floor().long().clamp(max=self.curve_table.shape[0] - 2)
|
||||
embedding = torch.lerp(
|
||||
self.curve_table[lower],
|
||||
self.curve_table[lower + 1],
|
||||
(position - lower).unsqueeze(1),
|
||||
)
|
||||
shift, scale = functional.linear(embedding, self.adaln_weight, self.adaln_bias).chunk(2, dim=-1)
|
||||
|
||||
outputs = []
|
||||
for segment, weight, bias in (
|
||||
(video_segment, self.video_weight, self.video_bias),
|
||||
(audio_segment, self.audio_weight, self.audio_bias),
|
||||
):
|
||||
global_start, global_stop, row = segment
|
||||
local_start, local_stop = context.target_intersection(global_start, global_stop)
|
||||
local_target = local_hidden[local_start:local_stop]
|
||||
if local_target.shape[0]:
|
||||
local_target = (
|
||||
rms_norm(local_target, self.norm_weight, self.eps)
|
||||
* (1.0 + scale[row])
|
||||
+ shift[row]
|
||||
).to(torch.float32)
|
||||
local_output = functional.linear(local_target, weight, bias)
|
||||
else:
|
||||
local_output = torch.empty(
|
||||
0, weight.shape[0], dtype=torch.float32, device=local_hidden.device,
|
||||
)
|
||||
outputs.append(context.all_gather_target_rows(local_output, global_start, global_stop))
|
||||
return outputs[0], outputs[1]
|
||||
|
|
|
|||
163
src/h3_blackwell_runtime/latent_upscaler.py
Normal file
163
src/h3_blackwell_runtime/latent_upscaler.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
"""MiniMax H3 learned latent upscaler inference.
|
||||
|
||||
Adapted from LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler (Apache-2.0).
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
import re
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
LATENTS_MEAN = (
|
||||
0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075,
|
||||
-0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975,
|
||||
-0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923,
|
||||
-0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543,
|
||||
-0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279,
|
||||
-0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264,
|
||||
)
|
||||
LATENTS_STD = (
|
||||
1.2223774194717407, 1.2767263650894165, 1.6831774711608887, 1.7549455165863037,
|
||||
1.5636216402053833, 2.194143533706665, 0.9653137922286987, 1.0569885969161987,
|
||||
0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647,
|
||||
0.7996809482574463, 0.44988900423049927, 0.7197399735450745, 0.6936293244361877,
|
||||
2.961095094680786, 2.7694199085235596, 3.0496184825897217, 2.1088054180145264,
|
||||
3.276226282119751, 3.1627357006073, 2.2816812992095947, 2.6127843856811523,
|
||||
)
|
||||
|
||||
|
||||
def _normalization(channels: int) -> nn.GroupNorm:
|
||||
return nn.GroupNorm(32, channels)
|
||||
|
||||
|
||||
class ResBlockEmb3D(nn.Module):
|
||||
def __init__(self, channels: int, emb_channels: int, dropout: float = 0.0):
|
||||
super().__init__()
|
||||
self.in_layers = nn.Sequential(
|
||||
_normalization(channels),
|
||||
nn.SiLU(),
|
||||
nn.Conv3d(channels, channels, 3, padding=1),
|
||||
)
|
||||
self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(emb_channels, 2 * channels))
|
||||
self.out_norm = _normalization(channels)
|
||||
self.out_layers = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Dropout(p=dropout),
|
||||
nn.Conv3d(channels, channels, 3, padding=1),
|
||||
)
|
||||
nn.init.zeros_(self.out_layers[-1].weight)
|
||||
nn.init.zeros_(self.out_layers[-1].bias)
|
||||
|
||||
def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
|
||||
hidden = self.in_layers(x)
|
||||
scale, shift = self.emb_layers(emb).to(hidden.dtype).chunk(2, dim=1)
|
||||
hidden = self.out_norm(hidden) * (1 + scale[:, :, None, None, None]) + shift[:, :, None, None, None]
|
||||
return x + self.out_layers(hidden)
|
||||
|
||||
|
||||
class TemporalConv(nn.Module):
|
||||
def __init__(self, channels: int, kernel_size: int = 5):
|
||||
super().__init__()
|
||||
self.norm = _normalization(channels)
|
||||
self.dwconv = nn.Conv3d(
|
||||
channels, channels, kernel_size=(kernel_size, 1, 1),
|
||||
padding=(kernel_size // 2, 0, 0), groups=channels,
|
||||
)
|
||||
self.pwconv = nn.Conv3d(channels, channels, 1)
|
||||
nn.init.zeros_(self.pwconv.weight)
|
||||
nn.init.zeros_(self.pwconv.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
hidden = self.dwconv(F.silu(self.norm(x)))
|
||||
return x + self.pwconv(hidden)
|
||||
|
||||
|
||||
class H3LatentResizer3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 24,
|
||||
in_blocks: int = 12,
|
||||
out_blocks: int = 12,
|
||||
channels: int = 512,
|
||||
dropout: float = 0.1,
|
||||
temporal_every: int = 2,
|
||||
temporal_kernel: int = 5,
|
||||
):
|
||||
super().__init__()
|
||||
self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1)
|
||||
embed_dim = 64
|
||||
self.embed = nn.Sequential(nn.Linear(1, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim))
|
||||
self.in_blocks = self._make_blocks(in_blocks, channels, embed_dim, dropout, temporal_every, temporal_kernel)
|
||||
self.out_blocks = self._make_blocks(out_blocks, channels, embed_dim, dropout, temporal_every, temporal_kernel)
|
||||
self.norm_out = _normalization(channels)
|
||||
self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1)
|
||||
|
||||
@staticmethod
|
||||
def _make_blocks(count, channels, embed_dim, dropout, temporal_every, temporal_kernel):
|
||||
blocks = nn.ModuleList()
|
||||
for index in range(count):
|
||||
blocks.append(ResBlockEmb3D(channels, embed_dim, dropout))
|
||||
if temporal_every > 0 and index % temporal_every == 0:
|
||||
blocks.append(TemporalConv(channels, temporal_kernel))
|
||||
return blocks
|
||||
|
||||
def forward(self, x: torch.Tensor, *, scale: float, target_size: tuple[int, int, int]) -> torch.Tensor:
|
||||
emb = self.embed(torch.tensor([[scale - 1]], dtype=x.dtype, device=x.device)).expand(x.shape[0], -1)
|
||||
hidden = self.conv_in(x)
|
||||
for block in self.in_blocks:
|
||||
hidden = block(hidden, emb) if isinstance(block, ResBlockEmb3D) else block(hidden)
|
||||
hidden = F.interpolate(hidden, size=target_size, mode="trilinear", align_corners=False)
|
||||
for block in self.out_blocks:
|
||||
hidden = block(hidden, emb) if isinstance(block, ResBlockEmb3D) else block(hidden)
|
||||
return self.conv_out(F.silu(self.norm_out(hidden)))
|
||||
|
||||
|
||||
def _checkpoint_config(state: dict[str, torch.Tensor]) -> dict:
|
||||
in_ids = {int(match.group(1)) for key in state if (match := re.match(r"in_blocks\.(\d+)\.in_layers\.", key))}
|
||||
out_ids = {int(match.group(1)) for key in state if (match := re.match(r"out_blocks\.(\d+)\.in_layers\.", key))}
|
||||
temporal_keys = [key for key in state if key.endswith("dwconv.weight")]
|
||||
conv_in = state["conv_in.weight"]
|
||||
return {
|
||||
"in_channels": conv_in.shape[1],
|
||||
"in_blocks": len(in_ids),
|
||||
"out_blocks": len(out_ids),
|
||||
"channels": conv_in.shape[0],
|
||||
"temporal_every": 2 if temporal_keys else 0,
|
||||
"temporal_kernel": state[temporal_keys[0]].shape[2] if temporal_keys else 5,
|
||||
}
|
||||
|
||||
|
||||
def load_h3_latent_upscaler(
|
||||
path: str | Path,
|
||||
*,
|
||||
device: str | torch.device = "cuda",
|
||||
dtype: torch.dtype = torch.float16,
|
||||
) -> H3LatentResizer3D:
|
||||
state = load_file(str(path), device="cpu")
|
||||
if any(key.startswith("upscaler.") for key in state):
|
||||
state = {key.removeprefix("upscaler."): value for key, value in state.items() if key.startswith("upscaler.")}
|
||||
model = H3LatentResizer3D(**_checkpoint_config(state))
|
||||
model.load_state_dict(state, strict=True)
|
||||
return model.to(device=device, dtype=dtype).eval().requires_grad_(False)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def upscale_h3_latent(model: H3LatentResizer3D, latent: torch.Tensor, *, scale: float = 2.0) -> torch.Tensor:
|
||||
if latent.ndim != 5 or latent.shape[1] != 24:
|
||||
raise ValueError("H3 latent must have shape [B,24,T,H,W]")
|
||||
if not 1.0 <= scale <= 4.0:
|
||||
raise ValueError("scale must be between 1.0 and 4.0")
|
||||
dtype = next(model.parameters()).dtype
|
||||
device = next(model.parameters()).device
|
||||
source = latent.to(device=device, dtype=dtype)
|
||||
if scale == 1.0:
|
||||
return source
|
||||
mean = torch.tensor(LATENTS_MEAN, device=device, dtype=dtype).view(1, 24, 1, 1, 1)
|
||||
std = torch.tensor(LATENTS_STD, device=device, dtype=dtype).view(1, 24, 1, 1, 1)
|
||||
target_size = (source.shape[2], round(source.shape[3] * scale), round(source.shape[4] * scale))
|
||||
result = model((source - mean) / std, scale=scale, target_size=target_size)
|
||||
return result.mul_(std).add_(mean)
|
||||
|
|
@ -16,6 +16,7 @@ from .block import configure_mlp_chunking
|
|||
from .checkpoint import H3Checkpoint
|
||||
from .denoiser import H3PackedDenoiser
|
||||
from .lora import load_lora_adapter, set_active_lora
|
||||
from .latent_upscaler import load_h3_latent_upscaler, upscale_h3_latent
|
||||
from .packing import H3PromptPacker
|
||||
from .qwen3vl_text import Qwen3VLPromptConditioner
|
||||
from .qwen3vl_vision import build_fl2va_presentation, Qwen3VL32BVision
|
||||
|
|
@ -45,6 +46,7 @@ class RuntimeConfig:
|
|||
mlp_chunks: int = 1
|
||||
mlp_chunk_threshold: int = 4096
|
||||
turbo_loras: tuple[tuple[str, str], ...] = ()
|
||||
latent_upscaler_path: str | None = None
|
||||
device: str = "cuda"
|
||||
|
||||
|
||||
|
|
@ -64,6 +66,18 @@ def _refiner_attention_backend(attention: str) -> str:
|
|||
return fallback if fallback in AVAILABLE_BACKENDS and fallback != "sol_attn" else "sage2"
|
||||
|
||||
|
||||
def normalize_upscale(value) -> float | None:
|
||||
if value is None or value == "" or value == "none" or value == 1 or value == 1.0:
|
||||
return None
|
||||
try:
|
||||
scale = float(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("upscale must be a number between 1.0 and 4.0") from exc
|
||||
if not 1.0 < scale <= 4.0:
|
||||
raise ValueError("upscale must be greater than 1.0 and at most 4.0")
|
||||
return scale
|
||||
|
||||
|
||||
class H3HotRuntime:
|
||||
"""Keep all prompt-only H3 models resident for repeated requests."""
|
||||
|
||||
|
|
@ -113,6 +127,12 @@ class H3HotRuntime:
|
|||
"vision_tower_loaded",
|
||||
lambda: Qwen3VL32BVision(config.text_encoder_path, device=config.device, dtype=torch.float32),
|
||||
)
|
||||
self.latent_upscaler = None
|
||||
if config.latent_upscaler_path is not None:
|
||||
self.latent_upscaler = self._timed_load(
|
||||
"latent_upscaler_loaded",
|
||||
lambda: load_h3_latent_upscaler(config.latent_upscaler_path, device=config.device, dtype=torch.float16),
|
||||
)
|
||||
|
||||
def _timed_load(self, stage: str, fn):
|
||||
_sync()
|
||||
|
|
@ -143,6 +163,7 @@ class H3HotRuntime:
|
|||
"current_attention": self.attention,
|
||||
"available_turbos": sorted(self.turbo_loras),
|
||||
"current_turbo": self.turbo,
|
||||
"latent_upscaler_loaded": self.latent_upscaler is not None,
|
||||
"vae_dtype": self.config.vae_dtype,
|
||||
"vae_tile_size": self.config.vae_tile_size,
|
||||
"mlp_chunks": self.config.mlp_chunks,
|
||||
|
|
@ -224,9 +245,11 @@ class H3HotRuntime:
|
|||
seed: int,
|
||||
attention: str | None = None,
|
||||
turbo: str | None = None,
|
||||
upscale: float | None = None,
|
||||
first_frame: torch.Tensor | None = None,
|
||||
last_frame: torch.Tensor | None = None,
|
||||
mux_audio: bool = True,
|
||||
keep_intermediates: bool = False,
|
||||
ffmpeg_loglevel: str = "error",
|
||||
save_latent: str | Path | None = None,
|
||||
cache_mode: str | None = None,
|
||||
|
|
@ -249,6 +272,9 @@ class H3HotRuntime:
|
|||
output = Path(output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
requested_turbo = None if turbo in {None, "", "none"} else turbo
|
||||
upscale_scale = normalize_upscale(upscale)
|
||||
if upscale_scale is not None and self.latent_upscaler is None:
|
||||
raise ValueError("H3 latent upscaler is not loaded")
|
||||
if requested_turbo is not None:
|
||||
if requested_turbo not in self.turbo_loras:
|
||||
raise ValueError(f"Turbo variant {requested_turbo!r} is not loaded")
|
||||
|
|
@ -301,6 +327,14 @@ class H3HotRuntime:
|
|||
latent, audio_latent = sampled
|
||||
else:
|
||||
latent, audio_latent = sampled, None
|
||||
source_width, source_height = width, height
|
||||
if upscale_scale is not None:
|
||||
latent = timed(
|
||||
"latent_upscaled",
|
||||
lambda: upscale_h3_latent(self.latent_upscaler, latent, scale=upscale_scale),
|
||||
)
|
||||
width = latent.shape[-1] * 16
|
||||
height = latent.shape[-2] * 16
|
||||
if save_latent is not None:
|
||||
latent_path = Path(save_latent)
|
||||
latent_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
|
@ -351,7 +385,10 @@ class H3HotRuntime:
|
|||
timed("audio_encode", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-f", "f32le", "-ar", "32000", "-ac", "2", "-i", str(audio_raw), str(audio_output)), check=True))
|
||||
audio_raw.unlink()
|
||||
timed("mux", lambda: subprocess.run(_ffmpeg_command(ffmpeg_loglevel, "-y", "-i", str(video_output), "-i", str(audio_output), "-c:v", "copy", "-c:a", "aac", "-shortest", str(output)), check=True))
|
||||
video_output.unlink()
|
||||
if not keep_intermediates:
|
||||
video_output.unlink()
|
||||
audio_output.unlink()
|
||||
audio_output = None
|
||||
|
||||
return {
|
||||
"output": str(output),
|
||||
|
|
@ -359,9 +396,13 @@ class H3HotRuntime:
|
|||
"frames": aligned_frames,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"source_width": source_width,
|
||||
"source_height": source_height,
|
||||
"seed": seed,
|
||||
"attention": self.attention,
|
||||
"turbo": self.turbo,
|
||||
"upscale": upscale_scale,
|
||||
"keep_intermediates": keep_intermediates,
|
||||
"vae_dtype": self.config.vae_dtype,
|
||||
"vae_tile_size": self.config.vae_tile_size,
|
||||
"stages": stages,
|
||||
|
|
|
|||
|
|
@ -103,6 +103,7 @@ def sample_video_res_multistep(
|
|||
cache_end_percent: float = 1.0,
|
||||
cache_subsample_factor: int = 2,
|
||||
cache_stats: dict | None = None,
|
||||
audio_step_trace: list[dict] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
|
||||
sigmas = beta_sigmas(steps, device=video.device)
|
||||
|
|
@ -128,6 +129,7 @@ def sample_video_res_multistep(
|
|||
step_started = time.perf_counter()
|
||||
previous_index = index - 1
|
||||
sigma_down = sigmas[index]
|
||||
audio_before = audio_carried if audio_step_trace is None else audio_carried.detach().cpu()
|
||||
current_percent = previous_index / total_steps
|
||||
can_cache = cache_mode is not None and cache_threshold > 0 and cache_start_percent <= current_percent <= cache_end_percent and cache["video_diff"] is not None
|
||||
skipped = False
|
||||
|
|
@ -187,6 +189,16 @@ def sample_video_res_multistep(
|
|||
previous_sigma = sigmas[previous_index - 1] if previous_index else None
|
||||
video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
|
||||
audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)
|
||||
if audio_step_trace is not None:
|
||||
audio_step_trace.append({
|
||||
"step": index,
|
||||
"video_sigma": float(sigma),
|
||||
"audio_sigma": float(_audio_sigma(sigma)),
|
||||
"video_sigma_down": float(sigma_down),
|
||||
"audio_before": audio_before,
|
||||
"audio_denoised": audio_denoised.detach().cpu(),
|
||||
"audio_after": audio_carried.detach().cpu(),
|
||||
})
|
||||
video_history, audio_history = video_denoised, audio_denoised
|
||||
video_history_sigma = audio_history_sigma = sigma_down
|
||||
if progress:
|
||||
|
|
|
|||
101
src/h3_blackwell_runtime/tensor_parallel.py
Normal file
101
src/h3_blackwell_runtime/tensor_parallel.py
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
"""True NVFP4 tensor-parallel sharding for the H3 denoiser."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from .distributed import SequenceParallelContext, balanced_ranges
|
||||
from .nvfp4 import Nvfp4Linear, Nvfp4LinearTensors
|
||||
|
||||
|
||||
def aligned_balanced_ranges(total: int, parts: int, alignment: int) -> tuple[tuple[int, int], ...]:
|
||||
"""Balance ranges in indivisible alignment-sized groups."""
|
||||
if alignment < 1 or total % alignment:
|
||||
raise ValueError(f"total {total} must be divisible by alignment {alignment}")
|
||||
groups = balanced_ranges(total // alignment, parts)
|
||||
return tuple((start * alignment, stop * alignment) for start, stop in groups)
|
||||
|
||||
|
||||
def _reject_lora(linear: Nvfp4Linear) -> None:
|
||||
if len(linear.lora_branches):
|
||||
raise ValueError("tensor-parallel sharding must occur before loading Turbo LoRA branches")
|
||||
|
||||
|
||||
def select_nvfp4_outputs(linear: Nvfp4Linear, ranges: tuple[tuple[int, int], ...]) -> Nvfp4Linear:
|
||||
"""Create one column-parallel NVFP4 linear from selected output-row ranges."""
|
||||
_reject_lora(linear)
|
||||
indices = torch.cat([
|
||||
torch.arange(start, stop, device=linear.weight.device)
|
||||
for start, stop in ranges
|
||||
])
|
||||
bias = None if linear.bias is None else linear.bias.index_select(0, indices)
|
||||
tensors = Nvfp4LinearTensors(
|
||||
weight=linear.weight.index_select(0, indices).contiguous(),
|
||||
weight_scale=linear.weight_scale.index_select(0, indices).contiguous(),
|
||||
weight_scale_2=linear.weight_scale_2,
|
||||
bias=bias.contiguous() if bias is not None else None,
|
||||
pre_quant_scale=linear.pre_quant_scale,
|
||||
full_precision_matrix_mult=linear.full_precision_matrix_mult,
|
||||
in_features=linear.in_features,
|
||||
out_features=indices.numel(),
|
||||
)
|
||||
return Nvfp4Linear(tensors, output_dtype=linear.output_dtype)
|
||||
|
||||
|
||||
def slice_nvfp4_inputs(linear: Nvfp4Linear, start: int, stop: int) -> tuple[Nvfp4Linear, torch.Tensor | None]:
|
||||
"""Create one row-parallel NVFP4 linear and return its once-only output bias."""
|
||||
_reject_lora(linear)
|
||||
if start < 0 or stop > linear.in_features or start >= stop:
|
||||
raise ValueError(f"invalid input shard [{start}, {stop}) for width {linear.in_features}")
|
||||
if start % 32 or stop % 32:
|
||||
raise ValueError("NVFP4 input shards must align to 32 features")
|
||||
bias = linear.bias
|
||||
tensors = Nvfp4LinearTensors(
|
||||
weight=linear.weight[:, start // 2 : stop // 2].contiguous(),
|
||||
weight_scale=linear.weight_scale[:, start // 16 : stop // 16].contiguous(),
|
||||
weight_scale_2=linear.weight_scale_2,
|
||||
bias=None,
|
||||
pre_quant_scale=(
|
||||
None if linear.pre_quant_scale is None
|
||||
else linear.pre_quant_scale[start:stop].contiguous()
|
||||
),
|
||||
full_precision_matrix_mult=linear.full_precision_matrix_mult,
|
||||
in_features=stop - start,
|
||||
out_features=linear.out_features,
|
||||
)
|
||||
return Nvfp4Linear(tensors, output_dtype=linear.output_dtype), bias
|
||||
|
||||
|
||||
def configure_h3_tensor_parallel(model, context: SequenceParallelContext) -> None:
|
||||
"""Shard all denoiser attention and MLP linears in place across ranks."""
|
||||
for block in model.backbone.blocks:
|
||||
attention = block.attention
|
||||
global_heads = attention.heads
|
||||
if global_heads != context.heads or attention.head_dim != context.head_dim:
|
||||
raise ValueError("tensor-parallel context does not match H3 attention dimensions")
|
||||
head_start, head_stop = context.local_head_range
|
||||
inner = global_heads * attention.head_dim
|
||||
local_start = head_start * attention.head_dim
|
||||
local_stop = head_stop * attention.head_dim
|
||||
attention.qkv_proj = select_nvfp4_outputs(attention.qkv_proj, (
|
||||
(local_start, local_stop),
|
||||
(inner + local_start, inner + local_stop),
|
||||
(2 * inner + local_start, 2 * inner + local_stop),
|
||||
))
|
||||
attention.out_proj, output_bias = slice_nvfp4_inputs(
|
||||
attention.out_proj, local_start, local_stop,
|
||||
)
|
||||
attention.heads = context.local_head_count
|
||||
attention.register_buffer("tensor_parallel_output_bias", output_bias, persistent=False)
|
||||
|
||||
mlp = block.mlp
|
||||
intermediate = mlp.fc2.in_features
|
||||
mlp_ranges = aligned_balanced_ranges(intermediate, context.world_size, 32)
|
||||
mlp_start, mlp_stop = mlp_ranges[context.rank]
|
||||
mlp.fc1 = select_nvfp4_outputs(mlp.fc1, (
|
||||
(mlp_start, mlp_stop),
|
||||
(intermediate + mlp_start, intermediate + mlp_stop),
|
||||
))
|
||||
mlp.fc2, output_bias = slice_nvfp4_inputs(mlp.fc2, mlp_start, mlp_stop)
|
||||
mlp.register_buffer("tensor_parallel_output_bias", output_bias, persistent=False)
|
||||
mlp.tensor_parallel_intermediate_ranges = mlp_ranges
|
||||
262
tests/test_distributed.py
Normal file
262
tests/test_distributed.py
Normal file
|
|
@ -0,0 +1,262 @@
|
|||
import tempfile
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
from torch.nn import functional as F
|
||||
|
||||
from h3_blackwell_runtime.distributed import (
|
||||
SequenceParallelContext,
|
||||
balanced_ranges,
|
||||
localize_segments,
|
||||
range_lengths,
|
||||
)
|
||||
from h3_blackwell_runtime.final import H3FinalLayer
|
||||
from h3_blackwell_runtime.nvfp4 import Nvfp4Linear, Nvfp4LinearTensors
|
||||
from h3_blackwell_runtime.tensor_parallel import (
|
||||
aligned_balanced_ranges,
|
||||
select_nvfp4_outputs,
|
||||
slice_nvfp4_inputs,
|
||||
)
|
||||
|
||||
|
||||
def _init_gloo(rank: int, world_size: int, init_file: str) -> None:
|
||||
dist.init_process_group(
|
||||
"gloo",
|
||||
init_method=f"file://{init_file}",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
|
||||
|
||||
def _transport_identity_worker(rank: int, world_size: int, init_file: str) -> None:
|
||||
_init_gloo(rank, world_size, init_file)
|
||||
try:
|
||||
sequence, heads, head_dim = 17, 56, 3
|
||||
context = SequenceParallelContext.create(sequence, heads, head_dim)
|
||||
start, stop = context.local_token_range
|
||||
full = torch.arange(sequence * heads * head_dim, dtype=torch.float32).reshape(1, sequence, heads, head_dim)
|
||||
q = full[:, start:stop].contiguous()
|
||||
k = q + 1_000_000
|
||||
v = q + 2_000_000
|
||||
full_q, full_k, full_v = context.seq_to_heads(q, k, v)
|
||||
torch.testing.assert_close(context.heads_to_seq(full_q), q[0], rtol=0, atol=0)
|
||||
torch.testing.assert_close(context.heads_to_seq(full_k), k[0], rtol=0, atol=0)
|
||||
torch.testing.assert_close(context.heads_to_seq(full_v), v[0], rtol=0, atol=0)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _attention_parity_worker(rank: int, world_size: int, init_file: str) -> None:
|
||||
_init_gloo(rank, world_size, init_file)
|
||||
try:
|
||||
torch.manual_seed(440420)
|
||||
sequence, heads, head_dim = 19, 56, 8
|
||||
context = SequenceParallelContext.create(sequence, heads, head_dim)
|
||||
q = torch.randn(1, sequence, heads, head_dim)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
start, stop = context.local_token_range
|
||||
local_q, local_k, local_v = context.seq_to_heads(
|
||||
q[:, start:stop].contiguous(),
|
||||
k[:, start:stop].contiguous(),
|
||||
v[:, start:stop].contiguous(),
|
||||
)
|
||||
local_heads = F.scaled_dot_product_attention(
|
||||
local_q.transpose(1, 2),
|
||||
local_k.transpose(1, 2),
|
||||
local_v.transpose(1, 2),
|
||||
).transpose(1, 2)
|
||||
actual = context.heads_to_seq(local_heads)
|
||||
expected = F.scaled_dot_product_attention(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
|
||||
).transpose(1, 2)[0, start:stop]
|
||||
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _final_projection_worker(rank: int, world_size: int, init_file: str) -> None:
|
||||
_init_gloo(rank, world_size, init_file)
|
||||
try:
|
||||
sequence, hidden_size = 9, 4
|
||||
context = SequenceParallelContext.create(sequence, heads=4, head_dim=2)
|
||||
hidden = torch.arange(sequence * hidden_size, dtype=torch.float32).reshape(sequence, hidden_size) / 10
|
||||
timesteps = torch.tensor([0.25, 0.75])
|
||||
layer = H3FinalLayer(
|
||||
torch.zeros(1025, 2),
|
||||
torch.ones(hidden_size),
|
||||
torch.zeros(2 * hidden_size, 2),
|
||||
torch.zeros(2 * hidden_size),
|
||||
torch.arange(3 * hidden_size, dtype=torch.float32).reshape(3, hidden_size) / 10,
|
||||
torch.tensor([0.1, 0.2, 0.3]),
|
||||
torch.arange(2 * hidden_size, dtype=torch.float32).reshape(2, hidden_size) / 20,
|
||||
torch.tensor([-0.1, 0.1]),
|
||||
hidden_size=hidden_size,
|
||||
)
|
||||
video_segment = (3, 9, 0)
|
||||
audio_segment = (0, 3, 1)
|
||||
expected_video, expected_audio = layer(hidden, timesteps, video_segment, audio_segment)
|
||||
start, stop = context.local_token_range
|
||||
actual_video, actual_audio = layer.forward_sequence_parallel(
|
||||
hidden[start:stop], timesteps, video_segment, audio_segment, context,
|
||||
)
|
||||
torch.testing.assert_close(actual_video, expected_video)
|
||||
torch.testing.assert_close(actual_audio, expected_audio)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _ragged_row_collectives_worker(rank: int, world_size: int, init_file: str) -> None:
|
||||
_init_gloo(rank, world_size, init_file)
|
||||
try:
|
||||
sequence, features = 17, 5
|
||||
context = SequenceParallelContext.create(sequence, heads=56, head_dim=2)
|
||||
full = torch.arange(sequence * features, dtype=torch.float32).reshape(sequence, features)
|
||||
start, stop = context.local_token_range
|
||||
gathered = context.all_gather_rows(full[start:stop].contiguous())
|
||||
torch.testing.assert_close(gathered, full, rtol=0, atol=0)
|
||||
partial = full * float(rank + 1)
|
||||
reduced = context.reduce_scatter_rows(partial)
|
||||
expected = full[start:stop] * sum(range(1, world_size + 1))
|
||||
torch.testing.assert_close(reduced, expected, rtol=0, atol=0)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _tensor_parallel_math_worker(rank: int, world_size: int, init_file: str) -> None:
|
||||
_init_gloo(rank, world_size, init_file)
|
||||
try:
|
||||
torch.manual_seed(440421)
|
||||
sequence, hidden, heads, head_dim, intermediate = 17, 16, 14, 4, 224
|
||||
context = SequenceParallelContext.create(sequence, heads=heads, head_dim=head_dim)
|
||||
full_x = torch.randn(sequence, hidden)
|
||||
start, stop = context.local_token_range
|
||||
gathered_x = context.all_gather_rows(full_x[start:stop].contiguous())
|
||||
|
||||
inner = heads * head_dim
|
||||
qkv_weight = torch.randn(3 * inner, hidden)
|
||||
output_weight = torch.randn(hidden, inner)
|
||||
output_bias = torch.randn(hidden)
|
||||
head_start, head_stop = context.local_head_range
|
||||
feature_start, feature_stop = head_start * head_dim, head_stop * head_dim
|
||||
indices = torch.cat((
|
||||
torch.arange(feature_start, feature_stop),
|
||||
torch.arange(inner + feature_start, inner + feature_stop),
|
||||
torch.arange(2 * inner + feature_start, 2 * inner + feature_stop),
|
||||
))
|
||||
local_qkv = F.linear(gathered_x, qkv_weight.index_select(0, indices))
|
||||
local_inner = context.local_head_count * head_dim
|
||||
local_q, local_k, local_v = local_qkv.split(local_inner, dim=-1)
|
||||
local_q = local_q.view(1, sequence, context.local_head_count, head_dim)
|
||||
local_k = local_k.view_as(local_q)
|
||||
local_v = local_v.view_as(local_q)
|
||||
local_attention = F.scaled_dot_product_attention(
|
||||
local_q.transpose(1, 2), local_k.transpose(1, 2), local_v.transpose(1, 2),
|
||||
).transpose(1, 2).reshape(sequence, local_inner)
|
||||
partial_attention = F.linear(
|
||||
local_attention, output_weight[:, feature_start:feature_stop],
|
||||
)
|
||||
actual_attention = context.reduce_scatter_rows(partial_attention) + output_bias
|
||||
|
||||
full_q, full_k, full_v = F.linear(full_x, qkv_weight).split(inner, dim=-1)
|
||||
full_q = full_q.view(1, sequence, heads, head_dim)
|
||||
full_k = full_k.view_as(full_q)
|
||||
full_v = full_v.view_as(full_q)
|
||||
expected_attention = F.linear(
|
||||
F.scaled_dot_product_attention(
|
||||
full_q.transpose(1, 2), full_k.transpose(1, 2), full_v.transpose(1, 2),
|
||||
).transpose(1, 2).reshape(sequence, inner),
|
||||
output_weight,
|
||||
output_bias,
|
||||
)[start:stop]
|
||||
torch.testing.assert_close(actual_attention, expected_attention, rtol=2e-5, atol=2e-5)
|
||||
|
||||
fc1_weight = torch.randn(2 * intermediate, hidden)
|
||||
fc2_weight = torch.randn(hidden, intermediate)
|
||||
fc2_bias = torch.randn(hidden)
|
||||
mlp_ranges = aligned_balanced_ranges(intermediate, world_size, 32)
|
||||
mlp_start, mlp_stop = mlp_ranges[rank]
|
||||
local_fc1_weight = torch.cat((
|
||||
fc1_weight[mlp_start:mlp_stop],
|
||||
fc1_weight[intermediate + mlp_start:intermediate + mlp_stop],
|
||||
))
|
||||
gate, up = F.linear(gathered_x, local_fc1_weight).chunk(2, dim=-1)
|
||||
partial_mlp = F.linear(F.silu(gate) * up, fc2_weight[:, mlp_start:mlp_stop])
|
||||
actual_mlp = context.reduce_scatter_rows(partial_mlp) + fc2_bias
|
||||
full_gate, full_up = F.linear(full_x, fc1_weight).chunk(2, dim=-1)
|
||||
expected_mlp = F.linear(F.silu(full_gate) * full_up, fc2_weight, fc2_bias)[start:stop]
|
||||
# TP reduction changes FP32 accumulation order across rank partials.
|
||||
torch.testing.assert_close(actual_mlp, expected_mlp, rtol=1e-4, atol=2e-4)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _run_distributed(worker, world_size: int) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
init_file = str(Path(directory) / "process-group")
|
||||
mp.spawn(worker, args=(world_size, init_file), nprocs=world_size, join=True)
|
||||
|
||||
|
||||
class DistributedPartitionContracts(unittest.TestCase):
|
||||
def test_balanced_ragged_ranges(self):
|
||||
ranges = balanced_ranges(56, 6)
|
||||
self.assertEqual(range_lengths(ranges), (10, 10, 9, 9, 9, 9))
|
||||
self.assertEqual(ranges[0], (0, 10))
|
||||
self.assertEqual(ranges[-1], (47, 56))
|
||||
|
||||
def test_ranges_reject_empty_partitions(self):
|
||||
with self.assertRaisesRegex(ValueError, "non-empty"):
|
||||
balanced_ranges(3, 4)
|
||||
|
||||
def test_segments_are_clipped_and_rebased(self):
|
||||
segments = [(0, 4, 1), (4, 10, 2), (10, 15, 3)]
|
||||
self.assertEqual(localize_segments(segments, 3, 12), [(0, 1, 1), (1, 7, 2), (7, 9, 3)])
|
||||
|
||||
def test_transport_identity_for_planned_world_sizes(self):
|
||||
for world_size in (2, 4, 6, 8):
|
||||
with self.subTest(world_size=world_size):
|
||||
_run_distributed(_transport_identity_worker, world_size)
|
||||
|
||||
def test_two_rank_sdpa_matches_single_process(self):
|
||||
_run_distributed(_attention_parity_worker, 2)
|
||||
|
||||
def test_distributed_final_projection_matches_single_process(self):
|
||||
_run_distributed(_final_projection_worker, 2)
|
||||
|
||||
def test_ragged_all_gather_and_reduce_scatter(self):
|
||||
_run_distributed(_ragged_row_collectives_worker, 6)
|
||||
|
||||
def test_tensor_parallel_attention_and_mlp_match_dense_math(self):
|
||||
for world_size in (2, 6):
|
||||
with self.subTest(world_size=world_size):
|
||||
_run_distributed(_tensor_parallel_math_worker, world_size)
|
||||
|
||||
def test_nvfp4_column_and_row_shards_preserve_layout(self):
|
||||
tensors = Nvfp4LinearTensors(
|
||||
weight=torch.arange(96 * 32, dtype=torch.int32).to(torch.uint8).reshape(96, 32),
|
||||
weight_scale=torch.arange(96 * 4, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(96, 4),
|
||||
weight_scale_2=torch.tensor(0.5),
|
||||
bias=torch.arange(96, dtype=torch.bfloat16),
|
||||
pre_quant_scale=torch.arange(64, dtype=torch.bfloat16),
|
||||
full_precision_matrix_mult=False,
|
||||
in_features=64,
|
||||
out_features=96,
|
||||
)
|
||||
linear = Nvfp4Linear(tensors)
|
||||
column = select_nvfp4_outputs(linear, ((0, 16), (32, 48)))
|
||||
self.assertEqual(tuple(column.weight.shape), (32, 32))
|
||||
self.assertEqual(tuple(column.weight_scale.shape), (32, 4))
|
||||
self.assertEqual(column.out_features, 32)
|
||||
row, bias = slice_nvfp4_inputs(linear, 32, 64)
|
||||
self.assertEqual(tuple(row.weight.shape), (96, 16))
|
||||
self.assertEqual(tuple(row.weight_scale.shape), (96, 2))
|
||||
self.assertEqual(row.in_features, 32)
|
||||
self.assertIsNone(row.bias)
|
||||
torch.testing.assert_close(bias, tensors.bias)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -63,7 +63,26 @@ class Fl2vaVAEContracts(unittest.TestCase):
|
|||
|
||||
class AttentionBackendContracts(unittest.TestCase):
|
||||
def test_hot_backends_include_benchmark_candidates(self):
|
||||
self.assertTrue({"sage2", "cudnn_sdpa", "ck_int8"}.issubset(AVAILABLE_BACKENDS))
|
||||
self.assertTrue({"sage2", "cudnn_sdpa", "ck_int8", "flash4"}.issubset(AVAILABLE_BACKENDS))
|
||||
|
||||
def test_flash4_dispatches_bshd_and_restores_hnd(self):
|
||||
q = torch.randn(1, 2, 3, 128, dtype=torch.bfloat16)
|
||||
expected_bshd = torch.randn(1, 3, 2, 128, dtype=torch.bfloat16)
|
||||
flash = unittest.mock.MagicMock(return_value=(expected_bshd, torch.empty(0)))
|
||||
package = SimpleNamespace(cute=SimpleNamespace(flash_attn_func=flash))
|
||||
with (
|
||||
patch.dict(sys.modules, {"flash_attn": package, "flash_attn.cute": package.cute}),
|
||||
patch.object(torch.Tensor, "is_cuda", new_callable=unittest.mock.PropertyMock, return_value=True),
|
||||
):
|
||||
actual = run_attention(q, q, q, backend="flash4", is_causal=False)
|
||||
torch.testing.assert_close(actual, expected_bshd.transpose(1, 2))
|
||||
flash.assert_called_once()
|
||||
called_q, called_k, called_v = flash.call_args.args
|
||||
self.assertEqual(called_q.shape, (1, 3, 2, 128))
|
||||
self.assertTrue(called_q.is_contiguous())
|
||||
self.assertTrue(called_k.is_contiguous())
|
||||
self.assertTrue(called_v.is_contiguous())
|
||||
self.assertFalse(flash.call_args.kwargs["causal"])
|
||||
|
||||
def test_cudnn_backend_is_forced_without_fallback(self):
|
||||
q = torch.randn(1, 2, 3, 4)
|
||||
|
|
|
|||
52
tests/test_latent_upscaler.py
Normal file
52
tests/test_latent_upscaler.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.latent_upscaler import H3LatentResizer3D, _checkpoint_config, upscale_h3_latent
|
||||
from h3_blackwell_runtime.runtime import normalize_upscale
|
||||
|
||||
|
||||
class LatentUpscalerContracts(unittest.TestCase):
|
||||
def test_request_upscale_normalization(self):
|
||||
self.assertIsNone(normalize_upscale(None))
|
||||
self.assertIsNone(normalize_upscale("none"))
|
||||
self.assertIsNone(normalize_upscale(1))
|
||||
self.assertEqual(normalize_upscale("2"), 2.0)
|
||||
with self.assertRaisesRegex(ValueError, "greater than 1.0"):
|
||||
normalize_upscale(0.5)
|
||||
with self.assertRaisesRegex(ValueError, "at most 4.0"):
|
||||
normalize_upscale(4.1)
|
||||
with self.assertRaisesRegex(ValueError, "must be a number"):
|
||||
normalize_upscale({"scale": 2})
|
||||
|
||||
def test_3d_model_preserves_time_and_scales_space(self):
|
||||
model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, dropout=0, temporal_every=0).eval()
|
||||
latent = torch.randn(1, 24, 3, 2, 4)
|
||||
output = upscale_h3_latent(model, latent, scale=2)
|
||||
self.assertEqual(output.shape, (1, 24, 3, 4, 8))
|
||||
|
||||
def test_invalid_latent_shape_is_rejected(self):
|
||||
model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, temporal_every=0).eval()
|
||||
with self.assertRaisesRegex(ValueError, r"\[B,24,T,H,W\]"):
|
||||
upscale_h3_latent(model, torch.randn(1, 16, 3, 2, 4))
|
||||
|
||||
def test_scale_one_preserves_latent_exactly(self):
|
||||
model = H3LatentResizer3D(in_blocks=1, out_blocks=1, channels=32, temporal_every=0).eval()
|
||||
latent = torch.randn(1, 24, 3, 2, 4)
|
||||
self.assertTrue(torch.equal(upscale_h3_latent(model, latent, scale=1), latent))
|
||||
|
||||
def test_checkpoint_architecture_detection_matches_module_layout(self):
|
||||
model = H3LatentResizer3D(in_blocks=3, out_blocks=2, channels=32, temporal_every=2, temporal_kernel=3)
|
||||
config = _checkpoint_config(model.state_dict())
|
||||
self.assertEqual(config, {
|
||||
"in_channels": 24,
|
||||
"in_blocks": 3,
|
||||
"out_blocks": 2,
|
||||
"channels": 32,
|
||||
"temporal_every": 2,
|
||||
"temporal_kernel": 3,
|
||||
})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -5,7 +5,7 @@ import torch
|
|||
from torch import nn
|
||||
|
||||
from h3_blackwell_runtime.lora import DynamicLoraMixin
|
||||
from h3_blackwell_runtime.sampler import sample_video_turbo, turbo_sigmas
|
||||
from h3_blackwell_runtime.sampler import sample_video_res_multistep, sample_video_turbo, turbo_sigmas
|
||||
|
||||
|
||||
class _Linear(DynamicLoraMixin, nn.Module):
|
||||
|
|
@ -87,6 +87,34 @@ class TurboLoraContracts(unittest.TestCase):
|
|||
torch.testing.assert_close(calls[1][1], torch.full_like(audio, 0.75))
|
||||
torch.testing.assert_close(calls[1][2], torch.tensor([1.0 / 7.0, 0.25]))
|
||||
|
||||
def test_base_sampler_records_opt_in_audio_step_trace(self):
|
||||
video = torch.zeros(1, 1, 1, 1, 1)
|
||||
audio = torch.zeros(1, 32, 2, 1)
|
||||
trace = []
|
||||
|
||||
def packer(*args, **kwargs):
|
||||
return (None, None, None, None, None, None)
|
||||
|
||||
def model(*args):
|
||||
return torch.ones(1), torch.ones(1)
|
||||
|
||||
with (
|
||||
patch("h3_blackwell_runtime.sampler.beta_sigmas", return_value=torch.tensor([1.0, 0.0])),
|
||||
patch("h3_blackwell_runtime.sampler.unpatchify_video", return_value=torch.zeros_like(video)),
|
||||
patch("h3_blackwell_runtime.sampler._unpack_audio", return_value=torch.ones_like(audio)),
|
||||
):
|
||||
_video, final_audio = sample_video_res_multistep(
|
||||
model, packer, torch.empty(0), video, audio,
|
||||
steps=1, return_audio=True, audio_step_trace=trace,
|
||||
)
|
||||
|
||||
self.assertEqual(len(trace), 1)
|
||||
self.assertEqual(trace[0]["step"], 1)
|
||||
torch.testing.assert_close(trace[0]["audio_before"], torch.zeros_like(audio))
|
||||
torch.testing.assert_close(trace[0]["audio_denoised"], torch.full_like(audio, 4.0))
|
||||
torch.testing.assert_close(trace[0]["audio_after"], torch.full_like(audio, 4.0))
|
||||
torch.testing.assert_close(final_audio, torch.ones_like(audio))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
|
|||
123
tools/analyze_audio_boundary.py
Normal file
123
tools/analyze_audio_boundary.py
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
"""Compare H3 audio-latent and lossless-waveform boundaries."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def dbfs(value: float) -> float:
|
||||
return 20.0 * math.log10(max(value, 1e-20))
|
||||
|
||||
|
||||
def load_audio_latent(path: Path) -> torch.Tensor:
|
||||
state = torch.load(path, map_location="cpu", weights_only=False)
|
||||
latent = state.get("audio_latent") if isinstance(state, dict) else state
|
||||
if latent is None or latent.ndim != 4:
|
||||
raise ValueError(f"{path} does not contain a [B,C,S,T] audio latent")
|
||||
return latent.float()
|
||||
|
||||
|
||||
def latent_metrics(latent: torch.Tensor) -> dict:
|
||||
frames = latent.movedim(-1, 0).flatten(1)
|
||||
frame_rms = frames.square().mean(1).sqrt()
|
||||
frame_mean = frames.mean(1)
|
||||
frame_max = frames.abs().amax(1)
|
||||
deltas = frames[1:] - frames[:-1]
|
||||
delta_rms = deltas.square().mean(1).sqrt()
|
||||
adjacent_cosine = torch.nn.functional.cosine_similarity(frames[:-1], frames[1:], dim=1)
|
||||
|
||||
block_frames = min(4, frames.shape[0] // 2)
|
||||
first = frames[:block_frames].flatten()
|
||||
last = frames[-block_frames:].flatten()
|
||||
first_last_cosine = torch.nn.functional.cosine_similarity(first, last, dim=0)
|
||||
|
||||
first_count = min(20, frames.shape[0])
|
||||
return {
|
||||
"shape": list(latent.shape),
|
||||
"dtype": str(latent.dtype),
|
||||
"first_20_frame_rms": frame_rms[:first_count].tolist(),
|
||||
"first_20_frame_mean": frame_mean[:first_count].tolist(),
|
||||
"first_20_frame_max_abs": frame_max[:first_count].tolist(),
|
||||
"first_19_delta_rms": delta_rms[: max(0, first_count - 1)].tolist(),
|
||||
"first_19_adjacent_cosine": adjacent_cosine[: max(0, first_count - 1)].tolist(),
|
||||
"first_4_rms": float(frames[:block_frames].square().mean().sqrt()),
|
||||
"frames_4_20_rms": float(frames[block_frames:first_count].square().mean().sqrt()),
|
||||
"remaining_rms": float(frames[first_count:].square().mean().sqrt()),
|
||||
"first_4_vs_last_4_cosine": float(first_last_cosine),
|
||||
"largest_delta_frame": int(delta_rms.argmax().item() + 1),
|
||||
"largest_delta_rms": float(delta_rms.max()),
|
||||
}
|
||||
|
||||
|
||||
def load_wav(path: Path) -> tuple[np.ndarray, int]:
|
||||
with wave.open(str(path), "rb") as source:
|
||||
if source.getsampwidth() != 2:
|
||||
raise ValueError(f"{path} must be PCM S16")
|
||||
channels = source.getnchannels()
|
||||
sample_rate = source.getframerate()
|
||||
samples = np.frombuffer(source.readframes(source.getnframes()), dtype="<i2")
|
||||
return samples.reshape(-1, channels).astype(np.float32) / 32768.0, sample_rate
|
||||
|
||||
|
||||
def waveform_metrics(samples: np.ndarray, sample_rate: int) -> dict:
|
||||
first_half_second = samples[: sample_rate // 2]
|
||||
mono = first_half_second.mean(1)
|
||||
window_samples = sample_rate // 100
|
||||
windows = []
|
||||
for start in range(0, len(first_half_second), window_samples):
|
||||
block = first_half_second[start : start + window_samples]
|
||||
if len(block) == 0:
|
||||
continue
|
||||
windows.append({
|
||||
"start_ms": start * 1000.0 / sample_rate,
|
||||
"peak_dbfs": dbfs(float(np.max(np.abs(block)))),
|
||||
"rms_dbfs": dbfs(float(np.sqrt(np.mean(block * block)))),
|
||||
"mean": float(block.mean()),
|
||||
})
|
||||
|
||||
spectrum_samples = min(sample_rate // 10, len(mono))
|
||||
windowed = mono[:spectrum_samples] * np.hanning(spectrum_samples)
|
||||
magnitudes = np.abs(np.fft.rfft(windowed))
|
||||
frequencies = np.fft.rfftfreq(spectrum_samples, 1.0 / sample_rate)
|
||||
dominant = np.argsort(magnitudes[1:])[-8:][::-1] + 1
|
||||
derivatives = np.max(np.abs(np.diff(first_half_second, axis=0)), axis=1)
|
||||
return {
|
||||
"sample_rate": sample_rate,
|
||||
"samples": len(samples),
|
||||
"first_sample": samples[0].tolist(),
|
||||
"first_sample_dbfs": [dbfs(float(abs(value))) for value in samples[0]],
|
||||
"first_500ms_peak_dbfs": dbfs(float(np.max(np.abs(first_half_second)))),
|
||||
"first_500ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_half_second**2)))),
|
||||
"largest_derivative": float(derivatives.max()),
|
||||
"largest_derivative_ms": float((derivatives.argmax() + 1) * 1000.0 / sample_rate),
|
||||
"dominant_first_100ms_hz": [float(frequencies[index]) for index in dominant],
|
||||
"windows_10ms": windows,
|
||||
}
|
||||
|
||||
|
||||
def analyze(latent_path: Path, wav_path: Path) -> dict:
|
||||
samples, sample_rate = load_wav(wav_path)
|
||||
return {
|
||||
"latent_path": str(latent_path),
|
||||
"wav_path": str(wav_path),
|
||||
"latent": latent_metrics(load_audio_latent(latent_path)),
|
||||
"waveform": waveform_metrics(samples, sample_rate),
|
||||
}
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--affected-latent", type=Path, required=True)
|
||||
parser.add_argument("--affected-wav", type=Path, required=True)
|
||||
parser.add_argument("--clean-latent", type=Path, required=True)
|
||||
parser.add_argument("--clean-wav", type=Path, required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
print(json.dumps({
|
||||
"affected": analyze(args.affected_latent, args.affected_wav),
|
||||
"clean": analyze(args.clean_latent, args.clean_wav),
|
||||
}, indent=2))
|
||||
143
tools/benchmark_ulysses.py
Normal file
143
tools/benchmark_ulysses.py
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
"""Benchmark ragged H3 Ulysses transport and attention under torchrun."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from h3_blackwell_runtime.attention import run_attention, run_flash4_attention_bshd, run_sol_attention_bshd
|
||||
from h3_blackwell_runtime.distributed import SequenceParallelContext
|
||||
|
||||
|
||||
def synchronize(device: torch.device) -> None:
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
|
||||
|
||||
def run_backend(backend: str, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
|
||||
if backend == "flash4":
|
||||
return run_flash4_attention_bshd(q, k, v, is_causal=False)
|
||||
if backend == "sol_attn":
|
||||
return run_sol_attention_bshd(q, k, v, is_causal=False)
|
||||
return run_attention(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
backend=backend,
|
||||
is_causal=False,
|
||||
).transpose(1, 2)
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--sequence", type=int, default=20000)
|
||||
parser.add_argument("--heads", type=int, default=56)
|
||||
parser.add_argument("--head-dim", type=int, default=128)
|
||||
parser.add_argument("--backend", default="sdpa")
|
||||
parser.add_argument("--warmup", type=int, default=3)
|
||||
parser.add_argument("--iterations", type=int, default=10)
|
||||
parser.add_argument("--output", type=Path)
|
||||
args = parser.parse_args()
|
||||
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
||||
use_cuda = torch.cuda.is_available()
|
||||
device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
|
||||
if use_cuda:
|
||||
torch.cuda.set_device(device)
|
||||
dist.init_process_group(backend="nccl" if use_cuda else "gloo", device_id=device if use_cuda else None)
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
context = SequenceParallelContext.create(args.sequence, args.heads, args.head_dim)
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(440420 + rank)
|
||||
shape = (1, context.local_token_length, args.heads, args.head_dim)
|
||||
dtype = torch.bfloat16 if use_cuda else torch.float32
|
||||
q = torch.randn(shape, generator=generator, dtype=dtype, device=device)
|
||||
k = torch.randn(shape, generator=generator, dtype=dtype, device=device)
|
||||
v = torch.randn(shape, generator=generator, dtype=dtype, device=device)
|
||||
|
||||
|
||||
def iteration() -> tuple[float, float, float, torch.Tensor]:
|
||||
dist.barrier()
|
||||
synchronize(device)
|
||||
started = time.perf_counter()
|
||||
full_q, full_k, full_v = context.seq_to_heads(q, k, v)
|
||||
synchronize(device)
|
||||
after_forward = time.perf_counter()
|
||||
head_output = run_backend(args.backend, full_q, full_k, full_v)
|
||||
synchronize(device)
|
||||
after_attention = time.perf_counter()
|
||||
local_output = context.heads_to_seq(head_output)
|
||||
synchronize(device)
|
||||
finished = time.perf_counter()
|
||||
return (
|
||||
after_forward - started,
|
||||
after_attention - after_forward,
|
||||
finished - after_attention,
|
||||
local_output,
|
||||
)
|
||||
|
||||
|
||||
for _ in range(args.warmup):
|
||||
*_timings, output = iteration()
|
||||
del output
|
||||
|
||||
rank_timings = []
|
||||
for _ in range(args.iterations):
|
||||
forward, attention, inverse, output = iteration()
|
||||
rank_timings.append((forward, attention, inverse, forward + attention + inverse))
|
||||
del output
|
||||
|
||||
timings = torch.tensor(rank_timings, dtype=torch.float64, device=device)
|
||||
gathered = [torch.empty_like(timings) for _ in range(world_size)]
|
||||
dist.all_gather(gathered, timings)
|
||||
|
||||
if use_cuda:
|
||||
peak_memory = torch.tensor([torch.cuda.max_memory_allocated(device)], dtype=torch.int64, device=device)
|
||||
else:
|
||||
peak_memory = torch.tensor([0], dtype=torch.int64, device=device)
|
||||
memory_by_rank = [torch.empty_like(peak_memory) for _ in range(world_size)]
|
||||
dist.all_gather(memory_by_rank, peak_memory)
|
||||
|
||||
if rank == 0:
|
||||
stacked = torch.stack(gathered).cpu()
|
||||
stage_names = ("forward_all_to_all", "attention", "inverse_all_to_all", "total")
|
||||
stages = {}
|
||||
for index, name in enumerate(stage_names):
|
||||
maximum_rank = stacked[:, :, index].amax(dim=0).tolist()
|
||||
stages[name] = {
|
||||
"median_seconds": statistics.median(maximum_rank),
|
||||
"minimum_seconds": min(maximum_rank),
|
||||
"maximum_seconds": max(maximum_rank),
|
||||
}
|
||||
element_size = q.element_size()
|
||||
report = {
|
||||
"world_size": world_size,
|
||||
"backend": args.backend,
|
||||
"device": torch.cuda.get_device_name(device) if use_cuda else "cpu",
|
||||
"torch": torch.__version__,
|
||||
"sequence": args.sequence,
|
||||
"heads": args.heads,
|
||||
"head_dim": args.head_dim,
|
||||
"token_lengths": list(context.token_lengths),
|
||||
"head_lengths": list(context.head_lengths),
|
||||
"dtype": str(dtype),
|
||||
"iterations": args.iterations,
|
||||
"aggregate_transport_bytes_per_iteration": {
|
||||
"forward_qkv": 3 * args.sequence * args.heads * args.head_dim * element_size,
|
||||
"inverse_output": args.sequence * args.heads * args.head_dim * element_size,
|
||||
},
|
||||
"stages": stages,
|
||||
"peak_allocated_bytes_by_rank": [int(value.item()) for value in memory_by_rank],
|
||||
}
|
||||
serialized = json.dumps(report, indent=2)
|
||||
if args.output is not None:
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(serialized)
|
||||
|
||||
dist.destroy_process_group()
|
||||
64
tools/compare_audio_pcm.py
Normal file
64
tools/compare_audio_pcm.py
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
"""Compare decoded audio streams using aligned float PCM arrays."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def decode(path: Path) -> np.ndarray:
|
||||
raw = subprocess.check_output([
|
||||
"ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(path),
|
||||
"-map", "0:a:0", "-f", "f32le", "-acodec", "pcm_f32le", "-",
|
||||
])
|
||||
return np.frombuffer(raw, dtype="<f4").reshape(-1, 2)
|
||||
|
||||
|
||||
def metrics(reference: np.ndarray, candidate: np.ndarray) -> dict:
|
||||
count = min(len(reference), len(candidate))
|
||||
reference = reference[:count]
|
||||
candidate = candidate[:count]
|
||||
error = candidate - reference
|
||||
signal_power = np.maximum(np.mean(reference**2, axis=0), 1e-30)
|
||||
noise_power = np.maximum(np.mean(error**2, axis=0), 1e-30)
|
||||
first_250ms = error[:8000]
|
||||
first_signal_power = np.maximum(np.mean(reference[:8000] ** 2, axis=0), 1e-30)
|
||||
first_noise_power = np.maximum(np.mean(first_250ms**2, axis=0), 1e-30)
|
||||
return {
|
||||
"reference_samples": len(reference),
|
||||
"candidate_samples": len(candidate),
|
||||
"compared_samples": count,
|
||||
"snr_db_by_channel": (10.0 * np.log10(signal_power / noise_power)).tolist(),
|
||||
"rmse": float(np.sqrt(np.mean(error**2))),
|
||||
"first_250ms_snr_db_by_channel": (
|
||||
10.0 * np.log10(first_signal_power / first_noise_power)
|
||||
).tolist(),
|
||||
"first_250ms_rmse": float(np.sqrt(np.mean(first_250ms**2))),
|
||||
"max_abs_error": float(np.max(np.abs(error))),
|
||||
}
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--reference", type=Path, required=True)
|
||||
parser.add_argument("--candidate", action="append", default=[], metavar="NAME=PATH")
|
||||
parser.add_argument("--output", type=Path)
|
||||
args = parser.parse_args()
|
||||
|
||||
reference = decode(args.reference)
|
||||
report = {"reference": str(args.reference), "candidates": {}}
|
||||
for value in args.candidate:
|
||||
if "=" not in value:
|
||||
raise ValueError(f"candidate must be NAME=PATH, got {value!r}")
|
||||
name, raw_path = value.split("=", 1)
|
||||
report["candidates"][name] = {
|
||||
"path": raw_path,
|
||||
**metrics(reference, decode(Path(raw_path))),
|
||||
}
|
||||
|
||||
serialized = json.dumps(report, indent=2)
|
||||
if args.output is not None:
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(serialized)
|
||||
|
|
@ -18,9 +18,9 @@ args = parser.parse_args()
|
|||
|
||||
state = torch.load(args.latent, map_location="cuda", weights_only=False)
|
||||
if isinstance(state, dict):
|
||||
latent = state.get("audio_latent", state.get("latent"))
|
||||
latent = state.get("audio_latent", state.get("final_audio", state.get("latent")))
|
||||
if latent is None:
|
||||
raise ValueError("saved state does not contain 'audio_latent' or 'latent'")
|
||||
raise ValueError("saved state does not contain 'audio_latent', 'final_audio', or 'latent'")
|
||||
else:
|
||||
latent = state
|
||||
latent = latent.to("cuda")
|
||||
|
|
|
|||
185
tools/distributed_t2va.py
Normal file
185
tools/distributed_t2va.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
"""Run prompt-only H3 T2VA with Ulysses or TP+sequence parallelism."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||||
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
||||
from h3_blackwell_runtime.distributed import SequenceParallelContext
|
||||
from h3_blackwell_runtime.packing import H3PromptPacker
|
||||
from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
|
||||
from h3_blackwell_runtime.sampler import sample_video_res_multistep
|
||||
from h3_blackwell_runtime.t2v import random_av_latents
|
||||
from h3_blackwell_runtime.tensor_parallel import configure_h3_tensor_parallel
|
||||
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
|
||||
|
||||
|
||||
class ParallelDenoiser:
|
||||
def __init__(self, model: H3PackedDenoiser, mode: str):
|
||||
self.model = model
|
||||
self.mode = mode
|
||||
|
||||
def __call__(self, hidden, timesteps, positions, segments, video_segment, audio_segment):
|
||||
context = SequenceParallelContext.create(hidden.shape[0], heads=56, head_dim=128)
|
||||
if self.mode == "ulysses":
|
||||
return self.model.forward_sequence_parallel(
|
||||
hidden, timesteps, positions, segments, video_segment, audio_segment, context,
|
||||
)
|
||||
return self.model.forward_tensor_parallel(
|
||||
hidden, timesteps, positions, segments, video_segment, audio_segment, context,
|
||||
)
|
||||
|
||||
|
||||
def synchronize(device: torch.device) -> None:
|
||||
torch.cuda.synchronize(device)
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--benchmark", type=Path, required=True)
|
||||
parser.add_argument("--mode", choices=("ulysses", "tensor"), required=True)
|
||||
parser.add_argument("--attention", default="sdpa")
|
||||
parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
parser.add_argument("--text-encoder", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
|
||||
parser.add_argument("--save-latent", type=Path)
|
||||
parser.add_argument("--report", type=Path)
|
||||
args = parser.parse_args()
|
||||
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.cuda.set_device(local_rank)
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
dist.init_process_group("nccl", device_id=device)
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
benchmark = json.loads(args.benchmark.read_text(encoding="utf-8"))
|
||||
|
||||
load_started = time.perf_counter()
|
||||
checkpoint = H3Checkpoint(args.model, device=device)
|
||||
model = H3PackedDenoiser.from_checkpoint(
|
||||
checkpoint, output_dtype=torch.bfloat16, attention_backend=args.attention,
|
||||
).eval()
|
||||
packer = H3PromptPacker(checkpoint)
|
||||
if args.mode == "tensor":
|
||||
partition_context = SequenceParallelContext.create(world_size, heads=56, head_dim=128)
|
||||
configure_h3_tensor_parallel(model, partition_context)
|
||||
synchronize(device)
|
||||
model_load_seconds = time.perf_counter() - load_started
|
||||
model_load_peak = torch.cuda.max_memory_allocated(device)
|
||||
torch.cuda.reset_peak_memory_stats(device)
|
||||
|
||||
conditioning_started = time.perf_counter()
|
||||
if rank == 0:
|
||||
conditioner = Qwen3VLPromptConditioner(args.text_encoder, device=device, dtype=torch.float32)
|
||||
refiner = H3TokenRefiner(checkpoint, attention_backend="sdpa").eval()
|
||||
text = refiner(conditioner(benchmark["prompt"])).to(torch.bfloat16)
|
||||
text_length = torch.tensor([text.shape[1]], dtype=torch.int64, device=device)
|
||||
del conditioner, refiner
|
||||
else:
|
||||
text = None
|
||||
text_length = torch.zeros(1, dtype=torch.int64, device=device)
|
||||
checkpoint.release_cache()
|
||||
torch.cuda.empty_cache()
|
||||
dist.broadcast(text_length, src=0)
|
||||
if rank != 0:
|
||||
text = torch.empty(1, int(text_length.item()), 5376, dtype=torch.bfloat16, device=device)
|
||||
dist.broadcast(text, src=0)
|
||||
synchronize(device)
|
||||
conditioning_seconds = time.perf_counter() - conditioning_started
|
||||
conditioning_peak = torch.cuda.max_memory_allocated(device)
|
||||
torch.cuda.reset_peak_memory_stats(device)
|
||||
|
||||
width, height = benchmark["resolution"]
|
||||
video, audio, aligned_frames = random_av_latents(
|
||||
width, height, benchmark["frames"], benchmark["seed"], device=device,
|
||||
)
|
||||
dist.barrier()
|
||||
synchronize(device)
|
||||
sampling_started = time.perf_counter()
|
||||
video, audio = sample_video_res_multistep(
|
||||
ParallelDenoiser(model, args.mode),
|
||||
packer,
|
||||
text,
|
||||
video,
|
||||
audio,
|
||||
steps=benchmark["steps"],
|
||||
seed=benchmark["seed"],
|
||||
return_audio=True,
|
||||
progress=rank == 0,
|
||||
)
|
||||
synchronize(device)
|
||||
sampling_seconds = time.perf_counter() - sampling_started
|
||||
|
||||
timing = torch.tensor(
|
||||
[model_load_seconds, conditioning_seconds, sampling_seconds],
|
||||
dtype=torch.float64,
|
||||
device=device,
|
||||
)
|
||||
timings = [torch.empty_like(timing) for _ in range(world_size)]
|
||||
dist.all_gather(timings, timing)
|
||||
checksums = torch.stack((video.float().sum(), audio.float().sum())).to(torch.float64)
|
||||
all_checksums = [torch.empty_like(checksums) for _ in range(world_size)]
|
||||
dist.all_gather(all_checksums, checksums)
|
||||
checksum_stack = torch.stack(all_checksums)
|
||||
if not torch.allclose(checksum_stack, checksum_stack[0].expand_as(checksum_stack), rtol=0, atol=1e-5):
|
||||
raise RuntimeError(f"rank outputs diverged: {checksum_stack.cpu().tolist()}")
|
||||
|
||||
peak_memory = torch.tensor(
|
||||
[model_load_peak, conditioning_peak, torch.cuda.max_memory_allocated(device)],
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)
|
||||
memory = [torch.empty_like(peak_memory) for _ in range(world_size)]
|
||||
dist.all_gather(memory, peak_memory)
|
||||
|
||||
if rank == 0:
|
||||
if args.save_latent is not None:
|
||||
args.save_latent.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save({
|
||||
"latent": video.cpu(),
|
||||
"audio_latent": audio.cpu(),
|
||||
"frames": aligned_frames,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"prompt": benchmark["prompt"],
|
||||
"seed": benchmark["seed"],
|
||||
"distributed_mode": args.mode,
|
||||
"world_size": world_size,
|
||||
"attention": args.attention,
|
||||
}, args.save_latent)
|
||||
timing_stack = torch.stack(timings).cpu()
|
||||
report = {
|
||||
"mode": args.mode,
|
||||
"world_size": world_size,
|
||||
"attention": args.attention,
|
||||
"device": torch.cuda.get_device_name(device),
|
||||
"torch": torch.__version__,
|
||||
"benchmark": str(args.benchmark),
|
||||
"resolution": [width, height],
|
||||
"frames": aligned_frames,
|
||||
"steps": benchmark["steps"],
|
||||
"seed": benchmark["seed"],
|
||||
"timings_max_rank_seconds": {
|
||||
"model_load": float(timing_stack[:, 0].max()),
|
||||
"conditioning": float(timing_stack[:, 1].max()),
|
||||
"sampling": float(timing_stack[:, 2].max()),
|
||||
},
|
||||
"peak_allocated_bytes_by_rank": {
|
||||
"model_load": [int(value[0].item()) for value in memory],
|
||||
"conditioning": [int(value[1].item()) for value in memory],
|
||||
"sampling": [int(value[2].item()) for value in memory],
|
||||
},
|
||||
"checksums": checksum_stack[0].cpu().tolist(),
|
||||
"latent": str(args.save_latent) if args.save_latent is not None else None,
|
||||
}
|
||||
serialized = json.dumps(report, indent=2)
|
||||
if args.report is not None:
|
||||
args.report.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.report.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(serialized)
|
||||
|
||||
dist.destroy_process_group()
|
||||
|
|
@ -1,10 +1,18 @@
|
|||
"""List selected H3 checkpoint tensor shapes."""
|
||||
|
||||
import argparse
|
||||
|
||||
from safetensors import safe_open
|
||||
|
||||
|
||||
path = "/models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
|
||||
with safe_open(path, framework="pt", device="cpu") as checkpoint:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
parser.add_argument("--prefix", action="append", default=[])
|
||||
args = parser.parse_args()
|
||||
prefixes = tuple(args.prefix) or ("adaln_t_table", "blocks.0.adaln_proj", "final_layer.adaln_proj")
|
||||
|
||||
with safe_open(args.checkpoint, framework="pt", device="cpu") as checkpoint:
|
||||
for name in checkpoint.keys():
|
||||
if name == "adaln_t_table" or name.startswith("blocks.0.adaln_proj") or name.startswith("final_layer.adaln_proj"):
|
||||
print(name, tuple(checkpoint.get_tensor(name).shape))
|
||||
if any(name == prefix or name.startswith(prefix) for prefix in prefixes):
|
||||
tensor = checkpoint.get_tensor(name)
|
||||
print(name, tuple(tensor.shape), tensor.dtype)
|
||||
|
|
|
|||
27
tools/poll_qwen38_health_spark.sh
Normal file
27
tools/poll_qwen38_health_spark.sh
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
#!/usr/bin/env bash
|
||||
set -u
|
||||
|
||||
for i in $(seq 0 39); do
|
||||
status=$(docker ps -a --filter name=qwen38-vllm --format '{{.Status}}')
|
||||
listening=$(ss -ltnp | grep -c ':8002 ' || true)
|
||||
|
||||
if curl -fsS --max-time 5 http://127.0.0.1:8002/health >/tmp/qwen38-health.out 2>/tmp/qwen38-health.err; then
|
||||
echo "HEALTHY poll=$i status=$status listening=$listening"
|
||||
docker logs --tail 60 qwen38-vllm
|
||||
exit 0
|
||||
fi
|
||||
|
||||
if [[ $status != Up* ]]; then
|
||||
echo "NOT_RUNNING poll=$i status=$status listening=$listening"
|
||||
docker logs --tail 160 qwen38-vllm
|
||||
exit 2
|
||||
fi
|
||||
|
||||
echo "poll=$i status=$status listening=$listening"
|
||||
docker logs --tail 10 qwen38-vllm
|
||||
sleep 30
|
||||
done
|
||||
|
||||
echo "TIMEOUT"
|
||||
docker logs --tail 200 qwen38-vllm
|
||||
exit 1
|
||||
182
tools/probe_audio_vae_boundary.py
Normal file
182
tools/probe_audio_vae_boundary.py
Normal file
|
|
@ -0,0 +1,182 @@
|
|||
"""Decode controlled H3 audio-latent boundary variants for diagnosis."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE
|
||||
|
||||
|
||||
SAMPLE_RATE = 32000
|
||||
SAMPLES_PER_LATENT = 800
|
||||
|
||||
|
||||
def load_latent(path: Path) -> torch.Tensor:
|
||||
state = torch.load(path, map_location="cpu", weights_only=False)
|
||||
latent = state.get("audio_latent") if isinstance(state, dict) else state
|
||||
if latent is None or latent.ndim != 4:
|
||||
raise ValueError(f"{path} does not contain a [B,C,S,T] audio latent")
|
||||
return latent
|
||||
|
||||
|
||||
def dbfs(value: float) -> float:
|
||||
return 20.0 * math.log10(max(value, 1e-20))
|
||||
|
||||
|
||||
def waveform_metrics(waveform: torch.Tensor) -> dict:
|
||||
samples = waveform.float().numpy().T
|
||||
first_100ms = samples[: SAMPLE_RATE // 10]
|
||||
first_500ms = samples[: SAMPLE_RATE // 2]
|
||||
derivatives = np.max(np.abs(np.diff(first_500ms, axis=0)), axis=1)
|
||||
return {
|
||||
"samples": len(samples),
|
||||
"first_sample": samples[0].tolist(),
|
||||
"first_sample_dbfs": [dbfs(float(abs(value))) for value in samples[0]],
|
||||
"first_100ms_peak_dbfs": dbfs(float(np.max(np.abs(first_100ms)))),
|
||||
"first_100ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_100ms**2)))),
|
||||
"first_500ms_peak_dbfs": dbfs(float(np.max(np.abs(first_500ms)))),
|
||||
"first_500ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_500ms**2)))),
|
||||
"largest_derivative": float(derivatives.max()),
|
||||
"largest_derivative_ms": float((derivatives.argmax() + 1) * 1000.0 / SAMPLE_RATE),
|
||||
}
|
||||
|
||||
|
||||
def comparison(reference: torch.Tensor, candidate: torch.Tensor) -> dict:
|
||||
count = min(reference.shape[-1], candidate.shape[-1])
|
||||
reference = reference[..., :count].float()
|
||||
candidate = candidate[..., :count].float()
|
||||
|
||||
def region_metrics(samples: int) -> dict:
|
||||
ref = reference[..., :samples]
|
||||
test = candidate[..., :samples]
|
||||
error = test - ref
|
||||
signal_power = ref.square().mean(dim=-1)
|
||||
noise_power = error.square().mean(dim=-1)
|
||||
psnr = 10.0 * torch.log10(signal_power.clamp_min(1e-30) / noise_power.clamp_min(1e-30))
|
||||
return {
|
||||
"rmse": float(error.square().mean().sqrt()),
|
||||
"max_abs": float(error.abs().max()),
|
||||
"psnr_db_by_channel": psnr.flatten().tolist(),
|
||||
}
|
||||
|
||||
return {
|
||||
"first_100ms": region_metrics(SAMPLE_RATE // 10),
|
||||
"first_500ms": region_metrics(SAMPLE_RATE // 2),
|
||||
"full": region_metrics(count),
|
||||
}
|
||||
|
||||
|
||||
def write_waveform(path: Path, waveform: torch.Tensor) -> None:
|
||||
raw = path.with_suffix(".f32le")
|
||||
waveform.transpose(0, 1).contiguous().numpy().tofile(raw)
|
||||
subprocess.run([
|
||||
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
||||
"-f", "f32le", "-ar", str(SAMPLE_RATE), "-ac", "2", "-i", str(raw),
|
||||
"-c:a", "pcm_f32le", str(path),
|
||||
], check=True)
|
||||
raw.unlink()
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--affected-latent", type=Path, required=True)
|
||||
parser.add_argument("--clean-latent", type=Path, required=True)
|
||||
parser.add_argument("--vae", type=Path, default=Path("/vae/minimax_h3_audio_vae_fp32.safetensors"))
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument("--report", type=Path)
|
||||
args = parser.parse_args()
|
||||
|
||||
affected = load_latent(args.affected_latent)
|
||||
clean = load_latent(args.clean_latent)
|
||||
if affected.shape != clean.shape:
|
||||
raise ValueError(f"latent shapes differ: {tuple(affected.shape)} != {tuple(clean.shape)}")
|
||||
|
||||
boundary_frames = 4
|
||||
variants = {
|
||||
"affected-original": (affected, 0),
|
||||
"clean-original": (clean, 0),
|
||||
"zero-normalized-latent": (torch.zeros_like(affected), 0),
|
||||
"affected-repeat-frame0": (affected[..., :1].expand_as(affected).clone(), 0),
|
||||
"affected-repeat-frame4": (affected[..., 4:5].expand_as(affected).clone(), 0),
|
||||
}
|
||||
|
||||
silent_carrier = affected[..., 4:5].expand_as(affected).clone()
|
||||
carrier_start = silent_carrier.clone()
|
||||
carrier_start[..., :boundary_frames] = affected[..., :boundary_frames]
|
||||
variants["carrier-affected-first4-at-start"] = (carrier_start, 0)
|
||||
|
||||
interior_frame = 40
|
||||
carrier_interior = silent_carrier.clone()
|
||||
carrier_interior[..., interior_frame : interior_frame + boundary_frames] = affected[..., :boundary_frames]
|
||||
variants["carrier-affected-first4-at-frame40"] = (carrier_interior, 0)
|
||||
|
||||
replaced_with_frame4 = affected.clone()
|
||||
replaced_with_frame4[..., :boundary_frames] = affected[..., 4:5]
|
||||
variants["affected-first4-repeat-frame4"] = (replaced_with_frame4, 0)
|
||||
|
||||
affected_with_clean = affected.clone()
|
||||
affected_with_clean[..., :boundary_frames] = clean[..., :boundary_frames]
|
||||
variants["affected-first4-from-clean"] = (affected_with_clean, 0)
|
||||
|
||||
clean_with_affected = clean.clone()
|
||||
clean_with_affected[..., :boundary_frames] = affected[..., :boundary_frames]
|
||||
variants["clean-first4-from-affected"] = (clean_with_affected, 0)
|
||||
|
||||
prefix_repeat = affected[..., :1].expand(*affected.shape[:-1], boundary_frames)
|
||||
variants["affected-prefix-repeat-frame0"] = (
|
||||
torch.cat((prefix_repeat, affected), dim=-1),
|
||||
boundary_frames * SAMPLES_PER_LATENT,
|
||||
)
|
||||
variants["affected-prefix-own-first4"] = (
|
||||
torch.cat((affected[..., :boundary_frames], affected), dim=-1),
|
||||
boundary_frames * SAMPLES_PER_LATENT,
|
||||
)
|
||||
|
||||
vae = MiniMaxH3AudioVAE.from_safetensors(args.vae, device="cuda").eval()
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
decoded = {}
|
||||
report = {"boundary_frames": boundary_frames, "variants": {}}
|
||||
with torch.inference_mode():
|
||||
for name, (latent, crop_start) in variants.items():
|
||||
waveform = vae.decode(latent.to("cuda", dtype=next(vae.parameters()).dtype)).cpu()[0]
|
||||
waveform = waveform[:, crop_start : crop_start + affected.shape[-1] * SAMPLES_PER_LATENT]
|
||||
decoded[name] = waveform
|
||||
output = args.output_dir / f"{name}.wav"
|
||||
write_waveform(output, waveform)
|
||||
report["variants"][name] = {
|
||||
"output": str(output),
|
||||
"crop_start_samples": crop_start,
|
||||
"metrics": waveform_metrics(waveform),
|
||||
}
|
||||
|
||||
affected_reference = decoded["affected-original"]
|
||||
for name, waveform in decoded.items():
|
||||
if name != "affected-original":
|
||||
report["variants"][name]["difference_from_affected_original"] = comparison(
|
||||
affected_reference, waveform,
|
||||
)
|
||||
|
||||
interior_start = interior_frame * SAMPLES_PER_LATENT
|
||||
segment_samples = boundary_frames * SAMPLES_PER_LATENT
|
||||
report["interior_placement"] = {
|
||||
"frame": interior_frame,
|
||||
"start_seconds": interior_start / SAMPLE_RATE,
|
||||
"affected_onset_vs_carrier_interior_event": comparison(
|
||||
affected_reference[..., :segment_samples],
|
||||
decoded["carrier-affected-first4-at-frame40"][..., interior_start : interior_start + segment_samples],
|
||||
),
|
||||
"carrier_start_event_vs_carrier_interior_event": comparison(
|
||||
decoded["carrier-affected-first4-at-start"][..., :segment_samples],
|
||||
decoded["carrier-affected-first4-at-frame40"][..., interior_start : interior_start + segment_samples],
|
||||
),
|
||||
}
|
||||
|
||||
serialized = json.dumps(report, indent=2)
|
||||
if args.report is not None:
|
||||
args.report.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.report.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(serialized)
|
||||
20
tools/run_distributed_matrix.sh
Normal file
20
tools/run_distributed_matrix.sh
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
available="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
||||
if [[ "$available" -lt 1 ]]; then
|
||||
echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
read -ra counts <<< "${H3_GPU_COUNTS:-1 2 4 6 8}"
|
||||
read -ra modes <<< "${H3_DISTRIBUTED_MODES:-ulysses tensor}"
|
||||
for count in "${counts[@]}"; do
|
||||
if [[ "$count" -gt "$available" ]]; then
|
||||
echo "Skipping ${count} GPUs; only ${available} are visible." >&2
|
||||
continue
|
||||
fi
|
||||
for mode in "${modes[@]}"; do
|
||||
tools/run_distributed_t2va.sh "$mode" "$count" "${H3_ATTENTION:-sdpa}"
|
||||
done
|
||||
done
|
||||
35
tools/run_distributed_t2va.sh
Normal file
35
tools/run_distributed_t2va.sh
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
mode="${1:?usage: run_distributed_t2va.sh ulysses|tensor [WORLD_SIZE] [ATTENTION]}"
|
||||
world_size="${2:-${H3_WORLD_SIZE:-}}"
|
||||
if [[ -z "$world_size" ]]; then
|
||||
world_size="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
||||
fi
|
||||
if [[ "$world_size" -lt 1 ]]; then
|
||||
echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
|
||||
exit 1
|
||||
fi
|
||||
attention="${3:-sdpa}"
|
||||
benchmark="${H3_DISTRIBUTED_BENCHMARK:-benchmarks/t2va-dialogue-quoted-864x480-141f-base12-sage2-seed440420.json}"
|
||||
output_root="${H3_DISTRIBUTED_OUTPUT:-/output/h3-baselines}"
|
||||
stem="distributed-${mode}-${world_size}gpu-${attention}"
|
||||
model="${H3_MODEL_PATH:-/models/minimax_h3_fl2va_pruned_nvfp4.safetensors}"
|
||||
text_encoder="${H3_TEXT_ENCODER_PATH:-/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors}"
|
||||
extra_args=()
|
||||
if [[ "${H3_SAVE_LATENTS:-1}" == "1" ]]; then
|
||||
extra_args+=(--save-latent "$output_root/$stem.latent.pt")
|
||||
fi
|
||||
|
||||
torchrun \
|
||||
--standalone \
|
||||
--nnodes=1 \
|
||||
--nproc-per-node="$world_size" \
|
||||
tools/distributed_t2va.py \
|
||||
--benchmark "$benchmark" \
|
||||
--mode "$mode" \
|
||||
--attention "$attention" \
|
||||
--model "$model" \
|
||||
--text-encoder "$text_encoder" \
|
||||
--report "$output_root/$stem.json" \
|
||||
"${extra_args[@]}"
|
||||
23
tools/run_ulysses_benchmark.sh
Normal file
23
tools/run_ulysses_benchmark.sh
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
world_size="${1:-${H3_WORLD_SIZE:-}}"
|
||||
if [[ -z "$world_size" ]]; then
|
||||
world_size="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
||||
fi
|
||||
if [[ "$world_size" -lt 1 ]]; then
|
||||
echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
|
||||
exit 1
|
||||
fi
|
||||
backend="${2:-sdpa}"
|
||||
sequence="${3:-20000}"
|
||||
output="${H3_BENCHMARK_OUTPUT:-/output/h3-baselines/ulysses-${world_size}gpu-${backend}-${sequence}t.json}"
|
||||
|
||||
torchrun \
|
||||
--standalone \
|
||||
--nnodes=1 \
|
||||
--nproc-per-node="$world_size" \
|
||||
tools/benchmark_ulysses.py \
|
||||
--sequence "$sequence" \
|
||||
--backend "$backend" \
|
||||
--output "$output"
|
||||
100
tools/runpod_api.py
Normal file
100
tools/runpod_api.py
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
"""Minimal RunPod API v2 client for the H3 single-node benchmark pod."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
|
||||
|
||||
API = "https://api.runpod.io/v2"
|
||||
BLACKWELL_GPUS = (
|
||||
"NVIDIA RTX PRO 6000 Blackwell Server Edition",
|
||||
"NVIDIA RTX PRO 6000 Blackwell Workstation Edition",
|
||||
"NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition",
|
||||
)
|
||||
DEFAULT_IMAGE = "runpod/pytorch:1.1.0-cu1300-torch291-ubuntu2404"
|
||||
|
||||
|
||||
def request(method: str, path: str, body=None, query=None):
|
||||
key = os.environ.get("RUNPOD_API_KEY")
|
||||
if not key:
|
||||
raise SystemExit("RUNPOD_API_KEY is required")
|
||||
url = f"{API}{path}"
|
||||
if query:
|
||||
url += "?" + urllib.parse.urlencode(query)
|
||||
data = None if body is None else json.dumps(body).encode("utf-8")
|
||||
call = urllib.request.Request(url, data=data, method=method)
|
||||
call.add_header("Authorization", f"Bearer {key}")
|
||||
call.add_header("Accept", "application/json")
|
||||
if data is not None:
|
||||
call.add_header("Content-Type", "application/json")
|
||||
try:
|
||||
with urllib.request.urlopen(call, timeout=60) as response:
|
||||
return json.load(response)
|
||||
except urllib.error.HTTPError as error:
|
||||
detail = error.read().decode("utf-8", errors="replace")
|
||||
raise SystemExit(f"RunPod API returned HTTP {error.code}: {detail}") from error
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
commands = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
catalog_parser = commands.add_parser("catalog")
|
||||
catalog_parser.add_argument("--count", type=int, default=8)
|
||||
catalog_parser.add_argument("--cloud", choices=("SECURE", "COMMUNITY"), default="SECURE")
|
||||
|
||||
create_parser = commands.add_parser("create")
|
||||
create_parser.add_argument("--gpu", choices=BLACKWELL_GPUS, default=BLACKWELL_GPUS[0])
|
||||
create_parser.add_argument("--count", type=int, default=8)
|
||||
create_parser.add_argument("--cloud", choices=("SECURE", "COMMUNITY"), default="SECURE")
|
||||
create_parser.add_argument("--datacenter")
|
||||
create_parser.add_argument("--image", default=DEFAULT_IMAGE)
|
||||
create_parser.add_argument("--disk", type=int, default=100)
|
||||
create_parser.add_argument("--volume", type=int, default=100)
|
||||
create_parser.add_argument("--yes", action="store_true")
|
||||
|
||||
get_parser = commands.add_parser("get")
|
||||
get_parser.add_argument("pod_id")
|
||||
|
||||
terminate_parser = commands.add_parser("terminate")
|
||||
terminate_parser.add_argument("pod_id")
|
||||
terminate_parser.add_argument("--yes", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
if args.command == "catalog":
|
||||
response = request("GET", "/catalog/gpus", query={
|
||||
"include": "AVAILABILITY",
|
||||
"product": "POD",
|
||||
"count": args.count,
|
||||
"cloud": args.cloud,
|
||||
"minCudaVersion": "12.8",
|
||||
})
|
||||
response["gpus"] = [gpu for gpu in response["gpus"] if gpu["id"] in BLACKWELL_GPUS]
|
||||
elif args.command == "create":
|
||||
if not args.yes:
|
||||
raise SystemExit("create rents billable GPUs; repeat with --yes after checking catalog")
|
||||
body = {
|
||||
"name": "h3-blackwell-distributed",
|
||||
"image": args.image,
|
||||
"gpu": {"id": args.gpu, "count": args.count, "minCudaVersion": "12.8"},
|
||||
"cloud": args.cloud,
|
||||
"disk": args.disk,
|
||||
"ports": ["22/tcp"],
|
||||
"mounts": {"persistent": {"size": args.volume, "path": "/workspace"}},
|
||||
"startSsh": True,
|
||||
}
|
||||
if args.datacenter:
|
||||
body["dataCenterIds"] = [args.datacenter]
|
||||
response = request("POST", "/pods", body=body)
|
||||
elif args.command == "get":
|
||||
response = request("GET", f"/pods/{args.pod_id}")
|
||||
else:
|
||||
if not args.yes:
|
||||
raise SystemExit("termination is irreversible; repeat with --yes")
|
||||
response = request("POST", f"/pods/{args.pod_id}/actions", body={"action": "terminate"})
|
||||
|
||||
json.dump(response, sys.stdout, indent=2)
|
||||
sys.stdout.write("\n")
|
||||
|
|
@ -14,7 +14,7 @@ import torch
|
|||
from PIL import Image
|
||||
|
||||
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND, attention_backend_status
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig, TURBO_VARIANTS
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig, TURBO_VARIANTS, normalize_upscale
|
||||
|
||||
|
||||
def _load_image(value) -> torch.Tensor | None:
|
||||
|
|
@ -57,6 +57,7 @@ parser.add_argument("--vae-tile-size", type=int, default=256)
|
|||
parser.add_argument("--mlp-chunks", type=int, default=1)
|
||||
parser.add_argument("--mlp-chunk-threshold", type=int, default=4096)
|
||||
parser.add_argument("--turbo-lora", action="append", default=[], metavar="NAME=PATH", help="Load a resident 4step or 8step Turbo adapter.")
|
||||
parser.add_argument("--latent-upscaler", help="Load the optional H3 3D latent upscaler for request-level spatial upscaling.")
|
||||
parser.add_argument("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.")
|
||||
parser.add_argument("--warmup-output", type=Path, default=Path("/output/h3-blackwell-runtime/hot-runtime-warmup.mp4"))
|
||||
args = parser.parse_args()
|
||||
|
|
@ -68,7 +69,7 @@ for value in args.turbo_lora:
|
|||
name, path = value.split("=", 1)
|
||||
turbo_loras.append((name, path))
|
||||
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold, turbo_loras=tuple(turbo_loras)))
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold, turbo_loras=tuple(turbo_loras), latent_upscaler_path=args.latent_upscaler))
|
||||
runtime_lock = threading.Lock()
|
||||
warmup_result = None
|
||||
if args.warmup:
|
||||
|
|
@ -142,6 +143,15 @@ class Handler(BaseHTTPRequestHandler):
|
|||
write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": list(AVAILABLE_BACKENDS)})
|
||||
return
|
||||
mux_audio = bool(payload.get("mux_audio", True))
|
||||
keep_intermediates = bool(payload.get("keep_intermediates", False))
|
||||
try:
|
||||
upscale = normalize_upscale(payload.get("upscale"))
|
||||
except ValueError as exc:
|
||||
write_json(self, 400, {"error": "invalid upscale", "message": str(exc)})
|
||||
return
|
||||
if upscale is not None and runtime.latent_upscaler is None:
|
||||
write_json(self, 400, {"error": "H3 latent upscaler is not loaded"})
|
||||
return
|
||||
ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error")
|
||||
first_frame = _load_image(payload.get("first_frame"))
|
||||
last_frame = _load_image(payload.get("last_frame"))
|
||||
|
|
@ -166,9 +176,11 @@ class Handler(BaseHTTPRequestHandler):
|
|||
seed=seed,
|
||||
attention=attention,
|
||||
turbo=turbo,
|
||||
upscale=upscale,
|
||||
first_frame=first_frame,
|
||||
last_frame=last_frame,
|
||||
mux_audio=mux_audio,
|
||||
keep_intermediates=keep_intermediates,
|
||||
ffmpeg_loglevel=ffmpeg_loglevel,
|
||||
save_latent=save_latent,
|
||||
cache_mode=cache_mode,
|
||||
|
|
|
|||
50
tools/smoke_flash4.py
Normal file
50
tools/smoke_flash4.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
"""Compare FlashAttention-4 against PyTorch SDPA on an H3-shaped operation."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--sequence", type=int, default=257)
|
||||
parser.add_argument("--heads", type=int, default=8)
|
||||
parser.add_argument("--head-dim", type=int, default=128)
|
||||
parser.add_argument("--iterations", type=int, default=5)
|
||||
args = parser.parse_args()
|
||||
|
||||
from flash_attn.cute import flash_attn_func
|
||||
|
||||
torch.manual_seed(440411)
|
||||
shape = (1, args.sequence, args.heads, args.head_dim)
|
||||
q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
with torch.inference_mode():
|
||||
result = flash_attn_func(q, k, v, causal=False)
|
||||
actual = result[0] if isinstance(result, tuple) else result
|
||||
expected = torch.nn.functional.scaled_dot_product_attention(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=False,
|
||||
).transpose(1, 2)
|
||||
torch.cuda.synchronize()
|
||||
started = torch.cuda.Event(enable_timing=True)
|
||||
finished = torch.cuda.Event(enable_timing=True)
|
||||
started.record()
|
||||
for _ in range(args.iterations):
|
||||
result = flash_attn_func(q, k, v, causal=False)
|
||||
actual = result[0] if isinstance(result, tuple) else result
|
||||
finished.record()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
error = (actual.float() - expected.float()).abs()
|
||||
print(json.dumps({
|
||||
"device": torch.cuda.get_device_name(),
|
||||
"capability": torch.cuda.get_device_capability(),
|
||||
"shape": tuple(actual.shape),
|
||||
"dtype": str(actual.dtype),
|
||||
"contiguous": actual.is_contiguous(),
|
||||
"max_abs_error": error.max().item(),
|
||||
"mean_abs_error": error.mean().item(),
|
||||
"milliseconds": started.elapsed_time(finished) / args.iterations,
|
||||
}, indent=2))
|
||||
8
tools/smoke_qwen38_vllm_spark.sh
Normal file
8
tools/smoke_qwen38_vllm_spark.sh
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
curl -fsS --max-time 120 http://127.0.0.1:8002/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
--data-binary @- <<'JSON'
|
||||
{"model":"qwen38","messages":[{"role":"user","content":"Reply with exactly: ready"}],"max_tokens":64,"temperature":0,"chat_template_kwargs":{"enable_thinking":false}}
|
||||
JSON
|
||||
61
tools/spark_stack_spark.sh
Normal file
61
tools/spark_stack_spark.sh
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
ROOT=/home/daniel/aeon-spark-test/h3
|
||||
STACK_DIR="$ROOT/spark-stack"
|
||||
COMPOSE_FILE="$STACK_DIR/compose.spark-stack.yml"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
spark_stack_spark.sh start h3|qwen|flux
|
||||
spark_stack_spark.sh stop
|
||||
spark_stack_spark.sh status
|
||||
|
||||
Ports:
|
||||
h3 -> http://192.168.1.162:8001
|
||||
qwen -> http://192.168.1.162:8002
|
||||
flux -> http://192.168.1.162:8192
|
||||
EOF
|
||||
}
|
||||
|
||||
managed_containers=(
|
||||
h3-hot-runtime
|
||||
h3-blackwell-runtime-h3-hot-runtime-1
|
||||
qwen38-vllm
|
||||
flux2-klein-comfy
|
||||
flux-nvfp4-test
|
||||
)
|
||||
|
||||
stop_managed() {
|
||||
for name in "${managed_containers[@]}"; do
|
||||
if docker ps -a --format '{{.Names}}' | grep -qx "$name"; then
|
||||
docker stop "$name" >/dev/null 2>&1 || true
|
||||
docker rm "$name" >/dev/null 2>&1 || true
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
case "${1:-}" in
|
||||
start)
|
||||
target="${2:-}"
|
||||
case "$target" in
|
||||
h3|qwen|flux) ;;
|
||||
*) usage; exit 2 ;;
|
||||
esac
|
||||
mkdir -p "$ROOT/vllm-cache"
|
||||
stop_managed
|
||||
docker compose -f "$COMPOSE_FILE" --profile "$target" up -d
|
||||
;;
|
||||
stop)
|
||||
stop_managed
|
||||
;;
|
||||
status)
|
||||
docker ps -a --format 'table {{.Names}}\t{{.Status}}\t{{.Ports}}' | grep -E 'NAMES|h3-hot-runtime|qwen38-vllm|flux2-klein-comfy|flux-nvfp4-test' || true
|
||||
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits || true
|
||||
;;
|
||||
*)
|
||||
usage
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
27
tools/start_qwen38_vllm_spark.sh
Normal file
27
tools/start_qwen38_vllm_spark.sh
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
ROOT=/home/daniel/aeon-spark-test/h3
|
||||
MODEL_DIR="$ROOT/models/Qwen3.8-27B-NVFP4"
|
||||
COMPOSE_FILE="$ROOT/qwen38-vllm/compose.qwen38-vllm.yml"
|
||||
|
||||
if [[ ! -f "$MODEL_DIR/model.safetensors" ]]; then
|
||||
echo "model.safetensors is not present yet: $MODEL_DIR/model.safetensors" >&2
|
||||
echo "Wait for qwen38-27b-nvfp4-download to finish before starting vLLM." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 - <<'PY'
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
tokenizer = Path('/home/daniel/aeon-spark-test/h3/models/Qwen3.8-27B-NVFP4/tokenizer.json')
|
||||
truncation = json.loads(tokenizer.read_text(encoding='utf-8')).get('truncation')
|
||||
if truncation is not None:
|
||||
raise SystemExit(f'tokenizer truncation must be null, got: {truncation!r}')
|
||||
print('tokenizer truncation: null')
|
||||
PY
|
||||
|
||||
mkdir -p "$ROOT/vllm-cache"
|
||||
docker compose -f "$COMPOSE_FILE" up -d
|
||||
docker logs -f qwen38-vllm
|
||||
91
tools/submit_h3_comfy_t2v.py
Normal file
91
tools/submit_h3_comfy_t2v.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
"""Submit a prompt-only MiniMax H3 T2V graph to ComfyUI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict) -> dict:
|
||||
request = Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
|
||||
with urlopen(request, timeout=30) as response:
|
||||
return json.loads(response.read().decode())
|
||||
|
||||
|
||||
def get_json(url: str) -> dict:
|
||||
with urlopen(url, timeout=30) as response:
|
||||
return json.loads(response.read().decode())
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--server", default="http://127.0.0.1:8188")
|
||||
parser.add_argument("--prompt", required=True)
|
||||
parser.add_argument("--filename-prefix", required=True)
|
||||
parser.add_argument("--seed", type=int, default=440407)
|
||||
parser.add_argument("--width", type=int, default=960)
|
||||
parser.add_argument("--height", type=int, default=544)
|
||||
parser.add_argument("--frames", type=int, default=124)
|
||||
parser.add_argument("--steps", type=int, default=12)
|
||||
parser.add_argument("--upscale", type=float)
|
||||
args = parser.parse_args()
|
||||
|
||||
graph = {
|
||||
"1": {"class_type": "UNETLoader", "inputs": {"unet_name": "minimax_h3_fl2va_pruned_nvfp4.safetensors", "weight_dtype": "default"}},
|
||||
"3": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}},
|
||||
"4": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}},
|
||||
"5": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_audio_vae_fp32.safetensors"}},
|
||||
"8": {"class_type": "MiniMaxH3ImageToVideo", "inputs": {"clip": ["3", 0], "vae": ["4", 0], "prompt": args.prompt, "width": args.width, "height": args.height, "length": args.frames}},
|
||||
"9": {"class_type": "BasicGuider", "inputs": {"model": ["1", 0], "conditioning": ["8", 0]}},
|
||||
"10": {"class_type": "RandomNoise", "inputs": {"noise_seed": args.seed}},
|
||||
"11": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "res_multistep"}},
|
||||
"12": {"class_type": "BasicScheduler", "inputs": {"model": ["1", 0], "scheduler": "beta", "steps": args.steps, "denoise": 1.0}},
|
||||
"13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["10", 0], "guider": ["9", 0], "sampler": ["11", 0], "sigmas": ["12", 0], "latent_image": ["8", 1]}},
|
||||
}
|
||||
decoded_latent = ["13", 0]
|
||||
if args.upscale is not None:
|
||||
graph.update({
|
||||
"18": {"class_type": "LTXVSeparateAVLatent", "inputs": {"av_latent": ["13", 0]}},
|
||||
"19": {"class_type": "MinimaxH3LatentUpscaler3D", "inputs": {
|
||||
"latent": ["18", 0],
|
||||
"model_name": "minimax_h3_latent_upscaler_3d_fp16.safetensors",
|
||||
"mode": "scale by multiplier",
|
||||
"mode.scale": args.upscale,
|
||||
"align": 32,
|
||||
"keep_proportion": True,
|
||||
"device": "cuda",
|
||||
"precision": "fp16",
|
||||
}},
|
||||
"20": {"class_type": "LTXVConcatAVLatent", "inputs": {"video_latent": ["19", 0], "audio_latent": ["18", 1]}},
|
||||
})
|
||||
decoded_latent = ["20", 0]
|
||||
graph.update({
|
||||
"14": {"class_type": "VAEDecode", "inputs": {"samples": decoded_latent, "vae": ["4", 0]}},
|
||||
"15": {"class_type": "VAEDecodeAudio", "inputs": {"samples": decoded_latent, "vae": ["5", 0]}},
|
||||
"16": {"class_type": "CreateVideo", "inputs": {"images": ["14", 0], "audio": ["15", 0], "bit_depth": 8, "fps": 24.0}},
|
||||
"17": {"class_type": "SaveVideo", "inputs": {"video": ["16", 0], "filename_prefix": args.filename_prefix, "format": "mp4", "codec": "auto"}},
|
||||
})
|
||||
|
||||
started = time.perf_counter()
|
||||
response = post_json(f"{args.server}/prompt", {"prompt": graph})
|
||||
prompt_id = response["prompt_id"]
|
||||
while True:
|
||||
time.sleep(1)
|
||||
history = get_json(f"{args.server}/history/{prompt_id}").get(prompt_id)
|
||||
if history is None:
|
||||
continue
|
||||
status = history.get("status", {})
|
||||
if status.get("completed") or status.get("status_str") in {"success", "error"}:
|
||||
break
|
||||
print(json.dumps({
|
||||
"prompt_id": prompt_id,
|
||||
"wall_seconds": time.perf_counter() - started,
|
||||
"status": history.get("status", {}),
|
||||
"outputs": history.get("outputs", {}),
|
||||
}, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
186
tools/sweep_dialogue_audio_boundary.py
Normal file
186
tools/sweep_dialogue_audio_boundary.py
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
"""Run a paired tagged-versus-quoted H3 dialogue audio sweep."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
|
||||
from h3_blackwell_runtime.sampler import sample_video_res_multistep
|
||||
from h3_blackwell_runtime.t2v import random_av_latents
|
||||
|
||||
|
||||
SAMPLE_RATE = 32000
|
||||
|
||||
|
||||
def dbfs(value: float) -> float:
|
||||
return 20.0 * math.log10(max(value, 1e-20))
|
||||
|
||||
|
||||
def waveform_metrics(waveform: torch.Tensor) -> dict:
|
||||
waveform = waveform.float()
|
||||
first_100ms = waveform[..., :3200]
|
||||
next_400ms = waveform[..., 3200:16000]
|
||||
first_500ms = waveform[..., :16000]
|
||||
derivatives = (first_500ms[..., 1:] - first_500ms[..., :-1]).abs()
|
||||
windows = waveform.unfold(-1, 320, 320)
|
||||
window_rms = windows.square().mean(dim=(0, 2)).sqrt()
|
||||
active = (20.0 * torch.log10(window_rms.clamp_min(1e-20)) > -40.0).nonzero()
|
||||
first_active_ms = None if active.numel() == 0 else int(active[0, 0]) * 10
|
||||
|
||||
first_rms = dbfs(float(first_100ms.square().mean().sqrt()))
|
||||
next_rms = dbfs(float(next_400ms.square().mean().sqrt()))
|
||||
return {
|
||||
"first_sample": waveform[..., 0].flatten().tolist(),
|
||||
"first_100ms_peak_dbfs": dbfs(float(first_100ms.abs().max())),
|
||||
"first_100ms_rms_dbfs": first_rms,
|
||||
"next_400ms_peak_dbfs": dbfs(float(next_400ms.abs().max())),
|
||||
"next_400ms_rms_dbfs": next_rms,
|
||||
"boundary_decay_db": first_rms - next_rms,
|
||||
"first_500ms_peak_dbfs": dbfs(float(first_500ms.abs().max())),
|
||||
"first_500ms_rms_dbfs": dbfs(float(first_500ms.square().mean().sqrt())),
|
||||
"full_peak_dbfs": dbfs(float(waveform.abs().max())),
|
||||
"full_rms_dbfs": dbfs(float(waveform.square().mean().sqrt())),
|
||||
"largest_first_500ms_derivative": float(derivatives.max()),
|
||||
"first_10ms_window_above_minus_40_dbfs_ms": first_active_ms,
|
||||
}
|
||||
|
||||
|
||||
def latent_metrics(latent: torch.Tensor) -> dict:
|
||||
frames = latent.float().movedim(-1, 0).flatten(1)
|
||||
return {
|
||||
"shape": list(latent.shape),
|
||||
"first_4_rms": float(frames[:4].square().mean().sqrt()),
|
||||
"frames_4_20_rms": float(frames[4:20].square().mean().sqrt()),
|
||||
"first_frame_rms": float(frames[0].square().mean().sqrt()),
|
||||
"frame_0_to_1_delta_rms": float((frames[1] - frames[0]).square().mean().sqrt()),
|
||||
}
|
||||
|
||||
|
||||
def write_waveform(path: Path, waveform: torch.Tensor) -> None:
|
||||
raw = path.with_suffix(".f32le")
|
||||
waveform.transpose(0, 1).contiguous().numpy().tofile(raw)
|
||||
subprocess.run([
|
||||
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
||||
"-f", "f32le", "-ar", str(SAMPLE_RATE), "-ac", "2", "-i", str(raw),
|
||||
"-c:a", "pcm_f32le", str(path),
|
||||
], check=True)
|
||||
raw.unlink()
|
||||
|
||||
|
||||
def save_report(path: Path, report: dict) -> None:
|
||||
temporary = path.with_suffix(path.suffix + ".tmp")
|
||||
temporary.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
||||
temporary.replace(path)
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--tagged-benchmark", type=Path, required=True)
|
||||
parser.add_argument("--quoted-benchmark", type=Path, required=True)
|
||||
parser.add_argument("--seed-start", type=int, default=440420)
|
||||
parser.add_argument("--seed-count", type=int, default=10)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument("--report", type=Path, required=True)
|
||||
parser.add_argument("--attention", default="sage2")
|
||||
args = parser.parse_args()
|
||||
|
||||
tagged = json.loads(args.tagged_benchmark.read_text(encoding="utf-8"))
|
||||
quoted = json.loads(args.quoted_benchmark.read_text(encoding="utf-8"))
|
||||
for field in ("resolution", "frames", "steps"):
|
||||
if tagged[field] != quoted[field]:
|
||||
raise ValueError(f"benchmark {field} differs: {tagged[field]} != {quoted[field]}")
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
args.report.parent.mkdir(parents=True, exist_ok=True)
|
||||
if args.report.exists():
|
||||
report = json.loads(args.report.read_text(encoding="utf-8"))
|
||||
else:
|
||||
report = {
|
||||
"tagged_benchmark": str(args.tagged_benchmark),
|
||||
"quoted_benchmark": str(args.quoted_benchmark),
|
||||
"attention": args.attention,
|
||||
"seed_start": args.seed_start,
|
||||
"seed_count": args.seed_count,
|
||||
"cases": {},
|
||||
"pairs": {},
|
||||
}
|
||||
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention))
|
||||
conditioned = {
|
||||
"tagged": runtime.refiner(runtime.conditioner(tagged["prompt"])),
|
||||
"quoted": runtime.refiner(runtime.conditioner(quoted["prompt"])),
|
||||
}
|
||||
|
||||
width, height = tagged["resolution"]
|
||||
for seed in range(args.seed_start, args.seed_start + args.seed_count):
|
||||
for prompt_format, benchmark in (("tagged", tagged), ("quoted", quoted)):
|
||||
key = f"{seed}:{prompt_format}"
|
||||
if key in report["cases"]:
|
||||
print(f"skip completed {key}", flush=True)
|
||||
continue
|
||||
|
||||
started = time.perf_counter()
|
||||
video, audio, aligned_frames = random_av_latents(
|
||||
width, height, benchmark["frames"], seed, device=runtime.config.device,
|
||||
)
|
||||
sampled_video, audio_latent = sample_video_res_multistep(
|
||||
runtime.model,
|
||||
runtime.packer,
|
||||
conditioned[prompt_format],
|
||||
video,
|
||||
audio,
|
||||
steps=benchmark["steps"],
|
||||
seed=seed,
|
||||
return_audio=True,
|
||||
)
|
||||
with torch.inference_mode():
|
||||
waveform = runtime.audio_vae.decode(
|
||||
audio_latent.to("cuda", dtype=next(runtime.audio_vae.parameters()).dtype),
|
||||
).cpu()[0]
|
||||
|
||||
stem = f"dialogue-{prompt_format}-base12-sage2-seed{seed}"
|
||||
wav_path = args.output_dir / f"{stem}.wav"
|
||||
latent_path = args.output_dir / f"{stem}.audio-latent.pt"
|
||||
write_waveform(wav_path, waveform)
|
||||
torch.save({
|
||||
"audio_latent": audio_latent.detach().cpu(),
|
||||
"prompt_format": prompt_format,
|
||||
"prompt": benchmark["prompt"],
|
||||
"seed": seed,
|
||||
}, latent_path)
|
||||
|
||||
report["cases"][key] = {
|
||||
"seed": seed,
|
||||
"prompt_format": prompt_format,
|
||||
"wav": str(wav_path),
|
||||
"audio_latent": str(latent_path),
|
||||
"frames": aligned_frames,
|
||||
"seconds": time.perf_counter() - started,
|
||||
"waveform": waveform_metrics(waveform),
|
||||
"latent": latent_metrics(audio_latent.cpu()),
|
||||
}
|
||||
del sampled_video, audio_latent, waveform, video, audio
|
||||
save_report(args.report, report)
|
||||
print(json.dumps(report["cases"][key]), flush=True)
|
||||
|
||||
tagged_case = report["cases"][f"{seed}:tagged"]
|
||||
quoted_case = report["cases"][f"{seed}:quoted"]
|
||||
report["pairs"][str(seed)] = {
|
||||
"quoted_peak_reduction_db": (
|
||||
tagged_case["waveform"]["first_100ms_peak_dbfs"]
|
||||
- quoted_case["waveform"]["first_100ms_peak_dbfs"]
|
||||
),
|
||||
"quoted_rms_reduction_db": (
|
||||
tagged_case["waveform"]["first_100ms_rms_dbfs"]
|
||||
- quoted_case["waveform"]["first_100ms_rms_dbfs"]
|
||||
),
|
||||
"tagged_boundary_decay_db": tagged_case["waveform"]["boundary_decay_db"],
|
||||
"quoted_boundary_decay_db": quoted_case["waveform"]["boundary_decay_db"],
|
||||
}
|
||||
save_report(args.report, report)
|
||||
|
||||
print(json.dumps(report["pairs"], indent=2))
|
||||
106
tools/trace_audio_denoising.py
Normal file
106
tools/trace_audio_denoising.py
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
"""Trace when an H3 audio-boundary artifact emerges during base sampling."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
|
||||
from h3_blackwell_runtime.sampler import _decode_audio_latent, sample_video_res_multistep
|
||||
from h3_blackwell_runtime.t2v import random_av_latents
|
||||
|
||||
|
||||
def dbfs(value: float) -> float:
|
||||
return 20.0 * math.log10(max(value, 1e-20))
|
||||
|
||||
|
||||
def latent_metrics(latent: torch.Tensor) -> dict:
|
||||
frames = latent.float().movedim(-1, 0).flatten(1)
|
||||
return {
|
||||
"first_4_rms": float(frames[:4].square().mean().sqrt()),
|
||||
"frames_4_20_rms": float(frames[4:20].square().mean().sqrt()),
|
||||
"first_frame_rms": float(frames[0].square().mean().sqrt()),
|
||||
"frame_0_to_1_delta_rms": float((frames[1] - frames[0]).square().mean().sqrt()),
|
||||
}
|
||||
|
||||
|
||||
def waveform_metrics(waveform: torch.Tensor) -> dict:
|
||||
waveform = waveform.float()
|
||||
first_100ms = waveform[..., :3200]
|
||||
first_500ms = waveform[..., :16000]
|
||||
derivative = (first_500ms[..., 1:] - first_500ms[..., :-1]).abs()
|
||||
return {
|
||||
"first_sample": waveform[..., 0].flatten().tolist(),
|
||||
"first_100ms_peak_dbfs": dbfs(float(first_100ms.abs().max())),
|
||||
"first_100ms_rms_dbfs": dbfs(float(first_100ms.square().mean().sqrt())),
|
||||
"first_500ms_peak_dbfs": dbfs(float(first_500ms.abs().max())),
|
||||
"first_500ms_rms_dbfs": dbfs(float(first_500ms.square().mean().sqrt())),
|
||||
"largest_derivative": float(derivative.max()),
|
||||
}
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--benchmark", type=Path, required=True)
|
||||
parser.add_argument("--trace", type=Path, required=True)
|
||||
parser.add_argument("--report", type=Path, required=True)
|
||||
parser.add_argument("--attention", default="sage2")
|
||||
args = parser.parse_args()
|
||||
|
||||
benchmark = json.loads(args.benchmark.read_text(encoding="utf-8"))
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention))
|
||||
video, initial_audio, aligned_frames = random_av_latents(
|
||||
benchmark["resolution"][0],
|
||||
benchmark["resolution"][1],
|
||||
benchmark["frames"],
|
||||
benchmark["seed"],
|
||||
device=runtime.config.device,
|
||||
)
|
||||
text = runtime.refiner(runtime.conditioner(benchmark["prompt"]))
|
||||
trace = []
|
||||
video, final_audio = sample_video_res_multistep(
|
||||
runtime.model,
|
||||
runtime.packer,
|
||||
text,
|
||||
video,
|
||||
initial_audio,
|
||||
steps=benchmark["steps"],
|
||||
seed=benchmark["seed"],
|
||||
return_audio=True,
|
||||
audio_step_trace=trace,
|
||||
)
|
||||
|
||||
report = {
|
||||
"benchmark": str(args.benchmark),
|
||||
"attention": args.attention,
|
||||
"seed": benchmark["seed"],
|
||||
"frames": aligned_frames,
|
||||
"steps": [],
|
||||
}
|
||||
with torch.inference_mode():
|
||||
for entry in trace:
|
||||
denoised = _decode_audio_latent(entry["audio_denoised"]).to(
|
||||
"cuda", dtype=next(runtime.audio_vae.parameters()).dtype,
|
||||
)
|
||||
waveform = runtime.audio_vae.decode(denoised).cpu()[0]
|
||||
report["steps"].append({
|
||||
"step": entry["step"],
|
||||
"video_sigma": entry["video_sigma"],
|
||||
"audio_sigma": entry["audio_sigma"],
|
||||
"latent": latent_metrics(denoised.cpu()),
|
||||
"denoised_waveform": waveform_metrics(waveform),
|
||||
})
|
||||
|
||||
args.trace.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save({
|
||||
"initial_audio": initial_audio.detach().cpu(),
|
||||
"final_audio": final_audio.detach().cpu(),
|
||||
"steps": trace,
|
||||
"prompt": benchmark["prompt"],
|
||||
"seed": benchmark["seed"],
|
||||
}, args.trace)
|
||||
args.report.parent.mkdir(parents=True, exist_ok=True)
|
||||
serialized = json.dumps(report, indent=2)
|
||||
args.report.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(serialized)
|
||||
55
tools/upscale_video_latent.py
Normal file
55
tools/upscale_video_latent.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
"""Upscale a saved direct-runtime H3 latent with the learned 3D model."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from h3_blackwell_runtime.latent_upscaler import load_h3_latent_upscaler, upscale_h3_latent
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--latent", type=Path, required=True)
|
||||
parser.add_argument("--model", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--scale", type=float, default=2.0)
|
||||
parser.add_argument("--precision", choices=("float16", "bfloat16", "float32"), default="float16")
|
||||
args = parser.parse_args()
|
||||
|
||||
dtype = getattr(torch, args.precision)
|
||||
state = torch.load(args.latent, map_location="cpu", weights_only=False)
|
||||
if not isinstance(state, dict) or "latent" not in state:
|
||||
state = {"latent": state}
|
||||
source = state["latent"]
|
||||
|
||||
torch.cuda.synchronize()
|
||||
started = time.perf_counter()
|
||||
model = load_h3_latent_upscaler(args.model, dtype=dtype)
|
||||
torch.cuda.synchronize()
|
||||
loaded = time.perf_counter()
|
||||
upscaled = upscale_h3_latent(model, source, scale=args.scale).cpu()
|
||||
torch.cuda.synchronize()
|
||||
finished = time.perf_counter()
|
||||
|
||||
result = dict(state)
|
||||
result["latent"] = upscaled
|
||||
result["width"] = upscaled.shape[-1] * 16
|
||||
result["height"] = upscaled.shape[-2] * 16
|
||||
result["upscale"] = {
|
||||
"model": args.model.name,
|
||||
"scale": args.scale,
|
||||
"precision": args.precision,
|
||||
"source_shape": tuple(source.shape),
|
||||
"output_shape": tuple(upscaled.shape),
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(result, args.output)
|
||||
print(json.dumps({
|
||||
"output": str(args.output),
|
||||
"source_shape": tuple(source.shape),
|
||||
"output_shape": tuple(upscaled.shape),
|
||||
"model_load_seconds": loaded - started,
|
||||
"upscale_seconds": finished - loaded,
|
||||
}, indent=2))
|
||||
Loading…
Add table
Reference in a new issue