From b9cf1563de979458e5c831ccfe8f5bfd47a55d01 Mon Sep 17 00:00:00 2001 From: "Chen, Zhentao" Date: Wed, 25 Feb 2026 09:01:14 +0800 Subject: [PATCH] [ROCm] Optimize Deepseek R1 on MI300X (#18242) Co-authored-by: Chen, Todd --- python/sglang/srt/layers/attention/aiter_backend.py | 1 + python/sglang/srt/layers/quantization/fp8_utils.py | 1 + python/sglang/srt/models/deepseek_v2.py | 7 +++++-- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index d851040cf..97d6ee573 100644 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -1560,6 +1560,7 @@ class AiterAttnBackend(AttentionBackend): o = torch.empty_like(q, dtype=self.input_dtype) if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( layer, forward_batch.out_cache_loc, k, v ) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 3b9616e27..cc889b580 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -61,6 +61,7 @@ if _use_aiter: aiter_per1x128_quant = get_hip_quant(aiter.QuantType.per_1x128) + if _is_cuda: from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 27afcf74e..f11d665c2 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -180,6 +180,9 @@ if _use_aiter_gfx95: get_dsv3_gemm_output_zero_allocator_size, ) +if _use_aiter: + from sglang.srt.layers.rocm_linear_utils import fused_qk_rope_cat_and_cache_mla + if _is_cuda: from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm, dsv3_router_gemm elif _is_cpu and _is_cpu_amx_available: @@ -1724,7 +1727,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): if ( self.rotary_emb is not None and (not self._fuse_rope_for_trtllm_mla(forward_batch)) - and (not _use_aiter or not _is_gfx95_supported or self.use_nsa) + and (not _use_aiter or not _is_hip or self.use_nsa) ): q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) @@ -1780,7 +1783,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): **(dict(topk_indices=topk_indices) if topk_indices is not None else {}), ) else: - if _use_aiter_gfx95: + if _use_aiter: cos = self.rotary_emb.cos_cache sin = self.rotary_emb.sin_cache