@@ -549,6 +615,16 @@ end-to-end latency on the target model and resolution before deployment.
| ❌ |
CUDA-only. Requires vsa. Configure sparsity via --attention-backend-config. |
+
+ | `video_sparse_attn_h3` |
+ Yes |
+ No |
+ No |
+ No |
+ ❌ |
+ ❌ |
+ CUDA-only (SM90 / SM100 / SM103). In-tree Triton kernel, no external dependency. Configure via --attention-backend-config. |
+
sla_attn |
Yes |
diff --git a/docs/docs/sglang-diffusion/compatibility_matrix.mdx b/docs/docs/sglang-diffusion/compatibility_matrix.mdx
index ffdb56e1a..4f3ed83fd 100644
--- a/docs/docs/sglang-diffusion/compatibility_matrix.mdx
+++ b/docs/docs/sglang-diffusion/compatibility_matrix.mdx
@@ -137,6 +137,12 @@ Rows are grouped when a family shares the same runtime path or optimization supp
T2VA / FL2VA / Ref2VA, 768p at 24 fps with synchronized audio |
Cache-DiTSageOnline FP8GGUF |
+
+ | FastH3 |
+ FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree
|
+ T2VA only, 4-step distilled, 768p at 24 fps with synchronized audio |
+ VSA-H3 |
+
| Wan2.1 Fun |
weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers
|
@@ -621,6 +627,21 @@ Optimization columns are abbreviated to keep the matrix readable:
❌ |
❌ |
+
+ | FastH3 4-step (T2VA only) |
+ FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree |
+ 768p · 24 fps |
+ ❌ |
+ ❌ |
+ ✅ |
+ ✅ |
+ ❌ |
+ ❌ |
+ ❌ |
+ ❌ |
+ ❌ |
+ ❌ |
+
| LTX-2.3 (one/two-stage/TI2V/HQ) |
Lightricks/LTX-2.3 |
@@ -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"`.
diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx
index 8bad3ff07..18116f4cb 100644
--- a/docs/docs/sglang-diffusion/quantization.mdx
+++ b/docs/docs/sglang-diffusion/quantization.mdx
@@ -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.
`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`.
diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md
index b6095a058..efef0fe80 100644
--- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md
+++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md
@@ -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. |
diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py
index b9269b5b9..6133d5f97 100755
--- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py
+++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py
@@ -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",
diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md
index 5ee2aef96..c6e97c4dd 100644
--- a/python/sglang/multimodal_gen/README.md
+++ b/python/sglang/multimodal_gen/README.md
@@ -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:
diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py
index 49b434003..63b5dd035 100644
--- a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py
@@ -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__()
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py
index 71e3e959a..b3847a44b 100644
--- a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py
+++ b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_audio.py
@@ -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,
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py
index d23b6740e..4fcc6540f 100644
--- a/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py
+++ b/python/sglang/multimodal_gen/configs/models/vaes/minimax_h3_video.py
@@ -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(
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
index 5621aff2e..9f35fa9c7 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
@@ -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",
]
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
index 1cacfe3ef..20b8e829c 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
@@ -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"]
diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
index 603500183..85934ec90 100644
--- a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
+++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
@@ -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"]
diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py
index 9d09359ad..cf76d3422 100644
--- a/python/sglang/multimodal_gen/registry.py
+++ b/python/sglang/multimodal_gen/registry.py
@@ -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,
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py
new file mode 100644
index 000000000..8367dce0f
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn_h3.py
@@ -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
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vsa_h3_kernels.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vsa_h3_kernels.py
new file mode 100644
index 000000000..e809cc8f1
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vsa_h3_kernels.py
@@ -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,
+ )
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
index 19d2bf39c..50f572b3d 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -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,
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py
index 8aee42fcf..046fc4648 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines/minimax_h3_pipeline.py
@@ -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]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
index 99f61de19..6b48354ae 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
@@ -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],
diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py
index 2f6b95df5..c23038695 100644
--- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py
+++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py
@@ -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,
diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py
index 17627e3f8..c6efec9ca 100644
--- a/python/sglang/multimodal_gen/runtime/platforms/interface.py
+++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py
@@ -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,
diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py
index 7373a8301..408332cc0 100644
--- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py
+++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py
@@ -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",
+ },
}
diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py
index 775a9531c..d12a1fff8 100644
--- a/python/sglang/multimodal_gen/test/server/gpu_cases.py
+++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py
@@ -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 = [
diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_benchmark_skill.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_benchmark_skill.py
index 795626b82..1b4c1bd50 100644
--- a/python/sglang/multimodal_gen/test/unit/test_diffusion_benchmark_skill.py
+++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_benchmark_skill.py
@@ -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",
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py
new file mode 100644
index 000000000..269b8c81e
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py
@@ -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
diff --git a/python/sglang/multimodal_gen/test/unit/test_vsa_h3_attention.py b/python/sglang/multimodal_gen/test/unit/test_vsa_h3_attention.py
new file mode 100644
index 000000000..999eb2daf
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_vsa_h3_attention.py
@@ -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