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

[Diffusion] Fuse FLUX.2 ModelOpt FP8 producers and QKV packing (#37162)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xiaoyu Zhang
2026-09-01 16:14:29 +08:00
committed by GitHub
co-authored by Cursor
parent 379e33d87e
commit 1c3ad92438
16 changed files with 1571 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/diffusion/test_flux2_fp8_norm_quant.py
+52
View File
@@ -0,0 +1,52 @@
import sys
import pytest
import torch
from sglang.kernels.ops.diffusion import (
fused_layernorm_modulate_fp8_quant_raw,
fused_layernorm_modulate_raw,
)
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
DEVICE = "cuda"
DTYPE = torch.bfloat16
HIDDEN = 6144
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)
scale = torch.randn((1, 1, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
shift = torch.randn((1, 1, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
return x, scale, shift
@pytest.mark.parametrize("rows", [1, 127, 512])
@pytest.mark.parametrize(
"input_scale_value", [0.005, 0.03125, 0.25, 0.4754464328289032, 1.0]
)
def test_flux2_layernorm_modulate_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)
normalized = fused_layernorm_modulate_raw(
x, scale.squeeze(1), shift.squeeze(1), EPS
)
expected, _ = static_quant_fp8(normalized, input_scale)
actual = fused_layernorm_modulate_fp8_quant_raw(
x, scale.squeeze(1), shift.squeeze(1), input_scale, EPS
)
assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_flux2_qkv_epilogue.py
+158
View File
@@ -0,0 +1,158 @@
import sys
from unittest.mock import patch
import pytest
import torch
from sglang.kernels.ops.diffusion import (
try_fused_flux2_qkv_epilogue,
)
from sglang.multimodal_gen.runtime.layers.layernorm import (
RMSNorm,
apply_qk_norm_with_optional_rope,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
DEVICE = "cuda"
DTYPE = torch.bfloat16
HEAD_DIM = 128
def _packed_qkv(tokens: int, heads: int, generator: torch.Generator):
source = [
torch.randn(
(1, tokens, heads, HEAD_DIM),
dtype=DTYPE,
device=DEVICE,
generator=generator,
)
for _ in range(3)
]
packed = torch.cat([tensor.flatten(2) for tensor in source], dim=-1)
views = [
tensor.unflatten(-1, (heads, HEAD_DIM)) for tensor in packed.chunk(3, dim=-1)
]
assert all(not tensor.is_contiguous() for tensor in views)
return views
@pytest.mark.parametrize("img_tokens,txt_tokens,heads", [(17, 7, 4), (256, 64, 8)])
def test_flux2_qkv_epilogue_is_bit_exact(
img_tokens: int, txt_tokens: int, heads: int
) -> None:
generator = torch.Generator(device=DEVICE)
generator.manual_seed(20260831 + img_tokens)
img_qkv = _packed_qkv(img_tokens, heads, generator)
txt_qkv = _packed_qkv(txt_tokens, heads, generator)
norms = [
RMSNorm(HEAD_DIM, eps=1e-6).to(device=DEVICE, dtype=DTYPE) for _ in range(4)
]
for norm in norms:
norm.weight.data.normal_(generator=generator)
angles = torch.randn(
(img_tokens + txt_tokens, HEAD_DIM // 2),
device=DEVICE,
generator=generator,
)
cache = torch.cat([angles.cos(), angles.sin()], dim=-1).contiguous()
img_reference = [tensor.contiguous() for tensor in img_qkv]
txt_reference = [tensor.contiguous() for tensor in txt_qkv]
txt_reference[0], txt_reference[1] = apply_qk_norm_with_optional_rope(
txt_reference[0],
txt_reference[1],
norms[2],
norms[3],
HEAD_DIM,
cache,
is_neox=False,
)
img_reference[0], img_reference[1] = apply_qk_norm_with_optional_rope(
img_reference[0],
img_reference[1],
norms[0],
norms[1],
HEAD_DIM,
cache,
is_neox=False,
position_offset=txt_tokens,
)
expected = tuple(
torch.cat([txt_reference[index], img_reference[index]], dim=1)
for index in range(3)
)
actual = try_fused_flux2_qkv_epilogue(
*img_qkv,
*txt_qkv,
norms[0].weight,
norms[1].weight,
norms[2].weight,
norms[3].weight,
cache,
1e-6,
1e-6,
)
assert actual is not None
assert all(
torch.equal(result, reference)
for result, reference in zip(actual, expected, strict=True)
)
def test_flux2_qkv_epilogue_rejects_compile() -> None:
tensor = torch.empty((1, 1, 1, HEAD_DIM), device=DEVICE, dtype=DTYPE)
weight = torch.empty((HEAD_DIM,), device=DEVICE, dtype=DTYPE)
cache = torch.empty((2, HEAD_DIM), device=DEVICE, dtype=torch.float32)
with patch("torch.compiler.is_compiling", return_value=True):
assert (
try_fused_flux2_qkv_epilogue(
tensor,
tensor,
tensor,
tensor,
tensor,
tensor,
weight,
weight,
weight,
weight,
cache,
1e-6,
1e-6,
)
is None
)
def test_flux2_qkv_epilogue_rejects_cuda_graph_capture() -> None:
tensor = torch.empty((1, 1, 1, HEAD_DIM), device=DEVICE, dtype=DTYPE)
weight = torch.empty((HEAD_DIM,), device=DEVICE, dtype=DTYPE)
cache = torch.empty((2, HEAD_DIM), device=DEVICE, dtype=torch.float32)
with patch("torch.cuda.is_current_stream_capturing", return_value=True):
assert (
try_fused_flux2_qkv_epilogue(
tensor,
tensor,
tensor,
tensor,
tensor,
tensor,
weight,
weight,
weight,
weight,
cache,
1e-6,
1e-6,
)
is None
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/kernels/ops/diffusion/test_flux2_token_cat_fp8.py
+56
View File
@@ -0,0 +1,56 @@
import sys
from unittest.mock import patch
import pytest
import torch
from sglang.kernels.ops.diffusion import try_flux2_token_cat_fp8
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
@pytest.mark.parametrize("tokens", [1, 127, 4096])
def test_flux2_token_cat_fp8_is_bit_exact(tokens: int) -> None:
generator = torch.Generator(device="cuda")
generator.manual_seed(20260831 + tokens)
attention = torch.randn(
(1, tokens, 6144),
dtype=torch.bfloat16,
device="cuda",
generator=generator,
)
mlp = torch.randn(
(1, tokens, 18432),
dtype=torch.bfloat16,
device="cuda",
generator=generator,
)
scale = torch.tensor([0.013], dtype=torch.float32, device="cuda")
expected, _ = static_quant_fp8(torch.cat([attention, mlp], dim=-1), scale)
actual = try_flux2_token_cat_fp8(attention, mlp, scale)
assert actual is not None
assert torch.equal(actual, expected)
def test_flux2_token_cat_fp8_rejects_compile() -> None:
attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16)
mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16)
scale = torch.ones((1,), device="cuda", dtype=torch.float32)
with patch("torch.compiler.is_compiling", return_value=True):
assert try_flux2_token_cat_fp8(attention, mlp, scale) is None
def test_flux2_token_cat_fp8_rejects_cuda_graph_capture() -> None:
attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16)
mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16)
scale = torch.ones((1,), device="cuda", dtype=torch.float32)
with patch("torch.cuda.is_current_stream_capturing", return_value=True):
assert try_flux2_token_cat_fp8(attention, mlp, scale) is None
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))