[diffusion] feat: support request-scoped skip-softmax attention (#37959)

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -59,6 +59,57 @@ QUALITY_LEVELS: tuple[str, ...] = ("lossless", "extra-high", "high")
 KERNEL_FUSION_QUALITY_LEVELS = frozenset({"extra-high", "high"})
 
 
+@dataclass(frozen=True)
+class SkipSoftmaxParams:
+    """Validated request-scoped BLASST/Skip-Softmax controls."""
+
+    threshold_scale_factor: float
+    start_step: int = 0
+
+
+def resolve_skip_softmax_params(
+    params: dict[str, Any] | None,
+) -> SkipSoftmaxParams | None:
+    if params is None:
+        return None
+    if not isinstance(params, dict):
+        raise ValueError(f"skip_softmax_params must be a dict, got {params!r}")
+
+    valid_keys = {"threshold_scale_factor", "start_step"}
+    unknown = sorted(set(params) - valid_keys)
+    if unknown:
+        raise ValueError(
+            f"Unknown skip_softmax_params keys: {unknown}. "
+            f"Valid keys: {sorted(valid_keys)}."
+        )
+    if "threshold_scale_factor" not in params:
+        raise ValueError("skip_softmax_params requires 'threshold_scale_factor'.")
+
+    threshold = params["threshold_scale_factor"]
+    if (
+        isinstance(threshold, bool)
+        or not isinstance(threshold, (int, float))
+        or not math.isfinite(float(threshold))
+        or float(threshold) <= 0
+    ):
+        raise ValueError(
+            "skip_softmax_params.threshold_scale_factor must be a finite "
+            f"positive number, got {threshold!r}"
+        )
+
+    start_step = params.get("start_step", 0)
+    if (
+        isinstance(start_step, bool)
+        or not isinstance(start_step, int)
+        or start_step < 0
+    ):
+        raise ValueError(
+            "skip_softmax_params.start_step must be a non-negative int, "
+            f"got {start_step!r}"
+        )
+    return SkipSoftmaxParams(float(threshold), start_step)
+
+
 def quality_allows_kernel_fusions(quality: str) -> bool:
     """Return whether a quality level includes request-gated kernel fusions."""
     return quality in KERNEL_FUSION_QUALITY_LEVELS
@@ -232,6 +283,11 @@ class SamplingParams:
     # request; see DenoisingStage._maybe_override_attention_backend.
     attention_backend_override: str | None = None
 
+    # Request-scoped BLASST/Skip-Softmax sparse attention. This is an explicit
+    # lossy opt-in; compatible self-attention layers are dispatched through
+    # FlashInfer while cross-attention remains on its normal backend.
+    skip_softmax_params: dict[str, Any] | None = None
+
     # Spectrum parameters
     enable_spectrum: bool = False
     spectrum_params: Any = None  # SpectrumParams
@@ -485,6 +541,8 @@ class SamplingParams:
                 f"quality must be one of {list(QUALITY_LEVELS)}, got {self.quality!r}"
             )
 
+        resolve_skip_softmax_params(self.skip_softmax_params)
+
         # These are always required to be sane regardless of pipeline.
         if (
             not isinstance(self.num_outputs_per_prompt, int)

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -309,6 +309,7 @@ async def generations(
             attention_backend_override=_get_extra_field(
                 request, "attention_backend_override"
             ),
+            skip_softmax_params=_get_extra_field(request, "skip_softmax_params"),
             quality=_runtime_sampling_quality(request.quality),
             output_compression=request.output_compression,
             output_quality=request.output_quality,

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -133,6 +133,7 @@ _MULTIPART_EXTRA_FORM_FIELDS = (
     "cfg_gate_step",
     "enable_cache_dit",
     "quality",
+    "skip_softmax_params",
 )
 
 
@@ -291,6 +292,7 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
         "attention_backend_override": _extra_value(
             request, "attention_backend_override"
         ),
+        "skip_softmax_params": _extra_value(request, "skip_softmax_params"),
         "enable_frame_interpolation": request.enable_frame_interpolation,
         "frame_interpolation_exp": request.frame_interpolation_exp,
         "frame_interpolation_scale": request.frame_interpolation_scale,

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -23,6 +23,20 @@ class AttentionRequirements:
     packed_varlen: bool = False
 
 
+def trailing_padding_used_len(
+    total_tokens: int,
+    max_seqlen: int,
+    bounds: tuple[int, ...],
+) -> int | None:
+    """Return the live prefix length for a packed, padded single sequence."""
+    if len(bounds) != 3:
+        return None
+    start, used, total = bounds
+    if start != 0 or used >= total or total != total_tokens or used != max_seqlen:
+        return None
+    return used
+
+
 class AttentionBackend(ABC):
     """Abstract class for attention backends."""
 

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

@@ -1,6 +1,7 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
+from numbers import Integral
 from typing import Any, List, Optional, Tuple
 
 import torch
@@ -11,8 +12,16 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
     AttentionImpl,
     AttentionMetadata,
     AttentionMetadataBuilder,
+    trailing_padding_used_len,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.skip_softmax import (
+    get_fixed_sequence_metadata,
+    get_host_sequence_lengths,
+    get_request_skip_softmax_params,
+    run_skip_softmax,
 )
 from sglang.multimodal_gen.runtime.layers.utils import register_custom_op
+from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
 )
@@ -365,6 +374,7 @@ class FlashAttentionImpl(AttentionImpl):
         softmax_scale: float,
         num_kv_heads: int | None = None,
         prefix: str = "",
+        is_cross_attention: bool = False,
         **extra_impl_args,
     ) -> None:
         self.num_heads = num_heads
