diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 924a08fbb..7297ce43c 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -118,6 +118,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--num-inference-steps {STEPS}` and `--seed {SEED}` - `--num-outputs-per-prompt {N}` / `--num-outputs {N}`: generate multiple outputs for each prompt. A scalar seed expands as `seed + output_index`. - `--quality {lossless,extra-high,high}`: cumulative request-level optimization tier. `lossless` (default) keeps the selected deployment's reference path and all unconditional bit-exact replacements. `extra-high` adds only request-gated DiT/VAE kernel fusions; the tier does not itself enable sparse, caching, or other approximate paths. `high` includes the complete `extra-high` set and may also enable model-owned approximate optimizations. Separately configured quantization, attention, or caching options still apply. Support and validation constraints are model-specific. +- `skip_softmax_params` (online request only): explicitly enables lossy BLASST/Skip-Softmax attention for compatible FA self-attention layers. See [Attention Backends](../attention_backends#request-scoped-skip-softmax). - `--height {HEIGHT}`, `--width {WIDTH}`, `--num-frames {N}`, `--fps {FPS}` - `--output-path {PATH}`, `--output-file-name {NAME}`, `--save-output`, `--return-frames` diff --git a/docs/docs/sglang-diffusion/attention_backends.mdx b/docs/docs/sglang-diffusion/attention_backends.mdx index 3b4b08bcb..1f594a604 100644 --- a/docs/docs/sglang-diffusion/attention_backends.mdx +++ b/docs/docs/sglang-diffusion/attention_backends.mdx @@ -789,6 +789,53 @@ dense switching; under ring parallelism the target must be ring-capable. Note `sage_attn` / `sage_attn_3` are lossy (quantized attention) — validate quality on your workload. +### Request-scoped Skip Softmax + +Skip Softmax (BLASST) keeps the QK matmul, but skips the exponential, +softmax-state update, V load, and PV matmul for attention tiles whose estimated +softmax mass is below a threshold. It is an explicit lossy optimization backed +by the FlashInfer kernels shipped with SGLang's pinned dependencies. See the +[BLASST paper](https://arxiv.org/abs/2512.12087) and NVIDIA's +[video-generation study](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog28_Accelerating_Video_Generation_with_GEMM_Quantization_Attention_Quantization_and_Skip_Softmax_Attention_in_TensorRT-LLM.md) +for the algorithm and measured quality/performance trade-offs. + +Pass `skip_softmax_params` in one image or video request. The runtime routes +compatible self-attention layers through the FA/FlashInfer path for that batch; +cross-attention keeps its normal backend, and the next request restores the +server default. + +```python +client.videos.create( + model="", + prompt="...", + extra_body={ + "skip_softmax_params": { + "threshold_scale_factor": 500.0, + "start_step": 14, + } + }, +) +``` + +- `threshold_scale_factor` is required and must be positive. The kernel uses + `threshold_scale_factor / context_length`; larger values skip more work and + usually increase quality loss. +- `start_step` is the zero-based denoising step at which sparse execution + starts. It defaults to `0`; keeping early high-noise steps dense is generally + safer. + +There is intentionally no default threshold. Calibrate the threshold and start +step against output-quality metrics for each model, resolution, step count, and +deployment backend; values tuned for one workload are not portable. + +Current support is FP16/BF16, head dimension 128 or 256, and unmasked +self-attention on Hopper SM90 and Blackwell SM100/SM103/SM107. Ulysses sequence +parallelism is supported because the kernel runs after its all-to-all. Ring +Attention, attention masks, `torch.compile`, and breakable CUDA graphs reject +the request instead of silently running dense attention. Models that merely +offer an FA backend do not automatically qualify: the runtime still checks the +GPU, dtype, head dimension, attention role, and execution mode. + ### Using SpargeAttention Install the optional CUDA extension, then select the backend explicitly: diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 9dfc3462e..bdd68652c 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 2e5a13ba1..59dfd4f83 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index 17245ec5e..b9665e440 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 7f50900f3..69191b0ee 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/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.""" diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index b0bb625c9..630677789 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py index a137d8017..13187e1d7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/skip_softmax.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/skip_softmax.py new file mode 100644 index 000000000..5d6b57aba --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 5ecdb0644..b04fcc17f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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 ) diff --git a/python/sglang/multimodal_gen/runtime/managers/forward_context.py b/python/sglang/multimodal_gen/runtime/managers/forward_context.py index e506929c6..f59e1524c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/forward_context.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index fef29927a..e49aed982 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 69ec0cadd..5841ee462 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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), ) diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py index c44ad2d41..89222181b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py +++ b/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")) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_packed_sequence.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_packed_sequence.py index 302707562..e8f3afcdd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_packed_sequence.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py index 47862045e..d708908ea 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py +++ b/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])