From bda995237751dd065041ae6c9ddb979d1b6528bd Mon Sep 17 00:00:00 2001 From: karverma-amd Date: Thu, 20 Aug 2026 23:46:23 -0500 Subject: [PATCH] [AMD] DeepSeek-V4 MI355X: eliminate bpreshuffle fp8-scale copies at producer sites (MoE down, MLA o_proj bmm) (#33166) Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com> Co-authored-by: Thomas Wang --- .../srt/layers/quantization/fp8_utils.py | 39 ++++- .../forward_mla_rocm.py | 32 +++- python/sglang/srt/models/deepseek_v2.py | 11 +- .../test_fp8_bpreshuffle_producer_mi35x.py | 156 ++++++++++++++++++ .../unit/layers/test_fp8_bpreshuffle_scale.py | 62 +++++++ 5 files changed, 291 insertions(+), 9 deletions(-) create mode 100644 test/registered/unit/layers/test_fp8_bpreshuffle_producer_mi35x.py diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index a79b441e0..93e1811bc 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -115,10 +115,21 @@ def materialize_bpreshuffle_fp8_scale(scale: torch.Tensor) -> torch.Tensor: def view_aiter_fused_rms_transposed_fp8_scale(scale: torch.Tensor) -> torch.Tensor: - """Expose AITER fused-RMS ``transpose_scale=True`` storage logically. + """Zero-copy view of a ``transpose_scale=True`` fp8 group scale. - The fused-RMS op returns transposed physical bytes through a row-major-looking - view. Restore logical ``[M, G]`` indexing without copying those bytes. + Producer-neutral counterpart of ``materialize_bpreshuffle_fp8_scale``. When an + AITER quant/fused-RMS kernel is asked for ``transpose_scale=True`` it writes the + per-token group scale directly in physical ``[num_groups, tokens]`` byte order + behind a row-major-looking ``[tokens, num_groups]`` tensor. Swapping the strides + restores logical ``[M, G]`` indexing over those same bytes -- i.e. the + column-major layout the gfx95 bpreshuffle GEMM consumes -- with no copy. Callers + that instead take the row-major (``transpose_scale=False``) path relayout via + ``materialize_bpreshuffle_fp8_scale``; this is the bit-identical no-copy path. + + Only valid for M(tokens) >= 2. At M == 1 the ``[1, G]`` and ``[G, 1]`` byte + orders coincide, so producers keep ``transpose_scale=False`` and materialize; + the stride swap here would be a no-op on shape but is never taken at M == 1. + Non-2-D scales (e.g. per-tensor) pass through unchanged. """ if scale.dim() != 2: return scale @@ -136,6 +147,28 @@ def materialize_bpreshuffle_fp8_scale_tuple( ) +def view_aiter_fused_rms_transposed_fp8_scale_tuple( + value: Tuple[torch.Tensor, ...], +) -> Tuple[torch.Tensor, ...]: + """Zero-copy scale reinterpret for FP8 ``(q_input, x_scale, ...)`` tuples.""" + return (value[0], view_aiter_fused_rms_transposed_fp8_scale(value[1]), *value[2:]) + + +def emit_transposed_bpreshuffle_scale(m: int, *, on_bpreshuffle_gfx95: bool) -> bool: + """Whether a producer should emit its fp8 scale already transposed. + + Producer sites choose between two equivalent gfx95 bpreshuffle scale layouts: + ``transpose_scale=True`` + zero-copy ``view_aiter_fused_rms_transposed_fp8_scale`` (this + predicate True), or row-major ``transpose_scale=False`` + + ``materialize_bpreshuffle_fp8_scale`` (this predicate False). The transposed + zero-copy path is only taken on gfx95 bpreshuffle and only for M(tokens) >= 2: + at M == 1 the ``[1, G]`` and ``[G, 1]`` byte orders coincide, so the transposed + emit buys nothing and the materialize path is used. Centralizes the gate shared + by the MoE-down and MLA o_proj producer sites. + """ + return on_bpreshuffle_gfx95 and m >= 2 + + def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool: if _FORCE_CK_W8A8: return False diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 928eb6352..378f3cbf6 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -31,7 +31,9 @@ from sglang.srt.layers.dcp import ( ) from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run from sglang.srt.layers.quantization.fp8_utils import ( + emit_transposed_bpreshuffle_scale, materialize_bpreshuffle_fp8_scale_tuple, + view_aiter_fused_rms_transposed_fp8_scale_tuple, ) from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp from sglang.srt.lora.deepseek_mla_correction import ( @@ -233,13 +235,24 @@ def rocm_absorb_v_bmm( if attn.o_proj.weight.dtype == torch.uint8: attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf) elif _is_block_scale_fp8(attn.o_proj): + # No-copy fp8 scale: emit the bpreshuffle scale already transposed and + # reinterpret it with a stride swap, instead of relaying out a copy. + # Falls back to the materialize (copy) path at M == 1 / non-gfx95. + _emit_bpre = emit_transposed_bpreshuffle_scale( + _bmm_buf.shape[0], + on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95, + ) attn_bmm_output = fused_flatten_fp8_group_quant( _bmm_buf, group_size=128, dtype_quant=torch.float8_e4m3fn, - transpose_scale=False, + transpose_scale=_emit_bpre, ) - if _use_aiter_bpreshuffle_gfx95: + if _emit_bpre: + attn_bmm_output = view_aiter_fused_rms_transposed_fp8_scale_tuple( + attn_bmm_output + ) + elif _use_aiter_bpreshuffle_gfx95: attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple( attn_bmm_output ) @@ -250,13 +263,24 @@ def rocm_absorb_v_bmm( attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output) elif _is_block_scale_fp8(attn.o_proj): attn_bmm_output = attn_bmm_output.transpose(0, 1) + # No-copy fp8 scale: emit the bpreshuffle scale already transposed and + # reinterpret it with a stride swap, instead of relaying out a copy. + # Falls back to the materialize (copy) path at M == 1 / non-gfx95. + _emit_bpre = emit_transposed_bpreshuffle_scale( + attn_bmm_output.shape[0], + on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95, + ) attn_bmm_output = fused_flatten_fp8_group_quant( attn_bmm_output, group_size=128, dtype_quant=torch.float8_e4m3fn, - transpose_scale=False, + transpose_scale=_emit_bpre, ) - if _use_aiter_bpreshuffle_gfx95: + if _emit_bpre: + attn_bmm_output = view_aiter_fused_rms_transposed_fp8_scale_tuple( + attn_bmm_output + ) + elif _use_aiter_bpreshuffle_gfx95: attn_bmm_output = materialize_bpreshuffle_fp8_scale_tuple(attn_bmm_output) else: attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 31047701e..570fa2bba 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -123,7 +123,9 @@ from sglang.srt.layers.moe.utils import ( from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8 import Fp8Config 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.layers.quantization.mxfp4_flashinfer_trtllm_moe import ( maybe_fuse_routed_scale_and_shared_add, @@ -419,14 +421,19 @@ class DeepseekV2MLP(nn.Module): if self._fused_clamp_use_fp8: from aiter import dtypes + _emit_bpre = emit_transposed_bpreshuffle_scale( + gate_up.shape[0], on_bpreshuffle_gfx95=_use_aiter_bpreshuffle_gfx95 + ) x_fp8, x_scale = fused_clamp_act_mul( gate_up, swiglu_limit=self.swiglu_limit, activation="silu", dtype_quant=dtypes.fp8, - transpose_scale=False, + transpose_scale=_emit_bpre, ) - if _use_aiter_bpreshuffle_gfx95: + if _emit_bpre: + x_scale = view_aiter_fused_rms_transposed_fp8_scale(x_scale) + elif _use_aiter_bpreshuffle_gfx95: x_scale = materialize_bpreshuffle_fp8_scale(x_scale) x = (x_fp8, x_scale) else: diff --git a/test/registered/unit/layers/test_fp8_bpreshuffle_producer_mi35x.py b/test/registered/unit/layers/test_fp8_bpreshuffle_producer_mi35x.py new file mode 100644 index 000000000..5a730f349 --- /dev/null +++ b/test/registered/unit/layers/test_fp8_bpreshuffle_producer_mi35x.py @@ -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() diff --git a/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py b/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py index b3ab726cf..b13ee3384 100644 --- a/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py +++ b/test/registered/unit/layers/test_fp8_bpreshuffle_scale.py @@ -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()