@@ -372,8 +382,26 @@ class FlashAttentionImpl(AttentionImpl):
         self.head_size = head_size
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.is_cross_attention = is_cross_attention
+        self.packed_trailing_padding = extra_impl_args.get(
+            "packed_trailing_padding", False
+        )
         self.attention_metadata = FlashAttentionMetadata()
 
+    def _request_skip_softmax_threshold(self) -> tuple[bool, float | None]:
+        params = get_request_skip_softmax_params()
+        if params is None or self.is_cross_attention:
+            return False, None
+        current_step = get_forward_context().current_timestep
+        if not isinstance(current_step, Integral):
+            raise RuntimeError(
+                "Skip Softmax requires the denoising loop to publish an integer "
+                f"step index, got {current_step!r}."
+            )
+        if current_step < params.start_step:
+            return True, None
+        return True, params.threshold_scale_factor
+
     def forward(
         self,
         query: torch.Tensor,
@@ -383,6 +411,38 @@ class FlashAttentionImpl(AttentionImpl):
         *,
         return_softmax_lse: bool = False,
     ):
+        use_trtllm, skip_threshold = self._request_skip_softmax_threshold()
+        if use_trtllm:
+            if return_softmax_lse:
+                raise NotImplementedError(
+                    "Skip Softmax does not support the LSE output required by "
+                    "Ring Attention."
+                )
+            batch_size, query_length = query.shape[:2]
+            kv_length = key.shape[1]
+            metadata = get_fixed_sequence_metadata(
+                query.device,
+                batch_size,
+                query_length,
+                kv_length,
+            )
+            output = run_skip_softmax(
+                query.reshape(-1, query.shape[-2], query.shape[-1]),
+                key.reshape(-1, key.shape[-2], key.shape[-1]),
+                value.reshape(-1, value.shape[-2], value.shape[-1]),
+                seq_lens=metadata.seq_lens,
+                cu_seqlens_q=metadata.cu_seqlens_q,
+                cu_seqlens_kv=metadata.cu_seqlens_kv,
+                max_seqlen_q=query_length,
+                max_seqlen_kv=kv_length,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                threshold_scale_factor=skip_threshold,
+                q_seq_lens_cpu=metadata.q_seq_lens_cpu,
+                kv_seq_lens_cpu=metadata.kv_seq_lens_cpu,
+            )
+            return output.view(batch_size, query_length, *output.shape[1:])
+
         if attn_metadata is not None:
             if attn_metadata.max_seqlen_q is None:
                 attn_metadata.max_seqlen_q = query.shape[1]
@@ -457,7 +517,62 @@ class FlashAttentionImpl(AttentionImpl):
         max_seqlen: int,
         cu_seqlens_host: tuple[int, ...] | None = None,
     ) -> torch.Tensor:
