From d6e0a8cbf4e81006684d1ebe69923cf97daa1279 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 5 Sep 2026 10:29:09 +0800 Subject: [PATCH] [diffusion] add Helios per-token gated-residual fusion (quality-gated) (#38042) Co-authored-by: BBuf --- .../sglang/kernels/ops/diffusion/__init__.py | 5 ++ .../sites/helios_gated_residual_site.py | 72 +++++++++++++++++++ .../runtime/models/dits/helios.py | 20 ++++-- .../pipelines_core/stages/denoising.py | 7 ++ .../model_specific_stages/helios_denoising.py | 28 ++++++++ 5 files changed, 128 insertions(+), 4 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/sites/helios_gated_residual_site.py diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 1c8cea5fe..2b57ece41 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -581,6 +581,11 @@ _EXPORTS: dict[str, str] = { "mount_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site", "try_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site", "unmount_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site", + "helios_gated_residual_active": "sites.helios_gated_residual_site", + "mark_helios_gated_residual_site": "sites.helios_gated_residual_site", + "mount_helios_gated_residual": "sites.helios_gated_residual_site", + "try_helios_gated_residual": "sites.helios_gated_residual_site", + "unmount_helios_gated_residual": "sites.helios_gated_residual_site", "lingbot_video_gated_residual_active": "sites.lingbot_video_gated_residual_site", "mark_lingbot_video_gated_residual_site": "sites.lingbot_video_gated_residual_site", "mount_lingbot_video_gated_residual": "sites.lingbot_video_gated_residual_site", diff --git a/python/sglang/kernels/ops/diffusion/sites/helios_gated_residual_site.py b/python/sglang/kernels/ops/diffusion/sites/helios_gated_residual_site.py new file mode 100644 index 000000000..53d03f111 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/sites/helios_gated_residual_site.py @@ -0,0 +1,72 @@ +"""Helios per-token gated-residual fusion, gated by request quality. + +Each Helios block applies ``residual + (gate * update).to(residual.dtype)`` at +the self-attention and FFN updates, where ``update`` (post-attn / post-FFN) and +the per-token ``gate`` (``[B, S, 1]``) stay in FP32 while ``residual`` is BF16. +The shared ``residual_gate_add`` kernel computes ``residual + update * gate`` +in a single pass but requires one dtype, so the gate and update are first cast +to BF16. That reordering of the FP32 multiply is numerically equivalent only at +half-precision rounding level (not bit-exact), so the fusion is opt-in: +``quality="extra-high"`` and ``quality="high"`` mount it, while the default +``quality="lossless"`` keeps the reference FP32-multiply form bit-for-bit. +""" + +from __future__ import annotations + +import logging + +import torch +import torch.nn as nn + +from sglang.kernels.ops.diffusion.sites.quality_gate import QualityGatedFusion + +logger = logging.getLogger(__name__) + +_FUSION = QualityGatedFusion( + name="Helios per-token gated residual", + marker_attr="_sgl_helios_gated_residual_site", + enabled_attr="_sgl_helios_gated_residual_enabled", +) + + +def mark_helios_gated_residual_site(module: nn.Module) -> None: + """Mark a Helios block; it starts on the reference path.""" + _FUSION.mark(module) + + +def helios_gated_residual_active(module: nn.Module) -> bool: + return _FUSION.is_enabled(module) + + +def mount_helios_gated_residual(root: nn.Module) -> bool: + return _FUSION.mount(root, logger=logger) + + +def unmount_helios_gated_residual(root: nn.Module) -> None: + _FUSION.unmount(root) + + +def try_helios_gated_residual( + site: nn.Module, + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, +) -> torch.Tensor | None: + """Return the fused ``residual + update * gate`` when the site is enabled. + + Returns ``None`` (caller runs the reference path) when the site is off or + the tensors are not eligible for the per-token fast path. + """ + if not _FUSION.is_enabled(site): + return None + from sglang.kernels.ops.diffusion import ( + can_use_residual_gate_add_cuda, + residual_gate_add, + ) + + if residual.dtype != update.dtype or residual.dtype != gate.dtype: + update = update.to(residual.dtype) + gate = gate.to(residual.dtype) + if not can_use_residual_gate_add_cuda(residual, update, gate): + return None + return residual_gate_add(residual, update, gate) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/helios.py b/python/sglang/multimodal_gen/runtime/models/dits/helios.py index 048956af5..1f7221754 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/helios.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/helios.py @@ -19,6 +19,8 @@ import torch.nn.functional as F from sglang.kernels.ops.diffusion import ( can_use_helios_qk_rope, fused_inplace_helios_qk_rope, + mark_helios_gated_residual_site, + try_helios_gated_residual, ) from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.configs.models.fsdp import is_block @@ -497,6 +499,8 @@ class HeliosTransformerBlock(nn.Module): # 4. Guidance cross-attention flag self.guidance_cross_attn = guidance_cross_attn + mark_helios_gated_residual_site(self) + def forward( self, hidden_states, @@ -526,8 +530,11 @@ class HeliosTransformerBlock(nn.Module): attn_output = self.attn1( norm_hidden_states, rotary_emb, original_context_length ) - hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as( - hidden_states + fused = try_helios_gated_residual(self, hidden_states, attn_output, gate_msa) + hidden_states = ( + fused + if fused is not None + else (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) ) # 2. Cross-attention @@ -562,9 +569,14 @@ class HeliosTransformerBlock(nn.Module): # 3. Feed-forward norm_hidden_states = self.norm3(hidden_states, c_shift_msa, c_scale_msa) ff_output = self.ffn(norm_hidden_states) + fused = try_helios_gated_residual(self, hidden_states, ff_output, c_gate_msa) hidden_states = ( - hidden_states.float() + ff_output.float() * c_gate_msa - ).type_as(hidden_states) + fused + if fused is not None + else (hidden_states.float() + ff_output.float() * c_gate_msa).type_as( + hidden_states + ) + ) return hidden_states 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 f5ac2cd03..69ec0cadd 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -25,6 +25,7 @@ from sglang.kernels.ops.diffusion import ( mount_fused_gate_rmsnorm, mount_fused_linear_gelu, mount_fused_ln_modulate, + mount_helios_gated_residual, mount_hunyuan_qknorm, mount_lingbot_video_gated_residual, mount_lingbot_video_rmsnorm, @@ -36,6 +37,7 @@ from sglang.kernels.ops.diffusion import ( unmount_fused_gate_rmsnorm, unmount_fused_linear_gelu, unmount_fused_ln_modulate, + unmount_helios_gated_residual, unmount_hunyuan_qknorm, unmount_lingbot_video_gated_residual, unmount_lingbot_video_rmsnorm, @@ -222,6 +224,11 @@ _QUALITY_FUSION_HANDLERS: tuple[ mount_lingbot_video_gated_residual, unmount_lingbot_video_gated_residual, ), + ( + "Helios per-token gated residual", + mount_helios_gated_residual, + unmount_helios_gated_residual, + ), ( "SANA-Video BF16-input linear attention", mount_sana_video_linear_attention, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py index 1a8796e35..c8237318d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py @@ -13,6 +13,13 @@ import numpy as np import torch import torch.nn.functional as F +from sglang.kernels.ops.diffusion import ( + mount_helios_gated_residual, + unmount_helios_gated_residual, +) +from sglang.multimodal_gen.configs.sample.sampling_params import ( + quality_allows_kernel_fusions, +) from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( @@ -106,6 +113,26 @@ class HeliosChunkedDenoisingStage(PipelineStage): super().__init__() self.transformer = transformer self.scheduler = scheduler + self._quality_fusions_mounted = False + + def _maybe_toggle_quality_fusions(self, batch: Req) -> None: + # Mount the request-scoped Helios gated-residual fast path for + # quality="extra-high"/"high"; "lossless" keeps the reference + # FP32-multiply form bit-for-bit. + quality = getattr(batch.sampling_params, "quality", "lossless") + want = quality_allows_kernel_fusions(quality) + if want == self._quality_fusions_mounted: + return + self._quality_fusions_mounted = want + if self.transformer is None: + return + if want: + if mount_helios_gated_residual(self.transformer): + logger.info( + "Mounted Helios per-token gated residual for quality=%s", quality + ) + else: + unmount_helios_gated_residual(self.transformer) @property def role_affinity(self) -> RoleType: @@ -481,6 +508,7 @@ class HeliosChunkedDenoisingStage(PipelineStage): def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Run the Helios chunked denoising loop.""" + self._maybe_toggle_quality_fusions(batch) pipeline_config = server_args.pipeline_config scheduler = get_or_create_request_scheduler(batch, self.scheduler) device = (