[Diffusion] Fuse Qwen-Image FP8 norm and activation quantization (#37156)
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xiaoyu Zhang
@@ -0,0 +1,82 @@
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.benchmark import marker
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
fused_norm_scale_shift_fp8,
|
||||
fused_scale_residual_norm_scale_shift_fp8,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
|
||||
from sglang.multimodal_gen.runtime.layers.layernorm import (
|
||||
LayerNormScaleShift,
|
||||
ScaleResidualLayerNormScaleShift,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
HIDDEN = 3072
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
@marker.parametrize("rows", [128, 1024, 4096], [128])
|
||||
@marker.parametrize("residual_path", [False, True], [False, True])
|
||||
@marker.benchmark("impl", ["split", "fused"], unit="us")
|
||||
def benchmark(rows: int, residual_path: bool, impl: str):
|
||||
if impl == "fused" and torch.cuda.get_device_capability()[0] < 10:
|
||||
marker.skip("Fused Qwen-Image norm+FP8 quant requires NVIDIA Blackwell")
|
||||
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260831 + rows + int(residual_path))
|
||||
x = torch.randn((1, rows, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
residual = torch.randn_like(x)
|
||||
gate = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
scale = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
shift = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
input_scale = torch.tensor(0.03125, dtype=torch.float32, device=DEVICE)
|
||||
|
||||
if residual_path:
|
||||
layer = ScaleResidualLayerNormScaleShift(
|
||||
HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE
|
||||
).to(DEVICE)
|
||||
|
||||
if impl == "split":
|
||||
|
||||
def fn():
|
||||
normalized, residual_out = layer.forward_cuda(
|
||||
residual, x, gate, shift, scale
|
||||
)
|
||||
quantized, _ = static_quant_fp8(normalized, input_scale)
|
||||
return quantized, residual_out
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
return fused_scale_residual_norm_scale_shift_fp8(
|
||||
residual, x, gate, scale, shift, input_scale, EPS
|
||||
)
|
||||
|
||||
else:
|
||||
layer = LayerNormScaleShift(
|
||||
HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE
|
||||
).to(DEVICE)
|
||||
|
||||
if impl == "split":
|
||||
|
||||
def fn():
|
||||
normalized = layer.forward_cuda(x, shift, scale)
|
||||
return static_quant_fp8(normalized, input_scale)[0]
|
||||
|
||||
else:
|
||||
|
||||
def fn():
|
||||
return fused_norm_scale_shift_fp8(x, scale, shift, input_scale, EPS)
|
||||
|
||||
return marker.do_bench(fn, disable_log_bandwidth=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -0,0 +1,84 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
fused_norm_scale_shift_fp8,
|
||||
fused_scale_residual_norm_scale_shift_fp8,
|
||||
)
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
|
||||
from sglang.multimodal_gen.runtime.layers.layernorm import (
|
||||
LayerNormScaleShift,
|
||||
ScaleResidualLayerNormScaleShift,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
HIDDEN = 3072
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
def _make_inputs(rows: int):
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260831 + rows)
|
||||
x = torch.randn((1, rows, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
residual = torch.randn_like(x)
|
||||
gate = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
scale = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
shift = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator)
|
||||
return x, residual, gate, scale, shift
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rows", [1, 127, 1024])
|
||||
@pytest.mark.parametrize(
|
||||
"input_scale_value", [0.005, 0.03125, 0.4263392984867096, 0.4754464328289032, 1.0]
|
||||
)
|
||||
def test_norm_scale_shift_fp8_is_bit_exact(rows: int, input_scale_value: float) -> None:
|
||||
x, _, _, scale, shift = _make_inputs(rows)
|
||||
input_scale = torch.tensor(input_scale_value, dtype=torch.float32, device=DEVICE)
|
||||
layer = LayerNormScaleShift(
|
||||
HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE
|
||||
).to(DEVICE)
|
||||
|
||||
normalized = layer.forward_cuda(x, shift, scale)
|
||||
expected, _ = static_quant_fp8(normalized, input_scale)
|
||||
actual_normalized, actual = fused_norm_scale_shift_fp8(
|
||||
x, scale, shift, input_scale, EPS
|
||||
)
|
||||
|
||||
assert torch.equal(actual_normalized, normalized)
|
||||
assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rows", [1, 127, 1024])
|
||||
@pytest.mark.parametrize(
|
||||
"input_scale_value", [0.005, 0.03125, 0.4263392984867096, 0.4754464328289032, 1.0]
|
||||
)
|
||||
def test_residual_norm_scale_shift_fp8_is_bit_exact(
|
||||
rows: int, input_scale_value: float
|
||||
) -> None:
|
||||
x, residual, gate, scale, shift = _make_inputs(rows)
|
||||
input_scale = torch.tensor(input_scale_value, dtype=torch.float32, device=DEVICE)
|
||||
layer = ScaleResidualLayerNormScaleShift(
|
||||
HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE
|
||||
).to(DEVICE)
|
||||
|
||||
normalized, expected_residual = layer.forward_cuda(residual, x, gate, shift, scale)
|
||||
expected, _ = static_quant_fp8(normalized, input_scale)
|
||||
actual_normalized, actual, actual_residual = (
|
||||
fused_scale_residual_norm_scale_shift_fp8(
|
||||
residual, x, gate, scale, shift, input_scale, EPS
|
||||
)
|
||||
)
|
||||
|
||||
assert torch.equal(actual_normalized, normalized)
|
||||
assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
|
||||
assert torch.equal(actual_residual, expected_residual)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Unit tests for Qwen-Image ModelOpt FP8 norm+quant activation gates."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp8LinearMethod,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
|
||||
QwenImageTransformerBlock,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _fp8_linear(input_scale: float) -> nn.Module:
|
||||
linear = nn.Module()
|
||||
linear.quant_method = object.__new__(ModelOptFp8LinearMethod)
|
||||
linear.register_parameter(
|
||||
"input_scale",
|
||||
nn.Parameter(
|
||||
torch.tensor(input_scale, dtype=torch.float32, device="cuda"),
|
||||
requires_grad=False,
|
||||
),
|
||||
)
|
||||
return linear
|
||||
|
||||
|
||||
def _attention(*, fused: bool, scales: tuple[float, ...]) -> SimpleNamespace:
|
||||
if fused:
|
||||
return SimpleNamespace(
|
||||
use_fused_qkv=True,
|
||||
to_qkv=_fp8_linear(scales[0]),
|
||||
added_kv_proj_dim=None,
|
||||
)
|
||||
return SimpleNamespace(
|
||||
use_fused_qkv=False,
|
||||
to_q=_fp8_linear(scales[0]),
|
||||
to_k=_fp8_linear(scales[1]),
|
||||
to_v=_fp8_linear(scales[2]),
|
||||
added_kv_proj_dim=None,
|
||||
)
|
||||
|
||||
|
||||
def _block(attn: SimpleNamespace) -> QwenImageTransformerBlock:
|
||||
block = object.__new__(QwenImageTransformerBlock)
|
||||
nn.Module.__init__(block)
|
||||
block.dim = 3072
|
||||
block.zero_cond_t = False
|
||||
block.attn = attn
|
||||
block.img_mlp = None
|
||||
block.txt_mlp = None
|
||||
block._fp8_img_attn_norm_quant = False
|
||||
block._fp8_txt_attn_norm_quant = False
|
||||
block._fp8_img_mlp_norm_quant = False
|
||||
block._fp8_txt_mlp_norm_quant = False
|
||||
return block
|
||||
|
||||
|
||||
@patch("torch.cuda.get_device_capability", return_value=(10, 0))
|
||||
@patch("torch.cuda.is_available", return_value=True)
|
||||
class TestQwenImageFp8NormQuantGate(CustomTestCase):
|
||||
def test_separate_qkv_requires_identical_input_scales(
|
||||
self, _is_available, _capability
|
||||
) -> None:
|
||||
matching = _block(_attention(fused=False, scales=(0.25, 0.25, 0.25)))
|
||||
mismatched = _block(_attention(fused=False, scales=(0.25, 0.5, 0.25)))
|
||||
|
||||
matching.configure_fp8_norm_quant()
|
||||
mismatched.configure_fp8_norm_quant()
|
||||
|
||||
self.assertTrue(matching._fp8_img_attn_norm_quant)
|
||||
self.assertFalse(mismatched._fp8_img_attn_norm_quant)
|
||||
|
||||
def test_merged_qkv_uses_its_materialized_input_scale(
|
||||
self, _is_available, _capability
|
||||
) -> None:
|
||||
block = _block(_attention(fused=True, scales=(0.25,)))
|
||||
|
||||
block.configure_fp8_norm_quant()
|
||||
|
||||
self.assertTrue(block._fp8_img_attn_norm_quant)
|
||||
|
||||
def test_nonpositive_scale_keeps_fusion_disabled(
|
||||
self, _is_available, _capability
|
||||
) -> None:
|
||||
block = _block(_attention(fused=True, scales=(0.0,)))
|
||||
|
||||
block.configure_fp8_norm_quant()
|
||||
|
||||
self.assertFalse(block._fp8_img_attn_norm_quant)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user