[diffusion] feat: support FastH3 (4-step VSA-distilled MiniMax-H3) with a VSA-H3 attention backend (#37480)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Kevin Mi
co-authored by
Claude Fable 5
Claude Fable 5
parent
99b9109553
commit
f586654518
@@ -643,7 +643,71 @@ SGLang projects those adapter factors onto the pruned coordinates at load time.
|
||||
A structurally modified checkpoint without that metadata still fails closed,
|
||||
and packed GGUF weights remain incompatible with LoRA.
|
||||
|
||||
## 6. Sampling and output controls
|
||||
## 6. FastH3: 4-step distilled preview
|
||||
|
||||
[FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree](https://huggingface.co/FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree)
|
||||
is a 4-step DMD2 distillation of MiniMax-H3, trained data-free with Video
|
||||
Sparse Attention (VSA) at 0.9 sparsity and 64-token tiles. Only the T2VA
|
||||
capability was distilled: requests must use `task: "t2va"`, and `fl2va` /
|
||||
`ref2va` requests are rejected. The checkpoint inherits the MiniMax-H3
|
||||
Community License.
|
||||
|
||||
Pass the repository directly to `--model-path`. The flat native-Diffusers
|
||||
upload is materialized into the base-H3 layout through a registered model
|
||||
overlay; the only non-symlink step is a one-time re-serialization of the
|
||||
roughly 10 GB video VAE on first launch.
|
||||
|
||||
```bash 4×B300 VSA-H3
|
||||
sglang serve \
|
||||
--model-path FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree \
|
||||
--num-gpus 4 \
|
||||
--attention-backend video_sparse_attn_h3 \
|
||||
--attention-backend-config '{"VSA_sparsity": 0.9}' \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
Requests use the same asynchronous video endpoint as the base model, with
|
||||
`task: "t2va"`, `conditions: []`, and a target such as
|
||||
`{"short_edge": 768, "aspect_ratio": "16:9", "duration_seconds": 5.0}`. The
|
||||
request default is `num_inference_steps: 5`: five points on the standard
|
||||
shift-12/shift-3 sigma grid, i.e. the four distilled DiT evaluations. Any
|
||||
other step count is rejected.
|
||||
|
||||
`video_sparse_attn_h3` (VSA-H3) is the trained sparse policy: an in-tree
|
||||
Triton block-sparse kernel (SM90 / SM100 / SM103) over segment-pure prefix
|
||||
tiles and (4, 4, 4) video tiles, driven by the checkpoint's trained
|
||||
`to_gate_compress` compression branch. Only the DiT runs sparse; the token
|
||||
refiner, text encoder, and VAEs keep their dense defaults. Ulysses sequence
|
||||
parallelism is supported. See
|
||||
[Attention Backends](/docs/sglang-diffusion/attention_backends) for
|
||||
`VSA_sparsity`, `vsa_mode`, `vsa_dense_first_n_steps`, and
|
||||
`vsa_dense_layers`. Every dense backend that runs on base H3 (`fa`,
|
||||
`torch_sdpa`, ...) also runs on the FastH3 weights without VSA flags, and
|
||||
`sglang generate` takes the same flags as `sglang serve`.
|
||||
|
||||
Measured latencies for the 4× B300 recipe are in
|
||||
[FastH3 on B300](#fasth3-on-b300).
|
||||
|
||||
FastH3 rejects deployment options that do not apply to the distilled preview
|
||||
instead of silently ignoring them: `--model-variant`, `quality: "high"`,
|
||||
`fl2va` / `ref2va` requests, and, with VSA-H3, `--ring-degree` greater than 1,
|
||||
`torch.compile`, and breakable CUDA graph execution.
|
||||
|
||||
<Warning>
|
||||
The sibling `FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA` adapters carry
|
||||
full-rank `.diff` / `.diff_b` deltas and `set_weight` gate tensors beyond the
|
||||
LoRA contract; `--lora-path` rejects them with an explicit error. Serve the
|
||||
merged VSA-DataFree checkpoint above instead.
|
||||
</Warning>
|
||||
|
||||
<Note>
|
||||
Upstream labels this checkpoint a preview. Quality gaps versus base H3 on hard
|
||||
motion and fine detail are properties of the released distillation, not of the
|
||||
SGLang port. Use base MiniMax-H3 when output quality matters more than
|
||||
latency.
|
||||
</Note>
|
||||
|
||||
## 7. Sampling and output controls
|
||||
|
||||
MiniMax-H3 supports more than one output per prompt. The video API accepts
|
||||
`num_outputs_per_prompt` (or OpenAI-compatible `n`) from 1 through 10. Offline
|
||||
@@ -818,7 +882,7 @@ has completed, but the `quality: "high"` path above remains fail-closed
|
||||
to the audited 4×H200 workload.
|
||||
</Warning>
|
||||
|
||||
## 7. Feature contracts and advanced recipes
|
||||
## 8. Feature contracts and advanced recipes
|
||||
|
||||
The generated command already contains the recommended topology and encoder
|
||||
setting. Use the detailed reference below only when applying an optional
|
||||
@@ -1059,7 +1123,7 @@ fold decision is not node-boundary aware:
|
||||
|
||||
</Tabs>
|
||||
|
||||
## 8. Configuration notes
|
||||
## 9. Configuration notes
|
||||
|
||||
- MiniMax-H3 produces the canonical 24 fps output; request duration is expressed through `target.duration_seconds`.
|
||||
- `target.duration_seconds` must be between 4 and 15 seconds, inclusive. The command picker defaults to the verified 5-second profile.
|
||||
@@ -1078,7 +1142,7 @@ fold decision is not node-boundary aware:
|
||||
- `speed` keeps model components resident, while `auto` applies the model-aware 120 GiB residency threshold. `memory` prioritizes avoiding OOM and includes the executable VAE decoder in its default layerwise set. A measured recipe with sufficient headroom can opt into `--component-residency vae=resident`; the 2×H100 CI recipe does this because the VAE's 4.8 GiB/GPU cost avoids repeated decoder transfers during tiled decode. DiT residency and prefetch knobs remain scoped to the DiT. Use `speed` only after confirming that the complete target workload fits.
|
||||
- Breakable CUDA graph execution is an explicit opt-in, not part of the recommended `speed` preset. It requires `--enable-breakable-cuda-graph`, every served size in `--warmup-resolutions`, and `--bcg-text-buckets` that cover the live H3 condition sequence. The validated 1344×768 Ref2VA recipe uses 5504; other task profiles and reference sets may need a different value. It preserves eager output for matching captured signatures, but graph capture consumes additional GPU memory and may provide little latency benefit when Ulysses attention and collectives dominate, so benchmark it on the target topology before enabling it.
|
||||
|
||||
## 9. Benchmarks
|
||||
## 10. Benchmarks
|
||||
|
||||
The picker exposes resident and FSDP profiles on NVIDIA datacenter GPUs. GPU
|
||||
counts are properties of the selected recipes, not a claim that every platform
|
||||
@@ -1189,6 +1253,30 @@ python3 -m sglang.multimodal_gen.benchmarks.bench_serving \
|
||||
| Ref2VA | FP8 | fold | 112.0 s | 34.44 s | **27.12 s** | 52,816 MB |
|
||||
| Ref2VA | FP8 | replicate | 116.0 s | 33.42 s | **27.12 s** | 93,396 MB |
|
||||
|
||||
### FastH3 on B300
|
||||
|
||||
The same 4× B300 host served [FastH3](#6-fasth3-4-step-distilled-preview)
|
||||
with the VSA-H3 recipe above (1344×768 at 24 fps with audio, `task: "t2va"`,
|
||||
`num_inference_steps: 5`, seed 1000, eager BF16, Ulysses4, `VSA_sparsity` 0.9).
|
||||
E2E is the client wall clock of a `/v1/videos` request including decode,
|
||||
muxing, and file output, median of three requests after one warm request;
|
||||
the stage columns are the server timings of the same request. H3 aligns the
|
||||
requested durations to 124, 243, and 362 frames. Client RTF is E2E divided by
|
||||
the video duration:
|
||||
|
||||
| Requested / aligned | Encoder | Denoise (4 forwards) | Decode | Transport + MP4 | E2E | Client RTF | Peak/GPU |
|
||||
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
|
||||
| 5 s / 124 | 0.07 s | 2.18 s | 0.87 s | 0.9 s | **4.1 s** | 0.79 | 95,744 MB |
|
||||
| 10 s / 243 | 0.07 s | 4.83 s | 1.71 s | 1.3 s | **8.0 s** | 0.79 | 102,666 MB |
|
||||
| 15 s / 362 | 0.07 s | 8.80 s | 2.56 s | 1.8 s | **13.3 s** | 0.88 | 110,774 MB |
|
||||
|
||||
All three requests finish faster than playback. Dense FA on the same weights
|
||||
and topology takes 3.77 / 9.84 / 18.45 s (`sglang generate`, stage sum): it is
|
||||
competitive at 5 s, and VSA-H3 pulls ahead from 10 s on. At 5 s,
|
||||
TP2 + Ulysses2 (3.42 s, 62,290 MB), FSDP + Ulysses4 (3.42 s, 50,984 MB), and
|
||||
online `--quantization fp8` (2.93 s, 64,204 MB) trade a little latency for
|
||||
peak memory.
|
||||
|
||||
### H200 topology comparison
|
||||
|
||||
The same four-card H200 host completed both lossless resident placements with
|
||||
|
||||
@@ -79,6 +79,11 @@ For SGLang-native pipelines, the CLI accepts the lowercase names of `AttentionBa
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`VIDEO_SPARSE_ATTN`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Requires <code>vsa</code>. Configure <code>sparsity</code> via <code>--attention-backend-config</code>.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`video_sparse_attn_h3`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`VIDEO_SPARSE_ATTN_H3`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Video Sparse Attention for MiniMax-H3 / FastH3 (VSA-H3). In-tree Triton block-sparse kernel (SM90 / SM100 / SM103); no external package. Configure via <code>--attention-backend-config</code>.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`vmoba_attn`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)", whiteSpace: "nowrap"}}>`VMOBA_ATTN`</td>
|
||||
@@ -238,6 +243,67 @@ SpargeAttention is approximate even when `topk=1`: the recommended upstream
|
||||
kernel quantizes attention through SageAttention2. Validate output quality and
|
||||
end-to-end latency on the target model and resolution before deployment.
|
||||
|
||||
**Video Sparse Attention for H3 (`video_sparse_attn_h3`)**
|
||||
|
||||
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
|
||||
<colgroup>
|
||||
<col style={{width: "24%"}} />
|
||||
<col style={{width: "14%"}} />
|
||||
<col style={{width: "44%"}} />
|
||||
<col style={{width: "18%"}} />
|
||||
</colgroup>
|
||||
<thead>
|
||||
<tr style={{borderBottom: "2px solid #d55816"}}>
|
||||
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>Parameter</th>
|
||||
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.05)"}}>Type</th>
|
||||
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.02)"}}>Description</th>
|
||||
<th style={{textAlign: "left", padding: "10px 12px", fontWeight: 700, whiteSpace: "nowrap", backgroundColor: "rgba(255,255,255,0.05)"}}>Default</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`VSA_sparsity`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`float`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Fraction of video tiles excluded from the top-k selection (0.0 - 1.0). `0.9` is the FastH3 trained policy.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`0.9`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`vsa_mode`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`str`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`exempt`: non-video keys (text/audio prefix tiles) are always selected. `compete`: they compete with video tiles in the top-k.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`exempt`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`vsa_dense_first_n_steps`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Use dense attention for the first N denoising steps.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`0`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`vsa_dense_layers`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`list[int]`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Layer indices kept dense, e.g. `[0, 1]`.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`[]`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`vsa_tile_size`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Kernel tile size. Only `64` (the trained (4, 4, 4) geometry) is accepted.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`64`</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
VSA-H3 constraints:
|
||||
|
||||
- Only the DiT runs sparse; the token refiner, text encoder, and VAEs keep
|
||||
their dense defaults. An explicit `--component-attention-backends
|
||||
text_encoder=fa` is rejected because the H3 text encoder has SDPA-only layers.
|
||||
- Uses the checkpoint's trained `to_gate_compress` compression branch. Base
|
||||
MiniMax-H3 weights load zero gates and run pure sparse.
|
||||
- Ulysses sequence parallelism is supported; `--ring-degree` greater than 1,
|
||||
`torch.compile`, and breakable CUDA graph execution are rejected.
|
||||
|
||||
**V-MoBA (`vmoba_attn`)**
|
||||
|
||||
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
|
||||
@@ -549,6 +615,16 @@ end-to-end latency on the target model and resolution before deployment.
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA-only. Requires <code>vsa</code>. Configure <code>sparsity</code> via <code>--attention-backend-config</code>.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`video_sparse_attn_h3`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>No</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>No</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>No</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA-only (SM90 / SM100 / SM103). In-tree Triton kernel, no external dependency. Configure via <code>--attention-backend-config</code>.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>sla_attn</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
|
||||
|
||||
@@ -137,6 +137,12 @@ Rows are grouped when a family shares the same runtime path or optimization supp
|
||||
<td>T2VA / FL2VA / Ref2VA, 768p at 24 fps with synchronized audio</td>
|
||||
<td><span className="sgd-chip">Cache-DiT</span><span className="sgd-chip">Sage</span><span className="sgd-chip">Online FP8</span><span className="sgd-chip">GGUF</span></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>FastH3</td>
|
||||
<td><div className="sgd-id-list"><code>FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree</code></div></td>
|
||||
<td>T2VA only, 4-step distilled, 768p at 24 fps with synchronized audio</td>
|
||||
<td><span className="sgd-chip">VSA-H3</span></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Wan2.1 Fun</td>
|
||||
<td><div className="sgd-id-list"><code>weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers</code></div></td>
|
||||
@@ -621,6 +627,21 @@ Optimization columns are abbreviated to keep the matrix readable:
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>FastH3 4-step (T2VA only)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree</code></td>
|
||||
<td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768p · 24 fps</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>✅</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3 (one/two-stage/TI2V/HQ)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Lightricks/LTX-2.3</code></td>
|
||||
@@ -704,6 +725,11 @@ Optimization columns are abbreviated to keep the matrix readable:
|
||||
difference is transformer depth and width, picked up from
|
||||
`transformer/config.json` at load time. A single checkpoint serves T2V,
|
||||
I2V (`--image-path`), and T2I (`--num-frames 1`).
|
||||
6. FastH3's VSA column refers to the dedicated `video_sparse_attn_h3` (VSA-H3)
|
||||
backend. Dense backends that run on base MiniMax-H3 also run on the FastH3
|
||||
weights.
|
||||
FastH3 serves `t2va` only and rejects `--model-variant` and
|
||||
`quality: "high"`.
|
||||
|
||||
</Accordion>
|
||||
|
||||
|
||||
@@ -373,7 +373,7 @@ For a pipeline whose primary DiT is named `transformer`, the shorter
|
||||
|
||||
MiniMax-H3 supports this path while preserving its required FP32 patch,
|
||||
timestep, and output projections. See the
|
||||
[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#7-feature-contracts-and-advanced-recipes)
|
||||
[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#8-feature-contracts-and-advanced-recipes)
|
||||
for its distributed serving recipe.
|
||||
|
||||
### MXFP4 Online Quantization
|
||||
@@ -432,7 +432,7 @@ projections take that path.
|
||||
<Warning>
|
||||
`kitchen_int8` is approximate and is not a consistency ground-truth mode.
|
||||
The BF16 path is unchanged when `comfy-kitchen` is not installed. See the
|
||||
[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#7-feature-contracts-and-advanced-recipes)
|
||||
[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#8-feature-contracts-and-advanced-recipes)
|
||||
for the 24 GB offload recipe, including why `vae` must stay out of
|
||||
`--layerwise-offload-components`.
|
||||
</Warning>
|
||||
|
||||
+1
@@ -314,6 +314,7 @@ Use the preset categories this way:
|
||||
| `cosmos3-super-t2v-cfg2tp2` | `nvidia/Cosmos3-Super` | No | Explicit four-GPU TP2 x CFG2 throughput comparator. On H200 it was 48.00% faster end to end than TP2, but the topology changed the deterministic output (SSIM 0.914244, PSNR 29.469771 dB), so do not treat it as lossless-equivalent or select it automatically. |
|
||||
| `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Nightly cat image and motion prompt, 1280x720, 81 frames, 4 GPUs, CFG parallel, Ulysses degree 2, text encoder CPU offload and pinned CPU memory |
|
||||
| `minimax-h3-t2va` | `MiniMaxAI/MiniMax-H3` | Yes: `minimax_h3_t2va_5s` | H3 FL2VA-partition T2VA baseline: 1344x768 resolved canvas, 5 seconds / 124 frames at 24 fps, 50 joint video-audio steps, 4 GPUs, TP2 + Ulysses2, eager BF16/FP32. The helper writes H3's request contract to a generated config. |
|
||||
| `fasth3-t2va-vsa` | `FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree` | No | FastH3 4-step distilled T2VA on the trained VSA-H3 backend: 1344x768, 10 seconds / 243 frames, five sigma points = four DiT forwards, 4 GPUs, Ulysses 4, eager, 2-step warmup request. Compare against `--attention-backend fa` on the same weights for the dense-fallback gap. |
|
||||
| `longcat-image` | `meituan-longcat/LongCat-Image` | No | Eager DiT baseline at 1024x1024, 50 steps, guidance 4.5; prompt rewrite is disabled so Qwen2.5-VL does not contaminate the DiT A/B. |
|
||||
| `longcat-image-edit` | `meituan-longcat/LongCat-Image-Edit` | No | Native edit baseline using the public SGLang edit fixture. Its 1536x1024 source resolves to 1264x848 under the checkpoint's roughly-one-megapixel aspect-ratio rule, and the BCG comparator captures that exact serving canvas; prompt rewrite is disabled to isolate the DiT. |
|
||||
| `longcat-image-edit-turbo` | `meituan-longcat/LongCat-Image-Edit-Turbo` | No | Matching distilled edit baseline using the same public fixture, prompt, and 1264x848 BCG canvas. Its registered sampling class owns the eight-step, guidance-1 schedule. |
|
||||
|
||||
+27
@@ -425,6 +425,33 @@ MODELS = {
|
||||
"num-inference-steps",
|
||||
},
|
||||
},
|
||||
# H3 rejects a 1-step warmup request, hence --warmup-steps=2.
|
||||
"fasth3-t2va-vsa": {
|
||||
"path": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree",
|
||||
"prompt": (
|
||||
"A curious raccoon peers through a vibrant field of yellow "
|
||||
"sunflowers, its eyes wide with interest."
|
||||
),
|
||||
"seed": 1000,
|
||||
"config_overrides": {
|
||||
"task": "t2va",
|
||||
"conditions": [],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 10.0,
|
||||
},
|
||||
"num_inference_steps": 5,
|
||||
},
|
||||
"extra_args": [
|
||||
"--num-gpus=4",
|
||||
"--attention-backend=video_sparse_attn_h3",
|
||||
'--attention-backend-config={"VSA_sparsity": 0.9}',
|
||||
"--enable-torch-compile=false",
|
||||
"--warmup-steps=2",
|
||||
],
|
||||
"force_eager": True,
|
||||
},
|
||||
# Source-tracked extras from current registry / GPU test coverage.
|
||||
"longcat-image": {
|
||||
"path": "meituan-longcat/LongCat-Image",
|
||||
|
||||
@@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus
|
||||
## Key Features
|
||||
|
||||
SGLang Diffusion has the following features:
|
||||
- Broad model support: Wan, FastWan, FLUX, Qwen-Image, LongCat-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3/LTX-2.5, MiniMax-H3, LingBot Video MoE, LingBot World, SANA-Video/SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
|
||||
- Broad model support: Wan, FastWan, FLUX, Qwen-Image, LongCat-Image, Z-Image, Ideogram 4, Krea-2, Cosmos3, LTX-2/LTX-2.3/LTX-2.5, MiniMax-H3, FastH3, LingBot Video MoE, LingBot World, SANA-Video/SANA-WM, JoyEcho, MOVA, GLM-Image, ERNIE-Image, Hunyuan3D, and more
|
||||
- Fast inference speed: empowered by optimized `sgl-kernel` kernels, scheduler/runtime improvements, caching acceleration, and native diffusion hot-path optimizations
|
||||
- Ease of use: OpenAI-compatible api, CLI, and python sdk support
|
||||
- Multi-platform support:
|
||||
|
||||
@@ -46,6 +46,7 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
3,
|
||||
),
|
||||
r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"blocks.\1.attn.out_proj.\2",
|
||||
r"^transformer_blocks\.(\d+)\.attn\.to_gate_compress\.(.*)$": r"blocks.\1.attn.to_gate_compress.\2",
|
||||
r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"blocks.\1.attn.q_norm.\2",
|
||||
r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"blocks.\1.attn.k_norm.\2",
|
||||
r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"blocks.\1.mlp.fc1.\2",
|
||||
@@ -99,6 +100,7 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
final_norm_eps: float = 1e-5
|
||||
checkpoint_uses_diffusers_layout: bool = False
|
||||
adaln_affine_input_dim: int | None = None
|
||||
has_gate_compress: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
|
||||
@@ -24,6 +24,14 @@ class MiniMaxH3AudioVAEConfig(VAEConfig):
|
||||
load_encoder: bool = True
|
||||
load_decoder: bool = True
|
||||
|
||||
def update_model_arch(self, source_model_dict: dict) -> None:
|
||||
# Native-Diffusers AutoencoderKLMiniMaxH3Audio config field name.
|
||||
aliases = {"sampling_rate": "sample_rate"}
|
||||
model_dict = {
|
||||
aliases.get(key, key): value for key, value in source_model_dict.items()
|
||||
}
|
||||
super().update_model_arch(model_dict)
|
||||
|
||||
def post_init(self) -> None:
|
||||
validate_minimax_h3_vae_latent_stats(
|
||||
self.arch_config,
|
||||
|
||||
@@ -56,6 +56,17 @@ class MiniMaxH3VideoVAEConfig(VAEConfig):
|
||||
f"{self.parallel_decode_mode!r}"
|
||||
)
|
||||
|
||||
def update_model_arch(self, source_model_dict: dict) -> None:
|
||||
# Native-Diffusers AutoencoderKLMiniMaxH3 config field names.
|
||||
aliases = {
|
||||
"clip_length": "vae_clip_length",
|
||||
"token_drop": "vae_token_drop",
|
||||
}
|
||||
model_dict = {
|
||||
aliases.get(key, key): value for key, value in source_model_dict.items()
|
||||
}
|
||||
super().update_model_arch(model_dict)
|
||||
|
||||
def post_init(self) -> None:
|
||||
self.resolved_parallel_decode_mode()
|
||||
validate_minimax_h3_vae_latent_stats(
|
||||
|
||||
@@ -44,6 +44,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
|
||||
LTX23PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
FastH3PipelineConfig,
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
|
||||
@@ -97,5 +98,6 @@ __all__ = [
|
||||
"LingBotWorldCausalDMDConfig",
|
||||
"LingBotWorldV2CausalDMDConfig",
|
||||
"LingBotVideoMoEPipelineConfig",
|
||||
"FastH3PipelineConfig",
|
||||
"MiniMaxH3PipelineConfig",
|
||||
]
|
||||
|
||||
@@ -260,6 +260,21 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
|
||||
)
|
||||
if selected_backend is None:
|
||||
return
|
||||
if selected_backend is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3:
|
||||
if server_args.ring_degree > 1:
|
||||
raise ValueError(
|
||||
"VSA-H3 does not support --ring-degree > 1; use Ulysses "
|
||||
"sequence parallelism."
|
||||
)
|
||||
if (
|
||||
server_args.enable_torch_compile
|
||||
or server_args.enable_breakable_cuda_graph
|
||||
):
|
||||
raise ValueError(
|
||||
"VSA-H3 builds per-step tile metadata eagerly and is not "
|
||||
"validated under torch.compile or the breakable CUDA "
|
||||
"graph; disable them or use --attention-backend fa."
|
||||
)
|
||||
get_attn_backend(
|
||||
self.dit_config.arch_config.attention_head_dim,
|
||||
torch.bfloat16,
|
||||
@@ -279,4 +294,28 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
|
||||
return safetensors_list
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3PipelineConfig"]
|
||||
@dataclass
|
||||
class FastH3PipelineConfig(MiniMaxH3PipelineConfig):
|
||||
"""FastH3: 4-step VSA-distilled MiniMax-H3, t2va only."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.dit_config.arch_config.has_gate_compress = True
|
||||
|
||||
def validate_quality_deployment(self, server_args) -> None:
|
||||
raise ValueError(
|
||||
'quality="high" is audited only for the base MiniMax-H3 50-step '
|
||||
"4xH200 deployment; the FastH3 4-step distilled checkpoint has no "
|
||||
'audited high-quality deployment. Use quality="lossless".'
|
||||
)
|
||||
|
||||
def validate_server_args(self, server_args) -> None:
|
||||
if server_args.model_variant is not None:
|
||||
raise ValueError(
|
||||
"FastH3 ships one t2va-distilled weight partition; "
|
||||
"--model-variant does not apply. FL2VA and Ref2VA tasks were "
|
||||
"not distilled; use MiniMaxAI/MiniMax-H3 for those."
|
||||
)
|
||||
super().validate_server_args(server_args)
|
||||
|
||||
|
||||
__all__ = ["FastH3PipelineConfig", "MiniMaxH3PipelineConfig"]
|
||||
|
||||
@@ -302,4 +302,26 @@ class MiniMaxH3SamplingParams(SamplingParams):
|
||||
req.extra.update(self.build_request_extra())
|
||||
|
||||
|
||||
__all__ = ["MiniMaxH3SamplingParams"]
|
||||
@dataclass
|
||||
class FastH3SamplingParams(MiniMaxH3SamplingParams):
|
||||
"""FastH3: five sigma grid points, i.e. the four distilled DiT forwards."""
|
||||
|
||||
num_inference_steps: int = 5
|
||||
|
||||
def _validate(self) -> None:
|
||||
super()._validate()
|
||||
if self.num_inference_steps != 5:
|
||||
raise ValueError(
|
||||
"FastH3 is distilled for exactly five sigma grid points (four DiT "
|
||||
f"forwards); got num_inference_steps={self.num_inference_steps}. "
|
||||
"Use MiniMaxAI/MiniMax-H3 for other schedules."
|
||||
)
|
||||
if self.task is not None and self.task.strip().lower() != "t2va":
|
||||
raise ValueError(
|
||||
"FastH3 is distilled for t2va only; fl2va and ref2va were not "
|
||||
f"distilled (got task={self.task!r}). Use MiniMaxAI/MiniMax-H3 "
|
||||
"for those tasks."
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["FastH3SamplingParams", "MiniMaxH3SamplingParams"]
|
||||
|
||||
@@ -29,6 +29,7 @@ if TYPE_CHECKING:
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs import (
|
||||
Cosmos3Config,
|
||||
FastH3PipelineConfig,
|
||||
FastHunyuanConfig,
|
||||
FluxPipelineConfig,
|
||||
HeliosDistilledConfig,
|
||||
@@ -162,7 +163,10 @@ from sglang.multimodal_gen.configs.sample.ltx_2 import (
|
||||
LTX23SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.ltx_2_5 import LTX25SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3 import (
|
||||
FastH3SamplingParams,
|
||||
MiniMaxH3SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.mova import (
|
||||
MOVA_360P_SamplingParams,
|
||||
MOVA_720P_SamplingParams,
|
||||
@@ -334,6 +338,7 @@ _MODEL_NAME_DETECTORS: List[Tuple[str, Callable[[str], bool]]] = []
|
||||
KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: Dict[str, str] = {
|
||||
"minimaxai/minimax-h3": "MiniMaxH3Pipeline",
|
||||
"minimax/minimax-h3": "MiniMaxH3Pipeline",
|
||||
"fastvideo/fastvideo-fasth3-4-step-preview-v1-vsa-datafree": "FastH3Pipeline",
|
||||
"lerobot/pi05": "Pi05Pipeline",
|
||||
"pi05": "Pi05Pipeline",
|
||||
"pi0.5": "Pi05Pipeline",
|
||||
@@ -972,6 +977,17 @@ def _register_configs():
|
||||
in model_id.lower().replace("-", "").replace("_", "")
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=FastH3SamplingParams,
|
||||
pipeline_config_cls=FastH3PipelineConfig,
|
||||
hf_model_paths=[
|
||||
"FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree",
|
||||
],
|
||||
model_detectors=[
|
||||
lambda model_id: "fasth3"
|
||||
in model_id.lower().replace("-", "").replace("_", "")
|
||||
],
|
||||
)
|
||||
# FLUX
|
||||
register_configs(
|
||||
sampling_param_cls=FluxSamplingParams,
|
||||
|
||||
+488
@@ -0,0 +1,488 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
# (the video_sparse_attn_h3 backend), rewritten for SGLang's packed-varlen
|
||||
# MiniMax-H3 attention contract.
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VSA for MiniMax-H3's packed mixed-modality self-attention.
|
||||
|
||||
H3 runs one joint bidirectional attention over
|
||||
``[text | condition keyframes | audio | generated video]``. Tiles are
|
||||
``[segment-pure prefix chunks] + [3D video tiles]``; prefix tiles never
|
||||
straddle segment boundaries. Selection is a top-k over pooled tile scores
|
||||
emitted directly as the per-query-tile index lists the vendored Triton
|
||||
tile-64 kernel consumes, with per-tile valid sizes so ragged interior tiles
|
||||
mask exactly.
|
||||
|
||||
Non-video queries are always dense. Non-video keys are either
|
||||
always-selected for every query ("exempt", default) or compete in top-k
|
||||
under a FLOP-matched budget ("compete"). The compression branch is gated by
|
||||
``to_gate_compress``: base H3 has no such weights and the gate loads as
|
||||
zeros (pure sparse); VSA-distilled students (FastH3) ship trained gates
|
||||
that activate the branch.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
|
||||
AttentionBackend,
|
||||
AttentionImpl,
|
||||
AttentionMetadata,
|
||||
AttentionMetadataBuilder,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import (
|
||||
construct_variable_block_sizes,
|
||||
get_non_pad_index,
|
||||
get_tile_partition_indices,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.vsa_h3_kernels import (
|
||||
vsa_h3_block_sparse_attn_forward,
|
||||
vsa_h3_pack_tiles,
|
||||
vsa_h3_untile,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
|
||||
# The FastH3 checkpoints are trained and served at the 64-token (4, 4, 4)
|
||||
# tile geometry; the kernel block size matches it exactly.
|
||||
VSA_H3_TILE_SHAPE = (4, 4, 4)
|
||||
VSA_H3_TILE_ELEMS = math.prod(VSA_H3_TILE_SHAPE)
|
||||
|
||||
_DIT_BLOCK_PREFIX = re.compile(r"^blocks\.(\d+)\.")
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
def _h3_tile_geometry(
|
||||
prefix_segments: tuple[int, ...],
|
||||
dit_seq_shape: tuple[int, int, int],
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
|
||||
"""Segment-pure prefix chunks, then video tiles.
|
||||
|
||||
Returns (variable_block_sizes int32 [n_tiles], pack_index int32 [S_pad]:
|
||||
padded position -> packed row or -1, unpack_index int32 [used]: packed row
|
||||
-> padded position, num_prefix_tiles, num_video_tiles).
|
||||
"""
|
||||
prefix_len = sum(prefix_segments)
|
||||
|
||||
prefix_sizes: list[int] = []
|
||||
for segment in prefix_segments:
|
||||
full, remainder = divmod(segment, VSA_H3_TILE_ELEMS)
|
||||
prefix_sizes.extend([VSA_H3_TILE_ELEMS] * full)
|
||||
if remainder:
|
||||
prefix_sizes.append(remainder)
|
||||
num_prefix_tiles = len(prefix_sizes)
|
||||
|
||||
num_tiles = (
|
||||
math.ceil(dit_seq_shape[0] / VSA_H3_TILE_SHAPE[0]),
|
||||
math.ceil(dit_seq_shape[1] / VSA_H3_TILE_SHAPE[1]),
|
||||
math.ceil(dit_seq_shape[2] / VSA_H3_TILE_SHAPE[2]),
|
||||
)
|
||||
video_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device)
|
||||
num_video_tiles = int(video_sizes.numel())
|
||||
|
||||
video_indices = (
|
||||
get_tile_partition_indices(dit_seq_shape, VSA_H3_TILE_SHAPE, device)
|
||||
+ prefix_len
|
||||
)
|
||||
tile_partition_indices = torch.cat(
|
||||
[
|
||||
torch.arange(prefix_len, device=device, dtype=torch.long),
|
||||
video_indices,
|
||||
]
|
||||
)
|
||||
variable_block_sizes = torch.cat(
|
||||
[
|
||||
torch.tensor(prefix_sizes, dtype=torch.long, device=device),
|
||||
video_sizes.to(torch.long),
|
||||
]
|
||||
)
|
||||
non_pad_index = get_non_pad_index(variable_block_sizes, VSA_H3_TILE_ELEMS)
|
||||
|
||||
total = prefix_len + math.prod(dit_seq_shape)
|
||||
sizes_sum = int(variable_block_sizes.sum())
|
||||
if sizes_sum != total or non_pad_index.numel() != total:
|
||||
raise ValueError(
|
||||
f"VSA-H3 tile geometry mismatch for prefix={prefix_segments}, "
|
||||
f"video={dit_seq_shape}: sizes sum {sizes_sum}, non-pad "
|
||||
f"{non_pad_index.numel()}, expected {total}."
|
||||
)
|
||||
|
||||
seq_pad = variable_block_sizes.numel() * VSA_H3_TILE_ELEMS
|
||||
pack_index = torch.full((seq_pad,), -1, dtype=torch.int32, device=device)
|
||||
pack_index[non_pad_index] = tile_partition_indices.to(torch.int32)
|
||||
unpack_index = torch.empty(total, dtype=torch.int32, device=device)
|
||||
unpack_index[tile_partition_indices] = non_pad_index.to(torch.int32)
|
||||
return (
|
||||
variable_block_sizes.to(torch.int32),
|
||||
pack_index,
|
||||
unpack_index,
|
||||
num_prefix_tiles,
|
||||
num_video_tiles,
|
||||
)
|
||||
|
||||
|
||||
class VideoSparseAttentionH3Backend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes() -> list[int]:
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_enum() -> AttentionBackendEnum:
|
||||
return AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls() -> type["VideoSparseAttentionH3Impl"]:
|
||||
return VideoSparseAttentionH3Impl
|
||||
|
||||
@staticmethod
|
||||
def get_metadata_cls() -> type["VideoSparseAttentionH3Metadata"]:
|
||||
return VideoSparseAttentionH3Metadata
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["VideoSparseAttentionH3MetadataBuilder"]:
|
||||
return VideoSparseAttentionH3MetadataBuilder
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoSparseAttentionH3Metadata(AttentionMetadata):
|
||||
VSA_sparsity: float
|
||||
total_seq_length: int
|
||||
num_prefix_tiles: int
|
||||
num_video_tiles: int
|
||||
exempt: bool
|
||||
variable_block_sizes: torch.Tensor
|
||||
pack_index: torch.Tensor
|
||||
unpack_index: torch.Tensor
|
||||
dense_layers: tuple[int, ...] = ()
|
||||
workspace_cache: dict = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def num_tiles(self) -> int:
|
||||
return self.num_prefix_tiles + self.num_video_tiles
|
||||
|
||||
|
||||
class VideoSparseAttentionH3MetadataBuilder(AttentionMetadataBuilder):
|
||||
def __init__(self) -> None:
|
||||
self._workspace_cache: dict = {}
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def build( # type: ignore[override]
|
||||
self,
|
||||
current_timestep: int,
|
||||
raw_latent_shape: tuple[int, int, int],
|
||||
patch_size: tuple[int, int, int],
|
||||
VSA_sparsity: float,
|
||||
prefix_segments: tuple[int, ...],
|
||||
device: torch.device,
|
||||
exempt: bool = True,
|
||||
dense_layers: tuple[int, ...] = (),
|
||||
dense_first_n_steps: int = 0,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> VideoSparseAttentionH3Metadata:
|
||||
dit_seq_shape = (
|
||||
raw_latent_shape[0] // patch_size[0],
|
||||
raw_latent_shape[1] // patch_size[1],
|
||||
raw_latent_shape[2] // patch_size[2],
|
||||
)
|
||||
prefix_segments = tuple(int(s) for s in prefix_segments if s > 0)
|
||||
if current_timestep < dense_first_n_steps:
|
||||
VSA_sparsity = 0.0
|
||||
|
||||
(
|
||||
variable_block_sizes,
|
||||
pack_index,
|
||||
unpack_index,
|
||||
num_prefix_tiles,
|
||||
num_video_tiles,
|
||||
) = _h3_tile_geometry(prefix_segments, dit_seq_shape, device)
|
||||
|
||||
return VideoSparseAttentionH3Metadata(
|
||||
current_timestep=current_timestep,
|
||||
VSA_sparsity=float(VSA_sparsity),
|
||||
total_seq_length=sum(prefix_segments) + math.prod(dit_seq_shape),
|
||||
num_prefix_tiles=num_prefix_tiles,
|
||||
num_video_tiles=num_video_tiles,
|
||||
exempt=exempt,
|
||||
variable_block_sizes=variable_block_sizes,
|
||||
pack_index=pack_index,
|
||||
unpack_index=unpack_index,
|
||||
dense_layers=tuple(int(layer) for layer in dense_layers),
|
||||
workspace_cache=self._workspace_cache,
|
||||
)
|
||||
|
||||
|
||||
def _compute_topk(sparsity: float, num_video_tiles: int) -> int:
|
||||
keep = math.ceil((1.0 - sparsity) * num_video_tiles)
|
||||
return max(1, min(keep, num_video_tiles))
|
||||
|
||||
|
||||
def _topk_tile_lists(
|
||||
scores: torch.Tensor,
|
||||
num_prefix_tiles: int,
|
||||
num_video_tiles: int,
|
||||
sparsity: float,
|
||||
exempt: bool,
|
||||
) -> torch.Tensor:
|
||||
"""scores [H, n_tiles, n_tiles] -> ascending int32 kv-tile lists
|
||||
[H, num_video_tiles, width]; width = num_prefix_tiles + keep (exempt) or
|
||||
min(keep + num_prefix_tiles, n_tiles) (compete)."""
|
||||
prefix = num_prefix_tiles
|
||||
keep = _compute_topk(sparsity, num_video_tiles)
|
||||
video_rows = scores[:, prefix:, :]
|
||||
if exempt or prefix == 0:
|
||||
picked = video_rows[:, :, prefix:].topk(keep, dim=-1).indices + prefix
|
||||
picked = picked.sort(dim=-1).values
|
||||
if prefix == 0:
|
||||
return picked.to(torch.int32)
|
||||
prefix_cols = torch.arange(prefix, device=scores.device).expand(
|
||||
*picked.shape[:-1], prefix
|
||||
)
|
||||
return torch.cat([prefix_cols, picked], dim=-1).to(torch.int32)
|
||||
keep_total = min(keep + prefix, scores.shape[-1])
|
||||
return (
|
||||
video_rows.topk(keep_total, dim=-1).indices.sort(dim=-1).values.to(torch.int32)
|
||||
)
|
||||
|
||||
|
||||
def _workspace_key(
|
||||
meta: VideoSparseAttentionH3Metadata,
|
||||
heads: int,
|
||||
head_dim: int,
|
||||
has_gate: bool,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> tuple:
|
||||
return (
|
||||
meta.num_tiles,
|
||||
meta.num_prefix_tiles,
|
||||
meta.exempt,
|
||||
heads,
|
||||
head_dim,
|
||||
has_gate,
|
||||
dtype,
|
||||
device,
|
||||
)
|
||||
|
||||
|
||||
class _Workspace:
|
||||
"""Per-geometry scratch: tiled q/k/v(/gate) [3|4, H, S_pad, D], pooled fp32
|
||||
tile means [3, H, n_tiles, D], and the kernel index lists (prefix rows and
|
||||
prefix columns are static; only the top-k video columns change per layer).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
meta: VideoSparseAttentionH3Metadata,
|
||||
heads: int,
|
||||
head_dim: int,
|
||||
has_gate: bool,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
n_tiles = meta.num_tiles
|
||||
seq_pad = n_tiles * VSA_H3_TILE_ELEMS
|
||||
self.key = _workspace_key(meta, heads, head_dim, has_gate, dtype, device)
|
||||
self.tiled = torch.empty(
|
||||
(3 + int(has_gate), heads, seq_pad, head_dim), dtype=dtype, device=device
|
||||
)
|
||||
self.pooled = torch.empty(
|
||||
(3, heads, n_tiles, head_dim), dtype=torch.float32, device=device
|
||||
)
|
||||
self.out_tiled = torch.empty(
|
||||
(heads, seq_pad, head_dim), dtype=dtype, device=device
|
||||
)
|
||||
all_tiles = torch.arange(n_tiles, dtype=torch.int32, device=device)
|
||||
self.dense_index = all_tiles.repeat(heads, n_tiles, 1)
|
||||
self.dense_num = torch.full(
|
||||
(heads, n_tiles), n_tiles, dtype=torch.int32, device=device
|
||||
)
|
||||
self.q2k_index = self.dense_index.clone()
|
||||
self.q2k_num = self.dense_num.clone()
|
||||
|
||||
def sparse_lists(
|
||||
self, video_lists: torch.Tensor, num_prefix_tiles: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
width = video_lists.shape[-1]
|
||||
self.q2k_index[:, num_prefix_tiles:, :width] = video_lists
|
||||
self.q2k_num[:, num_prefix_tiles:] = width
|
||||
return self.q2k_index, self.q2k_num
|
||||
|
||||
|
||||
def _get_workspace(
|
||||
meta: VideoSparseAttentionH3Metadata, query: torch.Tensor, has_gate: bool
|
||||
) -> _Workspace:
|
||||
heads, head_dim = query.shape[-2], query.shape[-1]
|
||||
key = _workspace_key(meta, heads, head_dim, has_gate, query.dtype, query.device)
|
||||
workspace = meta.workspace_cache.get("workspace")
|
||||
if workspace is None or workspace.key != key:
|
||||
workspace = _Workspace(
|
||||
meta, heads, head_dim, has_gate, query.dtype, query.device
|
||||
)
|
||||
meta.workspace_cache["workspace"] = workspace
|
||||
return workspace
|
||||
|
||||
|
||||
def _select_kv_lists(
|
||||
ws: _Workspace,
|
||||
meta: VideoSparseAttentionH3Metadata,
|
||||
scores: torch.Tensor | None,
|
||||
sparsity: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Kernel index lists for this layer: dense, or top-k video columns."""
|
||||
keep = _compute_topk(sparsity, meta.num_video_tiles)
|
||||
if sparsity <= 0.0 or keep >= meta.num_video_tiles:
|
||||
return ws.dense_index, ws.dense_num
|
||||
return ws.sparse_lists(
|
||||
_topk_tile_lists(
|
||||
scores,
|
||||
meta.num_prefix_tiles,
|
||||
meta.num_video_tiles,
|
||||
sparsity,
|
||||
meta.exempt,
|
||||
),
|
||||
meta.num_prefix_tiles,
|
||||
)
|
||||
|
||||
|
||||
class VideoSparseAttentionH3Impl(AttentionImpl):
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
head_size: int,
|
||||
causal: bool,
|
||||
softmax_scale: float,
|
||||
num_kv_heads: int | None = None,
|
||||
prefix: str = "",
|
||||
**extra_impl_args,
|
||||
) -> None:
|
||||
self.num_heads = num_heads
|
||||
self.head_size = head_size
|
||||
self.softmax_scale = softmax_scale
|
||||
self.prefix = prefix
|
||||
match = _DIT_BLOCK_PREFIX.match(prefix)
|
||||
self.layer_idx = int(match.group(1)) if match else None
|
||||
# The token refiner and any other non-packed caller resolve the same
|
||||
# backend object; they run the exact dense kernel instead.
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
|
||||
FlashAttentionImpl,
|
||||
)
|
||||
|
||||
self._dense_fallback = FlashAttentionImpl(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
causal=causal,
|
||||
softmax_scale=softmax_scale,
|
||||
num_kv_heads=num_kv_heads,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError(
|
||||
"VSA-H3 serves MiniMax-H3's packed varlen attention; use " "forward_varlen."
|
||||
)
|
||||
|
||||
def forward_varlen(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
cu_seqlens: torch.Tensor,
|
||||
max_seqlen: int,
|
||||
cu_seqlens_host: tuple[int, ...] | None = None,
|
||||
attn_metadata: VideoSparseAttentionH3Metadata | None = None,
|
||||
gate_compress: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""query/key/value: [T, H, D] packed rows (post-norm, post-RoPE)."""
|
||||
if self.layer_idx is None or attn_metadata is None:
|
||||
if attn_metadata is None and self.layer_idx is not None:
|
||||
raise RuntimeError(
|
||||
"VSA-H3 needs per-step attention metadata from the "
|
||||
"MiniMax-H3 denoising stage; none was set in the forward "
|
||||
"context."
|
||||
)
|
||||
return self._dense_fallback.forward_varlen(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
)
|
||||
|
||||
meta = attn_metadata
|
||||
bounds = (
|
||||
cu_seqlens_host
|
||||
if cu_seqlens_host is not None
|
||||
else tuple(int(item) for item in cu_seqlens.tolist())
|
||||
)
|
||||
used = int(bounds[1])
|
||||
if used != meta.total_seq_length:
|
||||
raise ValueError(
|
||||
f"VSA-H3 metadata was built for {meta.total_seq_length} packed "
|
||||
f"rows, got {used}. The step metadata and the packed sequence "
|
||||
"layout have diverged."
|
||||
)
|
||||
|
||||
sparsity = 0.0 if self.layer_idx in meta.dense_layers else meta.VSA_sparsity
|
||||
has_gate = gate_compress is not None
|
||||
ws = _get_workspace(meta, query, has_gate)
|
||||
|
||||
vsa_h3_pack_tiles(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
gate_compress,
|
||||
meta.pack_index,
|
||||
meta.variable_block_sizes,
|
||||
ws.tiled,
|
||||
ws.pooled,
|
||||
)
|
||||
q_pooled, k_pooled, v_pooled = ws.pooled
|
||||
|
||||
scores = None
|
||||
if sparsity > 0.0 or has_gate:
|
||||
scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) * (
|
||||
self.head_size**-0.5
|
||||
)
|
||||
|
||||
q2k_index, q2k_num = _select_kv_lists(ws, meta, scores, sparsity)
|
||||
vsa_h3_block_sparse_attn_forward(
|
||||
ws.tiled[0:1],
|
||||
ws.tiled[1:2],
|
||||
ws.tiled[2:3],
|
||||
q2k_index[None],
|
||||
q2k_num[None],
|
||||
meta.variable_block_sizes,
|
||||
out=ws.out_tiled[None],
|
||||
)
|
||||
|
||||
out_compress = None
|
||||
if has_gate:
|
||||
out_compress = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)
|
||||
|
||||
result = torch.empty(query.shape, dtype=query.dtype, device=query.device)
|
||||
vsa_h3_untile(
|
||||
ws.out_tiled,
|
||||
ws.tiled[3] if has_gate else None,
|
||||
out_compress,
|
||||
meta.unpack_index,
|
||||
used,
|
||||
result,
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,343 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
# (fastvideo-kernel triton_kernels/block_sparse_attn_triton.py).
|
||||
# Inference-only subset: the block-sparse forward. Backward stays upstream;
|
||||
# SGLang serves no-grad forwards. The tile pack/unpack kernels are SGLang's.
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Triton 64-token block-sparse attention for MiniMax-H3 VSA.
|
||||
|
||||
The attention kernel consumes an explicit per-query-block index list
|
||||
(``q2k_index`` / ``q2k_num``) plus per-key-block valid token counts
|
||||
(``variable_block_sizes``), so ragged interior tiles - segment-pure prefix
|
||||
chunks and 3D video tiles whose dimensions do not divide the tile shape - mask
|
||||
their pad columns exactly.
|
||||
|
||||
``vsa_h3_pack_tiles`` gathers packed ``[T, H, D]`` rows into the head-major
|
||||
padded tile layout the attention kernel reads and pools each tile in the same
|
||||
pass; ``vsa_h3_untile`` scatters the attention output back to packed rows and
|
||||
folds in the gated compression branch. Both replace chains of index copies
|
||||
and transposes that otherwise cost more than the attention kernel itself.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from triton.tools.tensor_descriptor import TensorDescriptor
|
||||
|
||||
# BLOCK_M / BLOCK_N are structural, not tunable: the kernel indexes the top-k
|
||||
# list per BLOCK_M q-tile and addresses keys as kv_idx * BLOCK_N, so both must
|
||||
# match the granularity q2k_index and variable_block_sizes were built at.
|
||||
VSA_H3_KERNEL_BLOCK = 64
|
||||
|
||||
# Pinned instead of autotuned: on B300 num_warps=4 wins at every sequence
|
||||
# length and num_stages=5 is within 0.5% of the best (7 spills); FastVideo
|
||||
# reports the same optimum for Blackwell.
|
||||
_ATTN_NUM_WARPS = 4
|
||||
_ATTN_NUM_STAGES = 5
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_fwd_sparse(
|
||||
desc_q,
|
||||
desc_k,
|
||||
desc_v,
|
||||
desc_o,
|
||||
sm_scale,
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
max_kv_blks,
|
||||
variable_block_sizes,
|
||||
H,
|
||||
N_CTX_Q,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
q_blk = tl.program_id(0)
|
||||
off_hz = tl.program_id(1)
|
||||
b = off_hz // H
|
||||
h = off_hz % H
|
||||
q_tiles = N_CTX_Q // BLOCK_M
|
||||
meta_base = off_hz * q_tiles + q_blk
|
||||
|
||||
kv_blocks = tl.load(q2k_num + meta_base)
|
||||
kv_ptr = q2k_index + meta_base.to(tl.int64) * max_kv_blks
|
||||
|
||||
q = desc_q.load([b, h, q_blk * BLOCK_M, 0]).reshape([BLOCK_M, HEAD_DIM])
|
||||
m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
|
||||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
|
||||
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
||||
qk_scale = sm_scale * 1.44269504 # 1/ln2
|
||||
|
||||
for i in range(0, kv_blocks):
|
||||
kv_idx = tl.load(kv_ptr + i).to(tl.int32)
|
||||
block_size = tl.load(variable_block_sizes + kv_idx)
|
||||
k = desc_k.load([b, h, kv_idx * BLOCK_N, 0]).reshape([BLOCK_N, HEAD_DIM])
|
||||
qk = tl.dot(q, tl.trans(k))
|
||||
mask = tl.arange(0, BLOCK_N) < block_size
|
||||
qk = tl.where(mask[None, :], qk, -float("inf"))
|
||||
|
||||
m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
|
||||
p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
|
||||
l_ij = tl.sum(p, 1)
|
||||
|
||||
alpha = tl.math.exp2(m_i - m_ij)
|
||||
l_i = l_i * alpha + l_ij
|
||||
acc = acc * alpha[:, None]
|
||||
|
||||
v = desc_v.load([b, h, kv_idx * BLOCK_N, 0]).reshape([BLOCK_N, HEAD_DIM])
|
||||
acc = tl.dot(p.to(tl.bfloat16), v, acc)
|
||||
m_i = m_ij
|
||||
|
||||
acc = acc / l_i[:, None]
|
||||
desc_o.store(
|
||||
[b, h, q_blk * BLOCK_M, 0],
|
||||
acc.to(desc_o.dtype).reshape([1, 1, BLOCK_M, HEAD_DIM]),
|
||||
)
|
||||
|
||||
|
||||
def vsa_h3_block_sparse_attn_forward(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
q2k_index: torch.Tensor,
|
||||
q2k_num: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""q/k/v: contiguous [B, H, S_pad, D] bf16 with S_pad = n_tiles * 64; pad
|
||||
rows zero. q2k_index/q2k_num: contiguous [B, H, n_tiles, max_kv] /
|
||||
[B, H, n_tiles] int32."""
|
||||
batch, heads, seq_q, head_dim = q.shape
|
||||
seq_kv = k.shape[2]
|
||||
if seq_q % VSA_H3_KERNEL_BLOCK or seq_kv % VSA_H3_KERNEL_BLOCK:
|
||||
raise ValueError(
|
||||
f"VSA-H3 kernel needs 64-multiple sequence lengths, got q={seq_q}, "
|
||||
f"kv={seq_kv}"
|
||||
)
|
||||
if variable_block_sizes.numel() != seq_kv // VSA_H3_KERNEL_BLOCK:
|
||||
raise ValueError(
|
||||
"variable_block_sizes must have one entry per 64-token key block: "
|
||||
f"{variable_block_sizes.numel()} vs {seq_kv // VSA_H3_KERNEL_BLOCK}"
|
||||
)
|
||||
if out is None:
|
||||
out = torch.empty_like(q)
|
||||
block = [1, 1, VSA_H3_KERNEL_BLOCK, head_dim]
|
||||
desc_q, desc_k, desc_v, desc_o = (
|
||||
TensorDescriptor.from_tensor(t, block_shape=block) for t in (q, k, v, out)
|
||||
)
|
||||
grid = (seq_q // VSA_H3_KERNEL_BLOCK, batch * heads, 1)
|
||||
_attn_fwd_sparse[grid](
|
||||
desc_q,
|
||||
desc_k,
|
||||
desc_v,
|
||||
desc_o,
|
||||
1.0 / math.sqrt(head_dim),
|
||||
q2k_index,
|
||||
q2k_num,
|
||||
q2k_index.shape[-1],
|
||||
variable_block_sizes,
|
||||
heads,
|
||||
seq_q,
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_M=VSA_H3_KERNEL_BLOCK,
|
||||
BLOCK_N=VSA_H3_KERNEL_BLOCK,
|
||||
num_warps=_ATTN_NUM_WARPS,
|
||||
num_stages=_ATTN_NUM_STAGES,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _pack_tiles_kernel(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
G,
|
||||
src_index,
|
||||
variable_block_sizes,
|
||||
Tiled,
|
||||
Pooled,
|
||||
stride_q_row,
|
||||
stride_q_head,
|
||||
stride_k_row,
|
||||
stride_k_head,
|
||||
stride_v_row,
|
||||
stride_v_head,
|
||||
stride_g_row,
|
||||
stride_g_head,
|
||||
H,
|
||||
S_PAD,
|
||||
N_TILES,
|
||||
HAS_GATE: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
tile = tl.program_id(0)
|
||||
h = tl.program_id(1)
|
||||
rows = tile * BLOCK + tl.arange(0, BLOCK)
|
||||
cols = tl.arange(0, HEAD_DIM)
|
||||
src = tl.load(src_index + rows)
|
||||
valid = src >= 0
|
||||
src = tl.where(valid, src, 0).to(tl.int64)
|
||||
mask = valid[:, None]
|
||||
size = tl.load(variable_block_sizes + tile).to(tl.float32)
|
||||
|
||||
tensor_stride = H.to(tl.int64) * S_PAD * HEAD_DIM
|
||||
out_off = (
|
||||
h.to(tl.int64) * S_PAD * HEAD_DIM + rows[:, None] * HEAD_DIM + cols[None, :]
|
||||
)
|
||||
pool_off = (h * N_TILES + tile).to(tl.int64) * HEAD_DIM + cols
|
||||
|
||||
x = tl.load(
|
||||
Q + src[:, None] * stride_q_row + h * stride_q_head + cols[None, :],
|
||||
mask=mask,
|
||||
other=0.0,
|
||||
)
|
||||
tl.store(Tiled + out_off, x)
|
||||
tl.store(Pooled + pool_off, tl.sum(x.to(tl.float32), 0) / size)
|
||||
|
||||
x = tl.load(
|
||||
K + src[:, None] * stride_k_row + h * stride_k_head + cols[None, :],
|
||||
mask=mask,
|
||||
other=0.0,
|
||||
)
|
||||
tl.store(Tiled + tensor_stride + out_off, x)
|
||||
tl.store(
|
||||
Pooled + H * N_TILES * HEAD_DIM + pool_off,
|
||||
tl.sum(x.to(tl.float32), 0) / size,
|
||||
)
|
||||
|
||||
x = tl.load(
|
||||
V + src[:, None] * stride_v_row + h * stride_v_head + cols[None, :],
|
||||
mask=mask,
|
||||
other=0.0,
|
||||
)
|
||||
tl.store(Tiled + 2 * tensor_stride + out_off, x)
|
||||
tl.store(
|
||||
Pooled + 2 * H * N_TILES * HEAD_DIM + pool_off,
|
||||
tl.sum(x.to(tl.float32), 0) / size,
|
||||
)
|
||||
|
||||
if HAS_GATE:
|
||||
x = tl.load(
|
||||
G + src[:, None] * stride_g_row + h * stride_g_head + cols[None, :],
|
||||
mask=mask,
|
||||
other=0.0,
|
||||
)
|
||||
tl.store(Tiled + 3 * tensor_stride + out_off, x)
|
||||
|
||||
|
||||
def vsa_h3_pack_tiles(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
gate: torch.Tensor | None,
|
||||
src_index: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
tiled: torch.Tensor,
|
||||
pooled: torch.Tensor,
|
||||
) -> None:
|
||||
"""Gather packed [T, H, D] rows into ``tiled`` [3|4, H, S_pad, D] and write
|
||||
fp32 per-tile means of q/k/v into ``pooled`` [3, H, n_tiles, D].
|
||||
``src_index`` maps each padded position to its packed row, or -1 (pad -> 0).
|
||||
"""
|
||||
_, heads, seq_pad, head_dim = tiled.shape
|
||||
n_tiles = seq_pad // VSA_H3_KERNEL_BLOCK
|
||||
has_gate = gate is not None
|
||||
g = gate if has_gate else q
|
||||
assert all(t.stride(-1) == 1 for t in (q, k, v, g))
|
||||
_pack_tiles_kernel[(n_tiles, heads)](
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
g,
|
||||
src_index,
|
||||
variable_block_sizes,
|
||||
tiled,
|
||||
pooled,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
k.stride(0),
|
||||
k.stride(1),
|
||||
v.stride(0),
|
||||
v.stride(1),
|
||||
g.stride(0),
|
||||
g.stride(1),
|
||||
heads,
|
||||
seq_pad,
|
||||
n_tiles,
|
||||
HAS_GATE=has_gate,
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK=VSA_H3_KERNEL_BLOCK,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _untile_kernel(
|
||||
OutTiled,
|
||||
Gate,
|
||||
OutC,
|
||||
dst_index,
|
||||
Res,
|
||||
used,
|
||||
total,
|
||||
S_PAD,
|
||||
N_TILES,
|
||||
HAS_GATE: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
row_block = tl.program_id(0)
|
||||
h = tl.program_id(1)
|
||||
H = tl.num_programs(1)
|
||||
rows = row_block * BLOCK + tl.arange(0, BLOCK)
|
||||
cols = tl.arange(0, HEAD_DIM)
|
||||
in_used = rows < used
|
||||
pos = tl.load(dst_index + rows, mask=in_used, other=0).to(tl.int64)
|
||||
head_base = h.to(tl.int64) * S_PAD * HEAD_DIM
|
||||
off = head_base + pos[:, None] * HEAD_DIM + cols[None, :]
|
||||
o = tl.load(OutTiled + off, mask=in_used[:, None], other=0.0).to(tl.float32)
|
||||
if HAS_GATE:
|
||||
g = tl.load(Gate + off, mask=in_used[:, None], other=0.0).to(tl.float32)
|
||||
tile = pos // BLOCK
|
||||
c = tl.load(
|
||||
OutC + (h * N_TILES + tile)[:, None] * HEAD_DIM + cols[None, :],
|
||||
mask=in_used[:, None],
|
||||
other=0.0,
|
||||
)
|
||||
o = o + c * g
|
||||
res_off = (rows[:, None] * H + h).to(tl.int64) * HEAD_DIM + cols[None, :]
|
||||
tl.store(Res + res_off, o.to(Res.type.element_ty), mask=(rows < total)[:, None])
|
||||
|
||||
|
||||
def vsa_h3_untile(
|
||||
out_tiled: torch.Tensor,
|
||||
gate_tiled: torch.Tensor | None,
|
||||
out_compress: torch.Tensor | None,
|
||||
dst_index: torch.Tensor,
|
||||
used: int,
|
||||
result: torch.Tensor,
|
||||
) -> None:
|
||||
"""Scatter ``out_tiled`` [H, S_pad, D] to packed rows of ``result`` [T, H, D]
|
||||
(rows past ``used`` are zero), adding ``out_compress`` [H, n_tiles, D] fp32
|
||||
scaled by ``gate_tiled`` when given."""
|
||||
heads, seq_pad, head_dim = out_tiled.shape
|
||||
total = result.shape[0]
|
||||
has_gate = gate_tiled is not None
|
||||
_untile_kernel[(triton.cdiv(total, VSA_H3_KERNEL_BLOCK), heads)](
|
||||
out_tiled,
|
||||
gate_tiled if has_gate else out_tiled,
|
||||
out_compress if has_gate else out_tiled,
|
||||
dst_index,
|
||||
result,
|
||||
used,
|
||||
total,
|
||||
seq_pad,
|
||||
seq_pad // VSA_H3_KERNEL_BLOCK,
|
||||
HAS_GATE=has_gate,
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK=VSA_H3_KERNEL_BLOCK,
|
||||
)
|
||||
@@ -67,6 +67,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.usp import _ring_attention_varlen
|
||||
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
is_layerwise_offloaded_module,
|
||||
@@ -86,6 +87,18 @@ logger = init_logger(__name__)
|
||||
|
||||
_ARCH_DEFAULTS = MiniMaxH3DiTArchConfig()
|
||||
|
||||
_NON_LORA_DELTA_SUFFIXES = (".diff", ".diff_b", ".set_weight")
|
||||
|
||||
|
||||
def _reject_non_lora_delta_tensors(adapter: dict[str, torch.Tensor]) -> None:
|
||||
offending = sorted(key for key in adapter if key.endswith(_NON_LORA_DELTA_SUFFIXES))
|
||||
if offending:
|
||||
raise ValueError(
|
||||
f"LoRA adapter carries {len(offending)} non-LoRA tensors "
|
||||
f"(.diff/.diff_b/.set_weight, e.g. {offending[0]}) that no MiniMax-H3 "
|
||||
"LoRA mapping rule applies; serve a checkpoint with them merged instead."
|
||||
)
|
||||
|
||||
|
||||
def _diffusers_h3_checkpoint(
|
||||
iterator: Iterable[tuple[str, torch.Tensor]],
|
||||
@@ -592,6 +605,7 @@ def _minimax_h3_attention_core_impl(
|
||||
ulysses_active: bool,
|
||||
subblock_sparse_query_block_mask: torch.Tensor | None = None,
|
||||
ring_active: bool = False,
|
||||
gate_compress: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Dynamic varlen attention and Ulysses/Ring collectives.
|
||||
|
||||
@@ -602,11 +616,14 @@ def _minimax_h3_attention_core_impl(
|
||||
|
||||
if ulysses_active:
|
||||
from sglang.multimodal_gen.runtime.layers.usp import (
|
||||
_usp_input_all_to_all,
|
||||
_usp_input_all_to_all_packed_qkv,
|
||||
_usp_output_all_to_all,
|
||||
)
|
||||
|
||||
q, k, v = _usp_input_all_to_all_packed_qkv(q, k, v)
|
||||
if gate_compress is not None:
|
||||
gate_compress = _usp_input_all_to_all(gate_compress[None], head_dim=2)[0]
|
||||
|
||||
if attention._attention_impl is None:
|
||||
attention._set_attention_backend(
|
||||
@@ -618,6 +635,26 @@ def _minimax_h3_attention_core_impl(
|
||||
)
|
||||
)
|
||||
|
||||
if attention._attention_backend_enum is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3:
|
||||
attn_metadata = (
|
||||
get_forward_context().attn_metadata
|
||||
if attention.prefix.startswith("blocks.")
|
||||
else None
|
||||
)
|
||||
out = attention._attention_impl.forward_varlen(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
cu_seqlens_host=cu_seqlens_host,
|
||||
attn_metadata=attn_metadata,
|
||||
gate_compress=gate_compress,
|
||||
)
|
||||
if ulysses_active:
|
||||
out = _usp_output_all_to_all(out[None], head_dim=2)[0]
|
||||
return out
|
||||
|
||||
if ring_active:
|
||||
ring_ws, _ = get_ring_ctx()
|
||||
if attention._attention_backend_enum is not AttentionBackendEnum.FA:
|
||||
@@ -767,6 +804,18 @@ class MiniMaxH3Attention(nn.Module):
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.out_proj",
|
||||
)
|
||||
# VSA compression gate; stays bf16 and unquantized (zero gate == pure sparse).
|
||||
self.to_gate_compress: ColumnParallelLinear | None = None
|
||||
if arch.has_gate_compress and prefix.startswith("blocks."):
|
||||
self.to_gate_compress = ColumnParallelLinear(
|
||||
arch.hidden_size,
|
||||
self.inner_dim,
|
||||
bias=False,
|
||||
gather_output=False,
|
||||
params_dtype=_BF16_DTYPE,
|
||||
quant_config=None,
|
||||
prefix=f"{prefix}.to_gate_compress",
|
||||
)
|
||||
|
||||
def _set_attention_backend(self, backend) -> None:
|
||||
if (
|
||||
@@ -1032,6 +1081,14 @@ class MiniMaxH3Attention(nn.Module):
|
||||
)
|
||||
q, k = _apply_rope_qk(q, k, cos_sin_cache, positions)
|
||||
|
||||
gate_compress = None
|
||||
if (
|
||||
self._attention_backend_enum is AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
|
||||
and self.to_gate_compress is not None
|
||||
):
|
||||
gate_flat, _ = self.to_gate_compress(x)
|
||||
gate_compress = gate_flat.view(total, self.num_heads, self.head_dim)
|
||||
|
||||
attention_core = (
|
||||
_minimax_h3_attention_core_bcg
|
||||
if self.bcg_breakpoint
|
||||
@@ -1048,6 +1105,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
subblock_sparse_query_block_mask=subblock_sparse_query_block_mask,
|
||||
ulysses_active=ulysses_active,
|
||||
ring_active=ring_active,
|
||||
gate_compress=gate_compress,
|
||||
)
|
||||
out = out.reshape(total, self.num_heads * self.head_dim)
|
||||
out, _ = self.out_proj(out)
|
||||
@@ -1780,6 +1838,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
|
||||
self, adapter: dict[str, torch.Tensor]
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Project released-checkpoint AdaLN LoRAs onto pruned coordinates."""
|
||||
_reject_non_lora_delta_tensors(adapter)
|
||||
full_width = self.arch.adaln_affine_input_dim
|
||||
if full_width is None:
|
||||
return adapter
|
||||
@@ -2102,6 +2161,16 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
|
||||
get_global_forced_attn_backend()
|
||||
or self._component_attention_backend_override
|
||||
)
|
||||
if selected_backend is None:
|
||||
selected_backend = next(
|
||||
(
|
||||
module._selected_attention_backend
|
||||
for module in self.modules()
|
||||
if isinstance(module, MiniMaxH3Attention)
|
||||
and module._selected_attention_backend is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
backend = get_attn_backend(
|
||||
self.arch.attention_head_dim,
|
||||
_BF16_DTYPE,
|
||||
|
||||
@@ -170,4 +170,17 @@ class MiniMaxH3Pipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
)
|
||||
|
||||
|
||||
EntryClass = MiniMaxH3Pipeline
|
||||
class FastH3Pipeline(MiniMaxH3Pipeline):
|
||||
"""FastH3: 4-step DMD2-distilled MiniMax-H3 (t2va only).
|
||||
|
||||
The flat single-partition repo is materialized into the base-H3 layout by
|
||||
the bundled model overlay (see model_overlays/), so every stage, loader,
|
||||
and admission path below is exactly the MiniMax-H3 one. There is no
|
||||
FL2VA/Ref2VA partition layout to default into.
|
||||
"""
|
||||
|
||||
pipeline_name = "FastH3Pipeline"
|
||||
default_model_subfolder = None
|
||||
|
||||
|
||||
EntryClass = [MiniMaxH3Pipeline, FastH3Pipeline]
|
||||
|
||||
+86
-3
@@ -5,7 +5,7 @@ single-branch execution, and payload validation.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from contextlib import contextmanager
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
@@ -39,7 +39,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
|
||||
VerificationResult,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.platforms import (
|
||||
AttentionBackendEnum,
|
||||
current_platform,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import maybe_nvtx_range
|
||||
@@ -731,6 +734,13 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
|
||||
device,
|
||||
placement_managed=placement_managed,
|
||||
)
|
||||
build_vsa_h3_step_metadata = _maybe_prepare_vsa_h3_step_metadata(
|
||||
model=model,
|
||||
packed=packed,
|
||||
ctx=ctx,
|
||||
server_args=server_args,
|
||||
device=device,
|
||||
)
|
||||
positive = MiniMaxH3DenoiseBranch(
|
||||
packed=packed,
|
||||
text_embeddings=emb["hidden_states"],
|
||||
@@ -770,6 +780,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
|
||||
self._forward_dit,
|
||||
batch=batch,
|
||||
attn_metadata=attn_metadata,
|
||||
build_vsa_h3_step_metadata=build_vsa_h3_step_metadata,
|
||||
),
|
||||
positive=positive,
|
||||
initial_video_rows=initial_video,
|
||||
@@ -823,6 +834,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
|
||||
*,
|
||||
batch: Req,
|
||||
attn_metadata: CubeSparseAttentionMetadata | None = None,
|
||||
build_vsa_h3_step_metadata: Callable[[int], Any] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Route the custom full loop through the native denoising runner."""
|
||||
|
||||
@@ -832,7 +844,11 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
|
||||
|
||||
with set_forward_context(
|
||||
current_timestep=step_index,
|
||||
attn_metadata=attn_metadata,
|
||||
attn_metadata=(
|
||||
build_vsa_h3_step_metadata(step_index)
|
||||
if build_vsa_h3_step_metadata is not None
|
||||
else attn_metadata
|
||||
),
|
||||
forward_batch=batch,
|
||||
):
|
||||
runner = self._maybe_get_bcg_runner(model)
|
||||
@@ -964,6 +980,73 @@ def _assemble_condition_rows(ctx: _FullLoopContext) -> None:
|
||||
ctx.keyframe_frame_count = int(ctx.keyframe["frame_count"])
|
||||
|
||||
|
||||
def _maybe_prepare_vsa_h3_step_metadata(
|
||||
*,
|
||||
model: Any,
|
||||
packed: Mapping[str, torch.Tensor],
|
||||
ctx: _FullLoopContext,
|
||||
server_args: ServerArgs,
|
||||
device: torch.device,
|
||||
) -> Callable[[int], Any] | None:
|
||||
"""Per-step VSA-H3 metadata builder over the request-static packed layout,
|
||||
or None off the VSA path."""
|
||||
model._resolve_attention_backend_once()
|
||||
if (
|
||||
model._resolved_attention_backend
|
||||
is not AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
|
||||
):
|
||||
return None
|
||||
if ctx.is_ref2va:
|
||||
raise NotImplementedError(
|
||||
"VSA-H3 supports the t2va/fl2va packed layout; the ref2va "
|
||||
"reference-block layout is not tiled yet. Use --attention-backend "
|
||||
"fa for ref2va."
|
||||
)
|
||||
|
||||
config = server_args.attention_backend_config or {}
|
||||
tile_size = int(config.get("vsa_tile_size", 64))
|
||||
if tile_size != 64:
|
||||
raise ValueError(
|
||||
"VSA-H3 in SGLang serves the trained 64-token (4, 4, 4) tile "
|
||||
f"geometry; got vsa_tile_size={tile_size}."
|
||||
)
|
||||
sparsity = float(config.get("VSA_sparsity", config.get("sparsity", 0.9)))
|
||||
if not 0.0 <= sparsity < 1.0:
|
||||
raise ValueError(f"VSA sparsity must be in [0, 1), got {sparsity}")
|
||||
mode = str(config.get("vsa_mode", "exempt"))
|
||||
if mode not in ("exempt", "compete"):
|
||||
raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {mode!r}")
|
||||
dense_first_n_steps = int(config.get("vsa_dense_first_n_steps", 0))
|
||||
dense_layers = tuple(int(layer) for layer in config.get("vsa_dense_layers", ()))
|
||||
|
||||
text_len = int(packed["text_pos"].numel())
|
||||
video_rows = int(packed["update_mask"].sum())
|
||||
cond_rows = int(packed["img_pos"].numel()) - video_rows
|
||||
audio_rows = int(packed["audio_pos"].numel())
|
||||
patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn_h3 import (
|
||||
VideoSparseAttentionH3MetadataBuilder,
|
||||
)
|
||||
|
||||
builder = VideoSparseAttentionH3MetadataBuilder()
|
||||
|
||||
def build(step_index: int):
|
||||
return builder.build(
|
||||
current_timestep=step_index,
|
||||
raw_latent_shape=(ctx.latent_t, ctx.latent_h, ctx.latent_w),
|
||||
patch_size=patch_size,
|
||||
VSA_sparsity=sparsity,
|
||||
prefix_segments=(text_len, cond_rows, audio_rows),
|
||||
device=device,
|
||||
exempt=mode == "exempt",
|
||||
dense_layers=dense_layers,
|
||||
dense_first_n_steps=dense_first_n_steps,
|
||||
)
|
||||
|
||||
return build
|
||||
|
||||
|
||||
def _build_packed_layout(
|
||||
ctx: _FullLoopContext,
|
||||
emb: Mapping[str, Any],
|
||||
|
||||
@@ -252,6 +252,40 @@ class _VideoSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
|
||||
raise ImportError("Video Sparse Attention backend is not installed.") from e
|
||||
|
||||
|
||||
class _VideoSparseAttentionH3BackendResolver(_CudaAttentionBackendResolver):
|
||||
backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3
|
||||
|
||||
# The vendored Triton tile-64 kernel is written against Hopper and
|
||||
# Blackwell block-sparse geometry; older architectures fail closed.
|
||||
supported_capabilities = {(9, 0), (10, 0), (10, 3)}
|
||||
|
||||
@classmethod
|
||||
def resolve(cls, platform) -> str:
|
||||
capability = platform.get_device_capability()
|
||||
capability_tuple = (
|
||||
(capability.major, capability.minor) if capability is not None else None
|
||||
)
|
||||
if capability_tuple not in cls.supported_capabilities:
|
||||
found = capability.as_version_str() if capability else "unknown"
|
||||
raise ValueError(
|
||||
"VSA-H3 (video_sparse_attn_h3) needs compute capability 9.0 "
|
||||
"(Hopper), 10.0 (B200 / GB200) or 10.3 (B300 / GB300); "
|
||||
f"this device reports {found}."
|
||||
)
|
||||
try:
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn_h3 import ( # noqa: F401
|
||||
VideoSparseAttentionH3Backend,
|
||||
)
|
||||
|
||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn_h3.VideoSparseAttentionH3Backend"
|
||||
except Exception as e:
|
||||
logger.error("Failed to import VSA-H3 attention backend: %s", str(e))
|
||||
raise ImportError(
|
||||
"VSA-H3 attention needs Triton and the in-tree tile-64 "
|
||||
"block-sparse kernel."
|
||||
) from e
|
||||
|
||||
|
||||
class _CubeSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
|
||||
backend = AttentionBackendEnum.CUBE_SPARSE_ATTN
|
||||
|
||||
@@ -435,6 +469,7 @@ _CUDA_ATTENTION_BACKEND_RESOLVERS = {
|
||||
_SageAttention3BackendResolver,
|
||||
_SpargeAttentionBackendResolver,
|
||||
_VideoSparseAttentionBackendResolver,
|
||||
_VideoSparseAttentionH3BackendResolver,
|
||||
_CubeSparseAttentionBackendResolver,
|
||||
_SparseVideoGen2AttentionBackendResolver,
|
||||
_SolAttnBackendResolver,
|
||||
|
||||
@@ -35,6 +35,7 @@ class AttentionBackendEnum(enum.Enum):
|
||||
SAGE_ATTN_3 = enum.auto()
|
||||
SPARGE_ATTN = enum.auto()
|
||||
VIDEO_SPARSE_ATTN = enum.auto()
|
||||
VIDEO_SPARSE_ATTN_H3 = enum.auto()
|
||||
SPARSE_VIDEO_GEN_2_ATTN = enum.auto()
|
||||
VMOBA_ATTN = enum.auto()
|
||||
AITER = enum.auto()
|
||||
@@ -57,6 +58,7 @@ class AttentionBackendEnum(enum.Enum):
|
||||
return self in {
|
||||
AttentionBackendEnum.SLIDING_TILE_ATTN,
|
||||
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
|
||||
AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,
|
||||
AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
|
||||
AttentionBackendEnum.VMOBA_ATTN,
|
||||
AttentionBackendEnum.SLA_ATTN,
|
||||
|
||||
@@ -42,6 +42,10 @@ BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
|
||||
"overlay_repo_id": "AgainstEntropy/SANA-WM_streaming-overlay",
|
||||
"overlay_revision": "62c6840871ecc3559189047513ba0670e1bf62e7",
|
||||
},
|
||||
"FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree": {
|
||||
"overlay_repo_id": "kevin-mi/FastH3-4step-Preview-overlay",
|
||||
"overlay_revision": "f769cb8001dae335089de7b250364335bc7cb183",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -688,6 +688,49 @@ MINIMAX_H3_FOUR_GPU_H100_CASES = [
|
||||
run_models_api_check=False,
|
||||
run_t2v_input_reference_check=False,
|
||||
),
|
||||
DiffusionTestCase(
|
||||
"fasth3_t2va_vsa_4gpu_h100",
|
||||
DiffusionServerArgs(
|
||||
model_path="FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree",
|
||||
modality="video",
|
||||
num_gpus=4,
|
||||
extras=[
|
||||
"--attention-backend",
|
||||
"video_sparse_attn_h3",
|
||||
"--attention-backend-config",
|
||||
'{"VSA_sparsity": 0.9}',
|
||||
"--enable-torch-compile",
|
||||
"false",
|
||||
],
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=(
|
||||
"A curious raccoon peers through a vibrant field of yellow "
|
||||
"sunflowers, its eyes wide with interest."
|
||||
),
|
||||
output_size="1344x768",
|
||||
seconds=5,
|
||||
output_format="mp4",
|
||||
expect_audio_output=True,
|
||||
num_outputs_per_prompt=1,
|
||||
extras={
|
||||
"task": "t2va",
|
||||
"conditions": [],
|
||||
"target": {
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 5.0,
|
||||
},
|
||||
"num_inference_steps": 5,
|
||||
"seed": 42,
|
||||
},
|
||||
),
|
||||
run_perf_check=False,
|
||||
run_consistency_check=False,
|
||||
run_component_accuracy_check=False,
|
||||
run_models_api_check=False,
|
||||
run_t2v_input_reference_check=False,
|
||||
),
|
||||
]
|
||||
|
||||
TWO_GPU_CASES = [
|
||||
|
||||
@@ -99,6 +99,7 @@ class TestDiffusionBenchmarkSkill(unittest.TestCase):
|
||||
"lingbot-world",
|
||||
"lingbot-world-v2",
|
||||
"fastwan21-t2v-1.3b",
|
||||
"fasth3-t2va-vsa",
|
||||
"wan22-t2v-nvfp4",
|
||||
"krea2-turbo",
|
||||
"krea2-raw",
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""FastH3 (4-step VSA-distilled MiniMax-H3) registration and admission contracts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
FastH3PipelineConfig,
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.minimax_h3 import FastH3SamplingParams
|
||||
from sglang.multimodal_gen.registry import (
|
||||
get_model_info,
|
||||
get_non_diffusers_pipeline_name,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
maybe_init_distributed_environment_and_model_parallel,
|
||||
model_parallel_is_initialized,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
|
||||
from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3DiTModel
|
||||
from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
|
||||
ensure_distributed_env_defaults,
|
||||
)
|
||||
|
||||
FASTH3_MODEL_ID = "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree"
|
||||
|
||||
|
||||
def _ensure_single_process_parallel_runtime() -> None:
|
||||
if model_parallel_is_initialized():
|
||||
return
|
||||
ensure_distributed_env_defaults()
|
||||
maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
|
||||
|
||||
|
||||
def test_registry_resolves_fasth3_configs() -> None:
|
||||
info = get_model_info(FASTH3_MODEL_ID)
|
||||
assert info.sampling_param_cls is FastH3SamplingParams
|
||||
assert info.pipeline_config_cls is FastH3PipelineConfig
|
||||
assert get_non_diffusers_pipeline_name(FASTH3_MODEL_ID) == "FastH3Pipeline"
|
||||
materialized = (
|
||||
"/cache/materialized_models/"
|
||||
"FastVideo__FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree-0123abcd"
|
||||
)
|
||||
assert get_model_info(materialized).sampling_param_cls is FastH3SamplingParams
|
||||
|
||||
|
||||
def test_fasth3_sampling_defaults_and_task_rejection() -> None:
|
||||
params = FastH3SamplingParams(prompt="p")
|
||||
assert params.num_inference_steps == 5
|
||||
assert params.guidance_scale == 1.0
|
||||
|
||||
with pytest.raises(ValueError, match="exactly five sigma grid points"):
|
||||
FastH3SamplingParams(prompt="p", num_inference_steps=50)
|
||||
|
||||
with pytest.raises(ValueError, match="distilled for t2va only"):
|
||||
FastH3SamplingParams(
|
||||
prompt="p",
|
||||
task="fl2va",
|
||||
conditions=[{"type": "image", "uri": "x.png", "role": "first_frame"}],
|
||||
target={
|
||||
"short_edge": 768,
|
||||
"aspect_ratio": "16:9",
|
||||
"duration_seconds": 5.0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_fasth3_pipeline_config_gates_and_rejections() -> None:
|
||||
config = FastH3PipelineConfig()
|
||||
assert config.dit_config.arch_config.has_gate_compress
|
||||
assert not MiniMaxH3PipelineConfig().dit_config.arch_config.has_gate_compress
|
||||
mapping = config.dit_config.arch_config.param_names_mapping
|
||||
source = "transformer_blocks.7.attn.to_gate_compress.weight"
|
||||
targets = [
|
||||
re.sub(pattern, target if isinstance(target, str) else target[0], source)
|
||||
for pattern, target in mapping.items()
|
||||
if re.match(pattern, source)
|
||||
]
|
||||
assert targets == ["blocks.7.attn.to_gate_compress.weight"]
|
||||
|
||||
with pytest.raises(ValueError, match="--model-variant does not apply"):
|
||||
config.validate_server_args(SimpleNamespace(model_variant="ref2va"))
|
||||
with pytest.raises(ValueError, match="no.*audited high-quality deployment"):
|
||||
config.validate_quality_deployment(server_args=None)
|
||||
|
||||
|
||||
def test_fasth3_lora_bundle_is_rejected_loudly() -> None:
|
||||
model = SimpleNamespace(arch=SimpleNamespace(adaln_affine_input_dim=None))
|
||||
plain = {
|
||||
"blocks.0.attn.qkv_proj.lora_A": torch.zeros(3, 64, 8),
|
||||
"blocks.0.attn.qkv_proj.lora_B": torch.zeros(3, 8, 64),
|
||||
}
|
||||
assert MiniMaxH3DiTModel.prepare_lora_adapter(model, dict(plain)) == plain
|
||||
|
||||
bundle = dict(plain)
|
||||
bundle["blocks.0.attn.qkv_proj.diff"] = torch.zeros(3, 64, 64)
|
||||
bundle["audio_patch_proj.diff_b"] = torch.zeros(64)
|
||||
bundle["blocks.0.attn.to_gate_compress.set_weight"] = torch.zeros(64, 64)
|
||||
with pytest.raises(ValueError, match="3 non-LoRA tensors.*set_weight"):
|
||||
MiniMaxH3DiTModel.prepare_lora_adapter(model, bundle)
|
||||
|
||||
|
||||
def test_fasth3_gates_stay_bf16_under_runtime_quantization() -> None:
|
||||
_ensure_single_process_parallel_runtime()
|
||||
with torch.device("meta"):
|
||||
model = MiniMaxH3DiTModel(
|
||||
config=FastH3PipelineConfig().dit_config,
|
||||
hf_config={},
|
||||
quant_config=Fp8Config(),
|
||||
)
|
||||
|
||||
attn = model.blocks[0].attn
|
||||
assert not isinstance(attn.qkv_proj.quant_method, UnquantizedLinearMethod)
|
||||
assert isinstance(attn.to_gate_compress.quant_method, UnquantizedLinearMethod)
|
||||
assert attn.to_gate_compress.weight.dtype == torch.bfloat16
|
||||
assert attn.to_gate_compress.weight.missing_param_init == "error"
|
||||
assert model.token_refiner.blocks[0].attn.to_gate_compress is None
|
||||
@@ -0,0 +1,241 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""VSA-H3 backend contracts.
|
||||
|
||||
The load-bearing check is sparsity -> 0: every tile is inside the budget, so
|
||||
the block-sparse kernel must reproduce dense attention over the packed rows to
|
||||
bf16 rounding. That single assertion pins the tile routing indices, ragged
|
||||
tile masking, the untile permutation, and the softmax scale.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn_h3 import (
|
||||
VSA_H3_TILE_ELEMS,
|
||||
VideoSparseAttentionH3Impl,
|
||||
VideoSparseAttentionH3MetadataBuilder,
|
||||
_topk_tile_lists,
|
||||
)
|
||||
|
||||
requires_cuda = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="VSA-H3 kernels need CUDA"
|
||||
)
|
||||
|
||||
# Ragged on purpose: text 70 and audio 100 are not tile multiples, and the
|
||||
# video canvas (5, 6, 10) is ragged in every tile dimension.
|
||||
PREFIX_SEGMENTS = (70, 0, 100)
|
||||
VIDEO_SHAPE = (5, 6, 10)
|
||||
HEADS = 4
|
||||
HEAD_DIM = 128
|
||||
|
||||
|
||||
def _build_metadata(sparsity: float, device, **kwargs):
|
||||
return VideoSparseAttentionH3MetadataBuilder().build(
|
||||
current_timestep=0,
|
||||
raw_latent_shape=VIDEO_SHAPE,
|
||||
patch_size=(1, 1, 1),
|
||||
VSA_sparsity=sparsity,
|
||||
prefix_segments=PREFIX_SEGMENTS,
|
||||
device=device,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _packed_qkv(device, seed: int = 7):
|
||||
used = sum(PREFIX_SEGMENTS) + math.prod(VIDEO_SHAPE)
|
||||
total = (used + 63) // 64 * 64
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
tensors = [
|
||||
torch.randn(
|
||||
(total, HEADS, HEAD_DIM), generator=generator, dtype=torch.float32
|
||||
).to(device=device, dtype=torch.bfloat16)
|
||||
for _ in range(3)
|
||||
]
|
||||
return used, total, tensors
|
||||
|
||||
|
||||
def _dense_reference(q, k, v, used: int) -> torch.Tensor:
|
||||
qf = q[:used].float().permute(1, 0, 2)
|
||||
kf = k[:used].float().permute(1, 0, 2)
|
||||
vf = v[:used].float().permute(1, 0, 2)
|
||||
scores = qf @ kf.transpose(-2, -1) / math.sqrt(HEAD_DIM)
|
||||
return (torch.softmax(scores, dim=-1) @ vf).permute(1, 0, 2)
|
||||
|
||||
|
||||
def _impl():
|
||||
impl = VideoSparseAttentionH3Impl(
|
||||
num_heads=HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
causal=False,
|
||||
softmax_scale=HEAD_DIM**-0.5,
|
||||
num_kv_heads=HEADS,
|
||||
prefix="blocks.3.attn",
|
||||
)
|
||||
assert impl.layer_idx == 3
|
||||
return impl
|
||||
|
||||
|
||||
def _run(impl, meta, used, total, q, k, v, gate=None):
|
||||
cu = torch.tensor([0, used, total], dtype=torch.int32, device=q.device)
|
||||
return impl.forward_varlen(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens=cu,
|
||||
max_seqlen=used,
|
||||
cu_seqlens_host=(0, used, total),
|
||||
attn_metadata=meta,
|
||||
gate_compress=gate,
|
||||
)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_zero_sparsity_matches_dense() -> None:
|
||||
device = torch.device("cuda")
|
||||
meta = _build_metadata(0.0, device)
|
||||
used, total, (q, k, v) = _packed_qkv(device)
|
||||
assert meta.total_seq_length == used
|
||||
|
||||
out = _run(_impl(), meta, used, total, q, k, v)
|
||||
reference = _dense_reference(q, k, v, used)
|
||||
diff = (out[:used].float() - reference).abs().max().item()
|
||||
assert diff < 2e-2, f"sparse(0) vs dense max diff {diff}"
|
||||
assert torch.all(out[used:] == 0)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_zero_gate_is_noop_and_trained_gate_activates() -> None:
|
||||
device = torch.device("cuda")
|
||||
meta = _build_metadata(0.5, device)
|
||||
used, total, (q, k, v) = _packed_qkv(device)
|
||||
impl = _impl()
|
||||
|
||||
base = _run(impl, meta, used, total, q, k, v).clone()
|
||||
zero_gate = torch.zeros_like(q)
|
||||
gated_zero = _run(impl, meta, used, total, q, k, v, gate=zero_gate)
|
||||
assert torch.equal(base, gated_zero)
|
||||
|
||||
gate = torch.randn_like(q) * 0.1
|
||||
gated = _run(impl, meta, used, total, q, k, v, gate=gate)
|
||||
assert not torch.equal(base, gated)
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_dense_overrides() -> None:
|
||||
device = torch.device("cuda")
|
||||
meta = _build_metadata(0.9, device, dense_layers=(3,))
|
||||
used, total, (q, k, v) = _packed_qkv(device)
|
||||
out = _run(_impl(), meta, used, total, q, k, v)
|
||||
reference = _dense_reference(q, k, v, used)
|
||||
diff = (out[:used].float() - reference).abs().max().item()
|
||||
assert diff < 2e-2, f"dense-layer opt-out vs dense max diff {diff}"
|
||||
|
||||
assert _build_metadata(0.9, device, dense_first_n_steps=0).VSA_sparsity == 0.9
|
||||
assert _build_metadata(0.9, device, dense_first_n_steps=1).VSA_sparsity == 0.0
|
||||
|
||||
|
||||
def _lists_to_mask(lists: torch.Tensor, n_tiles: int) -> torch.Tensor:
|
||||
mask = torch.zeros(*lists.shape[:-1], n_tiles, dtype=torch.bool)
|
||||
mask.scatter_(-1, lists.long(), True)
|
||||
return mask
|
||||
|
||||
|
||||
def test_topk_tile_list_semantics() -> None:
|
||||
num_prefix, num_video = 3, 10
|
||||
n_tiles = num_prefix + num_video
|
||||
scores = torch.randn(2, n_tiles, n_tiles)
|
||||
keep = math.ceil(0.5 * num_video)
|
||||
|
||||
exempt = _topk_tile_lists(scores, num_prefix, num_video, 0.5, True)
|
||||
assert exempt.shape == (2, num_video, num_prefix + keep)
|
||||
assert exempt.dtype == torch.int32
|
||||
assert torch.equal(exempt, exempt.sort(dim=-1).values)
|
||||
mask = _lists_to_mask(exempt, n_tiles)
|
||||
assert mask[..., :num_prefix].all()
|
||||
assert (mask[..., num_prefix:].sum(dim=-1) == keep).all()
|
||||
|
||||
compete = _topk_tile_lists(scores, num_prefix, num_video, 0.5, False)
|
||||
assert compete.shape == (2, num_video, min(keep + num_prefix, n_tiles))
|
||||
assert (_lists_to_mask(compete, n_tiles).sum(dim=-1) == keep + num_prefix).all()
|
||||
|
||||
|
||||
def _masked_dense_reference(meta, used, q, k, v, gate, sparsity):
|
||||
"""fp32 reference over the padded tile layout with the top-k tile mask."""
|
||||
n_tiles = meta.num_tiles
|
||||
seq_pad = n_tiles * VSA_H3_TILE_ELEMS
|
||||
valid = meta.pack_index >= 0
|
||||
src = meta.pack_index.clamp(min=0).long()
|
||||
|
||||
def tile(x):
|
||||
t = x[src].float().permute(1, 0, 2) # [H, S_pad, D]
|
||||
return t * valid.to(t.dtype)[None, :, None]
|
||||
|
||||
qt, kt, vt, gt = tile(q), tile(k), tile(v), tile(gate)
|
||||
sizes = meta.variable_block_sizes.float()
|
||||
pooled = [
|
||||
t.view(HEADS, n_tiles, VSA_H3_TILE_ELEMS, HEAD_DIM).sum(2)
|
||||
/ sizes[None, :, None]
|
||||
for t in (qt, kt, vt)
|
||||
]
|
||||
scores = pooled[0] @ pooled[1].transpose(-1, -2) / math.sqrt(HEAD_DIM)
|
||||
lists = _topk_tile_lists(
|
||||
scores, meta.num_prefix_tiles, meta.num_video_tiles, sparsity, meta.exempt
|
||||
)
|
||||
tile_mask = torch.ones(HEADS, n_tiles, n_tiles, dtype=torch.bool, device=q.device)
|
||||
tile_mask[:, meta.num_prefix_tiles :] = _lists_to_mask(lists.cpu(), n_tiles).to(
|
||||
q.device
|
||||
)
|
||||
row_mask = tile_mask.repeat_interleave(VSA_H3_TILE_ELEMS, 1).repeat_interleave(
|
||||
VSA_H3_TILE_ELEMS, 2
|
||||
)
|
||||
row_mask &= valid[None, None, :]
|
||||
logits = qt @ kt.transpose(-1, -2) / math.sqrt(HEAD_DIM)
|
||||
logits = logits.masked_fill(~row_mask, float("-inf"))
|
||||
out = torch.softmax(logits, dim=-1) @ vt
|
||||
compress = torch.softmax(scores, dim=-1) @ pooled[2] # [H, n_tiles, D]
|
||||
out = out + compress.repeat_interleave(VSA_H3_TILE_ELEMS, 1) * gt
|
||||
return out[:, meta.unpack_index.long()].permute(1, 0, 2) # [used, H, D]
|
||||
|
||||
|
||||
@requires_cuda
|
||||
def test_sparse_gated_matches_masked_dense_reference() -> None:
|
||||
device = torch.device("cuda")
|
||||
for exempt in (True, False):
|
||||
meta = _build_metadata(0.5, device, exempt=exempt)
|
||||
used, total, (q, k, v) = _packed_qkv(device)
|
||||
gate = (torch.randn_like(q) * 0.1).to(torch.bfloat16)
|
||||
out = _run(_impl(), meta, used, total, q, k, v, gate=gate)
|
||||
reference = _masked_dense_reference(meta, used, q, k, v, gate, 0.5)
|
||||
diff = (out[:used].float() - reference).abs().max().item()
|
||||
assert diff < 2e-2, f"exempt={exempt}: sparse+gate vs reference {diff}"
|
||||
assert torch.all(out[used:] == 0)
|
||||
|
||||
|
||||
def test_metadata_tile_geometry_accounts_every_row() -> None:
|
||||
device = torch.device("cpu")
|
||||
meta = _build_metadata(0.9, device)
|
||||
used = sum(PREFIX_SEGMENTS) + math.prod(VIDEO_SHAPE)
|
||||
assert int(meta.variable_block_sizes.sum()) == used
|
||||
assert meta.unpack_index.numel() == used
|
||||
assert int((meta.pack_index >= 0).sum()) == used
|
||||
assert meta.pack_index.numel() == meta.num_tiles * VSA_H3_TILE_ELEMS
|
||||
# Prefix chunks never straddle segment boundaries: 70 -> 64+6, 100 -> 64+36.
|
||||
assert meta.num_prefix_tiles == 4
|
||||
assert meta.variable_block_sizes[: meta.num_prefix_tiles].tolist() == [
|
||||
64,
|
||||
6,
|
||||
64,
|
||||
36,
|
||||
]
|
||||
video_tiles = (
|
||||
math.ceil(VIDEO_SHAPE[0] / 4)
|
||||
* math.ceil(VIDEO_SHAPE[1] / 4)
|
||||
* math.ceil(VIDEO_SHAPE[2] / 4)
|
||||
)
|
||||
assert meta.num_video_tiles == video_tiles
|
||||
assert meta.variable_block_sizes.numel() == meta.num_prefix_tiles + video_tiles
|
||||
assert VSA_H3_TILE_ELEMS == 64
|
||||
Reference in New Issue
Block a user