[Diffusion][minimax-h3] Restrict MiniMax-H3 SubBlock sparsity to video queries (#35850)

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 {}
+        selected_backend = self.resolve_transformer_attention_backend(server_args)
-        attention_backend = component_backends.get(
+        if (
-            "transformer", self._server_arg_value(server_args.attention_backend)
+            int(server_args.ring_degree or 1) > 1
-        )
+            and selected_backend is not AttentionBackendEnum.FA
-        if attention_backend is None:
+        ):
+            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,

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
+        # One pass over the score matrix instead of torch.topk's several. The
-        # accepts the blocks in any order, so nothing sorts them.
+        # 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)
 

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
+    # The caller sorts each active sparse prefix for SM90. Dense rows are the
-    # pipeline consumes each list from high slot to low slot and applies
+    # already-sorted complete range; entries beyond each row's count are ignored.
-    # sequence-tail masking to the first block. Sort explicitly so the largest
+    ordered_index = q2k_block_index
-    # block id -- the possible ragged tail -- occupies the highest slot without
+    if block_counts is None:
-    # depending on the fused top-k kernel's current ascending output order.
+        block_counts = torch.full(
-    ordered_index = q2k_block_index.sort(dim=-1).values
+            ordered_index.shape[:-1],
-    block_counts = torch.full(
+            topk,
-        ordered_index.shape[:-1],
+            dtype=torch.int32,
-        topk,
+            device=ordered_index.device,
-        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 --
+        expected_q_blocks = -(-q.shape[1] // SUBBLOCK_SPARSE_BLOCK_SIZE)
-        # the construction-time log above only says the layer was eligible.
+        if plan.index.shape[2] != expected_q_blocks:
-        logger.info_once(
+            raise ValueError(
-            f"SubBlock sparse attention active: S={k.shape[1]} heads={q.shape[2]} "
+                "SubBlock routing/kernel query-block mismatch: "
-            f"keeping {plan.topk}/{plan.num_blocks} key blocks per query block "
+                f"plan has {plan.index.shape[2]}, kernel needs {expected_q_blocks}"
-            f"(sparsity {1 - plan.density:.4f})"
+            )
+        # 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,
+            block_index,
-            plan.topk,
+            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
+        Each packed document keeps its own full K/V context. The optional
-        document is routed on its own. Documents shorter than ``min_seq_len``
+        first-segment mask selects sparse Q blocks; unselected blocks stay
-        -- in MiniMax H3 the padding tail -- go through the dense kernel.
+        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
+            # strided; the attention kernels handle those views directly, and
-            # anyway, and forcing contiguity here measured as a wasted
+            # 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]

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(
+        if (
-            q,
+            attention._attention_backend_enum
-            k,
+            is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
-            v,
+        ):
-            cu_seqlens=cu_seqlens,
+            impl = attention._attention_impl
-            max_seqlen=max_seqlen,
+            sparse_will_run = (
-            cu_seqlens_host=cu_seqlens_host,
+                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=(

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,

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__ = [

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]:
+    def build(
-        return (
+        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: ``<Picture i>: `` label followed
     by the vision block; audio j: ``<Audio j>: `` label only (audio content
-    never enters Qwen) — then the verbatim prompt. Returns (ids, token_tags)
+    never enters Qwen) — then the verbatim prompt. Returns ``(ids, token_tags)``
-    with the vision block tagged VIDEO(0) and everything else TEXT(1).
+    with the vision block tagged visual(0) and everything else text(1).
 
     condition_labels: [("image", 1), ("audio", 1), ...] with 1-based ordinals
     per type.
@@ -196,7 +220,10 @@ def minimax_h3_ref2va_video_presentation(
     image_token_count: int | list[int] | None,
     video_block_token_counts: list[int] | list[list[int]] | None,
     video_block_timestamps: list[float] | list[list[float]] | None,
-) -> tuple[torch.Tensor, torch.Tensor]:
+    return_video_mask: bool = False,
+) -> (
+    tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+):
     """ref2va (optionally with video refs) positive presentation:
 
     per condition in request order —
@@ -209,12 +236,16 @@ def minimax_h3_ref2va_video_presentation(
       counts repeat the last frame), emitting the
       ``<0.2 seconds>`` ..
       ``<4.0 seconds>`` sequence — note Python bankers-rounding at .1f.
-    then the verbatim prompt. Vision blocks are tagged VIDEO(0), everything
+    then the verbatim prompt. Image and video blocks both retain visual
-    else TEXT(1).
+    modality tag 0. When requested, the third return marks only VIDEO_PAD
+    content tokens; vision delimiters remain dense.
     """
     if not prompt:
         raise ValueError("prompt must be non-empty")
-    presentation = _Presentation()
+    presentation = _Presentation(track_video_mask=return_video_mask)
+    video_token_id = (
+        tokenizer.convert_tokens_to_ids(VIDEO_PAD) if return_video_mask else None
+    )
     image_token_counts = _as_int_list(image_token_count, name="image_token_count")
     video_counts_by_ref = _as_nested_int_list(
         video_block_token_counts,
@@ -259,6 +290,7 @@ def minimax_h3_ref2va_video_presentation(
                 counts=counts,
                 timestamps=timestamps,
                 context="",
+                video_token_id=video_token_id,
             )
         else:
             raise ValueError(f"unsupported ref2va condition type {cond_type!r}")
@@ -267,7 +299,7 @@ def minimax_h3_ref2va_video_presentation(
     if video_seen != len(video_counts_by_ref):
         raise ValueError("unused video block token count entries")
     presentation.text(_text_ids(tokenizer, prompt))
-    return presentation.build()
+    return presentation.build(return_video_mask=return_video_mask)
 
 
 __all__ = [

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py

@@ -615,6 +615,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
         """
         from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import (
             MiniMaxH3DenoiseBranch,
+            _minimax_h3_subblock_video_query_indices,
             minimax_h3_denoise_loop,
         )
 
@@ -639,11 +640,28 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
         _assemble_condition_rows(ctx)
 
         emb = ctx.embeddings["positive"]
-        packed = _build_packed_layout(ctx, emb)
+        subblock_enabled = server_args.pipeline_config.uses_subblock_attention(
+            server_args
+        )
+        packed = _build_packed_layout(
+            ctx,
+            emb,
+            include_video_pos=subblock_enabled,
+        )
         tags = packed["token_tags"]
         tags[packed["text_pos"].view(-1)] = (
             emb["text_token_tags"].view(-1).to(torch.long)
         )
+        video_query_indices = None
+        if subblock_enabled:
+            text_video_token_mask = emb.get("text_video_token_mask")
+            # Legacy/precomputed presentations may lack this optional
+            # provenance. The helper then keeps their Qwen rows dense while
+            # retaining packed reference/target video rows as sparse.
+            video_query_indices = _minimax_h3_subblock_video_query_indices(
+                packed,
+                text_video_token_mask,
+            )
 
         sampling = batch.sampling_params
         imgvid_noise_aug, audio_noise_aug = minimax_h3_condition_noise_aug(sampling)
@@ -667,6 +685,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
                 packed=packed,
                 text_embeddings=emb["hidden_states"],
                 token_tags=tags,
+                video_query_indices=video_query_indices,
                 device=device,
             )
             _precompute_refined_prompt_embeds(
@@ -892,6 +911,8 @@ def _assemble_condition_rows(ctx: _FullLoopContext) -> None:
 def _build_packed_layout(
     ctx: _FullLoopContext,
     emb: Mapping[str, Any],
+    *,
+    include_video_pos: bool = False,
 ) -> dict[str, torch.Tensor]:
     """Build the per-task packed layout for the positive branch."""
 
@@ -912,6 +933,7 @@ def _build_packed_layout(
             ref_blocks=ctx.ref2va_positive_blocks,
             keyframe_frame_indices=ctx.keyframe_frame_indices,
             frame_count=ctx.keyframe_frame_count,
+            include_video_pos=include_video_pos,
         )
     else:
         packed = minimax_h3_packed_sequence(
@@ -925,6 +947,7 @@ def _build_packed_layout(
                 ctx.keyframe_frame_indices if ctx.include_cond else None
             ),
             frame_count=ctx.keyframe_frame_count,
+            include_video_pos=include_video_pos,
         )
     return packed
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py

@@ -54,7 +54,13 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
         plan = minimax_h3_plan_from_batch(batch)
         if plan is not None:
             try:
-                self._encode_from_plan(batch, plan)
+                self._encode_from_plan(
+                    batch,
+                    plan,
+                    include_video_token_mask=(
+                        server_args.pipeline_config.uses_subblock_attention(server_args)
+                    ),
+                )
                 self._publish_native_text_conditioning(batch)
                 if current_platform.is_mps():
                     self._finish_active_component_use()
@@ -229,7 +235,13 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
         batch.prompt_embeds = [hidden_states]
         batch.prompt_seq_lens = [[text_len]]
 
-    def _encode_from_plan(self, batch: Req, plan) -> None:
+    def _encode_from_plan(
+        self,
+        batch: Req,
+        plan,
+        *,
+        include_video_token_mask: bool = False,
+    ) -> None:
         """Encode the positive Qwen3VL presentation into layer-50 states.
 
         MiniMax H3 only supports the CFG-distilled model path, so every task
@@ -276,7 +288,12 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
             )
         with set_forward_context(current_timestep=0, attn_metadata=None):
             if plan.task == "ref2va":
-                embeddings = self._encode_ref2va(batch, plan, encode_ids)
+                embeddings = self._encode_ref2va(
+                    batch,
+                    plan,
+                    encode_ids,
+                    include_video_token_mask=include_video_token_mask,
+                )
             elif keyframes:
                 embeddings = self._encode_fl2va_keyframes(
                     batch,
@@ -352,10 +369,17 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
                 "hidden_states": pos_hidden,
                 "text_len": int(pos_ids.shape[0]),
                 "text_token_tags": pos_tags,
-            },
+            }
         }
 
-    def _encode_ref2va(self, batch: Req, plan, encode_ids) -> dict:
+    def _encode_ref2va(
+        self,
+        batch: Req,
+        plan,
+        encode_ids,
+        *,
+        include_video_token_mask: bool = False,
+    ) -> dict:
         """Encode the positive ref2va presentation.
 
         Per condition in order — image i: '<Picture i>: ' label +
@@ -521,14 +545,20 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
                 video_block_timestamps.append(timestamps)
 
         if has_video:
-            pos_ids, pos_tags = minimax_h3_ref2va_video_presentation(
+            presentation = minimax_h3_ref2va_video_presentation(
                 self.tokenizer,
                 prompt=plan.prompt,
                 condition_labels=condition_labels,
                 image_token_count=n_image_tokens,
                 video_block_token_counts=video_block_token_counts,
                 video_block_timestamps=video_block_timestamps,
+                return_video_mask=include_video_token_mask,
             )
+            if include_video_token_mask:
+                pos_ids, pos_tags, pos_video_mask = presentation
+            else:
+                pos_ids, pos_tags = presentation
+                pos_video_mask = None
         else:
             pos_ids, pos_tags = minimax_h3_ref2va_presentation(
                 self.tokenizer,
@@ -536,6 +566,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
                 condition_labels=condition_labels,
                 image_token_count=n_image_tokens,
             )
+            pos_video_mask = None
         pos_hidden = encode_ids(
             pos_ids,
             pixel_values=pixel_values,
@@ -545,13 +576,14 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
         )
         if batch.extra.get(_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY):
             batch.extra.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, None)
-        return {
+        positive = {
-            "positive": {
+            "hidden_states": pos_hidden,
-                "hidden_states": pos_hidden,
+            "text_len": int(pos_ids.shape[0]),
-                "text_len": int(pos_ids.shape[0]),
+            "text_token_tags": pos_tags,
-                "text_token_tags": pos_tags,
-            },
         }
+        if pos_video_mask is not None:
+            positive["text_video_token_mask"] = pos_video_mask
+        return {"positive": positive}
 
 
 __all__ = ["MiniMaxH3TextEncodingStage"]

python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py

@@ -412,13 +412,12 @@ def test_quality_admission_fails_closed_outside_validated_request():
 
 
 def test_validate_server_args_requires_packed_varlen_backend():
-    config = SimpleNamespace(
+    config = MiniMaxH3PipelineConfig()
-        vae_config=SimpleNamespace(resolved_parallel_decode_mode=lambda: None),
-        dit_config=SimpleNamespace(arch_config=SimpleNamespace(attention_head_dim=128)),
-        _server_arg_value=MiniMaxH3PipelineConfig._server_arg_value,
-    )
     server_args = SimpleNamespace(
-        component_attention_backends={}, attention_backend="sage_attn"
+        component_attention_backends={},
+        attention_backend="sage_attn",
+        ring_degree=1,
+        resolve_component_attention_backend=lambda *_names: (None, None),
     )
     with patch(
         "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3.get_attn_backend"
@@ -438,12 +437,63 @@ def test_validate_server_args_requires_packed_varlen_backend():
             MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
 
 
-def test_mps_admission_requires_layerwise_residency_for_every_h3_component():
+def test_validate_server_args_accepts_transformer_backend_override():
-    config = SimpleNamespace(
+    config = MiniMaxH3PipelineConfig()
-        vae_config=SimpleNamespace(resolved_parallel_decode_mode=lambda: None),
+    server_args = SimpleNamespace(
-        dit_config=SimpleNamespace(arch_config=SimpleNamespace(attention_head_dim=128)),
+        component_attention_backends={"transformer": "subblock_sparse_attn"},
-        _server_arg_value=MiniMaxH3PipelineConfig._server_arg_value,
+        attention_backend="fa",
+        ring_degree=1,
+        resolve_component_attention_backend=lambda *_names: (
+            AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
+            "transformer",
+        ),
     )
+
+    with patch(
+        "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3.get_attn_backend"
+    ) as get_attn_backend:
+        MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
+    get_attn_backend.assert_called_once_with(
+        128,
+        torch.bfloat16,
+        selected_attention_backend=AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
+        attention_requirements=AttentionRequirements(packed_varlen=True),
+    )
+
+
+def test_resolve_transformer_attention_backend_uses_selector_precedence():
+    config = MiniMaxH3PipelineConfig()
+    subblock = AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
+    fa = AttentionBackendEnum.FA
+    sdpa = AttentionBackendEnum.TORCH_SDPA
+    cases = (
+        ("fa", subblock, None, subblock),
+        ("subblock_sparse_attn", fa, None, fa),
+        (subblock, None, None, subblock),
+        ("fa", subblock, sdpa, sdpa),
+    )
+    for global_backend, component_backend, forced_backend, expected in cases:
+        server_args = SimpleNamespace(
+            attention_backend=global_backend,
+            resolve_component_attention_backend=lambda *_names: (
+                component_backend,
+                "transformer" if component_backend is not None else None,
+            ),
+        )
+        with patch(
+            "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3."
+            "get_global_forced_attn_backend",
+            return_value=forced_backend,
+        ):
+            resolved = config.resolve_transformer_attention_backend(server_args)
+            assert resolved is expected
+            assert config.uses_subblock_attention(server_args) is (
+                expected is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
+            )
+
+
+def test_mps_admission_requires_layerwise_residency_for_every_h3_component():
+    config = MiniMaxH3PipelineConfig()
     modes = {
         "transformer": LAYERWISE_OFFLOAD,
         "text_encoder": LAYERWISE_OFFLOAD,
@@ -454,7 +504,9 @@ def test_mps_admission_requires_layerwise_residency_for_every_h3_component():
         component_attention_backends={},
         attention_backend=None,
         enable_torch_compile=False,
+        ring_degree=1,
         residency_mode=modes.get,
+        resolve_component_attention_backend=lambda *_names: (None, None),
     )
 
     with patch.object(current_platform, "is_mps", return_value=True):

python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py

@@ -2,7 +2,7 @@
 """Mixed-precision weight and TP/Ulysses numerical contracts for H3 DiT."""
 
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 import torch
@@ -15,7 +15,13 @@ 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.attention.backends.attention_backend import (
+    AttentionRequirements,
+)
 from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPAImpl
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    component_attn_backend_context_manager,
+)
 from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
 from sglang.multimodal_gen.runtime.layers.quantization.fp8 import (
     Fp8Config,
@@ -33,6 +39,7 @@ from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
     _modulate_gate,
     _reorder_grouped_qkv_to_qkv,
 )
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
     ensure_distributed_env_defaults,
 )
@@ -272,6 +279,48 @@ def test_cache_dit_input_preservation_toggles_every_block():
     assert not any(block.preserve_input_for_cache_dit for block in model.blocks)
 
 
+@pytest.mark.parametrize(
+    ("global_backend", "expected_backend"),
+    [
+        (None, AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN),
+        (AttentionBackendEnum.FA, AttentionBackendEnum.FA),
+    ],
+)
+def test_lazy_attention_resolution_preserves_backend_precedence(
+    global_backend, expected_backend
+):
+    model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel)
+    torch.nn.Module.__init__(model)
+    model.arch = SimpleNamespace(attention_head_dim=128)
+    model._component_attention_backend_override = (
+        AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
+    )
+    model._resolved_attention_backend = None
+    backend = Mock()
+    backend.get_enum.return_value = expected_backend
+
+    with (
+        patch(
+            "sglang.multimodal_gen.runtime.models.dits.minimax_h3."
+            "get_global_forced_attn_backend",
+            return_value=global_backend,
+        ),
+        patch(
+            "sglang.multimodal_gen.runtime.models.dits.minimax_h3.get_attn_backend",
+            return_value=backend,
+        ) as get_attn_backend,
+    ):
+        model._resolve_attention_backend_once()
+
+    get_attn_backend.assert_called_once_with(
+        128,
+        torch.bfloat16,
+        selected_attention_backend=expected_backend,
+        attention_requirements=AttentionRequirements(packed_varlen=True),
+    )
+    assert model._resolved_attention_backend is expected_backend
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_cache_dit_out_of_place_gate_preserves_cuda_input():
     x = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16)
@@ -356,6 +405,27 @@ def test_tp_and_ulysses_admission_uses_tp_local_shapes():
         )
 
 
+def test_meta_model_captures_component_attention_override():
+    _ensure_single_process_parallel_runtime()
+    with (
+        component_attn_backend_context_manager(
+            AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
+            component_name="transformer",
+        ),
+        torch.device("meta"),
+    ):
+        model = MiniMaxH3DiTModel(
+            config=MiniMaxH3DiTConfig(),
+            hf_config={},
+            quant_config=None,
+        )
+
+    assert (
+        model._component_attention_backend_override
+        is AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN
+    )
+
+
 def test_meta_model_enforces_mixed_precision_weight_contract():
     expected_fp32 = set(MINIMAX_H3_FP32_PARAM_NAMES) | set(MINIMAX_H3_FP32_BUFFER_NAMES)
     _ensure_single_process_parallel_runtime()

python/sglang/multimodal_gen/test/unit/test_subblock_sparse_attention.py

@@ -15,7 +15,7 @@ none of which an accuracy-only comparison at real sparsity would pin down.
 from __future__ import annotations
 
 import unittest
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import torch
 
@@ -190,7 +190,7 @@ class TestSubBlockSparseBackend(unittest.TestCase):
         )
         self.assertEqual(metadata.current_timestep, 7)
 
-    def test_sm90_adapter_sorts_indices_and_uses_64x64_blocks(self):
+    def test_sm90_adapter_uses_presorted_indices_and_64x64_blocks(self):
         captured = {}
 
         class _FakeBlockSparseTensors:
@@ -202,7 +202,7 @@ class TestSubBlockSparseBackend(unittest.TestCase):
             captured.update(kwargs)
             return q, None
 
-        index = torch.tensor([[[[5, 1, 7, 3]]]], dtype=torch.int32)
+        index = torch.tensor([[[[1, 3, 5, 7]]]], dtype=torch.int32)
         q = torch.empty(1, 64, 1, HEAD_DIM, dtype=torch.bfloat16)
         with patch(
             "sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn._load_sm90_block_sparse_attention",
@@ -229,8 +229,11 @@ class TestSubBlockGating(unittest.TestCase):
     """The schedule must decide sparsity from the layer and the step alone."""
 
     def _impl(self, prefix: str, **config) -> SubBlockSparseAttentionImpl:
-        with _patch_schedule(config), patch.object(
+        with (
-            SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
+            _patch_schedule(config),
+            patch.object(
+                SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
+            ),
         ):
             return SubBlockSparseAttentionImpl(
                 num_heads=NUM_HEADS,
@@ -254,8 +257,11 @@ class TestSubBlockGating(unittest.TestCase):
         self.assertFalse(self._impl("token_refiner.blocks.0.attn").layer_enabled)
 
     def test_head_dim_other_than_128_is_dense(self):
-        with _patch_schedule({}), patch.object(
+        with (
-            SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
+            _patch_schedule({}),
+            patch.object(
+                SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
+            ),
         ):
             impl = SubBlockSparseAttentionImpl(
                 num_heads=NUM_HEADS,
@@ -319,7 +325,7 @@ class TestSubBlockNumerics(unittest.TestCase):
         ref = _dense_reference(q, k, v, HEAD_DIM**-0.5)
         self.assertGreater(_cosine(out, ref), 0.999)
 
-    def test_unsorted_ragged_tail_oversubscribes_sms(self):
+    def test_presorted_ragged_tail_oversubscribes_sms(self):
         """Make tail-mask ordering observable across multiple SM waves."""
         if torch.cuda.get_device_capability() != (9, 0):
             self.skipTest("the reverse-consumption constraint is specific to SM90")
@@ -337,24 +343,21 @@ class TestSubBlockNumerics(unittest.TestCase):
         self.assertGreater(num_tiles, 2 * num_sms)
 
         # The SM90 consumer visits slots from high to low and applies the tail
-        # mask to the first block. Put the ragged block in the lowest slot, so
+        # mask to the first block, so the caller places the ragged block last.
-        # removing the adapter's sort leaves its 27 padded rows unmasked. With
-        # zero Q/K and unit V that changes the output magnitude from 1 to
-        # (7 * 64 + 37) / (8 * 64), which an assert_close cannot overlook.
         topk = 8
         tail_block = num_blocks - 1
-        unsorted_blocks = torch.tensor(
+        sorted_blocks = torch.tensor(
-            [tail_block, 0, 1, 2, 3, 4, 5, 6],
+            [0, 1, 2, 3, 4, 5, 6, tail_block],
             device=device,
             dtype=torch.int32,
         )
-        unsorted_index = (
+        sorted_index = (
-            unsorted_blocks.view(1, 1, 1, topk)
+            sorted_blocks.view(1, 1, 1, topk)
             .expand(1, NUM_HEADS, num_blocks, topk)
             .clone()
         )
         out = _run_subblock_sparse_attention(
-            q, k, v, unsorted_index, topk, HEAD_DIM**-0.5
+            q, k, v, sorted_index, topk, HEAD_DIM**-0.5
         )
 
         torch.testing.assert_close(out, torch.ones_like(out), rtol=0, atol=2e-3)
@@ -378,14 +381,14 @@ class TestSubBlockNumerics(unittest.TestCase):
 
         num_blocks = (self.seq_len + 63) // 64
         topk = impl.router.route(q, k, sparsity=0.75, softmax_scale=HEAD_DIM**-0.5).topk
-        # A random permutation per row, not `randint`: sampling with replacement
+        # Draw a random subset per row, not `randint`: sampling with replacement
-        # would leave the control holding duplicate blocks, so it would attend
+        # would attend fewer distinct blocks and distort the softmax mass. Sort
-        # fewer distinct blocks than the router at the same budget, and the
+        # the selected subset to honor the SM90 kernel's direct-call contract.
-        # repeats would distort the softmax mass on top of that.
         random_index = (
             torch.rand(1, NUM_HEADS, num_blocks, num_blocks, device=device)
             .argsort(dim=-1)[..., :topk]
-            .to(torch.int32)
+            .sort(dim=-1)
+            .values.to(torch.int32)
         )
         random_out = _run_subblock_sparse_attention(
             q,
@@ -397,6 +400,97 @@ class TestSubBlockNumerics(unittest.TestCase):
         )
         self.assertLess(_cosine(random_out, ref), 0.9)
 
+    def _assert_kernel_backed_mixed_query_mask(self):
+        device = torch.device("cuda")
+        scale = HEAD_DIM**-0.5
+        generator = torch.Generator(device=device).manual_seed(7)
+        q = torch.randn(
+            1,
+            3 * 64,
+            NUM_HEADS,
+            HEAD_DIM,
+            device=device,
+            dtype=torch.bfloat16,
+            generator=generator,
+        )
+        k = torch.randn(
+            1,
+            16 * 64,
+            NUM_HEADS,
+            HEAD_DIM,
+            device=device,
+            dtype=torch.bfloat16,
+            generator=generator,
+        )
+        v = torch.randn(
+            k.shape,
+            device=device,
+            dtype=torch.bfloat16,
+            generator=generator,
+        )
+
+        topk = 8
+        sparse_blocks = torch.tensor(
+            [
+                [0, 2, 4, 6, 8, 10, 12, 14],
+                [1, 3, 5, 7, 9, 11, 13, 15],
+                [0, 1, 4, 5, 8, 9, 12, 13],
+            ],
+            device=device,
+            dtype=torch.int32,
+        )
+        routed_index = (
+            sparse_blocks.view(1, 1, 3, topk).expand(1, NUM_HEADS, 3, topk).clone()
+        )
+        plan = Mock(
+            index=routed_index,
+            topk=topk,
+            num_blocks=16,
+            density=0.5,
+        )
+        impl = self._impl(sparsity=0.5)
+        impl.router = Mock(route=Mock(return_value=plan))
+        sparse_query_block_mask = torch.tensor([True, False, True], device=device)
+
+        out = impl._sparse_attention(
+            q,
+            k,
+            v,
+            sparse_query_block_mask=sparse_query_block_mask,
+        )
+
+        reference = torch.empty_like(q)
+        all_blocks = torch.arange(16, device=device)
+        token_offsets = torch.arange(64, device=device)
+        for query_block in range(3):
+            selected_blocks = (
+                sparse_blocks[query_block].long()
+                if sparse_query_block_mask[query_block]
+                else all_blocks
+            )
+            selected_tokens = (
+                selected_blocks[:, None] * 64 + token_offsets[None, :]
+            ).reshape(-1)
+            query_slice = slice(query_block * 64, (query_block + 1) * 64)
+            reference[:, query_slice] = _dense_reference(
+                q[:, query_slice],
+                k.index_select(1, selected_tokens),
+                v.index_select(1, selected_tokens),
+                scale,
+            )
+
+        self.assertGreater(_cosine(out, reference), 0.999)
+
+    def test_sm90_kernel_backed_mixed_query_mask(self):
+        if torch.cuda.get_device_capability() != (9, 0):
+            self.skipTest("requires the SM90 SubBlock kernel")
+        self._assert_kernel_backed_mixed_query_mask()
+
+    def test_sm100_kernel_backed_mixed_query_mask(self):
+        if torch.cuda.get_device_capability() != (10, 0):
+            self.skipTest("requires the SM100 SubBlock kernel")
+        self._assert_kernel_backed_mixed_query_mask()
+
     def test_skipped_step_is_bitwise_dense(self):
         device = torch.device("cuda")
         q, k, v = _structured_qkv(self.seq_len, device)

test/registered/cpu/test_subblock_sparse_attention.py

@@ -1,15 +1,41 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import unittest
-from unittest.mock import patch
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
 
 import torch
 
+from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
+    MiniMaxH3PipelineConfig,
+)
 from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
+    SubBlockSparseAttentionImpl,
     _get_subblock_sparse_attention_runner,
     _sm90_sparse_attention,
     _sm100_sparse_attention,
 )
+from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
+    _minimax_h3_attention_core_impl,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import (
+    MiniMaxH3DenoiseBranch,
+    _minimax_h3_subblock_sparse_query_block_mask,
+    _minimax_h3_subblock_video_query_indices,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.packed_sequence import (
+    minimax_h3_packed_sequence,
+    minimax_h3_packed_sequence_ref2va_blocks,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.presentation import (
+    IMAGE_PAD,
+    VIDEO_PAD,
+    minimax_h3_ref2va_video_presentation,
+)
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -50,5 +76,320 @@ class TestSubBlockSparseAttentionDispatch(CustomTestCase):
                 _get_subblock_sparse_attention_runner(device)
 
 
+class TestSubBlockSparseAttentionModalities(CustomTestCase):
+    def test_transformer_subblock_with_ring_fails_admission(self):
+        config = MiniMaxH3PipelineConfig()
+        server_args = SimpleNamespace(
+            attention_backend="fa",
+            ring_degree=2,
+            resolve_component_attention_backend=lambda *_names: (
+                AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
+                "transformer",
+            ),
+        )
+
+        with (
+            patch.object(current_platform, "is_mps", return_value=False),
+            self.assertRaisesRegex(ValueError, "ring parallelism requires"),
+        ):
+            config.validate_server_args(server_args)
+
+    @staticmethod
+    def _run_attention_core_without_query_mask(
+        *,
+        sparse_ready: bool,
+        min_seq_len: int,
+    ) -> Mock:
+        impl = Mock()
+        impl._sparse_ready.return_value = sparse_ready
+        impl.schedule = SimpleNamespace(min_seq_len=min_seq_len)
+        q = torch.zeros(4, 1, 2)
+        impl.forward_varlen.return_value = torch.zeros_like(q)
+        attention = SimpleNamespace(
+            _attention_impl=impl,
+            _attention_backend_enum=AttentionBackendEnum.SUBBLOCK_SPARSE_ATTN,
+        )
+        _minimax_h3_attention_core_impl(
+            attention,
+            q,
+            q,
+            q,
+            cu_seqlens=torch.tensor([0, 4], dtype=torch.int32),
+            cu_seqlens_host=(0, 4),
+            max_seqlen=4,
+            ulysses_active=False,
+            subblock_sparse_query_block_mask=None,
+        )
+        return impl
+
+    def test_missing_query_mask_is_allowed_for_dense_fallback(self):
+        impl = self._run_attention_core_without_query_mask(
+            sparse_ready=False,
+            min_seq_len=4,
+        )
+
+        impl.forward_varlen.assert_called_once()
+
+    def test_missing_query_mask_is_allowed_when_segments_are_short(self):
+        impl = self._run_attention_core_without_query_mask(
+            sparse_ready=True,
+            min_seq_len=8,
+        )
+
+        impl.forward_varlen.assert_called_once()
+
+    def test_missing_query_mask_fails_only_when_sparse_attention_will_run(self):
+        with self.assertRaisesRegex(
+            ValueError,
+            "when SubBlock sparse attention is active",
+        ):
+            self._run_attention_core_without_query_mask(
+                sparse_ready=True,
+                min_seq_len=4,
+            )
+
+    def test_fl2va_keyframe_images_remain_dense(self):
+        packed = minimax_h3_packed_sequence(
+            text_len=5,
+            latent_t=2,
+            latent_h=4,
+            latent_w=4,
+            audio_t=5,
+            include_keyframe_cond=True,
+            keyframe_frame_indices=[0, -1],
+            frame_count=5,
+            include_video_pos=True,
+        )
+        video_indices = _minimax_h3_subblock_video_query_indices(
+            packed,
+            None,
+        )
+        condition_image_indices = set(
+            packed["img_pos"][~packed["update_mask"]].tolist()
+        )
+
+        torch.testing.assert_close(video_indices, packed["video_pos"])
+        self.assertTrue(condition_image_indices.isdisjoint(video_indices.tolist()))
+
+    def test_ref2va_images_are_dense_but_reference_and_target_video_are_sparse(self):
+        packed = minimax_h3_packed_sequence_ref2va_blocks(
+            text_len=5,
+            latent_t=2,
+            latent_h=4,
+            latent_w=4,
+            audio_t=5,
+            ref_blocks=[
+                {"kind": "image", "latent_h": 4, "latent_w": 4},
+                {
+                    "kind": "video_audio",
+                    "ref_audio_t": 3,
+                    "latent_t": 2,
+                    "latent_h": 4,
+                    "latent_w": 4,
+                },
+            ],
+            include_video_pos=True,
+        )
+        text_video_mask = torch.tensor([False, True, False, True, False])
+        video_indices = _minimax_h3_subblock_video_query_indices(
+            packed,
+            text_video_mask,
+        )
+
+        image_indices = set(packed["img_pos"].tolist()) - set(
+            packed["video_pos"].tolist()
+        )
+        text_video_indices = set(packed["text_pos"][text_video_mask].tolist())
+        text_non_video_indices = set(packed["text_pos"][~text_video_mask].tolist())
+
+        video_index_set = set(video_indices.tolist())
+        self.assertTrue(image_indices.isdisjoint(video_index_set))
+        self.assertTrue(set(packed["audio_pos"].tolist()).isdisjoint(video_index_set))
+        self.assertTrue(text_non_video_indices.isdisjoint(video_index_set))
+        self.assertTrue(text_video_indices.issubset(video_index_set))
+        self.assertTrue(set(packed["video_pos"].tolist()).issubset(video_index_set))
+
+    def test_ref2va_presentation_marks_only_video_vision_blocks_sparse(self):
+        class FakeTokenizer:
+            _special_ids = {
+                "<|vision_start|>": 10,
+                "<|vision_end|>": 11,
+                IMAGE_PAD: 12,
+                VIDEO_PAD: 13,
+            }
+
+            def __call__(self, text, *, add_special_tokens):
+                del add_special_tokens
+                return {"input_ids": [100 + len(text)]}
+
+            def convert_tokens_to_ids(self, token):
+                return self._special_ids[token]
+
+        ids, tags, video_mask = minimax_h3_ref2va_video_presentation(
+            FakeTokenizer(),
+            prompt="prompt",
+            condition_labels=[("image", 1), ("video", 1)],
+            image_token_count=2,
+            video_block_token_counts=[[2]],
+            video_block_timestamps=[[0.0]],
+            return_video_mask=True,
+        )
+
+        self.assertFalse(video_mask[ids == 12].any())
+        self.assertTrue(video_mask[ids == 13].all())
+        self.assertFalse(video_mask[ids == 10].any())
+        self.assertFalse(video_mask[ids == 11].any())
+        self.assertEqual(int(video_mask.sum()), 2)
+        self.assertEqual(tags[ids == 12].unique().tolist(), [0])
+        self.assertEqual(tags[ids == 13].unique().tolist(), [0])
+
+        default_result = minimax_h3_ref2va_video_presentation(
+            FakeTokenizer(),
+            prompt="prompt",
+            condition_labels=[("video", 1)],
+            image_token_count=None,
+            video_block_token_counts=[[1]],
+            video_block_timestamps=[[0.0]],
+        )
+        self.assertEqual(len(default_result), 2)
+
+    def test_video_query_indices_validate_first_segment_bounds(self):
+        for invalid in (
+            torch.tensor([-1]),
+            torch.tensor([5]),
+            torch.tensor([2, 2]),
+        ):
+            with self.subTest(indices=invalid.tolist()), self.assertRaises(ValueError):
+                _minimax_h3_subblock_sparse_query_block_mask(invalid, used_len=5)
+
+    def test_ref2va_video_positions_are_subblock_only_metadata(self):
+        kwargs = dict(
+            text_len=3,
+            latent_t=2,
+            latent_h=4,
+            latent_w=4,
+            audio_t=3,
+            ref_blocks=[
+                {
+                    "kind": "video",
+                    "ref_audio_t": 0,
+                    "latent_t": 2,
+                    "latent_h": 4,
+                    "latent_w": 4,
+                }
+            ],
+        )
+
+        ordinary = minimax_h3_packed_sequence_ref2va_blocks(**kwargs)
+        subblock = minimax_h3_packed_sequence_ref2va_blocks(
+            **kwargs,
+            include_video_pos=True,
+        )
+
+        self.assertNotIn("video_pos", ordinary)
+        self.assertIn("video_pos", subblock)
+        self.assertTrue(
+            set(subblock["video_pos"].tolist()).issubset(
+                set(subblock["img_pos"].tolist())
+            )
+        )
+
+    def test_non_subblock_branch_does_not_retain_dense_query_metadata(self):
+        packed = minimax_h3_packed_sequence(
+            text_len=3,
+            latent_t=2,
+            latent_h=4,
+            latent_w=4,
+            audio_t=3,
+            include_keyframe_cond=False,
+        )
+
+        self.assertNotIn("video_pos", packed)
+        branch = MiniMaxH3DenoiseBranch(
+            packed=packed,
+            text_embeddings=torch.zeros(3, 5120),
+            token_tags=packed["token_tags"],
+            video_query_indices=None,
+            device=torch.device("cpu"),
+        )
+
+        self.assertNotIn(
+            "subblock_sparse_query_block_mask",
+            branch.static_kwargs,
+        )
+
+    def test_query_mask_marks_only_pure_video_blocks_sparse(self):
+        sparse_query_block_mask = _minimax_h3_subblock_sparse_query_block_mask(
+            torch.cat([torch.arange(64), torch.arange(128, 140)]),
+            used_len=140,
+        )
+        torch.testing.assert_close(
+            sparse_query_block_mask,
+            torch.tensor([True, False, True]),
+        )
+
+    def test_hybrid_query_routing_uses_one_heterogeneous_bsa_call(self):
+        impl = object.__new__(SubBlockSparseAttentionImpl)
+        impl.softmax_scale = 2**-0.5
+        impl.causal = False
+        impl.schedule = SimpleNamespace(sparsity=0.75)
+        plan = SimpleNamespace(
+            index=torch.tensor(
+                [[[[7, 1, 4], [6, 2, 0], [5, 0, 3]]]], dtype=torch.int32
+            ),
+            topk=3,
+            num_blocks=8,
+            density=3 / 8,
+        )
+        impl.router = Mock(route=Mock(return_value=plan))
+        q = torch.zeros(1, 3 * 64, 1, 2)
+        k = torch.zeros(1, 8 * 64, 1, 2)
+        v = torch.zeros_like(k)
+        sparse_query_block_mask = torch.tensor([True, False, True])
+        impl.dense_impl = Mock()
+        sparse_out = torch.ones_like(q)
+
+        for runner, sparse_rows in (
+            (_sm90_sparse_attention, ([1, 4, 7], [0, 3, 5])),
+            (_sm100_sparse_attention, ([7, 1, 4], [5, 0, 3])),
+        ):
+            with (
+                self.subTest(runner=runner.__name__),
+                patch(
+                    "sglang.multimodal_gen.runtime.layers.attention.backends."
+                    "subblock_sparse_attn._run_subblock_sparse_attention",
+                    return_value=sparse_out,
+                ) as run_sparse,
+                patch(
+                    "sglang.multimodal_gen.runtime.layers.attention.backends."
+                    "subblock_sparse_attn._get_subblock_sparse_attention_runner",
+                    return_value=runner,
+                ),
+            ):
+                out = impl._sparse_attention(
+                    q,
+                    k,
+                    v,
+                    sparse_query_block_mask=sparse_query_block_mask,
+                )
+
+            impl.dense_impl.forward.assert_not_called()
+            torch.testing.assert_close(out, sparse_out)
+            routing_q = impl.router.route.call_args.args[0]
+            self.assertEqual(routing_q.shape[1], 3 * 64)
+            sparse_call = run_sparse.call_args.args
+            self.assertEqual(sparse_call[0].shape[1], 3 * 64)
+            self.assertIs(sparse_call[1], k)
+            self.assertIs(sparse_call[2], v)
+            self.assertEqual(sparse_call[4], 8)
+            torch.testing.assert_close(
+                sparse_call[6], torch.tensor([[[3, 8, 3]]], dtype=torch.int32)
+            )
+            block_index = sparse_call[3]
+            self.assertEqual(block_index[0, 0, 0, :3].tolist(), sparse_rows[0])
+            self.assertEqual(block_index[0, 0, 1].tolist(), list(range(8)))
+            self.assertEqual(block_index[0, 0, 2, :3].tolist(), sparse_rows[1])
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=3)