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 6b0fac19d..ab956dcf1 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py @@ -24,7 +24,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config impo from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionRequirements, ) -from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend +from sglang.multimodal_gen.runtime.layers.attention.selector import ( + get_attn_backend, + get_global_forced_attn_backend, +) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( LAYERWISE_OFFLOAD, ) @@ -101,6 +104,32 @@ class MiniMaxH3PipelineConfig(PipelineConfig): def _server_arg_value(value): return getattr(value, "value", value) + def resolve_transformer_attention_backend( + self, server_args + ) -> AttentionBackendEnum | None: + """Resolve the H3 DiT backend using the selector's precedence.""" + selected_backend = get_global_forced_attn_backend() + if selected_backend is None: + selected_backend, _ = server_args.resolve_component_attention_backend( + "transformer" + ) + if selected_backend is not None: + return selected_backend + attention_backend = server_args.attention_backend + if attention_backend is None or isinstance( + attention_backend, AttentionBackendEnum + ): + return attention_backend + attention_backend = self._server_arg_value(attention_backend) + return AttentionBackendEnum[str(attention_backend).strip().upper()] + + def uses_subblock_attention(self, server_args) -> bool: + """Return whether H3 must build SubBlock-only request metadata.""" + return ( + self.resolve_transformer_attention_backend(server_args) + is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN + ) + def validate_quality_deployment(self, server_args) -> None: """Fail closed unless the resident server matches the deployment audited for quality="high".""" @@ -220,17 +249,17 @@ class MiniMaxH3PipelineConfig(PipelineConfig): "MiniMax-H3 MPS execution does not support torch.compile; " "pass --enable-torch-compile false" ) - component_backends = server_args.component_attention_backends or {} - attention_backend = component_backends.get( - "transformer", self._server_arg_value(server_args.attention_backend) - ) - if attention_backend is None: + selected_backend = self.resolve_transformer_attention_backend(server_args) + if ( + int(server_args.ring_degree or 1) > 1 + and selected_backend is not AttentionBackendEnum.FA + ): + raise ValueError( + "MiniMax-H3 ring parallelism requires the FlashAttention " + "backend for the transformer" + ) + if selected_backend is None: return - selected_backend = ( - attention_backend - if isinstance(attention_backend, AttentionBackendEnum) - else AttentionBackendEnum[str(attention_backend).strip().upper()] - ) get_attn_backend( self.dit_config.arch_config.attention_head_dim, torch.bfloat16, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py index f701b2640..5ebffa33e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py @@ -243,8 +243,9 @@ class SubBlockRouter: scores = self.scores(q, k, softmax_scale) # [B, H, Gq, Gk] gq = scores.shape[2] topk = _snap_up_to_8(math.ceil((1.0 - sparsity) * gk), gk) - # One pass over the score matrix instead of torch.topk's several; the kernel - # accepts the blocks in any order, so nothing sorts them. + # One pass over the score matrix instead of torch.topk's several. The + # output order is unspecified: SM100 consumes it directly, while the + # SM90 backend sorts compact active prefixes before heterogeneous expansion. index = fused_topk(scores.reshape(-1, gk), topk).view(b, h, gq, topk) return RoutingPlan(index=index, topk=topk, num_blocks=gk) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py index b1c10fd5c..f2bb780e1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py @@ -137,22 +137,21 @@ def _sm90_sparse_attention( q2k_block_index: torch.Tensor, topk: int, softmax_scale: float, + block_counts: torch.Tensor | None = None, ) -> torch.Tensor: """Run a SubBlock routing plan through the existing SM90 CuTe kernel.""" BlockSparseTensorsTorch, flash_attn_func = _load_sm90_block_sparse_attention() - # The router contract permits indices in any order, while the SM90 sparse - # pipeline consumes each list from high slot to low slot and applies - # sequence-tail masking to the first block. Sort explicitly so the largest - # block id -- the possible ragged tail -- occupies the highest slot without - # depending on the fused top-k kernel's current ascending output order. - ordered_index = q2k_block_index.sort(dim=-1).values - block_counts = torch.full( - ordered_index.shape[:-1], - topk, - dtype=torch.int32, - device=ordered_index.device, - ) + # The caller sorts each active sparse prefix for SM90. Dense rows are the + # already-sorted complete range; entries beyond each row's count are ignored. + ordered_index = q2k_block_index + if block_counts is None: + block_counts = torch.full( + ordered_index.shape[:-1], + topk, + dtype=torch.int32, + device=ordered_index.device, + ) sparse_tensors = BlockSparseTensorsTorch( mask_block_cnt=block_counts, mask_block_idx=ordered_index, @@ -183,6 +182,7 @@ def _sm100_sparse_attention( q2k_block_index: torch.Tensor, topk: int, softmax_scale: float, + block_counts: torch.Tensor | None = None, ) -> torch.Tensor: """Run a SubBlock routing plan through FlashInfer's SM100 kernel.""" out = load_bsa_attn_blk64_fwd()( @@ -192,7 +192,7 @@ def _sm100_sparse_attention( q2k_block_index, topk, block_sizes=_cached_block_sizes(k.shape[1], k.device), - q2k_block_nums=None, # the budget is uniform across rows + q2k_block_nums=block_counts, softmax_scale=softmax_scale, ) return out[0] if isinstance(out, tuple) else out @@ -219,8 +219,14 @@ def _run_subblock_sparse_attention( q2k_block_index: torch.Tensor, topk: int, softmax_scale: float, + block_counts: torch.Tensor | None = None, ) -> torch.Tensor: - """Dispatch the same 64x64 routing plan to Hopper or Blackwell.""" + """Dispatch a prepared 64x64 routing plan to Hopper or Blackwell. + + SM90 requires every active index prefix to be sorted in ascending order; + SM100 accepts the router's original order. Heterogeneous callers must sort + compact sparse prefixes before expanding them to full-width dense rows. + """ runner = _get_subblock_sparse_attention_runner(q.device) return runner( q, @@ -229,6 +235,7 @@ def _run_subblock_sparse_attention( q2k_block_index, topk, softmax_scale, + block_counts, ) @@ -401,26 +408,90 @@ class SubBlockSparseAttentionImpl(AttentionImpl): ) def _sparse_attention( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + sparse_query_block_mask: torch.Tensor | None = None, ) -> torch.Tensor: - """q, k, v: ``[1, S, H, 128]`` bf16 -> same shape.""" + """Q ``[1, Sq, H, 128]`` against K/V ``[1, Sk, H, 128]``.""" plan = self.router.route( - q, k, sparsity=self.schedule.sparsity, softmax_scale=self.softmax_scale + q, + k, + sparsity=self.schedule.sparsity, + softmax_scale=self.softmax_scale, ) - # Proof that the sparse path actually ran, with the shape it ran on -- - # the construction-time log above only says the layer was eligible. - logger.info_once( - f"SubBlock sparse attention active: S={k.shape[1]} heads={q.shape[2]} " - f"keeping {plan.topk}/{plan.num_blocks} key blocks per query block " - f"(sparsity {1 - plan.density:.4f})" + expected_q_blocks = -(-q.shape[1] // SUBBLOCK_SPARSE_BLOCK_SIZE) + if plan.index.shape[2] != expected_q_blocks: + raise ValueError( + "SubBlock routing/kernel query-block mismatch: " + f"plan has {plan.index.shape[2]}, kernel needs {expected_q_blocks}" + ) + # Proof that the sparse path actually ran -- the construction-time log + # above only says the layer was eligible. + if sparse_query_block_mask is None: + logger.info_once( + f"SubBlock sparse attention active: Sq={q.shape[1]} " + f"Sk={k.shape[1]} heads={q.shape[2]} " + f"keeping {plan.topk}/{plan.num_blocks} key blocks per query " + f"block (sparsity {1 - plan.density:.4f})" + ) + else: + logger.info_once( + f"SubBlock heterogeneous BSA active: Sq={q.shape[1]} " + f"Sk={k.shape[1]} heads={q.shape[2]}; selected query blocks " + f"keep {plan.topk}/{plan.num_blocks} key blocks and unselected " + "query blocks are dense" + ) + block_counts = None + runner = _get_subblock_sparse_attention_runner(q.device) + block_index = ( + plan.index.sort(dim=-1).values + if runner is _sm90_sparse_attention + else plan.index ) + kernel_topk = plan.topk + if sparse_query_block_mask is not None: + sparse_query_block_mask = sparse_query_block_mask.to( + device=q.device, dtype=torch.bool + ).view(-1) + if sparse_query_block_mask.numel() != expected_q_blocks: + raise ValueError( + "SubBlock sparse query-block mask length does not match Q" + ) + num_k_blocks = plan.num_blocks + full_index = torch.arange( + num_k_blocks, device=q.device, dtype=block_index.dtype + ).view(1, 1, 1, num_k_blocks) + heterogeneous_index = full_index.expand( + *block_index.shape[:-1], num_k_blocks + ).clone() + sparse_rows = sparse_query_block_mask.view(1, 1, -1, 1) + heterogeneous_index[..., : plan.topk] = torch.where( + sparse_rows, + block_index, + heterogeneous_index[..., : plan.topk], + ) + block_counts = ( + torch.where( + sparse_query_block_mask.view(1, 1, -1), + plan.topk, + num_k_blocks, + ) + .expand(*block_index.shape[:-1]) + .to(torch.int32) + ) + block_index = heterogeneous_index + kernel_topk = num_k_blocks return _run_subblock_sparse_attention( q, k, v, - plan.index, - plan.topk, + block_index, + kernel_topk, self.softmax_scale, + block_counts, ) def forward( @@ -444,12 +515,15 @@ class SubBlockSparseAttentionImpl(AttentionImpl): cu_seqlens: torch.Tensor, max_seqlen: int, cu_seqlens_host: tuple[int, ...] | None = None, + first_segment_sparse_query_block_mask: torch.Tensor | None = None, ) -> torch.Tensor: """Packed ``[T, H, D]`` rows split into documents by ``cu_seqlens``. - The block-sparse kernel takes one contiguous sequence, so each packed - document is routed on its own. Documents shorter than ``min_seq_len`` - -- in MiniMax H3 the padding tail -- go through the dense kernel. + Each packed document keeps its own full K/V context. The optional + first-segment mask selects sparse Q blocks; unselected blocks stay + dense within the same heterogeneous BSA call. + Documents shorter than ``min_seq_len`` -- in H3, the padding tail -- + stay on the existing dense segment path. """ def all_dense() -> torch.Tensor: @@ -487,14 +561,24 @@ class SubBlockSparseAttentionImpl(AttentionImpl): for start, stop in segments: # Deliberately not `.contiguous()`. After the Ulysses all-to-all, # q/k/v are last-dim slices of one packed buffer, so they are - # strided; both the block-sparse kernel and SDPA permute them - # anyway, and forcing contiguity here measured as a wasted + # strided; the attention kernels handle those views directly, and + # forcing contiguity here measured as a wasted # full-tensor copy (0.46 ms per call at S=37.7k on B200). q_seg = query[start:stop].unsqueeze(0) k_seg = key[start:stop].unsqueeze(0) v_seg = value[start:stop].unsqueeze(0) if (start, stop) in sparse_segments: - seg_out = self._sparse_attention(q_seg, k_seg, v_seg) + sparse_query_block_mask = ( + first_segment_sparse_query_block_mask + if (start, stop) == segments[0] + else None + ) + seg_out = self._sparse_attention( + q_seg, + k_seg, + v_seg, + sparse_query_block_mask=sparse_query_block_mask, + ) else: seg_out = self._dense_segment(q_seg, k_seg, v_seg) out[start:stop] = seg_out[0] 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 ecd0d34b2..aaa51a8bb 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -51,7 +51,11 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionRequirements, ) -from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend +from sglang.multimodal_gen.runtime.layers.attention.selector import ( + get_attn_backend, + get_component_forced_attn_backend, + get_global_forced_attn_backend, +) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, @@ -176,6 +180,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( "token_tags", "block_token_tags", "block_combined_indices", + "subblock_sparse_query_block_mask", "skip_mask_out_condition", "prompt_embeds", "refined_prompt_embeds_length", @@ -559,6 +564,7 @@ def _minimax_h3_attention_core_impl( cu_seqlens_host: tuple[int, ...] | None, max_seqlen: int, ulysses_active: bool, + subblock_sparse_query_block_mask: torch.Tensor | None = None, ring_active: bool = False, ) -> torch.Tensor: """Dynamic varlen attention and Ulysses/Ring collectives. @@ -603,14 +609,47 @@ def _minimax_h3_attention_core_impl( ring_ws=ring_ws, ) else: - out = attention._attention_impl.forward_varlen( - q, - k, - v, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - cu_seqlens_host=cu_seqlens_host, - ) + if ( + attention._attention_backend_enum + is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN + ): + impl = attention._attention_impl + sparse_will_run = ( + cu_seqlens_host is not None + and impl._sparse_ready(q, k) + and any( + stop - start >= impl.schedule.min_seq_len + for start, stop in zip( + cu_seqlens_host[:-1], + cu_seqlens_host[1:], + ) + ) + ) + if sparse_will_run and subblock_sparse_query_block_mask is None: + raise ValueError( + "MiniMax H3 requires subblock_sparse_query_block_mask " + "when SubBlock sparse attention is active" + ) + out = attention._attention_impl.forward_varlen( + q, + k, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + first_segment_sparse_query_block_mask=( + subblock_sparse_query_block_mask + ), + ) + else: + out = attention._attention_impl.forward_varlen( + q, + k, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + cu_seqlens_host=cu_seqlens_host, + ) if ulysses_active: out = _usp_output_all_to_all(out[None], head_dim=2)[0] return out @@ -867,6 +906,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens: torch.Tensor, cu_seqlens_host: tuple[int, ...] | None = None, max_seqlen: int, + subblock_sparse_query_block_mask: torch.Tensor | None = None, ulysses_active: bool = False, ring_active: bool = False, ) -> torch.Tensor: @@ -943,6 +983,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen, + subblock_sparse_query_block_mask=subblock_sparse_query_block_mask, ulysses_active=ulysses_active, ring_active=ring_active, ) @@ -1485,6 +1526,7 @@ class MiniMaxH3DiTBlock(nn.Module): cu_seqlens: torch.Tensor, cu_seqlens_host: tuple[int, ...] | None = None, max_seqlen: int, + subblock_sparse_query_block_mask: torch.Tensor | None = None, ulysses_active: bool = False, ring_active: bool = False, adaln_params: tuple[torch.Tensor, ...] | None = None, @@ -1515,6 +1557,7 @@ class MiniMaxH3DiTBlock(nn.Module): cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen, + subblock_sparse_query_block_mask=subblock_sparse_query_block_mask, ulysses_active=ulysses_active, ring_active=ring_active, ) @@ -1965,6 +2008,9 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): if self._adaln_precomputed else None ) + # Component overrides disappear when the loader context exits. Preserve + # only that selection; process-wide overrides are resolved at first use. + self._component_attention_backend_override = get_component_forced_attn_backend() self._resolved_attention_backend: AttentionBackendEnum | None = None self._mark_missing_params_required() @@ -1987,9 +2033,14 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): def _resolve_attention_backend_once(self) -> None: if self._resolved_attention_backend is not None: return + selected_backend = ( + get_global_forced_attn_backend() + or self._component_attention_backend_override + ) backend = get_attn_backend( self.arch.attention_head_dim, _BF16_DTYPE, + selected_attention_backend=selected_backend, attention_requirements=AttentionRequirements(packed_varlen=True), ) for module in self.modules(): @@ -2333,6 +2384,9 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): _required_kwarg(kwargs, "inverse_indices").view(-1).to(torch.long) ) update_mask = _required_kwarg(kwargs, "update_mask") + subblock_sparse_query_block_mask = kwargs.get( + "subblock_sparse_query_block_mask" + ) block_token_tags = kwargs.get("block_token_tags") token_tags = kwargs.get("token_tags") if block_token_tags is None: @@ -2395,6 +2449,10 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}" ) device = x.device + if subblock_sparse_query_block_mask is not None and not isinstance( + subblock_sparse_query_block_mask, torch.Tensor + ): + raise ValueError("subblock_sparse_query_block_mask must be a tensor") self._resolve_attention_backend_once() # Row split is 2D: ring first (an outer, contiguous ring_chunk_len @@ -2530,6 +2588,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): cu_seqlens=cu_seqlens, cu_seqlens_host=cu_seqlens_host, max_seqlen=max_seqlen, + subblock_sparse_query_block_mask=subblock_sparse_query_block_mask, ulysses_active=ulysses_ws > 1, ring_active=ring_ws > 1, adaln_params=( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py index 8f914faf9..0ca97cb20 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py @@ -28,6 +28,69 @@ MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0 # (24 * 1 * 2 * 2 = 96); audio rows carry the 32-dim audio latent. MINIMAX_H3_VIDEO_ROW_WIDTH = 96 MINIMAX_H3_AUDIO_ROW_WIDTH = 32 +_MINIMAX_H3_SUBBLOCK_QUERY_BLOCK_SIZE = 64 + + +def _minimax_h3_subblock_video_query_indices( + packed: dict[str, torch.Tensor], + text_video_token_mask: torch.Tensor | None, +) -> torch.Tensor: + """Return packed video rows plus any Qwen rows marked as video.""" + latent_video_pos = packed["video_pos"].view(-1).to(dtype=torch.long) + if text_video_token_mask is None: + return latent_video_pos + text_pos = packed["text_pos"].view(-1).to(dtype=torch.long) + text_video_token_mask = text_video_token_mask.view(-1).to( + device=text_pos.device, + dtype=torch.bool, + ) + if text_video_token_mask.shape[0] != text_pos.shape[0]: + raise ValueError( + "text_video_token_mask must align with packed text rows " + f"({text_pos.shape[0]}), got {text_video_token_mask.shape[0]}" + ) + return torch.cat([text_pos[text_video_token_mask], latent_video_pos]) + + +def _minimax_h3_subblock_sparse_query_block_mask( + video_query_indices: torch.Tensor, + *, + used_len: int, +) -> torch.Tensor: + """Return True for pure-video 64-row Q blocks and False otherwise.""" + if used_len < 0: + raise ValueError(f"used_len must be non-negative, got {used_len}") + if video_query_indices.ndim != 1: + raise ValueError( + f"video_query_indices must be rank 1, got {list(video_query_indices.shape)}" + ) + video_query_indices = video_query_indices.to(dtype=torch.long) + if video_query_indices.numel(): + first = int(video_query_indices.min()) + last = int(video_query_indices.max()) + if first < 0 or last >= used_len: + raise ValueError( + "video_query_indices must be first-segment-relative and in " + f"[0, {used_len}), got min={first}, max={last}" + ) + if torch.unique(video_query_indices).numel() != video_query_indices.numel(): + raise ValueError("video_query_indices must not contain duplicates") + block_size = _MINIMAX_H3_SUBBLOCK_QUERY_BLOCK_SIZE + num_query_blocks = -(-used_len // block_size) + video_rows_per_block = torch.bincount( + torch.div(video_query_indices, block_size, rounding_mode="floor"), + minlength=num_query_blocks, + ) + block_ids = torch.arange( + num_query_blocks, device=video_query_indices.device, dtype=torch.long + ) + real_rows_per_block = (used_len - block_ids * block_size).clamp( + min=0, max=block_size + ) + # BSA is block-granular. Only blocks whose real rows are all video may use + # the sparse budget; mixed boundary blocks stay dense so their non-video + # rows retain exact attention in the single heterogeneous BSA call. + return video_rows_per_block == real_rows_per_block @torch.inference_mode() @@ -95,7 +158,9 @@ class MiniMaxH3DenoiseBranch: `packed` is a minimax_h3_packed_sequence(...) result (or equivalent layout dict); `text_embeddings` is the branch's [text_len, 5120] hidden states; - `token_tags` must already carry any fl2va vision-span overrides. + `token_tags` must already carry any fl2va vision-span overrides, and + `video_query_indices` explicitly identifies the only rows eligible for + SubBlock sparsity. """ def __init__( @@ -105,6 +170,7 @@ class MiniMaxH3DenoiseBranch: text_embeddings: torch.Tensor, token_tags: torch.Tensor, device: torch.device, + video_query_indices: torch.Tensor | None = None, ) -> None: seq_len = int(packed["seq_len"]) self.seq_len = seq_len @@ -183,6 +249,14 @@ class MiniMaxH3DenoiseBranch: sp_world_size = ulysses_world_size * ring_world_size sp_rank = ring_rank * ulysses_world_size + ulysses_rank token_tags_host = token_tags.view(-1).to(dtype=torch.long) + subblock_sparse_query_block_mask = ( + _minimax_h3_subblock_sparse_query_block_mask( + video_query_indices.view(-1).to(dtype=torch.long), + used_len=int(cu[1]), + ).to(device) + if video_query_indices is not None + else None + ) local_seq_len = seq_len // sp_world_size local_row_start = sp_rank * local_seq_len local_row_stop = local_row_start + local_seq_len @@ -228,6 +302,10 @@ class MiniMaxH3DenoiseBranch: "max_seqlen_q": text_len, }, } + if subblock_sparse_query_block_mask is not None: + self.static_kwargs["subblock_sparse_query_block_mask"] = ( + subblock_sparse_query_block_mask + ) def forward_kwargs( self, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py index 45a27926f..cd7fd66b5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/packed_sequence.py @@ -126,6 +126,7 @@ def minimax_h3_packed_sequence( include_keyframe_cond: bool, keyframe_frame_indices: list[int] | tuple[int, ...] | None = None, frame_count: int | None = None, + include_video_pos: bool = False, ) -> dict[str, Any]: """Build the packed-sequence structural fields for one CFG branch. @@ -204,10 +205,10 @@ def minimax_h3_packed_sequence( token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING token_tags[text_sl] = 1 # TEXT (fl2va image-segment override happens upstream) token_tags[audio_sl] = 2 # AUDIO - token_tags[img_pos] = 0 # VIDEO + token_tags[img_pos] = 0 # VISUAL (condition images + target video) cu = torch.tensor([0, used, seq_len], dtype=torch.int32) - return { + packed = { "seq_len": seq_len, "img_pos": img_pos, "audio_pos": audio_pos, @@ -217,6 +218,11 @@ def minimax_h3_packed_sequence( "token_tags": token_tags, "cu_seqlens": cu, } + if include_video_pos: + # Conditioning keyframes are images. Only generated video rows are + # eligible for SubBlock sparsity. + packed["video_pos"] = target_img_pos + return packed def _positive_int( @@ -283,6 +289,7 @@ def minimax_h3_packed_sequence_ref2va_blocks( frame_count: int | None = None, audio_channel: int = 2, seq_len: int | None = None, + include_video_pos: bool = False, ) -> dict[str, Any]: """General ref2va-family packed layout. @@ -398,6 +405,7 @@ def minimax_h3_packed_sequence_ref2va_blocks( audio_sl = slice(cursor, cursor + audio_rows) video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows) ref_img_pos_parts: list[torch.Tensor] = [] + ref_video_pos_parts: list[torch.Tensor] | None = [] if include_video_pos else None ref_audio_pos_parts: list[torch.Tensor] = [] g = torch.zeros(seq_len, 3, dtype=torch.float64) g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64) @@ -448,7 +456,10 @@ def minimax_h3_packed_sequence_ref2va_blocks( vh = int(item["latent_h"]) vw = int(item["latent_w"]) ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl)) - ref_img_pos_parts.append(_range_for_slice(visual_sl)) + visual_pos = _range_for_slice(visual_sl) + ref_img_pos_parts.append(visual_pos) + if ref_video_pos_parts is not None: + ref_video_pos_parts.append(visual_pos) ref_area = np.sqrt(vh * vw) rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area) @@ -512,10 +523,10 @@ def minimax_h3_packed_sequence_ref2va_blocks( token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING token_tags[text_sl] = 1 # TEXT token_tags[audio_pos] = 2 # AUDIO (refs + target) - token_tags[img_pos] = 0 # VIDEO (refs + target) + token_tags[img_pos] = 0 # VISUAL (reference images/videos + target video) cu = torch.tensor([0, used, seq_len], dtype=torch.int32) - return { + packed = { "seq_len": seq_len, "img_pos": img_pos, "audio_pos": audio_pos, @@ -526,6 +537,11 @@ def minimax_h3_packed_sequence_ref2va_blocks( "token_tags": token_tags, "cu_seqlens": cu, } + if ref_video_pos_parts is not None: + # Reference image blocks remain dense; reference videos and the + # generated target video are eligible for SubBlock sparsity. + packed["video_pos"] = _cat_ranges(ref_video_pos_parts + [target_img_pos]) + return packed __all__ = [ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py index 7473d6390..3f4b835e1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/presentation.py @@ -24,7 +24,7 @@ IMAGE_PAD = "<|image_pad|>" VIDEO_PAD = "<|video_pad|>" _TEXT_TAG = 1 -_VIDEO_TAG = 0 +_VISUAL_TAG = 0 def _text_ids(tokenizer: Any, text: str) -> list[int]: @@ -40,25 +40,45 @@ def _vision_block_ids(tokenizer: Any, pad_token: str, count: int) -> list[int]: class _Presentation: - """Accumulates aligned (ids, token_tags) presentation segments.""" + """Accumulates aligned ids, modality tags, and video-query metadata.""" - def __init__(self) -> None: + def __init__(self, *, track_video_mask: bool = False) -> None: self.ids: list[int] = [] self.tags: list[int] = [] + self.video_mask: list[bool] | None = [] if track_video_mask else None def text(self, token_ids: list[int]) -> None: self.ids += token_ids self.tags += [_TEXT_TAG] * len(token_ids) + if self.video_mask is not None: + self.video_mask += [False] * len(token_ids) - def vision(self, token_ids: list[int]) -> None: + def vision( + self, token_ids: list[int], *, video_token_id: int | None = None + ) -> None: self.ids += token_ids - self.tags += [_VIDEO_TAG] * len(token_ids) + self.tags += [_VISUAL_TAG] * len(token_ids) + if self.video_mask is not None: + self.video_mask += [ + video_token_id is not None and token_id == video_token_id + for token_id in token_ids + ] - def build(self) -> tuple[torch.Tensor, torch.Tensor]: - return ( + def build( + self, *, return_video_mask: bool = False + ) -> ( + tuple[torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ): + result = ( torch.tensor(self.ids, dtype=torch.long), torch.tensor(self.tags, dtype=torch.long), ) + if not return_video_mask: + return result + if self.video_mask is None: + raise ValueError("video mask was not tracked for this presentation") + return (*result, torch.tensor(self.video_mask, dtype=torch.bool)) def _timestamped_video_blocks( @@ -68,6 +88,7 @@ def _timestamped_video_blocks( counts: Sequence[int], timestamps: Sequence[float], context: str, + video_token_id: int | None, ) -> None: """Emit per-temporal-block ``<{t:.1f} seconds>`` text + VIDEO vision.""" @@ -79,7 +100,10 @@ def _timestamped_video_blocks( if count <= 0: raise ValueError(f"{context}video block token count must be positive") presentation.text(_text_ids(tokenizer, f"<{timestamp:.1f} seconds>")) - presentation.vision(_vision_block_ids(tokenizer, VIDEO_PAD, count)) + presentation.vision( + _vision_block_ids(tokenizer, VIDEO_PAD, count), + video_token_id=video_token_id, + ) def minimax_h3_text_only_ids(tokenizer: Any, prompt: str) -> torch.Tensor: @@ -118,8 +142,8 @@ def minimax_h3_ref2va_presentation( per condition in request order — image i: ``: `` label followed by the vision block; audio j: ``