diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 4e1060fd9..a79b441e0 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1163,15 +1163,24 @@ def aiter_w8a8_block_fp8_linear( # On ROCm >= 7.2, scale is in bpreshuffle's transposed layout. # Triton needs a row-major view, so adjust strides only. No copy. elif use_triton and _use_aiter_bpreshuffle_gfx95: - x_scale = torch.as_strided(x_scale, x_scale.shape, (1, x_scale.shape[0])) + x_scale = view_aiter_fused_rms_transposed_fp8_scale(x_scale) else: materialize_bpreshuffle_scale = _use_aiter_bpreshuffle_gfx95 and not use_triton + # No-copy bpreshuffle scale: emit it already transposed and stride-reinterpret + # to the column-major bpreshuffle layout, instead of a .t().contiguous().t() + # copy. Bit-identical for M>=2; M==1 keeps materialize (there the [1,G] and + # [G,1] byte orders coincide, so materialize is a no-op view anyway). + emit_bpreshuffle_scale = ( + materialize_bpreshuffle_scale and input_2d.shape[0] >= 2 + ) q_input, x_scale = aiter_per1x128_quant( input_2d, quant_dtype=aiter.dtypes.fp8, - transpose_scale=False, + transpose_scale=emit_bpreshuffle_scale, ) - if materialize_bpreshuffle_scale: + if emit_bpreshuffle_scale: + x_scale = view_aiter_fused_rms_transposed_fp8_scale(x_scale) + elif materialize_bpreshuffle_scale: x_scale = materialize_bpreshuffle_fp8_scale(x_scale) if use_triton: diff --git a/test/registered/unit/layers/test_fp8_bpreshuffle_dense_linear_mi35x.py b/test/registered/unit/layers/test_fp8_bpreshuffle_dense_linear_mi35x.py new file mode 100644 index 000000000..2fc0c4a59 --- /dev/null +++ b/test/registered/unit/layers/test_fp8_bpreshuffle_dense_linear_mi35x.py @@ -0,0 +1,188 @@ +"""Real-path validation of the dense w8a8 bpreshuffle fp8-scale no-copy on gfx95. + +The CPU tests in ``test_fp8_bpreshuffle_scale.py`` only prove the stride formula +against a fabricated layout. They cannot catch ``aiter_per1x128_quant( +transpose_scale=True)`` emitting the wrong layout, the ``emit_bpreshuffle_scale`` +gating being wrong, or an integration mismatch with the CK bpreshuffle GEMM. + +These tests exercise the real kernels ``aiter_w8a8_block_fp8_linear`` routes +through and prove the optimized (``transpose_scale=True`` + no-copy view) path is +equivalent to the original (``transpose_scale=False`` + materialize copy) path: + +- ``test_quant_producer_scale_equivalence`` -- at the quant-producer level: + identical quantized bytes, identical scale *values* after relayout, the + bpreshuffle ``(1, M)`` column-major stride, and zero-copy storage sharing. +- ``test_dense_linear_paths_bit_exact`` -- end-to-end through + ``aiter_w8a8_block_fp8_linear``: the new path (as shipped) vs the old path + (forced by patching the quant to row-major + the relayout to materialize) must + produce **bit-identical** GEMM output. Covers M == 1 (materialize fallback). + +Requires a gfx95 (MI35X) GPU with aiter and ROCm >= 7.2 (bpreshuffle); skips +otherwise. +""" + +import unittest +from unittest import mock + +import torch + +from sglang.srt.layers.quantization import fp8_utils +from sglang.srt.layers.quantization.fp8_utils import ( + materialize_bpreshuffle_fp8_scale, + view_aiter_fused_rms_transposed_fp8_scale, +) +from sglang.srt.utils.common import is_gfx95_supported, is_hip +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=90, suite="stage-b-test-1-gpu-small-amd-mi35x") + +# N, K chosen off the tuned-Triton list so aiter_w8a8_block_fp8_linear takes the +# CK bpreshuffle GEMM path (use_triton == False) -- the path this PR optimizes. +_N, _K = 512, 256 +_BLOCK = [128, 128] + + +@unittest.skipUnless( + is_hip() and is_gfx95_supported() and fp8_utils._use_aiter_bpreshuffle_gfx95, + "dense bpreshuffle scale no-copy is a gfx95 (MI35X) + aiter + ROCm>=7.2 path", +) +class TestDenseBpreshuffleScaleNoCopy(CustomTestCase): + @classmethod + def setUpClass(cls): + # These module globals only exist when aiter/gfx95 imports succeeded. + for attr in ("aiter", "aiter_per1x128_quant"): + if not hasattr(fp8_utils, attr): + raise unittest.SkipTest(f"fp8_utils.{attr} unavailable (no aiter)") + cls.device = "cuda" # torch maps "cuda" onto the ROCm HIP device + + def setUp(self): + torch.manual_seed(0) + + def _rand_input(self, m): + return torch.randn(m, _K, device=self.device, dtype=torch.bfloat16) + + # ---------------------------------------------------------- quant producer + + def test_quant_producer_scale_equivalence(self): + fp8 = fp8_utils.aiter.dtypes.fp8 + for m in (1, 2, 8, 16): + with self.subTest(m=m): + x = self._rand_input(m) + + # Original path: row-major emit + materialize. + q_f, s_f = fp8_utils.aiter_per1x128_quant( + x, quant_dtype=fp8, transpose_scale=False + ) + mat = materialize_bpreshuffle_fp8_scale(s_f) + + # M == 1 stays on the materialize path in production (>= 2 gate). + # There materialize is a no-op: the [1, G] row-major and [G, 1] + # column-major byte orders coincide, so it keeps the natural (G, 1) + # stride (NOT the (1, M) column-major stride taken for M >= 2) and + # shares storage; the scale values must survive intact. + if m < 2: + self.assertEqual(mat.stride(), (s_f.shape[1], 1)) # (G, 1) + self.assertTrue(torch.equal(mat, s_f)) + continue + + self.assertEqual(mat.stride(), (1, m)) # bpreshuffle column-major + + # Optimized path: transposed emit + zero-copy view. + q_t, s_t = fp8_utils.aiter_per1x128_quant( + x, quant_dtype=fp8, transpose_scale=True + ) + nocopy = view_aiter_fused_rms_transposed_fp8_scale(s_t) + + # Quantized bytes are layout-independent. + self.assertTrue( + torch.equal(q_t.view(torch.uint8), q_f.view(torch.uint8)), + "quantized output differs between transpose_scale paths", + ) + # Scale values match the materialized copy, on the (1, M) stride, + # with no allocation (view over the producer's buffer). + self.assertEqual(nocopy.shape, mat.shape) + self.assertTrue(torch.equal(nocopy, mat)) + self.assertEqual(nocopy.stride(), (1, m)) + self.assertEqual(nocopy.data_ptr(), s_t.data_ptr()) + + # ---------------------------------------------------------- end-to-end GEMM + + def _make_weight(self): + finfo = torch.finfo(torch.float8_e4m3fn) + w = (torch.rand(_N, _K, device=self.device, dtype=torch.float32) - 0.5) * 2 + weight = (w * finfo.max).clamp(finfo.min, finfo.max).to(torch.float8_e4m3fn) + weight_scale = ( + torch.rand( + _N // _BLOCK[0], + _K // _BLOCK[1], + device=self.device, + dtype=torch.float32, + ) + * 1e-2 + + 1e-3 + ) + return weight, weight_scale + + def test_dense_linear_paths_bit_exact(self): + weight, weight_scale = self._make_weight() + real_quant = fp8_utils.aiter_per1x128_quant + + def _row_major_quant(inp, **kwargs): + # Force the pre-optimization behavior: emit the scale row-major. + kwargs["transpose_scale"] = False + return real_quant(inp, **kwargs) + + for m in (1, 2, 8, 16): + with self.subTest(m=m): + x = self._rand_input(m) + + # New path, exactly as shipped. Also pin that it really takes the + # CK bpreshuffle GEMM (use_triton == False) -- the path this PR + # optimizes -- so a future tuned-shape-list change can't silently + # route this coverage through Triton and void the equivalence check. + with ( + mock.patch.object( + fp8_utils, + "gemm_a8w8_blockscale_bpreshuffle", + wraps=fp8_utils.gemm_a8w8_blockscale_bpreshuffle, + ) as spy_bpreshuffle, + mock.patch.object( + fp8_utils, + "triton_gemm_a8w8_blockscale", + wraps=fp8_utils.triton_gemm_a8w8_blockscale, + ) as spy_triton, + ): + out_new = fp8_utils.aiter_w8a8_block_fp8_linear( + x, weight, _BLOCK, weight_scale + ) + spy_bpreshuffle.assert_called_once() + spy_triton.assert_not_called() + + # Old path: row-major quant + materialize relayout. Patching both + # the quant flag and the relayout helper reconstructs the original + # `transpose_scale=False` + `materialize_bpreshuffle_fp8_scale` + # branch through the same public function and downstream GEMM. + with ( + mock.patch.object( + fp8_utils, "aiter_per1x128_quant", _row_major_quant + ), + mock.patch.object( + fp8_utils, + "view_aiter_fused_rms_transposed_fp8_scale", + materialize_bpreshuffle_fp8_scale, + ), + ): + out_old = fp8_utils.aiter_w8a8_block_fp8_linear( + x, weight, _BLOCK, weight_scale + ) + + self.assertEqual(out_new.shape, out_old.shape) + self.assertTrue( + torch.equal(out_new, out_old), + f"dense linear output differs between scale-layout paths (M={m})", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py b/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py index b9e12d884..b3ab726cf 100644 --- a/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py +++ b/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py @@ -14,6 +14,19 @@ from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") +def _simulate_transpose_scale_emit(values: torch.Tensor) -> torch.Tensor: + """Model the scale a quant kernel returns when called with + ``transpose_scale=True``: the per-group scale is written directly in + column-major (``[num_groups, tokens]``) byte order, exposed as a ``[M, G]`` + tensor. We reproduce that by laying the column-major bytes into contiguous + storage and reinterpreting it as ``[M, G]`` -- the logical row-major view is + scrambled, but the *storage* holds exactly the bytes the no-copy stride + reinterpret is meant to recover.""" + m, g = values.shape + colmajor_bytes = values.t().contiguous() # [G, M], storage == col-major of values + return colmajor_bytes.view(m, g) # [M, G] over the same (unchanged) storage + + class TestBpreshuffleScaleMaterialization(CustomTestCase): def test_materializes_transposed_physical_storage(self): scale = torch.arange(12, dtype=torch.float32).reshape(3, 4) @@ -102,5 +115,65 @@ class TestBpreshuffleScaleMaterialization(CustomTestCase): self.assertEqual(scale_out.stride(), (1, scale.shape[0])) +class TestBpreshuffleScaleFreshQuantNoCopy(CustomTestCase): + """The dense w8a8 fresh-quant path asks the quant kernel for the scale in + bpreshuffle byte-order (``transpose_scale=True``) and reinterprets its strides + via ``view_aiter_fused_rms_transposed_fp8_scale`` (the shared #31727 helper) + instead of relaying it out with ``materialize_bpreshuffle_fp8_scale`` (a + ``.t().contiguous().t()`` copy). These pin the PR's core claim: the reinterpret + is bit-identical to the copy path for M>=2, and allocates nothing. The real + quant/GEMM equivalence is validated on gfx95 in + ``test_fp8_bpreshuffle_dense_linear_mi35x.py``.""" + + def test_nocopy_matches_materialize(self): + for m, g in ((3, 4), (2, 2), (8, 5), (16, 128)): + with self.subTest(m=m, g=g): + values = torch.arange(m * g, dtype=torch.float32).reshape(m, g) + emitted = _simulate_transpose_scale_emit(values) + + nocopy = view_aiter_fused_rms_transposed_fp8_scale(emitted) + materialized = materialize_bpreshuffle_fp8_scale(values) + + self.assertTrue(torch.equal(nocopy, materialized)) + self.assertEqual(nocopy.shape, values.shape) + self.assertEqual(nocopy.stride(), (1, m)) + self.assertEqual(nocopy.stride(), materialized.stride()) + self.assertTrue(nocopy.t().is_contiguous()) + + def test_nocopy_shares_storage_no_allocation(self): + values = torch.arange(12, dtype=torch.float32).reshape(3, 4) + emitted = _simulate_transpose_scale_emit(values) + + nocopy = view_aiter_fused_rms_transposed_fp8_scale(emitted) + + # The reinterpret is a view over the producer's buffer -- no new storage. + self.assertEqual(nocopy.data_ptr(), emitted.data_ptr()) + # ...unlike the materialize path it replaces. + materialized = materialize_bpreshuffle_fp8_scale(values) + self.assertNotEqual(materialized.data_ptr(), values.data_ptr()) + + def test_m1_uses_materialize_path_values_and_layout(self): + """Production gates the no-copy emit on ``input_2d.shape[0] >= 2`` + (`emit_bpreshuffle_scale`), so a single row (M == 1) keeps the materialize + path. At M == 1 the ``[1, G]`` row-major and ``[G, 1]`` column-major byte + orders coincide, so ``materialize_bpreshuffle_fp8_scale`` is a no-op: the + ``[G, 1]`` transpose is already contiguous for the singleton dim, so + ``.contiguous()`` copies nothing and the result keeps the natural + ``(G, 1)`` stride (NOT the ``(1, M)`` column-major stride it produces for + M >= 2) while sharing the input's storage. Values must survive intact; the + downstream bpreshuffle GEMM consumes the same bytes either way. The actual + M==1 gating through aiter_w8a8_block_fp8_linear is exercised on gfx95 in + test_fp8_bpreshuffle_dense_linear_mi35x.py.""" + scale = torch.arange(4, dtype=torch.float32).reshape(1, 4) # [M=1, G=4] + + materialized = materialize_bpreshuffle_fp8_scale(scale) + + self.assertTrue(torch.equal(materialized, scale)) + self.assertEqual(materialized.shape, (1, 4)) + self.assertEqual(materialized.stride(), (scale.shape[1], 1)) # (G, 1) + self.assertEqual(materialized.data_ptr(), scale.data_ptr()) # no-op share + self.assertTrue(materialized.t().is_contiguous()) + + if __name__ == "__main__": unittest.main()