[AMD] Fix aiter page-size handling, DeepSeek MLA tuple inputs, and HiCache/FA3 decode-backend override (#16531)

python/sglang/srt/layers/attention/aiter_backend.py

@@ -279,7 +279,7 @@ class AiterAttnBackend(AttentionBackend):
     ):
 
         nhead_kv = 1
-        page_size = 1
+        page_size = self.page_size
         dtype = self.kv_cache_dtype
 
         meta = get_mla_metadata_v1(
@@ -1654,7 +1654,6 @@ class AiterMultiStepDraftBackend:
         # Cached variables for generate_draft_decode_kv_indices
         self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
         self.page_size = model_runner.server_args.page_size
-        assert self.page_size == 1, "Page size must be 1"
 
     def common_template(
         self, forward_batch: ForwardBatch, kv_indices_buffer: torch.Tensor, call_fn: int

python/sglang/srt/models/deepseek_v2.py

@@ -2036,8 +2036,15 @@ class DeepseekV2AttentionMLA(nn.Module):
         enable_rope_fusion = (
             os.getenv("SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION", "1") == "1"
         )
-        q_len = hidden_states.shape[0]
-        q_input = hidden_states.new_empty(
+        # NOTE: hidden_states can be a tuple for some quantization paths.
+        # For shape/device/dtype, use the first tensor; still pass the original
+        # hidden_states through linear ops which may accept tuple inputs.
+        hidden_states_tensor = (
+            hidden_states[0] if isinstance(hidden_states, tuple) else hidden_states
+        )
+
+        q_len = hidden_states_tensor.shape[0]
+        q_input = hidden_states_tensor.new_empty(
             q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim
         )
         if self.q_lora_rank is not None:

python/sglang/srt/server_args.py

@@ -1993,21 +1993,31 @@ class ServerArgs:
             or self.disaggregation_decode_enable_offload_kvcache
         ) and self.hicache_io_backend == "kernel":
             # fix for the compatibility issue with FlashAttention3 decoding and HiCache kernel backend
-            if self.decode_attention_backend is None:
-                if not self.use_mla_backend():
-                    self.decode_attention_backend = (
-                        "flashinfer" if is_flashinfer_available() else "triton"
-                    )
+            # Only override when the *effective* decode backend would be FA3.
+            # Otherwise, respect the user's chosen attention backend (e.g., aiter on ROCm).
+            effective_decode_backend = (
+                self.decode_attention_backend
+                if self.decode_attention_backend is not None
+                else self.attention_backend
+            )
+            if effective_decode_backend == "fa3":
+                if self.decode_attention_backend is None:
+                    # If decode backend wasn't explicitly set, pick a safe default that works with HiCache kernel IO.
+                    if not self.use_mla_backend():
+                        self.decode_attention_backend = (
+                            "flashinfer" if is_flashinfer_available() else "triton"
+                        )
+                    else:
+                        self.decode_attention_backend = (
+                            "flashinfer" if is_sm100_supported() else "triton"
+                        )
                 else:
-                    self.decode_attention_backend = (
-                        "flashinfer" if is_sm100_supported() else "triton"
+                    # If user explicitly requested FA3 decode, fall back to direct IO.
+                    self.hicache_io_backend = "direct"
+                    logger.warning(
+                        "FlashAttention3 decode backend is not compatible with hierarchical cache. "
+                        "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes."
                     )
-            elif self.decode_attention_backend == "fa3":
-                self.hicache_io_backend = "direct"
-                logger.warning(
-                    "FlashAttention3 decode backend is not compatible with hierarchical cache. "
-                    "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes."
-                )
 
     def _handle_speculative_decoding(self):
         if (