Repository: vllm-project/vllm
File: vllm/v1/attention/backends/rocm_aiter_fa.py (Python)

Code before:
                        num_decodes : num_decodes + num_extends
                    ],
                    slot_mapping=attn_metadata.slot_mapping[
                        num_decodes : num_decodes + num_extends
                    ],
                    k_scale=k_scale,
                    v_scale=v_scale,
                )

            # calculate for decodes
            if num_decodes > 0:
                assert attn_metadata.decode_metadata is not None
                if self.sliding_window[0] != -1:
                    assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                        "Sliding window with shuffle layout is not supported yet."
                    )

                # The ll4mi kernel in paged_attention_v1 requires
                # HEAD_SIZE >= 16 * NWARPS (= 64 on ROCm with NWARPS=4).
                # For smaller head sizes or sliding window attention,
                # fall back to the unified_attention triton kernel which
                # handles both correctly.
                _MIN_HEAD_SIZE_FOR_LL4MI = 64
                use_unified_attention = (
                    self.sliding_window[0] != -1
                    or self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI
                )

                if use_unified_attention:
                    assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                        "unified_attention fallback with shuffle layout "
                        "is not supported yet."
                    )
                    from aiter.ops.triton.unified_attention import (
                        unified_attention,
                    )

                    decode_cu_seqlens_q = attn_metadata.query_start_loc[
                        : num_decodes + 1
                    ]
                    descale_shape = (
                        num_decodes,
                        key_cache.shape[2],
                    )
                    unified_attention(
                        q=query[:num_decode_tokens],
                        k=key_cache,
                        v=value_cache,
                        out=output[:num_decode_tokens],
                        cu_seqlens_q=decode_cu_seqlens_q,
                        max_seqlen_q=1,

Review comment (bug):
If you are trying to revert the changes, why there is a need for the second conditions?

`head_size < 64 (below the ll4mi kernel's minimum supported size)` ?

Was the test failing before they remove  `unified_attention` (before https://github.com/vllm-project/vllm/pull/34378)?


Code after:
                        num_decodes : num_decodes + num_extends
                    ],
                    slot_mapping=attn_metadata.slot_mapping[
                        num_decodes : num_decodes + num_extends
                    ],
                    k_scale=k_scale,
                    v_scale=v_scale,
                )

            # calculate for decodes
            if num_decodes > 0:
                assert attn_metadata.decode_metadata is not None
                decode_max_query_len = attn_metadata.decode_metadata.max_query_len

                # Use unified_attention for speculative decoding (multi-token)
                # or when sliding window is enabled
                if self.sliding_window[0] != -1 or decode_max_query_len > 1:
                    assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
                        "Shuffle KV cache layout is not supported with sliding "
                        "window or speculative decoding (multi-token decode)."
                    )
                    from aiter.ops.triton.unified_attention import (
                        unified_attention,
                    )

                    descale_shape = (
                        attn_metadata.query_start_loc[:num_decodes].shape[0] - 1,
                        key_cache.shape[2],
                    )
                    unified_attention(
                        q=query[:num_decode_tokens],
                        k=key_cache,
                        v=value_cache,
                        out=output[:num_decode_tokens],
                        cu_seqlens_q=attn_metadata.query_start_loc[:num_decodes],
                        max_seqlen_q=decode_max_query_len,
                        seqused_k=attn_metadata.seq_lens[:num_decodes],
                        max_seqlen_k=attn_metadata.max_seq_len,
                        softmax_scale=self.scale,
                        causal=True,
                        alibi_slopes=self.alibi_slopes,
                        window_size=self.sliding_window,
                        block_table=attn_metadata.block_table[:num_decodes],
                        softcap=self.logits_soft_cap,
                        q_descale=None,
                        k_descale=layer._k_scale.expand(descale_shape),
                        v_descale=layer._v_scale.expand(descale_shape),
                    )
                    return

                # The ll4mi kernel in paged_attention_v1 requires