[kernels] Reorganize ops/diffusion by operator domain behind a lazy facade (#35114)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 5
Claude Opus 5
parent
7605529bdf
commit
ae6945e112
@@ -0,0 +1,141 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp8Config,
|
||||
ModelOptFp8LinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
cutlass_fp8_supported,
|
||||
input_to_float8,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
MAX_FP8_DIFF = 5e-4
|
||||
TEST_CASES = [
|
||||
pytest.param(19, 150, 80, id="misaligned_projection_shape"),
|
||||
pytest.param(512, 3072, 4096, id="flux2_added_kv_projection_shape"),
|
||||
]
|
||||
|
||||
|
||||
def _modelopt_fp8_supported() -> bool:
|
||||
return torch.cuda.is_available() and cutlass_fp8_supported()
|
||||
|
||||
|
||||
def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
if denominator == 0:
|
||||
return 0.0
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return (1 - sim).item()
|
||||
|
||||
|
||||
def _dequantize_fp8_input(qinput: torch.Tensor, x_scale: torch.Tensor) -> torch.Tensor:
|
||||
return qinput.to(torch.float32) * x_scale.to(torch.float32)
|
||||
|
||||
|
||||
def _dequantize_fp8_weight(
|
||||
weight: torch.Tensor, weight_scale: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
if weight_scale.ndim == 0 or weight_scale.numel() == 1:
|
||||
scale = weight_scale.to(torch.float32)
|
||||
else:
|
||||
scale = weight_scale.to(torch.float32).reshape(-1, 1).t()
|
||||
return weight.to(torch.float32) * scale
|
||||
|
||||
|
||||
def _build_layer(
|
||||
weight_q: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
input_scale: torch.Tensor,
|
||||
) -> tuple[torch.nn.Module, ModelOptFp8LinearMethod]:
|
||||
output_size, input_size = weight_q.shape
|
||||
method = ModelOptFp8LinearMethod(
|
||||
ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
|
||||
)
|
||||
layer = torch.nn.Module()
|
||||
method.create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size,
|
||||
output_partition_sizes=[output_size],
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=DTYPE,
|
||||
weight_loader=lambda *args, **kwargs: None,
|
||||
)
|
||||
layer = layer.to(device=DEVICE)
|
||||
|
||||
layer.weight.data.copy_(weight_q)
|
||||
layer.weight_scale.data.copy_(weight_scale.reshape_as(layer.weight_scale))
|
||||
layer.input_scale.data.copy_(input_scale.reshape_as(layer.input_scale))
|
||||
method.process_weights_after_loading(layer)
|
||||
return layer, method
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _modelopt_fp8_supported(),
|
||||
reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
|
||||
)
|
||||
@pytest.mark.parametrize("m,n,k", TEST_CASES)
|
||||
def test_checkpoint_processing(m: int, n: int, k: int) -> None:
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260410 + m + n + k)
|
||||
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight_q, weight_scale = input_to_float8(weight)
|
||||
input_scale = torch.tensor(1.0, device=DEVICE, dtype=torch.float32)
|
||||
|
||||
layer, _ = _build_layer(weight_q, weight_scale, input_scale)
|
||||
|
||||
assert tuple(layer.weight.shape) == (k, n)
|
||||
assert tuple(layer.weight.stride()) == (1, k)
|
||||
assert layer.weight.dtype == torch.float8_e4m3fn
|
||||
assert layer.input_scale.ndim == 0
|
||||
assert tuple(layer.weight_scale.shape) == (n, 1)
|
||||
|
||||
expected_weight = weight_q.t().to(torch.float32) * weight_scale.to(torch.float32)
|
||||
actual_weight = _dequantize_fp8_weight(layer.weight, layer.weight_scale)
|
||||
torch.testing.assert_close(actual_weight, expected_weight, atol=0.0, rtol=0.0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _modelopt_fp8_supported(),
|
||||
reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
|
||||
)
|
||||
@pytest.mark.parametrize("m,n,k", TEST_CASES)
|
||||
def test_shape_correctness(m: int, n: int, k: int) -> None:
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260410 + m + n + k)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight_q, weight_scale = input_to_float8(weight)
|
||||
_, input_scale = input_to_float8(x)
|
||||
|
||||
layer, method = _build_layer(weight_q, weight_scale, input_scale)
|
||||
|
||||
qinput, x_scale = static_quant_fp8(
|
||||
x.contiguous(),
|
||||
layer.input_scale,
|
||||
repeat_scale=method.cutlass_fp8_supported,
|
||||
)
|
||||
expected = torch.matmul(
|
||||
_dequantize_fp8_input(qinput, x_scale),
|
||||
_dequantize_fp8_weight(layer.weight, layer.weight_scale),
|
||||
)
|
||||
|
||||
actual = method.apply(layer, x)
|
||||
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
|
||||
assert diff < MAX_FP8_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,447 @@
|
||||
import sys
|
||||
|
||||
import flashinfer
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.quantization import (
|
||||
modelopt_quant as diffusion_modelopt_quant,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp4Config,
|
||||
ModelOptFp4LinearMethod,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
# B200-only correctness coverage for diffusion NVFP4 scaled mm.
|
||||
register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
BLOCK_SIZE = 16
|
||||
FLOAT4_E2M1_MAX = 6.0
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
FP4_VALUE_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
|
||||
DEEPGEMM_FP4_MAX_DIFF = 0.02
|
||||
TEST_CASES = [
|
||||
pytest.param(19, 150, 80, id="padding_regression"),
|
||||
pytest.param(512, 6144, 128, id="flux2_projection_shape"),
|
||||
]
|
||||
FLUX2_PROJECTION_SHAPE = (512, 6144, 128)
|
||||
|
||||
|
||||
def _nvfp4_supported() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
|
||||
|
||||
|
||||
def _make_global_scale(x: torch.Tensor) -> torch.Tensor:
|
||||
max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
|
||||
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)
|
||||
|
||||
|
||||
def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
if denominator == 0:
|
||||
return 0.0
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return (1 - sim).item()
|
||||
|
||||
|
||||
def _swap_fp4_nibbles(packed: torch.Tensor) -> torch.Tensor:
|
||||
return ((packed >> 4) | (packed << 4)).contiguous()
|
||||
|
||||
|
||||
def _fp4_lut(device: torch.device) -> torch.Tensor:
|
||||
return torch.tensor(FP4_VALUE_LUT, dtype=torch.float32, device=device)
|
||||
|
||||
|
||||
def _unpack_fp4_bytes(packed: torch.Tensor) -> torch.Tensor:
|
||||
assert packed.dtype == torch.uint8
|
||||
lut = _fp4_lut(packed.device)
|
||||
|
||||
def _decode(nibbles: torch.Tensor) -> torch.Tensor:
|
||||
values = lut[(nibbles & 0x7).to(torch.long)]
|
||||
return torch.where((nibbles & 0x8) != 0, -values, values)
|
||||
|
||||
low = _decode(packed & 0x0F)
|
||||
high = _decode((packed & 0xF0) >> 4)
|
||||
return torch.stack((low, high), dim=-1).reshape(
|
||||
packed.shape[0], packed.shape[1] * 2
|
||||
)
|
||||
|
||||
|
||||
def _swizzled_to_linear(
|
||||
scales_swizzled: torch.Tensor,
|
||||
rows: int,
|
||||
cols: int,
|
||||
) -> torch.Tensor:
|
||||
scales_swizzled = scales_swizzled.view(torch.float8_e4m3fn)
|
||||
row_tiles = (rows + 128 - 1) // 128
|
||||
tile_cols = BLOCK_SIZE * 4
|
||||
col_tiles = (cols + tile_cols - 1) // tile_cols
|
||||
tmp = scales_swizzled.reshape(1, row_tiles, col_tiles, 32, 4, 4)
|
||||
tmp = tmp.permute(0, 1, 4, 3, 2, 5)
|
||||
linear = tmp.reshape(row_tiles * 128, col_tiles * tile_cols // BLOCK_SIZE)
|
||||
return linear[:rows, : cols // BLOCK_SIZE]
|
||||
|
||||
|
||||
def _dequantize_nvfp4(
|
||||
packed: torch.Tensor,
|
||||
scales_swizzled: torch.Tensor,
|
||||
global_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
rows, packed_cols = packed.shape
|
||||
cols = packed_cols * 2
|
||||
unpacked = _unpack_fp4_bytes(packed).reshape(rows, cols // BLOCK_SIZE, BLOCK_SIZE)
|
||||
scales_linear = _swizzled_to_linear(scales_swizzled, rows, cols).to(torch.float32)
|
||||
return (unpacked * (scales_linear / global_scale).unsqueeze(-1)).reshape(rows, cols)
|
||||
|
||||
|
||||
def _quantize_weight_for_checkpoint(
|
||||
weight: torch.Tensor, weight_global_scale: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
weight_fp4, weight_scale_linear = flashinfer.fp4_quantize(
|
||||
weight,
|
||||
weight_global_scale,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
if weight_scale_linear.dtype == torch.uint8:
|
||||
weight_scale_linear = weight_scale_linear.view(torch.float8_e4m3fn)
|
||||
return weight_fp4, weight_scale_linear.contiguous()
|
||||
|
||||
|
||||
def _set_diffusion_fp4_backend(
|
||||
monkeypatch: pytest.MonkeyPatch, backend: str | None
|
||||
) -> None:
|
||||
if backend is None:
|
||||
monkeypatch.delenv(
|
||||
"SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", raising=False
|
||||
)
|
||||
else:
|
||||
monkeypatch.setenv("SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", backend)
|
||||
|
||||
current_platform.__class__.get_modelopt_flashinfer_fp4_backend.cache_clear()
|
||||
current_platform.__class__.get_modelopt_fp4_gemm_op.cache_clear()
|
||||
diffusion_modelopt_quant._get_fp4_gemm_op.cache_clear()
|
||||
|
||||
|
||||
def _build_layer(
|
||||
weight_fp4: torch.Tensor,
|
||||
weight_scale_linear: torch.Tensor,
|
||||
input_global_scale: torch.Tensor,
|
||||
weight_global_scale: torch.Tensor,
|
||||
*,
|
||||
weight_scale_device: torch.device | str | None = None,
|
||||
checkpoint_weight_scale_layout: str = "linear",
|
||||
) -> tuple[ModelOptFp4LinearMethod, torch.nn.Module]:
|
||||
output_size, input_size_half = weight_fp4.shape
|
||||
input_size = input_size_half * 2
|
||||
method = ModelOptFp4LinearMethod(
|
||||
ModelOptFp4Config(
|
||||
is_checkpoint_nvfp4_serialized=True,
|
||||
group_size=BLOCK_SIZE,
|
||||
swap_weight_nibbles=True,
|
||||
checkpoint_weight_scale_layout=checkpoint_weight_scale_layout,
|
||||
)
|
||||
)
|
||||
layer = torch.nn.Module()
|
||||
method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=input_size,
|
||||
output_partition_sizes=[output_size],
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=DTYPE,
|
||||
weight_loader=lambda *args, **kwargs: None,
|
||||
)
|
||||
layer = layer.to(device=DEVICE)
|
||||
|
||||
checkpoint_weight = _swap_fp4_nibbles(weight_fp4)
|
||||
layer.weight.data.copy_(checkpoint_weight)
|
||||
layer.input_scale.data.copy_(
|
||||
(1.0 / input_global_scale).reshape_as(layer.input_scale)
|
||||
)
|
||||
layer.weight_scale_2.data.copy_(
|
||||
(1.0 / weight_global_scale).reshape_as(layer.weight_scale_2)
|
||||
)
|
||||
layer.weight_scale.data.copy_(weight_scale_linear)
|
||||
if weight_scale_device is not None:
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
layer.weight_scale.detach().to(weight_scale_device), requires_grad=False
|
||||
)
|
||||
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
_, flashinfer_backend = current_platform.get_modelopt_fp4_gemm_op()
|
||||
if flashinfer_backend == "trtllm":
|
||||
expected_weight, _ = pad_nvfp4_weight(
|
||||
weight_fp4, n_alignment=128, k_alignment=0
|
||||
)
|
||||
expected_scale = (
|
||||
_swizzled_to_linear(weight_scale_linear, output_size, input_size)
|
||||
if checkpoint_weight_scale_layout == "swizzled"
|
||||
else weight_scale_linear
|
||||
)
|
||||
if expected_scale.shape[0] != expected_weight.shape[0]:
|
||||
pad_n = expected_weight.shape[0] - expected_scale.shape[0]
|
||||
expected_scale = torch.nn.functional.pad(expected_scale, (0, 0, 0, pad_n))
|
||||
|
||||
expected_padding_cols = 0
|
||||
if expected_scale.shape[1] % 4 != 0:
|
||||
padded_scale_k = ((expected_scale.shape[1] + 4 - 1) // 4) * 4
|
||||
pad_scale_k = padded_scale_k - expected_scale.shape[1]
|
||||
expected_scale = torch.nn.functional.pad(
|
||||
expected_scale, (0, pad_scale_k, 0, 0)
|
||||
)
|
||||
pad_weight_k = pad_scale_k * 8
|
||||
expected_weight = torch.nn.functional.pad(
|
||||
expected_weight, (0, pad_weight_k, 0, 0)
|
||||
)
|
||||
expected_padding_cols = pad_weight_k
|
||||
|
||||
expected_weight = flashinfer.shuffle_matrix_a(
|
||||
expected_weight.view(torch.uint8), 128
|
||||
)
|
||||
expected_scale = (
|
||||
flashinfer.shuffle_matrix_sf_a(expected_scale.view(torch.uint8), 128)
|
||||
.reshape(expected_scale.shape)
|
||||
.view(torch.float8_e4m3fn)
|
||||
)
|
||||
|
||||
assert torch.equal(layer.weight, expected_weight)
|
||||
assert torch.equal(
|
||||
layer.weight_scale_interleaved.view(torch.uint8),
|
||||
expected_scale.view(torch.uint8),
|
||||
)
|
||||
assert layer.weights_padding_cols == expected_padding_cols
|
||||
else:
|
||||
expected_weight, expected_padding_cols = pad_nvfp4_weight(weight_fp4)
|
||||
expected_scale_shape = (
|
||||
((output_size + 128 - 1) // 128) * 128,
|
||||
(((input_size // BLOCK_SIZE) + 4 - 1) // 4) * 4,
|
||||
)
|
||||
|
||||
assert torch.equal(layer.weight, expected_weight)
|
||||
assert layer.weight_scale_interleaved.shape == expected_scale_shape
|
||||
assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn
|
||||
assert layer.weights_padding_cols == expected_padding_cols
|
||||
torch.testing.assert_close(
|
||||
layer.alpha,
|
||||
(1.0 / (input_global_scale * weight_global_scale)).to(torch.float32),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
layer.input_scale_inv,
|
||||
input_global_scale.to(torch.float32),
|
||||
)
|
||||
return method, layer
|
||||
|
||||
|
||||
def _resolve_mode(mode: str):
|
||||
if mode == "flashinfer2":
|
||||
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn"
|
||||
if mode == "flashinfer_trtllm":
|
||||
return flashinfer.fp4_quantize, flashinfer.mm_fp4, "trtllm"
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"backend", [None, "flashinfer_trtllm"], ids=["default", "flashinfer_trtllm"]
|
||||
)
|
||||
@pytest.mark.parametrize("m,n,k", TEST_CASES)
|
||||
def test_checkpoint_processing(
|
||||
monkeypatch: pytest.MonkeyPatch, backend: str | None, m: int, n: int, k: int
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, backend)
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260404 + m + n + k)
|
||||
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
|
||||
_build_layer(
|
||||
weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
@pytest.mark.parametrize("mode", ["flashinfer2"])
|
||||
def test_flux2_shape_correctness(mode: str) -> None:
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
quantize_op, gemm_op, gemm_backend = _resolve_mode(mode)
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260404 + m + n + k)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = _make_global_scale(x)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
|
||||
|
||||
x_fp4, x_scale_swizzled = quantize_op(x, input_global_scale)
|
||||
weight_fp4, weight_scale_swizzled = quantize_op(weight, weight_global_scale)
|
||||
if x_scale_swizzled.dtype == torch.uint8:
|
||||
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
if weight_scale_swizzled.dtype == torch.uint8:
|
||||
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
|
||||
expected = torch.matmul(
|
||||
_dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
|
||||
_dequantize_nvfp4(weight_fp4, weight_scale_swizzled, weight_global_scale).t(),
|
||||
)
|
||||
|
||||
actual = gemm_op(
|
||||
x_fp4,
|
||||
weight_fp4.t(),
|
||||
x_scale_swizzled,
|
||||
weight_scale_swizzled.t(),
|
||||
alpha,
|
||||
DTYPE,
|
||||
backend=gemm_backend,
|
||||
)
|
||||
|
||||
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
|
||||
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{mode=}, {m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
def test_flux2_shape_correctness_flashinfer_trtllm(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
|
||||
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260404 + m + n + k + 17)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = _make_global_scale(x)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
|
||||
method, layer = _build_layer(
|
||||
weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
|
||||
)
|
||||
actual = method.apply(layer, x)
|
||||
|
||||
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
|
||||
weight_fp4_ref, weight_scale_swizzled = flashinfer.fp4_quantize(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
if x_scale_swizzled.dtype == torch.uint8:
|
||||
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
if weight_scale_swizzled.dtype == torch.uint8:
|
||||
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
|
||||
expected = torch.matmul(
|
||||
_dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
|
||||
_dequantize_nvfp4(
|
||||
weight_fp4_ref, weight_scale_swizzled, weight_global_scale
|
||||
).t(),
|
||||
)
|
||||
|
||||
diff = _calc_diff(actual, expected.to(dtype=DTYPE))
|
||||
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
def test_flux2_swizzled_scale_checkpoint_flashinfer_trtllm_matches_cudnn(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
|
||||
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260517 + m + n + k)
|
||||
|
||||
x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = _make_global_scale(x)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
|
||||
|
||||
x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
|
||||
weight_fp4, weight_scale_swizzled = flashinfer.fp4_quantize(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
if x_scale_swizzled.dtype == torch.uint8:
|
||||
x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
if weight_scale_swizzled.dtype == torch.uint8:
|
||||
weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
|
||||
|
||||
method, layer = _build_layer(
|
||||
weight_fp4,
|
||||
weight_scale_swizzled,
|
||||
input_global_scale,
|
||||
weight_global_scale,
|
||||
checkpoint_weight_scale_layout="swizzled",
|
||||
)
|
||||
actual = method.apply(layer, x)
|
||||
|
||||
expected = flashinfer.mm_fp4(
|
||||
x_fp4,
|
||||
weight_fp4.t(),
|
||||
x_scale_swizzled,
|
||||
weight_scale_swizzled.t(),
|
||||
alpha,
|
||||
DTYPE,
|
||||
backend="cudnn",
|
||||
)
|
||||
|
||||
diff = _calc_diff(actual, expected)
|
||||
assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _nvfp4_supported(),
|
||||
reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
|
||||
)
|
||||
def test_checkpoint_processing_flashinfer_trtllm_cpu_weight_scale(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
|
||||
|
||||
m, n, k = FLUX2_PROJECTION_SHAPE
|
||||
generator = torch.Generator(device=DEVICE)
|
||||
generator.manual_seed(20260413 + m + n + k)
|
||||
|
||||
weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
|
||||
input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
|
||||
weight_global_scale = _make_global_scale(weight)
|
||||
weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
|
||||
weight, weight_global_scale
|
||||
)
|
||||
|
||||
_build_layer(
|
||||
weight_fp4,
|
||||
weight_scale_linear,
|
||||
input_global_scale,
|
||||
weight_global_scale,
|
||||
weight_scale_device="cpu",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
Reference in New Issue
Block a user