diff --git a/.github/workflows/pr-test-jit-kernel.yml b/.github/workflows/pr-test-jit-kernel.yml index 001ace62d..772e69392 100644 --- a/.github/workflows/pr-test-jit-kernel.yml +++ b/.github/workflows/pr-test-jit-kernel.yml @@ -6,6 +6,9 @@ on: jit_kernel: required: true type: string + b200_runner: + required: true + type: string pr_head_sha: required: false type: string @@ -115,3 +118,30 @@ jobs: run: | cd test/ python3 run_suite.py --hw cuda --suite stage-b-kernel-benchmark-1-gpu-large + + jit-kernel-b200-test: + if: | + github.event_name != 'schedule' && + inputs.test_parallel_dispatch != 'true' && + !inputs.target_stage + runs-on: ${{ inputs.b200_runner }} + timeout-minutes: 240 + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} + + - uses: ./.github/actions/check-stage-health + + - uses: ./.github/actions/check-maintenance + + - name: Install dependencies + timeout-minutes: 20 + run: | + bash scripts/ci/cuda/ci_install_dependency.sh diffusion + + - name: Run B200 diffusion test + timeout-minutes: 30 + run: | + cd test/ + python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-b200 diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 588ebcfe3..494c8f767 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -546,6 +546,7 @@ jobs: uses: ./.github/workflows/pr-test-jit-kernel.yml with: jit_kernel: ${{ needs.check-changes.outputs.jit_kernel }} + b200_runner: ${{ needs.check-changes.outputs.b200_runner }} pr_head_sha: ${{ inputs.pr_head_sha || '' }} git_ref: ${{ inputs.git_ref || '' }} target_stage: ${{ inputs.target_stage || '' }} diff --git a/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py b/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py new file mode 100644 index 000000000..5ea8189ea --- /dev/null +++ b/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py @@ -0,0 +1,251 @@ +import sys + +import flashinfer +import pytest +import torch + +from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, + ModelOptFp4LinearMethod, +) +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, suite="stage-b-kernel-unit-1-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 _build_layer( + weight_fp4: torch.Tensor, + weight_scale_linear: torch.Tensor, + input_global_scale: torch.Tensor, + weight_global_scale: torch.Tensor, +) -> None: + 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) + ) + 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) + + method.process_weights_after_loading(layer) + + 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), + ) + + +def _resolve_mode(mode: str): + if mode == "jit_cutlass": + return scaled_fp4_quant, cutlass_scaled_fp4_mm, None + if mode == "flashinfer2": + return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn" + raise ValueError(f"Unknown mode: {mode}") + + +@pytest.mark.skipif( + not _nvfp4_supported(), + reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs", +) +@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(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", ["jit_cutlass", "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(), + ) + + if gemm_backend is None: + actual = gemm_op( + x_fp4, + weight_fp4, + x_scale_swizzled, + weight_scale_swizzled, + alpha, + DTYPE, + ) + else: + 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}" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/run_suite.py b/test/run_suite.py index b4344e9e4..7fe58ad91 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -43,6 +43,7 @@ PER_COMMIT_SUITES = { "stage-b-test-2-gpu-large", "stage-b-test-4-gpu-b200", "stage-b-kernel-unit-1-gpu-large", + "stage-b-kernel-unit-1-gpu-b200", "stage-b-kernel-unit-8-gpu-h200", "stage-b-kernel-benchmark-1-gpu-large", "stage-c-test-4-gpu-h100",