[diffusion] add Helios per-token gated-residual fusion (quality-gated) (#38042)

Co-authored-by: BBuf <bbuf@users.noreply.github.com>

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",

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)

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
 

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,

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 = (