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

[AMD] DeepSeek-V4 MI355X: eliminate bpreshuffle fp8-scale copies at producer sites (MoE down, MLA o_proj bmm) (#33166)

Browse Source 
Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
This commit is contained in:
karverma-amd
2026-08-20 21:46:23 -07:00
committed by GitHub
co-authored by kk Thomas Wang
parent 6127d1daee
commit bda9952377
5 changed files with 291 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/test_fp8_bpreshuffle_producer_mi35x.py
+156
View File
@@ -0,0 +1,156 @@
"""Producer-level validation of the gfx95 bpreshuffle fp8-scale no-copy path.
The CPU tests in ``test_fp8_bpreshuffle_scale.py`` only prove that the stride
reinterpret recovers a *fabricated* ``transpose_scale=True`` layout. They cannot
catch the real AITER producers ignoring or mis-implementing ``transpose_scale``.
These tests invoke the two producers the optimization actually routes through --
``fused_clamp_act_mul`` (MoE down) and ``fused_flatten_fp8_group_quant`` (MLA
o_proj) -- and prove, on real gfx95 kernels, that the two producer paths are
equivalent:
transpose_scale=True + view_aiter_fused_rms_transposed_fp8_scale (the optimized path)
transpose_scale=False + materialize_bpreshuffle_fp8_scale (the row-major path)
For M(tokens) >= 2 they must agree bit-for-bit on the quantized output and on the
scale *values* after relayout, with the no-copy path landing on the bpreshuffle
``(1, M)`` column-major stride and sharing the producer's storage. M == 1 is the
materialize-only fallback (``emit_transposed_bpreshuffle_scale`` gates the
transposed emit on M >= 2); we pin that its materialized scale is well-formed.
Requires a gfx95 (MI35X) GPU with aiter; skips otherwise.
"""
import unittest
import torch
from sglang.srt.layers.quantization.fp8_utils import (
emit_transposed_bpreshuffle_scale,
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=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
_GROUP_SIZE = 128
@unittest.skipUnless(
is_hip() and is_gfx95_supported(),
"bpreshuffle fp8-scale no-copy is a gfx95 (MI35X) + aiter optimization",
)
class TestBpreshuffleProducerScaleNoCopy(CustomTestCase):
@classmethod
def setUpClass(cls):
try:
from aiter import dtypes # noqa: F401
from aiter.ops.triton.fused_fp8_quant import ( # noqa: F401
fused_flatten_fp8_group_quant,
)
from aiter.ops.triton.fusions.fused_clamp_act_mul import ( # noqa: F401
fused_clamp_act_mul,
)
except Exception as err: # pragma: no cover - env-dependent
raise unittest.SkipTest(f"aiter producers unavailable: {err}")
cls.device = "cuda" # torch maps "cuda" onto the ROCm HIP device
def setUp(self):
torch.manual_seed(0)
# --- producer adapters: run the producer and normalize to (q, scale) ------
def _run_fused_clamp_act_mul(self, m, transpose_scale):
from aiter import dtypes
from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
inter = 4 * _GROUP_SIZE # G = 4 groups
gate_up = torch.randn(m, 2 * inter, device=self.device, dtype=torch.bfloat16)
q, scale = fused_clamp_act_mul(
gate_up,
swiglu_limit=7.0,
activation="silu",
dtype_quant=dtypes.fp8,
transpose_scale=transpose_scale,
)
return q, scale, gate_up
def _run_fused_flatten_fp8_group_quant(self, m, transpose_scale):
from aiter.ops.triton.fused_fp8_quant import fused_flatten_fp8_group_quant
heads, dim = 8, _GROUP_SIZE # heads*dim = 1024 -> G = 8 groups
buf = torch.randn(m, heads, dim, device=self.device, dtype=torch.bfloat16)
out = fused_flatten_fp8_group_quant(
buf,
group_size=_GROUP_SIZE,
dtype_quant=torch.float8_e4m3fn,
transpose_scale=transpose_scale,
)
return out[0], out[1], buf
def _assert_producer_paths_equivalent(self, run_producer, name):
for m in (1, 2, 8, 16):
with self.subTest(producer=name, m=m):
# Row-major path (transpose_scale=False) + materialize. This is
# the path M == 1 takes in production, so it must always be valid.
torch.manual_seed(m)
q_f, s_f, _ = run_producer(m, transpose_scale=False)
mat = materialize_bpreshuffle_fp8_scale(s_f)
self.assertEqual(mat.dim(), 2)
self.assertEqual(mat.shape[0], m)
if not emit_transposed_bpreshuffle_scale(m, on_bpreshuffle_gfx95=True):
# M == 1: the transposed emit is skipped by design. At M == 1 the
# [1, G] row-major and [G, 1] column-major byte orders coincide,
# so materialize is a no-op that keeps the natural (G, 1) stride
# (NOT the (1, M) column-major stride taken for M >= 2) while
# sharing storage; the scale values must survive intact.
self.assertEqual(m, 1)
self.assertEqual(mat.stride(), (s_f.shape[1], 1)) # (G, 1)
self.assertTrue(torch.equal(mat, s_f))
continue
# M >= 2: materialize lands on the bpreshuffle (1, M) column-major
# stride.
self.assertEqual(mat.stride(), (1, m))
# Optimized path: same input, transpose_scale=True + no-copy view.
torch.manual_seed(m)
q_t, s_t, _ = run_producer(m, transpose_scale=True)
nocopy = view_aiter_fused_rms_transposed_fp8_scale(s_t)
# Quantized output is layout-independent: transpose_scale only
# changes the *scale* storage, never the quantized bytes.
self.assertEqual(
q_t.dtype, q_f.dtype, "quant dtype differs between paths"
)
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 row-major + materialize path exactly...
self.assertEqual(nocopy.shape, mat.shape)
self.assertTrue(
torch.equal(nocopy, mat),
"no-copy scale values differ from materialized",
)
# ...on the bpreshuffle (1, M) stride, with no allocation.
self.assertEqual(nocopy.stride(), (1, m))
self.assertEqual(nocopy.data_ptr(), s_t.data_ptr())
def test_fused_clamp_act_mul_producer_paths_equivalent(self):
self._assert_producer_paths_equivalent(
self._run_fused_clamp_act_mul, "fused_clamp_act_mul"
)
def test_fused_flatten_fp8_group_quant_producer_paths_equivalent(self):
self._assert_producer_paths_equivalent(
self._run_fused_flatten_fp8_group_quant, "fused_flatten_fp8_group_quant"
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/test_fp8_bpreshuffle_scale.py
+62
View File
@@ -4,9 +4,11 @@ from unittest.mock import patch
import torch
from sglang.srt.layers.quantization.fp8_utils import (
emit_transposed_bpreshuffle_scale,
materialize_bpreshuffle_fp8_scale,
materialize_bpreshuffle_fp8_scale_tuple,
view_aiter_fused_rms_transposed_fp8_scale,
view_aiter_fused_rms_transposed_fp8_scale_tuple,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -175,5 +177,65 @@ class TestBpreshuffleScaleFreshQuantNoCopy(CustomTestCase):
self.assertTrue(materialized.t().is_contiguous())
class TestBpreshuffleScaleProducerNoCopy(CustomTestCase):
"""Producer-site (MoE down, MLA o_proj bmm) coverage for the shared no-copy
reinterpret that isn't exercised by the dense fresh-quant class above: the
guard that leaves non-2D scales untouched, and the tuple wrapper the producers
emit through (``view_aiter_fused_rms_transposed_fp8_scale_tuple``), which must
reinterpret only the scale slot and pass the rest through by identity."""
def test_nocopy_passthrough_for_non_2d_scale(self):
for scale in (
torch.arange(5, dtype=torch.float32), # 1-D (per-tensor scale)
torch.arange(24, dtype=torch.float32).reshape(2, 3, 4), # 3-D
):
with self.subTest(dim=scale.dim()):
self.assertIs(view_aiter_fused_rms_transposed_fp8_scale(scale), scale)
def test_tuple_helper_reinterprets_only_the_scale_slot(self):
q_input = torch.ones((3, 8), dtype=torch.float8_e4m3fn)
values = torch.arange(12, dtype=torch.float32).reshape(3, 4)
emitted = _simulate_transpose_scale_emit(values)
bf16_side = torch.ones((3, 8), dtype=torch.bfloat16)
q_out, scale_out, bf16_out = view_aiter_fused_rms_transposed_fp8_scale_tuple(
(q_input, emitted, bf16_side)
)
self.assertIs(q_out, q_input)
self.assertIs(bf16_out, bf16_side)
self.assertTrue(
torch.equal(scale_out, materialize_bpreshuffle_fp8_scale(values))
)
self.assertEqual(scale_out.stride(), (1, values.shape[0]))
self.assertEqual(scale_out.data_ptr(), emitted.data_ptr())
class TestEmitTransposedBpreshuffleScaleGate(CustomTestCase):
"""Pins the producer emit-gate shared by the MoE-down and MLA o_proj sites:
the transposed zero-copy path is taken only on gfx95 bpreshuffle and only for
M(tokens) >= 2; M == 1 must fall back to the materialize path. Guards the
``>= 2`` boundary against being widened to `M >= 1` (which would send a
degenerate single-token scale down the stride-swap path)."""
def test_gate_false_off_gfx95_regardless_of_m(self):
for m in (1, 2, 8):
with self.subTest(m=m):
self.assertFalse(
emit_transposed_bpreshuffle_scale(m, on_bpreshuffle_gfx95=False)
)
def test_gate_requires_m_ge_2_on_gfx95(self):
# M == 1 -> materialize fallback; M >= 2 -> transposed zero-copy path.
self.assertFalse(
emit_transposed_bpreshuffle_scale(1, on_bpreshuffle_gfx95=True)
)
for m in (2, 3, 16):
with self.subTest(m=m):
self.assertTrue(
emit_transposed_bpreshuffle_scale(m, on_bpreshuffle_gfx95=True)
)
if __name__ == "__main__":
unittest.main()