minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Tanmay patil
2026-08-10 23:12:05 -07:00
committed by GitHub
parent 667e18d99d
commit afa2d5570b
2 changed files with 207 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/krea2.py
+143 -21
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_krea2_fp8.py
+64
View File
@@ -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)