diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index ecde3b946..46a324e44 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -2880,7 +2880,9 @@ class DeepseekSparseAttnBackend( f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" ) - # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues + # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues. + # gfx950 reports device capability sm_(9,5), so it never enters this SM100+ + # branch and falls through to the aiter flash_attn_varlen_func path below. if self.device_sm_major >= 10: import flashinfer @@ -3323,11 +3325,13 @@ class DeepseekSparseAttnBackend( sum_seq_lens = sum(forward_batch.seq_lens_cpu) device_sm = get_device_sm() - # Requirements: H200/B200, short sequences, supported dtype, fits in chunk + # Requirements: H200/B200/MI355X, short sequences, supported dtype, fits in chunk self.use_mha = ( ( - device_sm == 90 or (device_sm >= 100 and device_sm < 110) - ) # SM90/SM100 only + device_sm == 90 + or (device_sm >= 100 and device_sm < 110) + or _IS_GFX95 + ) # SM90/SM100 (NVIDIA) or gfx95x (MI355X) and max_kv_len <= envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() # Short enough for MHA and self.token_to_kv_pool.dtype in [torch.bfloat16, torch.float8_e4m3fn] diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 69d90b600..4eeb8bf49 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -23,6 +23,7 @@ from sglang.srt.models.deepseek_common.utils import ( _is_cuda, _is_musa, _is_npu, + _use_aiter_gfx95, ) from sglang.srt.runtime_context import ( get_exec, @@ -483,6 +484,21 @@ class DeepseekMHAForwardMixin: kv_indices is not None ), "page_table_1_flattened should have been generated for FP8 MHA path" + if _use_aiter_gfx95: + # ROCm (gfx950) stores the FP8 MLA KV in the raw + # (kv_lora_rank + qk_rope_head_dim) layout, not the scaled 656-byte + # layout that dequantize_k_cache_paged expects (it asserts dim==656). + # Dequantize the raw layout via the pool's HIP-aware path instead — + # the same routine _get_mla_kv_buffer uses for the BF16 MHA path. + # Without this, a chunked-prefill split (extend_prefix_lens != 0) that + # reads cached prefix KV crashes with "576 != 656". + kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices) + kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer( + self.attn_mha, kv_indices, torch.bfloat16 + ) + kv_a = kv_a.squeeze(1).contiguous() + return kv_a, k_pe + kv_cache_fp8 = get_token_to_kv_pool().get_key_buffer(self.attn_mha.layer_id) kv_latent_bf16 = dequantize_k_cache_paged(kv_cache_fp8, kv_indices)