minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
HuangJi
2026-08-26 10:55:09 +08:00
committed by GitHub
parent 4382947b58
commit cc3b61873f
13 changed files with 1033 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
+40 -11
View File
@@ -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,
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/router.py
+3 -2
View File
@@ -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)
python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse_attn.py
+115 -31
View File
@@ -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]
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+68 -9
View File
@@ -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=(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py
+79 -1
View File
@@ -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
+21 -5
View File
@@ -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
+47 -15
View File
@@ -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: ``<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)
with the vision block tagged VIDEO(0) and everything else TEXT(1).
never enters Qwen) — then the verbatim prompt. Returns ``(ids, token_tags)``
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
else TEXT(1).
then the verbatim prompt. Image and video blocks both retain visual
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
+24 -1
View File
@@ -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
+44 -12
View File
@@ -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": {
"hidden_states": pos_hidden,
"text_len": int(pos_ids.shape[0]),
"text_token_tags": pos_tags,
},
positive = {
"hidden_states": pos_hidden,
"text_len": int(pos_ids.shape[0]),
"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
+63 -11
View File
@@ -412,13 +412,12 @@ def test_quality_admission_fails_closed_outside_validated_request():
def test_validate_server_args_requires_packed_varlen_backend():
config = SimpleNamespace(
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,
)
config = MiniMaxH3PipelineConfig()
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():
config = SimpleNamespace(
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,
def test_validate_server_args_accepts_transformer_backend_override():
config = MiniMaxH3PipelineConfig()
server_args = SimpleNamespace(
component_attention_backends={"transformer": "subblock_sparse_attn"},
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
+71 -1
View File
@@ -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
+116 -22
View File
@@ -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(
SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
with (
_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(
SubBlockSparseAttentionImpl, "_build_dense_impl", return_value=None
with (
_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
# 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.
# mask to the first block, so the caller places the ragged block last.
topk = 8
tail_block = num_blocks - 1
unsorted_blocks = torch.tensor(
[tail_block, 0, 1, 2, 3, 4, 5, 6],
sorted_blocks = torch.tensor(
[0, 1, 2, 3, 4, 5, 6, tail_block],
device=device,
dtype=torch.int32,
)
unsorted_index = (
unsorted_blocks.view(1, 1, 1, topk)
sorted_index = (
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
# would leave the control holding duplicate blocks, so it would attend
# fewer distinct blocks than the router at the same budget, and the
# repeats would distort the softmax mass on top of that.
# Draw a random subset per row, not `randint`: sampling with replacement
# would attend fewer distinct blocks and distort the softmax mass. Sort
# the selected subset to honor the SM90 kernel's direct-call contract.
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)