[AMD] [GLM5] Enable dense-MHA short-context prefill fallback on gfx950 (#30808)

Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>

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)
+                    device_sm == 90
-                )  # SM90/SM100 only
+                    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]

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)