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