-        del cu_seqlens_host
+        use_trtllm, skip_threshold = self._request_skip_softmax_threshold()
+        if use_trtllm:
+            bounds = cu_seqlens_host
+            used = (
+                trailing_padding_used_len(
+                    total_tokens=query.shape[0],
+                    max_seqlen=max_seqlen,
+                    bounds=bounds,
+                )
+                if self.packed_trailing_padding and bounds is not None
+                else None
+            )
+            if used is not None:
+                live_cu_seqlens = cu_seqlens[:2]
+                live_output = run_skip_softmax(
+                    query[:used],
+                    key[:used],
+                    value[:used],
+                    seq_lens=live_cu_seqlens[1:] - live_cu_seqlens[:-1],
+                    cu_seqlens_q=live_cu_seqlens,
+                    cu_seqlens_kv=live_cu_seqlens,
+                    max_seqlen_q=used,
+                    max_seqlen_kv=used,
+                    softmax_scale=self.softmax_scale,
+                    causal=self.causal,
+                    threshold_scale_factor=skip_threshold,
+                    q_seq_lens_cpu=get_host_sequence_lengths(bounds[:2]),
+                    kv_seq_lens_cpu=get_host_sequence_lengths(bounds[:2]),
+                )
+                output = live_output.new_zeros(
+                    query.shape[0], query.shape[1], value.shape[2]
+                )
+                output[:used] = live_output
+                return output
+
+            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            seq_lens_cpu = (
+                get_host_sequence_lengths(cu_seqlens_host)
+                if cu_seqlens_host is not None
+                else None
+            )
+            return run_skip_softmax(
+                query,
+                key,
+                value,
+                seq_lens=seq_lens,
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_kv=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_kv=max_seqlen,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                threshold_scale_factor=skip_threshold,
+                q_seq_lens_cpu=seq_lens_cpu,
+                kv_seq_lens_cpu=seq_lens_cpu,
+            )
         output = flash_attn_varlen_func(
             query,
             key,

python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py

@@ -10,6 +10,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
     AttentionBackend,
     AttentionImpl,
     AttentionMetadata,
+    trailing_padding_used_len,
 )
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -17,21 +18,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 
-def _trailing_padding_used_len(
-    *,
-    total_tokens: int,
-    max_seqlen: int,
-    bounds: tuple[int, ...],
-) -> int | None:
-    """Return live token count for H3-style [0, used, total] trailing padding."""
-    if len(bounds) != 3:
-        return None
-    start, used, total = bounds
-    if start != 0 or used >= total or total != total_tokens or used != max_seqlen:
-        return None
-    return used
-
-
 class SageAttentionBackend(AttentionBackend):
     @classmethod
     def supports_ring_rotation(cls) -> bool:
@@ -66,6 +52,9 @@ class SageAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
         self.dropout = extra_impl_args.get("dropout_p", 0.0)
+        self.packed_trailing_padding = extra_impl_args.get(
+            "packed_trailing_padding", False
+        )
 
     def forward(
         self,
@@ -125,10 +114,14 @@ class SageAttentionImpl(AttentionImpl):
     ) -> torch.Tensor:
         # MiniMax-H3 packs one live document as bounds=(0, used, total):
         # [0, used) are real tokens; [used, total) is 64-aligned tail padding.
