diff --git a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py index c3d986ee6..3d88d2a0c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py +++ b/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.ff = SwiGLU(features, multiplier, bias) + self.attn = Attention( + 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) - for _ in range(2) + TextFusionBlock( + 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) - for _ in range(2) + TextFusionBlock( + 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.ff = SwiGLU(features, multiplier, bias) + self.attn = Attention( + 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) diff --git a/python/sglang/multimodal_gen/test/unit/test_krea2_fp8.py b/python/sglang/multimodal_gen/test/unit/test_krea2_fp8.py new file mode 100644 index 000000000..0f8f14c62 --- /dev/null +++ b/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)