[Diffusion] Add online FP8 support for Krea-2 (#34136)

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

@@ -28,6 +28,9 @@ from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
     RowParallelLinear,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
 )
@@ -201,18 +204,42 @@ class RMSNorm(nn.Module):
 
 class SwiGLU(nn.Module):
     def __init__(
-        self, features: int, multiplier: int, bias: bool = False, multiple: int = 128
+        self,
+        features: int,
+        multiplier: int,
+        bias: bool = False,
+        multiple: int = 128,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         mlpdim = int(2 * features / 3) * multiplier
         mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
         # Tensor-parallel: gate/up shard the hidden dim by column, down all-reduces.
         self.gate = ColumnParallelLinear(
-            features, mlpdim, bias=bias, gather_output=False
+            features,
+            mlpdim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate",
         )
-        self.up = ColumnParallelLinear(features, mlpdim, bias=bias, gather_output=False)
+        self.up = ColumnParallelLinear(
+            features,
+            mlpdim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.up",
+        )
+
         self.down = RowParallelLinear(
-            mlpdim, features, bias=bias, input_is_parallel=True
+            mlpdim,
+            features,
+            bias=bias,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down",
         )
 
     def forward(self, x: Tensor) -> Tensor:
@@ -223,7 +250,15 @@ class SwiGLU(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self, dim: int, heads: int, kvheads: int = None, bias: bool = False):
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        kvheads: int = None,
+        bias: bool = False,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.heads = heads
         self.kvheads = kvheads if kvheads is not None else heads
@@ -240,21 +275,52 @@ class Attention(nn.Module):
         self.local_kvheads = self.kvheads // tp
 
         self.to_q = ColumnParallelLinear(
-            dim, self.headdim * self.heads, bias=bias, gather_output=False
+            dim,
+            self.headdim * self.heads,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q",
         )
         self.to_k = ColumnParallelLinear(
-            dim, self.headdim * self.kvheads, bias=bias, gather_output=False
+            dim,
+            self.headdim * self.kvheads,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k",
         )
         self.to_v = ColumnParallelLinear(
-            dim, self.headdim * self.kvheads, bias=bias, gather_output=False
+            dim,
+            self.headdim * self.kvheads,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v",
+        )
+        self.to_gate = ColumnParallelLinear(
+            dim,
+            dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_gate",
         )
-        self.to_gate = ColumnParallelLinear(dim, dim, bias=bias, gather_output=False)
         self.norm_q = RMSNorm(self.headdim)
         self.norm_k = RMSNorm(self.headdim)
         # to_out is a ModuleList ([linear]) so the param is to_out.0.weight, matching
         # the diffusers Attention layout in the released checkpoint.
         self.to_out = nn.ModuleList(
-            [RowParallelLinear(dim, dim, bias=bias, input_is_parallel=True)]
+            [
+                RowParallelLinear(
+                    dim,
+                    dim,
+                    bias=bias,
+                    input_is_parallel=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0",
+                )
+            ]
         )
         # Native GQA flash via the platform backend; parameterless.
         self.attn = USPAttention(
@@ -371,12 +437,27 @@ class TextFusionBlock(nn.Module):
         multiplier: int,
         bias: bool = False,
         kvheads: int = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.norm1 = RMSNorm(features)
         self.norm2 = RMSNorm(features)
-        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.attn = Attention(
-        self.ff = SwiGLU(features, multiplier, bias)
+            dim=features,
+            heads=heads,
+            bias=bias,
+            kvheads=kvheads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+        self.ff = SwiGLU(
+            features,
+            multiplier,
+            bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff",
+        )
 
     def forward(
         self,
@@ -410,19 +491,37 @@ class TextFusionTransformer(nn.Module):
         multiplier: int,
         bias: bool = False,
         kvheads: int = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.layerwise_blocks = nn.ModuleList(
             [
-                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                TextFusionBlock(
-                for _ in range(2)
+                    txt_dim,
+                    heads,
+                    multiplier,
+                    bias,
+                    kvheads,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layerwise_blocks.{i}",
+                )
+                for i in range(2)
             ]
         )
         self.projector = nn.Linear(num_txt_layers, 1, bias=False)
         self.refiner_blocks = nn.ModuleList(
             [
-                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                TextFusionBlock(
-                for _ in range(2)
+                    txt_dim,
+                    heads,
+                    multiplier,
+                    bias,
+                    kvheads,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.refiner_blocks.{i}",
+                )
+                for i in range(2)
             ]
         )
 
@@ -452,6 +551,8 @@ class SingleStreamBlock(nn.Module):
         multiplier: int,
         bias: bool = False,
         kvheads: int = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         # (6, features) modulation table added to the timestep projection (AdaLN-single),
@@ -459,8 +560,21 @@ class SingleStreamBlock(nn.Module):
         self.scale_shift_table = nn.Parameter(torch.zeros(6, features))
         self.norm1 = RMSNorm(features)
         self.norm2 = RMSNorm(features)
-        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.attn = Attention(
-        self.ff = SwiGLU(features, multiplier, bias)
+            dim=features,
+            heads=heads,
+            bias=bias,
+            kvheads=kvheads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+        self.ff = SwiGLU(
+            features,
+            multiplier,
+            bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff",
+        )
 
     def forward(
         self,
@@ -519,7 +633,7 @@ class Krea2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         self,
         config: Krea2DitConfig,
         hf_config: dict[str, Any],
-        quant_config: Optional[Any] = None,
+        quant_config: Optional[QuantizationConfig] = None,
     ) -> None:
         super().__init__(config=config, hf_config=hf_config)
         ac = self.config
@@ -542,9 +656,15 @@ class Krea2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         self.transformer_blocks = nn.ModuleList(
             [
                 SingleStreamBlock(
-                    ac.features, ac.heads, ac.multiplier, ac.bias, ac.kvheads
+                    ac.features,
+                    ac.heads,
+                    ac.multiplier,
+                    ac.bias,
+                    ac.kvheads,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(ac.layers)
+                for i in range(ac.layers)
             ]
         )
         self.time_embed = TimeEmbed(ac.tdim, ac.features)
@@ -555,6 +675,8 @@ class Krea2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
             ac.multiplier,
             ac.bias,
             ac.txtkvheads,
+            quant_config=quant_config,
+            prefix="text_fusion",
         )
         self.txt_in = TxtIn(ac.txtdim, ac.features)
         self.final_layer = LastLayer(ac.features, ac.patch, ac.channels)

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

@@ -0,0 +1,64 @@
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.krea2 import Krea2DitConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    maybe_init_distributed_environment_and_model_parallel,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import (
+    Fp8Config,
+    Fp8LinearMethod,
+)
+from sglang.multimodal_gen.runtime.models.dits.krea2 import (
+    Krea2Transformer2DModel,
+)
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
+    ensure_distributed_env_defaults,
+)
+
+
+def _ensure_single_process_parallel_runtime() -> None:
+    if model_parallel_is_initialized():
+        return
+    ensure_distributed_env_defaults()
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+
+
+def test_krea2_online_fp8_quantizes_transformer_linears():
+    _ensure_single_process_parallel_runtime()
+
+    quant_config = Fp8Config(
+        ignored_layers=[
+            "transformer_blocks.0.attn.to_q",
+            "transformer_blocks.0.ff.up",
+        ]
+    )
+
+    with torch.device("meta"):
+        model = Krea2Transformer2DModel(
+            config=Krea2DitConfig(),
+            hf_config={},
+            quant_config=quant_config,
+        )
+
+    block = model.transformer_blocks[0]
+
+    assert isinstance(block.attn.to_q.quant_method, UnquantizedLinearMethod)
+    assert isinstance(block.attn.to_k.quant_method, Fp8LinearMethod)
+    assert isinstance(block.attn.to_v.quant_method, Fp8LinearMethod)
+    assert isinstance(block.attn.to_out[0].quant_method, Fp8LinearMethod)
+
+    assert isinstance(block.ff.gate.quant_method, Fp8LinearMethod)
+    assert isinstance(block.ff.up.quant_method, UnquantizedLinearMethod)
+    assert isinstance(block.ff.down.quant_method, Fp8LinearMethod)
+
+    text_block = model.text_fusion.layerwise_blocks[0]
+    assert isinstance(text_block.attn.to_q.quant_method, Fp8LinearMethod)
+    assert isinstance(text_block.ff.gate.quant_method, Fp8LinearMethod)
+
+    # Numerically sensitive boundary layers remain BF16-native.
+    assert isinstance(model.img_in, torch.nn.Linear)
+    assert isinstance(model.time_mod_proj, torch.nn.Linear)
+    assert isinstance(model.final_layer.linear, torch.nn.Linear)
+    assert isinstance(model.text_fusion.projector, torch.nn.Linear)