From ba1d980b353b67b0b724512c6792416f9e4903a2 Mon Sep 17 00:00:00 2001 From: zijiec Date: Fri, 14 Aug 2026 02:42:25 -0400 Subject: [PATCH] [AMD] Accelerate AITER unified-attention decode with scaled FP8 Q (#31856) Co-authored-by: jacky.cheng Co-authored-by: YC Yen-Ching Tseng --- .../srt/layers/attention/aiter_backend.py | 13 +- .../test_aiter_fp8_q_unified_attention.py | 304 ++++++++++++++++++ 2 files changed, 315 insertions(+), 2 deletions(-) create mode 100644 test/registered/attention/test_aiter_fp8_q_unified_attention.py diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index dfcba404f..3943a5812 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -65,7 +65,10 @@ from sglang.kernels.ops.attention.utils import ( launch_reshape_and_cache_flash, pad_sequence_with_mask, ) -from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype +from sglang.kernels.ops.quantization.fp8_kernel import ( + fp8_dtype, + scaled_fp8_quant, +) from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.aiter_utils import ( forward_decode_vectorized_5d, @@ -2605,6 +2608,12 @@ class AiterAttnBackend(AttentionBackend): page_table = self.forward_metadata.swa_page_table max_kv_len = page_table.shape[1] * self.page_size + q_descale = None + if self.kv_cache_dtype == fp8_dtype: + q_descale = ( + layer.k_scale if layer.k_scale is not None else self.k_scale + ) + q, _ = scaled_fp8_quant(q, q_descale) unified_attention( q=q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), @@ -2624,7 +2633,7 @@ class AiterAttnBackend(AttentionBackend): window_size=window_size, block_table=page_table, softcap=0, - q_descale=None, + q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, sinks=sinks, diff --git a/test/registered/attention/test_aiter_fp8_q_unified_attention.py b/test/registered/attention/test_aiter_fp8_q_unified_attention.py new file mode 100644 index 000000000..bc1c608f8 --- /dev/null +++ b/test/registered/attention/test_aiter_fp8_q_unified_attention.py @@ -0,0 +1,304 @@ +"""Coverage for the AITER FP8-Q unified-attention decode path.""" + +import math +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=120, suite="stage-b-test-1-gpu-small-amd-mi35x") + +_RUNNABLE = is_hip() +if _RUNNABLE: + from aiter.ops.triton.attention.unified_attention import unified_attention + + import sglang.srt.layers.attention.aiter_backend as aiter_backend + from sglang.kernels.ops.quantization.fp8_kernel import ( + fp8_dtype, + scaled_fp8_quant, + ) + from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend + + +class _FakeKVPool: + def __init__(self, k_cache, v_cache): + self.k_cache = k_cache + self.v_cache = v_cache + + def get_kv_buffer(self, _layer_id): + return self.k_cache, self.v_cache + + def get_key_buffer(self, _layer_id): + return self.k_cache + + +@unittest.skipUnless(_RUNNABLE, "requires HIP with AITER unified attention") +class TestAiterFP8QUnifiedAttention(CustomTestCase): + def _make_backend_case(self, branch, kv_cache_dtype=None): + if kv_cache_dtype is None: + kv_cache_dtype = fp8_dtype + batch, num_q_heads, num_kv_heads, head_dim = 2, 2, 1, 8 + device = "cuda" + scale = torch.tensor([0.02], dtype=torch.float32, device=device) + k_cache = torch.zeros( + batch, + 1, + num_kv_heads, + head_dim, + dtype=kv_cache_dtype, + device=device, + ) + v_cache = torch.zeros_like(k_cache) + + backend = object.__new__(AiterAttnBackend) + backend.use_mla = branch == "mla" + backend.kv_cache_is_vectorized_5d = branch == "vectorized" + backend.use_triton_unified_attention = branch != "legacy" + backend.kv_cache_dtype = kv_cache_dtype + backend.input_dtype = torch.bfloat16 + backend.page_size = 1 + backend.scale = head_dim**-0.5 + backend.logits_soft_cap = 0.0 + backend.k_scale = scale + backend.v_scale = scale + backend.workspace_buffer = torch.empty(1, device=device) + backend.max_num_partitions = 1 + backend.kv_last_page_len = torch.ones(batch, dtype=torch.int32, device=device) + backend.token_to_kv_pool = _FakeKVPool(k_cache, v_cache) + backend.forward_metadata = SimpleNamespace( + kv_indices=torch.arange(batch, dtype=torch.int32, device=device).view( + batch, 1 + ), + swa_page_table=None, + qo_indptr=torch.arange(batch + 1, dtype=torch.int32, device=device), + kv_indptr=torch.arange(batch + 1, dtype=torch.int32, device=device), + kv_last_page_len=backend.kv_last_page_len, + max_q_len=1, + work_metadata=None, + work_indptr=None, + work_info_set=None, + reduce_indptr=None, + reduce_final_map=None, + reduce_partial_map=None, + num_kv_splits=None, + ) + backend._mla_decode_fwd_with_head_pad = mock.Mock( + return_value=torch.empty( + batch, + num_q_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + ) + backend._get_aiter_paged_ragged_kv_cache_dtype = mock.Mock( + return_value="fp8_e4m3" + ) + + layer = SimpleNamespace( + layer_id=0, + tp_q_head_num=num_q_heads, + tp_k_head_num=num_kv_heads, + tp_v_head_num=num_kv_heads, + qk_head_dim=head_dim, + v_head_dim=head_dim, + k_scale=scale, + v_scale=scale, + sliding_window_size=-1, + scaling=head_dim**-0.5, + logit_cap=0.0, + ) + forward_batch = SimpleNamespace( + batch_size=batch, + seq_lens=torch.ones(batch, dtype=torch.int32, device=device), + ) + q = torch.randn( + batch, + num_q_heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + return backend, layer, forward_batch, q + + def test_q_quantization_is_isolated_to_unified_attention(self): + for branch in ("mla", "vectorized", "unified", "legacy"): + with self.subTest(branch=branch): + backend, layer, forward_batch, q = self._make_backend_case(branch) + original_q = q.reshape(q.shape[0], -1).clone() + sentinel_q = torch.full( + original_q.shape, 7, dtype=fp8_dtype, device=original_q.device + ) + + with ( + mock.patch.object( + aiter_backend, + "scaled_fp8_quant", + return_value=(sentinel_q, layer.k_scale), + ) as quant, + mock.patch.object(aiter_backend, "unified_attention") as unified, + mock.patch.object( + aiter_backend, "forward_decode_vectorized_5d" + ) as vectorized, + mock.patch.object( + aiter_backend, "paged_attention_ragged" + ) as legacy, + ): + output = backend.forward_decode( + q, None, None, layer, forward_batch, save_kv_cache=False + ) + + self.assertEqual(output.numel(), q.numel()) + if branch == "unified": + quant.assert_called_once() + self.assertIs(quant.call_args.args[1], layer.k_scale) + self.assertIs(unified.call_args.kwargs["q_descale"], layer.k_scale) + torch.testing.assert_close( + unified.call_args.kwargs["q"].reshape(q.shape[0], -1), + sentinel_q, + ) + else: + quant.assert_not_called() + if branch == "mla": + observed_q = ( + backend._mla_decode_fwd_with_head_pad.call_args.args[0] + ) + elif branch == "vectorized": + observed_q = vectorized.call_args.args[1] + else: + observed_q = legacy.call_args.args[2] + torch.testing.assert_close( + observed_q.reshape(q.shape[0], -1), original_q + ) + + def test_bf16_kv_keeps_bf16_q(self): + backend, layer, forward_batch, q = self._make_backend_case( + "unified", torch.bfloat16 + ) + original_q = q.reshape(q.shape[0], -1).clone() + + with ( + mock.patch.object(aiter_backend, "scaled_fp8_quant") as quant, + mock.patch.object(aiter_backend, "unified_attention") as unified, + ): + output = backend.forward_decode( + q, None, None, layer, forward_batch, save_kv_cache=False + ) + + self.assertEqual(output.numel(), q.numel()) + quant.assert_not_called() + self.assertIsNone(unified.call_args.kwargs["q_descale"]) + observed_q = unified.call_args.kwargs["q"] + self.assertEqual(observed_q.dtype, torch.bfloat16) + torch.testing.assert_close( + observed_q.reshape(q.shape[0], -1), + original_q, + ) + + def test_fp8_q_kv_matches_bf16_reference_at_decode_shape(self): + # This is the per-TP-rank production shape used by the full trace. + batch, num_q_heads, num_kv_heads = 4, 16, 1 + seq_len, head_dim, page_size = 8192, 256, 16 + device = "cuda" + torch.manual_seed(0) + + base = torch.randn(batch, head_dim, device=device, dtype=torch.float32) + q = ( + base[:, None, :].expand(-1, num_q_heads, -1) + + 0.01 + * torch.randn( + batch, num_q_heads, head_dim, device=device, dtype=torch.float32 + ) + ).to(torch.bfloat16) + k = 0.1 * torch.randn( + batch, + seq_len, + num_kv_heads, + head_dim, + device=device, + dtype=torch.float32, + ) + k[:, 0, 0, :] = 2 * base + k = k.to(torch.bfloat16) + v = ( + 0.75 + + 0.25 + * torch.randn( + batch, + seq_len, + num_kv_heads, + head_dim, + device=device, + dtype=torch.float32, + ) + ).to(torch.bfloat16) + + fp8_max = torch.finfo(fp8_dtype).max + k_scale = (k.abs().float().amax() / fp8_max).clamp(min=1e-9).view(1) + v_scale = (v.abs().float().amax() / fp8_max).clamp(min=1e-9).view(1) + q_fp8, _ = scaled_fp8_quant(q.reshape(batch, -1), k_scale) + k_fp8, _ = scaled_fp8_quant(k.reshape(-1, head_dim), k_scale) + v_fp8, _ = scaled_fp8_quant(v.reshape(-1, head_dim), v_scale) + + q_fp8 = q_fp8.view(batch, num_q_heads, head_dim) + k_fp8 = k_fp8.view(-1, page_size, num_kv_heads, head_dim) + v_fp8 = v_fp8.view(-1, page_size, num_kv_heads, head_dim) + pages_per_seq = seq_len // page_size + block_table = torch.arange( + batch * pages_per_seq, dtype=torch.int32, device=device + ).view(batch, pages_per_seq) + output = torch.empty_like(q, dtype=torch.bfloat16) + + unified_attention( + q=q_fp8, + k=k_fp8, + v=v_fp8, + out=output, + cu_seqlens_q=torch.arange(batch + 1, dtype=torch.int32, device=device), + seqused_k=torch.full((batch,), seq_len, dtype=torch.int32, device=device), + max_seqlen_q=1, + max_seqlen_k=seq_len, + softmax_scale=1 / math.sqrt(head_dim), + causal=True, + window_size=(-1, -1), + block_table=block_table, + softcap=0, + q_descale=k_scale, + k_descale=k_scale, + v_descale=v_scale, + sinks=None, + ) + + scores = torch.einsum( + "bhd,btd->bht", q.float(), k[:, :, 0, :].float() + ) / math.sqrt(head_dim) + expected = torch.einsum( + "bht,btd->bhd", + torch.softmax(scores, dim=-1), + v[:, :, 0, :].float(), + ) + actual = output.float() + + self.assertTrue(bool(torch.isfinite(actual).all())) + self.assertGreater(expected.abs().mean().item(), 0.25) + mismatch = (actual - expected).abs() > 0.15 + 0.15 * expected.abs() + mismatch_fraction = mismatch.float().mean().item() + self.assertLess( + mismatch_fraction, + 0.005, + f"FP8 mismatch fraction {mismatch_fraction:.4%} exceeds 0.5%; " + f"max abs diff={(actual - expected).abs().max().item():.6f}", + ) + cosine = torch.nn.functional.cosine_similarity( + actual.flatten(), expected.flatten(), dim=0 + ).item() + self.assertGreater(cosine, 0.99) + + +if __name__ == "__main__": + unittest.main()