-        used = _trailing_padding_used_len(
-            total_tokens=query.shape[0],
-            max_seqlen=max_seqlen,
-            bounds=bounds,
+        used = (
+            trailing_padding_used_len(
+                total_tokens=query.shape[0],
+                max_seqlen=max_seqlen,
+                bounds=bounds,
+            )
+            if self.packed_trailing_padding
+            else None
         )
         if used is not None:
             live_out = self.forward(

python/sglang/multimodal_gen/runtime/layers/attention/backends/skip_softmax.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SkipSoftmaxParams
+from sglang.multimodal_gen.runtime.managers.forward_context import (
+    get_forward_context_or_none,
+)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core import Req
+
+_WORKSPACE_BYTES = 128 * 1024 * 1024
+_REQUEST_KEY = "_skip_softmax_params"
+_workspaces: dict[tuple[int, int], torch.Tensor] = {}
+
+
+@dataclass(frozen=True)
+class SkipSoftmaxSequenceMetadata:
+    seq_lens: torch.Tensor
+    cu_seqlens_q: torch.Tensor
+    cu_seqlens_kv: torch.Tensor
+    q_seq_lens_cpu: torch.Tensor
+    kv_seq_lens_cpu: torch.Tensor
+
+
+def set_request_skip_softmax_params(
+    batch: "Req", params: SkipSoftmaxParams | None
+) -> None:
+    batch.extra[_REQUEST_KEY] = params
+
+
+def get_request_skip_softmax_params() -> SkipSoftmaxParams | None:
+    context = get_forward_context_or_none()
+    if context is None:
+        return None
+    batch = context.forward_batch
+    if batch is None:
+        return None
+    params = batch.extra.get(_REQUEST_KEY)
+    assert params is None or isinstance(params, SkipSoftmaxParams)
+    return params
+
+
+def get_fixed_sequence_metadata(
+    device: torch.device,
+    batch_size: int,
+    query_length: int,
+    kv_length: int,
+) -> SkipSoftmaxSequenceMetadata:
+    device_index = device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    return _cached_fixed_sequence_metadata(
+        device_index, batch_size, query_length, kv_length
+    )
+
+
+@lru_cache(maxsize=32)
+def _cached_fixed_sequence_metadata(
+    device_index: int,
+    batch_size: int,
+    query_length: int,
+    kv_length: int,
+) -> SkipSoftmaxSequenceMetadata:
+    device = torch.device("cuda", device_index)
+    seq_lens = torch.full((batch_size,), kv_length, dtype=torch.int32, device=device)
+    cu_seqlens_q = torch.arange(
+        0,
+        (batch_size + 1) * query_length,
+        query_length,
+        dtype=torch.int32,
+        device=device,
+    )
+    cu_seqlens_kv = torch.arange(
+        0,
+        (batch_size + 1) * kv_length,
+        kv_length,
+        dtype=torch.int32,
+        device=device,
+    )
+    return SkipSoftmaxSequenceMetadata(
+        seq_lens=seq_lens,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_kv=cu_seqlens_kv,
+        q_seq_lens_cpu=torch.full((batch_size,), query_length, dtype=torch.int32),
+        kv_seq_lens_cpu=torch.full((batch_size,), kv_length, dtype=torch.int32),
+    )
+
+
+@lru_cache(maxsize=32)
+def get_host_sequence_lengths(cu_seqlens: tuple[int, ...]) -> torch.Tensor:
+    return torch.tensor(
+        [stop - start for start, stop in zip(cu_seqlens[:-1], cu_seqlens[1:])],
+        dtype=torch.int32,
+    )
+
+
+def run_skip_softmax(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    seq_lens: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_kv: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_kv: int,
+    softmax_scale: float,
+    causal: bool,
+    threshold_scale_factor: float | None,
+    q_seq_lens_cpu: torch.Tensor | None = None,
+    kv_seq_lens_cpu: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run FlashInfer TRTLLM attention, optionally with BLASST sparsity."""
+    _validate_inputs(query, key, value)
+    capability = torch.cuda.get_device_capability(query.device)
+    workspace = _get_workspace(query.device)
+    query, key, value = (tensor.contiguous() for tensor in (query, key, value))
+
+    common_kwargs = dict(
+        seq_lens=seq_lens,
+        max_q_len=max_seqlen_q,
+        max_kv_len=max_seqlen_kv,
+        bmm1_scale=softmax_scale,
+        bmm2_scale=1.0,
+        batch_size=cu_seqlens_q.numel() - 1,
+        cum_seq_lens_q=cu_seqlens_q,
+        cum_seq_lens_kv=cu_seqlens_kv,
+        skip_softmax_threshold_scale_factor=threshold_scale_factor,
+    )
+    if capability == (9, 0):
+        from flashinfer.prefill import trtllm_fmha_v2_prefill
+
+        return trtllm_fmha_v2_prefill(
+            (query, torch.stack((key, value), dim=1)),
+            "CONTIGUOUS_Q_KV",
+            workspace_buffer=workspace,
+            mask_mode="causal" if causal else "padding",
+            **common_kwargs,
+        )
+
+    if capability in ((10, 0), (10, 3), (10, 7)):
+        from flashinfer.prefill import trtllm_ragged_attention_deepseek
+
+        output = torch.empty(
+            (query.shape[0], query.shape[1], value.shape[2]),
+            dtype=query.dtype,
+            device=query.device,
+        )
+        return trtllm_ragged_attention_deepseek(
+            query,
+            key,
+            value,
+            workspace,
+            o_sf_scale=-1.0,
+            window_left=-1,
+            enable_pdl=None,
+            is_causal=causal,
+            return_lse=False,
+            out=output,
+            q_seq_lens_cpu=q_seq_lens_cpu,
+            kv_seq_lens_cpu=kv_seq_lens_cpu,
+            **common_kwargs,
+        )
+
+    raise ValueError(
+        "Skip Softmax requires Hopper SM90 or Blackwell SM100/SM103/SM107; "
+        f"found SM{capability[0]}{capability[1]}."
+    )
+
+
+def _validate_inputs(
+    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+) -> None:
+    if query.device.type != "cuda":
+        raise ValueError("Skip Softmax requires a CUDA device.")
+    if key.device != query.device or value.device != query.device:
+        raise ValueError("Skip Softmax requires Q/K/V on the same CUDA device.")
+    if query.ndim != 3 or key.ndim != 3 or value.ndim != 3:
+        raise ValueError("Skip Softmax expects packed [tokens, heads, head_dim] Q/K/V.")
+    if query.dtype not in (torch.float16, torch.bfloat16):
+        raise ValueError(
+            f"Skip Softmax requires FP16 or BF16 Q/K/V, got {query.dtype}."
+        )
+    if query.dtype != key.dtype or query.dtype != value.dtype:
+        raise ValueError("Skip Softmax requires Q/K/V to have the same dtype.")
+    if query.shape[-1] != key.shape[-1] or key.shape[-1] != value.shape[-1]:
+        raise ValueError("Skip Softmax requires matching Q/K/V head dimensions.")
+    if key.shape[:2] != value.shape[:2]:
+        raise ValueError("Skip Softmax requires matching K/V token and head counts.")
+    if query.shape[1] % key.shape[1] != 0:
+        raise ValueError(
+            "Skip Softmax requires the query head count to divide by KV heads."
+        )
+    if query.shape[-1] not in (128, 256):
+        raise ValueError(
+            f"Skip Softmax supports head dimensions 128 and 256; got {query.shape[-1]}."
+        )
+
+
+def _get_workspace(device: torch.device) -> torch.Tensor:
+    device_index = device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    key = (device_index, torch.cuda.current_stream(device).cuda_stream)
+    workspace = _workspaces.get(key)
+    if workspace is None:
+        workspace = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=device)
+        _workspaces[key] = workspace
+    return workspace

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -42,6 +42,9 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
     AttentionImpl,
     wrap_attention_impl_forward,
 )
+from sglang.multimodal_gen.runtime.layers.attention.backends.skip_softmax import (
+    get_request_skip_softmax_params,
+)
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
 from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import (
     async_a2a_communicate,
@@ -344,6 +347,18 @@ def apply_attention_backend_override(
     layer.backend = target
 
 
+def supports_skip_softmax(layer: nn.Module) -> bool:
+    return (
+        not layer.is_cross_attention
+        and layer.head_size in (128, 256)
+        and layer.dtype in (torch.float16, torch.bfloat16)
+        and (
+            layer._required_attention_backend is None
+            or layer.backend is AttentionBackendEnum.FA
+        )
+    )
+
+
 class UlyssesAttention(nn.Module):
     """Ulysses-style SequenceParallelism attention layer."""
 
@@ -357,6 +372,7 @@ class UlyssesAttention(nn.Module):
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
         required_attention_backend: AttentionBackendEnum | None = None,
         prefix: str = "",
+        is_cross_attention: bool = False,
         **extra_impl_args,
     ) -> None:
         super().__init__()
@@ -392,6 +408,7 @@ class UlyssesAttention(nn.Module):
             softmax_scale=self.softmax_scale,
             num_kv_heads=num_kv_heads,
             prefix=f"{prefix}.impl",
+            is_cross_attention=is_cross_attention,
             **extra_impl_args,
         )
         self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
@@ -404,6 +421,7 @@ class UlyssesAttention(nn.Module):
         self._attn_impl_by_backend = {self.backend: self.attn_impl}
         self._supported_attention_backends = supported_attention_backends
         self._required_attention_backend = required_attention_backend
+        self.is_cross_attention = is_cross_attention
         self.dtype = dtype
         self.causal = causal
         self.sp_attention_mode, self.sp_attention_mode_is_auto = (
@@ -658,6 +676,7 @@ class LocalAttention(nn.Module):
             softmax_scale=self.softmax_scale,
             num_kv_heads=num_kv_heads,
             causal=causal,
+            is_cross_attention=is_cross_attention,
             **extra_impl_args,
         )
         self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
@@ -670,6 +689,7 @@ class LocalAttention(nn.Module):
         self._attn_impl_by_backend = {self.backend: self.attn_impl}
         self._supported_attention_backends = supported_attention_backends
         self._required_attention_backend = required_attention_backend
+        self.is_cross_attention = is_cross_attention
         self.dtype = dtype
 
     def forward(
@@ -697,6 +717,13 @@ class LocalAttention(nn.Module):
         ctx_attn_metadata = forward_context.attn_metadata
 
         if attn_mask is not None:
+            if (
+                not self.is_cross_attention
+                and get_request_skip_softmax_params() is not None
+            ):
+                raise NotImplementedError(
+                    "Skip Softmax does not support LocalAttention masks."
+                )
             q_ = q.transpose(1, 2)
             k_ = k.transpose(1, 2)
             v_ = v.transpose(1, 2)
@@ -821,6 +848,7 @@ class USPAttention(nn.Module):
             softmax_scale=self.softmax_scale,
             num_kv_heads=num_kv_heads,
             prefix=f"{prefix}.impl",
+            is_cross_attention=is_cross_attention,
             **extra_impl_args,
         )
         self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
@@ -833,6 +861,7 @@ class USPAttention(nn.Module):
         self._attn_impl_by_backend = {self.backend: self.attn_impl}
         self._supported_attention_backends = supported_attention_backends
         self._required_attention_backend = required_attention_backend
+        self.is_cross_attention = is_cross_attention
         self.dtype = dtype
         self.causal = causal
         self.dropout_p = dropout_rate
@@ -897,6 +926,14 @@ class USPAttention(nn.Module):
         """
         forward_context: ForwardContext = get_forward_context()
         ctx_attn_metadata = forward_context.attn_metadata
+        if (
+            attn_mask is not None
+            and get_request_skip_softmax_params() is not None
+            and not self.is_cross_attention
+        ):
+            raise NotImplementedError(
+                "Skip Softmax does not support USPAttention masks."
+            )
         effective_skip_sp = (
             self.skip_sequence_parallel or skip_sequence_parallel_override
         )

python/sglang/multimodal_gen/runtime/managers/forward_context.py

@@ -61,6 +61,10 @@ def get_forward_context() -> "ForwardContext":
     return _forward_context
 
 
+def get_forward_context_or_none() -> "ForwardContext | None":
+    return _forward_context
+
+
 # TODO(will): finalize the interface
 @contextmanager
 def set_forward_context(

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -846,6 +846,7 @@ class MiniMaxH3Attention(nn.Module):
             softmax_scale=self.softmax_scale,
             num_kv_heads=self.num_heads,
             prefix=self.prefix,
+            packed_trailing_padding=True,
         )
         # Ring only supports FA (see _minimax_h3_attention_core_impl); keep
         # the resolved enum alongside the impl instance instead of a second

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

@@ -55,6 +55,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import (
 from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     quality_allows_kernel_fusions,
+    resolve_skip_softmax_params,
 )
 from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
     prompt_padding as bcg_utils,
@@ -96,12 +97,16 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_classifier_free_guidance_world_size,
     world_group_is_initialized,
 )
+from sglang.multimodal_gen.runtime.layers.attention.backends.skip_softmax import (
+    set_request_skip_softmax_params,
+)
 from sglang.multimodal_gen.runtime.layers.attention.layer import (
     LocalAttention,
     UlyssesAttention,
     USPAttention,
     apply_attention_backend_override,
     prepare_attention_backend_override,
+    supports_skip_softmax,
 )
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
 from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
@@ -383,6 +388,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self._attn_backend_default = self.attn_backend
         self._attn_metadata_head_size = attn_head_size
         self._attention_backend_active_override: AttentionBackendEnum | None = None
+        self._skip_softmax_forced_fa = False
 
         # cfg
         self.guidance = None
@@ -570,13 +576,37 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self, num_inference_steps: int | tuple[int, int], batch: Req
     ) -> None:
         """Apply request-dependent transformer acceleration in trace-safe order."""
-        self._maybe_override_attention_backend(batch)
+        skip_softmax_params = resolve_skip_softmax_params(
+            batch.sampling_params.skip_softmax_params
+        )
+        if skip_softmax_params is not None:
+            capability = current_platform.get_device_capability()
+            capability_tuple = (
+                (capability.major, capability.minor) if capability is not None else None
+            )
+            if capability_tuple not in ((9, 0), (10, 0), (10, 3), (10, 7)):
+                found = capability.as_version_str() if capability else "unknown"
+                raise ValueError(
+                    "skip_softmax_params requires Hopper SM90 or Blackwell "
+                    f"SM100/SM103/SM107; found {found}."
+                )
+            if (self.server_args.ring_degree or 1) > 1:
+                raise ValueError(
+                    "skip_softmax_params does not support Ring Attention because "
+                    "the ring merge requires dense per-hop softmax statistics."
+                )
+        set_request_skip_softmax_params(batch, skip_softmax_params)
+        self._maybe_override_attention_backend(
+            batch, force_fa_for_self_attention=skip_softmax_params is not None
+        )
         self._maybe_toggle_quality_fusions(batch)
         self._maybe_enable_cache_dit(num_inference_steps, batch)
         for transformer in filter(None, [self.transformer, self.transformer_2]):
             self._maybe_torch_compile(transformer)
 
-    def _maybe_override_attention_backend(self, batch: Req) -> None:
+    def _maybe_override_attention_backend(
+        self, batch: Req, *, force_fa_for_self_attention: bool = False
+    ) -> None:
         """Two-phase per-request backend switch: prepare all layers (may
         raise, mutates nothing), then flip all — a rejected request leaves the
         transformers untouched. Safe at this batch boundary because the field
@@ -584,21 +614,59 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         target = self._parse_attention_backend_override(
             batch.sampling_params.attention_backend_override
         )
-        if target == self._attention_backend_active_override:
+        if force_fa_for_self_attention and target not in (
+            None,
+            AttentionBackendEnum.FA,
+        ):
+            raise ValueError(
+                "skip_softmax_params requires the FA attention backend; "
+                f"attention_backend_override={target.name.lower()!r} is incompatible."
+            )
+        if (
+            target == self._attention_backend_active_override
+            and force_fa_for_self_attention == self._skip_softmax_forced_fa
+        ):
             return
         layers = self._request_switchable_attention_layers()
         stage_backend = self._attn_backend_default
+        layer_targets: list[tuple[nn.Module, AttentionBackendEnum | None]]
         if target is not None:
             stage_backend = self._validate_attention_backend_override(target, layers)
-            for layer in layers:
-                prepare_attention_backend_override(layer, target)
-        for layer in layers:
-            apply_attention_backend_override(layer, target)
+            layer_targets = [(layer, target) for layer in layers]
+        elif force_fa_for_self_attention:
+            self_attention_layers = [
+                layer for layer in layers if supports_skip_softmax(layer)
+            ]
+            if self_attention_layers:
+                stage_backend = self._validate_attention_backend_override(
+                    AttentionBackendEnum.FA, self_attention_layers
+                )
+            elif layers or self.attn_backend.get_enum() is not AttentionBackendEnum.FA:
+                self._validate_attention_backend_override(
+                    AttentionBackendEnum.FA, self_attention_layers
+                )
+            layer_targets = [
+                (
+                    layer,
+                    (AttentionBackendEnum.FA if supports_skip_softmax(layer) else None),
+                )
+                for layer in layers
+            ]
+        else:
+            layer_targets = [(layer, None) for layer in layers]
+
+        for layer, layer_target in layer_targets:
+            if layer_target is not None:
+                prepare_attention_backend_override(layer, layer_target)
+        for layer, layer_target in layer_targets:
+            apply_attention_backend_override(layer, layer_target)
         self.attn_backend = stage_backend
         self._attention_backend_active_override = target
+        self._skip_softmax_forced_fa = force_fa_for_self_attention
         logger.debug(
-            "Attention backend for this batch: %s (%d layers switched)",
+            "Attention backend for this batch: %s%s (%d layers considered)",
             target.name.lower() if target else "server default",
+            "; FA for self-attention" if force_fa_for_self_attention else "",
             len(layers),
         )
 

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

@@ -6,7 +6,15 @@ import unittest
 from types import SimpleNamespace
 from unittest.mock import patch
 
+import torch
+
 from sglang.multimodal_gen.runtime.layers.attention import layer as layer_module
+from sglang.multimodal_gen.runtime.layers.attention.backends import (
+    flash_attn as flash_attn_module,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+    FlashAttentionImpl,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     denoising as denoising_module,
 )
@@ -30,7 +38,9 @@ def _fake_backend_cls(enum, *, ring_capable=True):
     )
 
 
-def _fake_layer(default=AttentionBackendEnum.FA) -> SimpleNamespace:
+def _fake_layer(
+    default=AttentionBackendEnum.FA, *, is_cross_attention=False
+) -> SimpleNamespace:
     return SimpleNamespace(
         backend=default,
         _default_attn_backend=default,
@@ -39,11 +49,38 @@ def _fake_layer(default=AttentionBackendEnum.FA) -> SimpleNamespace:
         _required_attention_backend=None,
         _attn_impl_ctor_kwargs={"num_heads": 2},
         attn_impl=f"{default.name.lower()}_impl",
-        head_size=64,
-        dtype="bf16",
+        head_size=128,
+        dtype=torch.bfloat16,
+        is_cross_attention=is_cross_attention,
     )
 
 
+class TestSkipSoftmaxDispatch(unittest.TestCase):
+    def test_dense_trtllm_precedes_skip_threshold(self):
+        impl = FlashAttentionImpl(
+            num_heads=2,
+            head_size=128,
+            causal=False,
+            softmax_scale=128**-0.5,
+        )
+        params = SimpleNamespace(start_step=14, threshold_scale_factor=500.0)
+        with patch.object(
+            flash_attn_module, "get_request_skip_softmax_params", return_value=params
+        ):
+            with patch.object(
+                flash_attn_module,
+                "get_forward_context",
+                return_value=SimpleNamespace(current_timestep=13),
+            ):
+                self.assertEqual(impl._request_skip_softmax_threshold(), (True, None))
+            with patch.object(
+                flash_attn_module,
+                "get_forward_context",
+                return_value=SimpleNamespace(current_timestep=14),
+            ):
+                self.assertEqual(impl._request_skip_softmax_threshold(), (True, 500.0))
+
+
 class TestMaybeOverrideAttentionBackend(unittest.TestCase):
     def setUp(self):
         self.default_backend_cls = _fake_backend_cls(AttentionBackendEnum.FA)
@@ -59,6 +96,7 @@ class TestMaybeOverrideAttentionBackend(unittest.TestCase):
         self.stage._attn_backend_default = self.default_backend_cls
         self.stage._attn_metadata_head_size = 64
         self.stage._attention_backend_active_override = None
+        self.stage._skip_softmax_forced_fa = False
 
         self.layers = [_fake_layer(), _fake_layer()]
         self.prepare_calls = []
@@ -122,6 +160,38 @@ class TestMaybeOverrideAttentionBackend(unittest.TestCase):
         self.assertIs(self.stage.attn_backend, self.default_backend_cls)
         self.assertIsNone(self.stage._attention_backend_active_override)
 
+    def test_skip_softmax_forces_fa_only_for_self_attention(self):
+        self.layers[1] = _fake_layer(is_cross_attention=True)
+        self.stage._maybe_override_attention_backend(
+            _batch(None), force_fa_for_self_attention=True
+        )
+        self.assertEqual(
+            self.prepare_calls,
+            [(self.layers[0], AttentionBackendEnum.FA)],
+        )
+        self.assertEqual(
+            self.apply_calls,
+            [
+                (self.layers[0], AttentionBackendEnum.FA),
+                (self.layers[1], None),
+            ],
+        )
+
+    def test_skip_softmax_rejects_non_fa_override(self):
+        with self.assertRaisesRegex(ValueError, "requires the FA attention backend"):
+            self.stage._maybe_override_attention_backend(
+                _batch("sage_attn"), force_fa_for_self_attention=True
+            )
+
+    def test_skip_softmax_uses_custom_model_fa_without_shared_layers(self):
+        self.layers.clear()
+        self.stage._maybe_override_attention_backend(
+            _batch(None), force_fa_for_self_attention=True
+        )
+        self.assertEqual(self.prepare_calls, [])
+        self.assertEqual(self.apply_calls, [])
+        self.assertIs(self.stage.attn_backend, self.default_backend_cls)
+
     def test_unknown_backend_name_rejected(self):
         with self.assertRaisesRegex(ValueError, "Unknown attention_backend_override"):
             self.stage._maybe_override_attention_backend(_batch("bogus_attn"))

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

@@ -3,6 +3,9 @@
 
 import unittest
 
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    trailing_padding_used_len,
+)
 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,
@@ -10,6 +13,11 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.m
 
 
 class TestMiniMaxH3PackedSequence(unittest.TestCase):
+    def test_trailing_padding_layout(self):
+        self.assertEqual(trailing_padding_used_len(64, 61, (0, 61, 64)), 61)
+        self.assertIsNone(trailing_padding_used_len(64, 31, (0, 32, 64)))
+        self.assertIsNone(trailing_padding_used_len(64, 61, (0, 61)))
+
     def test_t2va_structure(self):
         built = minimax_h3_packed_sequence(
             text_len=97,

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

@@ -28,8 +28,10 @@ from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingPara
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     QUALITY_LEVELS,
     SamplingParams,
+    SkipSoftmaxParams,
     _json_safe,
     quality_allows_kernel_fusions,
+    resolve_skip_softmax_params,
 )
 from sglang.multimodal_gen.configs.sample.spectrum import SpectrumParams
 from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams
@@ -69,6 +71,29 @@ class TestSamplingParamsValidate(unittest.TestCase):
             with self.assertRaisesRegex(ValueError, r"quality must be one of"):
                 SamplingParams(quality=bad)  # type: ignore[arg-type]
 
+    def test_skip_softmax_params(self):
+        params = {"threshold_scale_factor": 500, "start_step": 14}
+        self.assertEqual(
+            resolve_skip_softmax_params(params),
+            SkipSoftmaxParams(threshold_scale_factor=500.0, start_step=14),
+        )
+        self.assertEqual(
+            SamplingParams(skip_softmax_params=params).skip_softmax_params, params
+        )
+
+    def test_skip_softmax_params_reject_invalid_values(self):
+        invalid = (
+            {},
+            {"threshold_scale_factor": 0},
+            {"threshold_scale_factor": math.inf},
+            {"threshold_scale_factor": 1, "start_step": -1},
+            {"threshold_scale_factor": 1, "unknown": True},
+        )
+        for params in invalid:
+            with self.subTest(params=params):
+                with self.assertRaisesRegex(ValueError, "skip_softmax_params"):
+                    SamplingParams(skip_softmax_params=params)
+
     def test_seed_accepts_int_or_non_empty_int_list(self):
         self.assertEqual(SamplingParams(seed=7).seed, 7)
         self.assertEqual(SamplingParams(seed=[7, 8]).seed, [7, 8])