[ROCm][CI] Add attention-sink support to ROCm AITER sparse MLA

tests/kernels/attention/test_rocm_aiter_mla_op_registration.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""ROCm custom op schema test for AITER MLA decode.
+"""ROCm custom op schema tests for AITER MLA decode.
 
-A single ``opcheck`` call on ``torch.ops.vllm.rocm_aiter_mla_decode_fwd``
-verifies that the custom op is registered and that its schema and fake
-implementation are consistent with the real kernel: fake-tensor support for
-torch.compile tracing and the ``mutates_args=["o"]`` in-place output aliasing.
+``opcheck`` verifies that the decode ops are registered and that their schemas
+and fake implementations are consistent with the real kernels: fake-tensor
+support for torch.compile tracing and ``mutates_args=["o"]`` in-place output
+aliasing.
 """
 
 import pytest
@@ -14,23 +14,20 @@
 from tests.kernels.utils import opcheck
 from vllm.platforms import current_platform
 
-_SKIP_NON_MI3XX = True
-if current_platform.is_rocm():
-    from vllm.platforms.rocm import on_mi3xx
-
-    _SKIP_NON_MI3XX = not on_mi3xx()
-
-pytestmark = [
-    pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
-    pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
-]
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-specific tests"
+)
 
 Q_HEAD_DIM = 576  # kv_lora_rank + qk_rope_head_dim
 V_HEAD_DIM = 512  # kv_lora_rank
 
 
 def _require_aiter():
     from vllm._aiter_ops import is_aiter_found_and_supported
+    from vllm.platforms.rocm import get_cdna_version
+
+    if get_cdna_version() not in (3, 4):
+        pytest.skip("AITER MLA requires CDNA 3 or 4")
 
     if not is_aiter_found_and_supported():
         pytest.skip("aiter is required on supported ROCm hardware for this test")
@@ -77,3 +74,34 @@ def test_mla_decode_fwd_op_schema() -> None:
             "reduce_partial_map": None,
         },
     )
+
+
+@torch.inference_mode()
+def test_mla_decode_fwd_lse_op_schema() -> None:
+    """Validate graph registration and mutation schema for LSE decode."""
+    _require_aiter()
+    # Import ensures the custom op is registered.
+    from vllm._aiter_ops import rocm_aiter_ops  # noqa: F401
+
+    batch_size, nhead = 2, 16
+    q = torch.randn(batch_size, nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    kv_buffer = torch.randn(32, 1, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    o = torch.zeros(batch_size, nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    qo_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda")
+    kv_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") * 16
+    kv_indices = torch.arange(32, dtype=torch.int32, device="cuda")
+    kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
+
+    opcheck(
+        torch.ops.vllm.rocm_aiter_mla_decode_fwd_lse,
+        (q, kv_buffer, o, qo_indptr, 1),
+        {
+            "kv_indptr": kv_indptr,
+            "kv_indices": kv_indices,
+            "kv_last_page_lens": kv_last_page_lens,
+            "sm_scale": Q_HEAD_DIM**-0.5,
+            "logit_cap": 0.0,
+            "q_scale": None,
+            "kv_scale": None,
+        },
+    )