diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 28ee0c1f0..fe112743a 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -349,6 +349,7 @@ def _fused_rope_cat_and_cache( q_nope_out.dtype if attn.kv_cache_dtype == "fp8_e4m3" and attn.current_attention_backend == "aiter" + and q_nope_out.shape[-2] == 12 else kv_cache_dtype ) kv_pool = get_token_to_kv_pool()