diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index a2c4a2e19..e71c051a4 100644 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 81a2058f5..96bb3812d 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 35ed102ed..a76486a3b 100644 --- a/python/sglang/srt/server_args.py +++ b/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 (