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

[AMD] DeepSeek-V4 MI355X: eliminate bpreshuffle fp8-scale relayout copy in dense w8a8 linear (#33165)

Browse Source
This commit is contained in:
karverma-amd
2026-08-19 03:02:40 -07:00
committed by GitHub
parent f22442d3a4
commit ce1830c59b
3 changed files with 273 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8_utils.py
+12 -3
View File
@@ -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:
test/registered/unit/layers/test_fp8_bpreshuffle_dense_linear_mi35x.py
+188
View File
@@ -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()
test/registered/unit/layers/test_fp8_bpreshuffle_scale.py
+73
View File
@@ -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()