[ROCM][Perf] Add ROCM_SEGMENTED_ATTN attention-backend works across RDNA3/3.5/4 (#59132)

Signed-off-by: root <jeffaw99@hotmail.com>
Signed-off-by: big-yellow-duck <jeffaw99@hotmail.com>
Co-authored-by: Tan Pin Siang <1716735+tanpinsiang@users.noreply.github.com>
Co-authored-by: vllmellm <190700713+vllmellm@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -61,6 +61,216 @@ def test_aiter_unified_attention_uses_dedicated_metadata_builder():
     )
 
 
+def test_segmented_attention_uses_dedicated_backend_components():
+    from vllm.v1.attention.backends.rocm_segmented_attn import (
+        RocmSegmentedAttentionBackend,
+        RocmSegmentedAttentionImpl,
+        RocmSegmentedAttentionMetadataBuilder,
+    )
+
+    assert (
+        RocmSegmentedAttentionBackend.get_builder_cls()
+        is RocmSegmentedAttentionMetadataBuilder
+    )
+    assert RocmSegmentedAttentionBackend.get_impl_cls() is RocmSegmentedAttentionImpl
+    assert RocmSegmentedAttentionBackend.get_name() == "ROCM_SEGMENTED_ATTN"
+    assert RocmSegmentedAttentionBackend.supports_non_causal()
+    assert RocmSegmentedAttentionBackend.supports_sliding_window()
+    assert RocmSegmentedAttentionBackend.supports_sink()
+    assert not RocmSegmentedAttentionBackend.supports_mm_prefix()
+
+
+@pytest.mark.parametrize(
+    "sliding_window,logits_soft_cap,has_sinks",
+    [(4096, None, False), (None, 50.0, False), (None, None, True)],
+)
+def test_segmented_attention_admits_unified_fallback_features(
+    monkeypatch, sliding_window, logits_soft_cap, has_sinks
+):
+    from vllm.platforms import rocm
+    from vllm.v1.attention.backend import AttentionType
+    from vllm.v1.attention.backends.rocm_segmented_attn import (
+        RocmSegmentedAttentionImpl,
+    )
+
+    monkeypatch.setattr(rocm, "on_gfx1x", lambda: True)
+    monkeypatch.setattr(rocm, "on_gfx12x", lambda: True)
+    sinks = torch.zeros(8) if has_sinks else None
+
+    impl = RocmSegmentedAttentionImpl(
+        num_heads=8,
+        head_size=128,
+        scale=128**-0.5,
+        num_kv_heads=2,
+        alibi_slopes=None,
+        sliding_window=sliding_window,
+        kv_cache_dtype="auto",
+        logits_soft_cap=logits_soft_cap,
+        attn_type=AttentionType.DECODER,
+        sinks=sinks,
+    )
+
+    assert isinstance(impl, RocmSegmentedAttentionImpl)
+
+
+@pytest.mark.parametrize(
+    "platform,rdna",
+    [("cuda", False), ("cpu", False), ("xpu", False), ("rocm", False), ("rocm", True)],
+)
+def test_segmented_attention_requires_rocm_rdna(monkeypatch, platform, rdna):
+    """Reject unsupported platforms before probing ROCm or constructing kernels."""
+    from vllm.platforms import rocm
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.attention.backends import rocm_segmented_attn as segmented
+
+    monkeypatch.setattr(
+        segmented,
+        "current_platform",
+        SimpleNamespace(is_rocm=lambda: platform == "rocm"),
+    )
+
+    def on_rdna():
+        assert platform == "rocm", "Must not probe ROCm on another platform"
+        return rdna
+
+    monkeypatch.setattr(rocm, "on_gfx1x", on_rdna)
+    reason = segmented.RocmSegmentedAttentionBackend.supports_combination(
+        128,
+        torch.bfloat16,
+        "auto",
+        32,
+        False,
+        False,
+        False,
+        False,
+        DeviceCapability(12, 0),
+    )
+    if rdna:
+        assert reason is None
+        segmented.RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto")
+    else:
+        assert reason == "ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm"
+        with pytest.raises(ValueError, match="requires AMD RDNA GPUs on ROCm"):
+            segmented.RocmSegmentedAttentionImpl(
+                8, 128, 128**-0.5, 2, None, None, "auto"
+            )
+
+
+@pytest.mark.parametrize("explicit_cache", [False, True])
+def test_segmented_triton_cache_survives_standalone_compilation(
+    monkeypatch, tmp_path, explicit_cache
+):
+    """Backend initialization preserves one cache across model compilations."""
+    import os
+
+    from vllm.compilation.compiler_interface import InductorStandaloneAdaptor
+    from vllm.platforms import rocm
+    from vllm.v1.attention.backends.rocm_segmented_attn import (
+        RocmSegmentedAttentionImpl,
+    )
+
+    monkeypatch.setattr(rocm, "on_gfx1x", lambda: True)
+    monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path))
+    monkeypatch.delenv("TRITON_CACHE_DIR", raising=False)
+    monkeypatch.delenv("TORCHINDUCTOR_CACHE_DIR", raising=False)
+    expected = tmp_path / "rocm_segmented_attention" / "triton"
+    if explicit_cache:
+        expected = tmp_path / "custom"
+        monkeypatch.setenv("TRITON_CACHE_DIR", str(expected))
+    for model in ("target", "draft", "restart_target"):
+        RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto")
+        assert os.environ["TRITON_CACHE_DIR"] == str(expected)
+        assert expected.is_dir()
+        InductorStandaloneAdaptor("binary").initialize_cache(str(tmp_path / model))
+        assert os.environ["TRITON_CACHE_DIR"] == str(expected)
+
+
+def test_segmented_attention_is_opt_in(monkeypatch):
+    from vllm.platforms import rocm
+    from vllm.platforms.rocm import RocmPlatform, _get_backend_priorities
+
+    assert AttentionBackendEnum.ROCM_SEGMENTED_ATTN not in _get_backend_priorities(
+        use_mla=False, use_sparse=False
+    )
+
+    monkeypatch.setattr(rocm, "on_gfx1x", lambda: True)
+    monkeypatch.setattr(rocm, "on_gfx12x", lambda: True)
+    config = AttentionSelectorConfig(
+        head_size=128,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="fp8",
+        block_size=32,
+    )
+
+    path = RocmPlatform.get_attn_backend_cls(
+        selected_backend=AttentionBackendEnum.ROCM_SEGMENTED_ATTN,
+        attn_selector_config=config,
+    )
+
+    assert path == AttentionBackendEnum.ROCM_SEGMENTED_ATTN.get_path()
+
+
+def test_segmented_attention_rejects_mm_prefix_at_selection(monkeypatch):
+    from vllm.platforms import rocm
+    from vllm.platforms.rocm import RocmPlatform
+
+    monkeypatch.setattr(rocm, "on_gfx1x", lambda: True)
+    config = AttentionSelectorConfig(
+        head_size=128,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        block_size=32,
+        use_mm_prefix=True,
+    )
+
+    with pytest.raises(ValueError, match="multimodal token full attention"):
+        RocmPlatform.get_attn_backend_cls(
+            selected_backend=AttentionBackendEnum.ROCM_SEGMENTED_ATTN,
+            attn_selector_config=config,
+        )
+
+
+def test_segmented_attention_forward_uses_dedicated_dispatch(monkeypatch):
+    from vllm.config import KernelConfig
+    from vllm.platforms import rocm
+    from vllm.v1.attention.backends.rocm_segmented_attn import (
+        RocmSegmentedAttentionImpl,
+    )
+
+    monkeypatch.setattr(rocm, "on_gfx1x", lambda: True)
+    impl = RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto")
+    impl._segmented_attention_config = SimpleNamespace(
+        kernel_config=KernelConfig(),
+        scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=3),
+        model_config=SimpleNamespace(max_model_len=16),
+    )
+    query = torch.empty(3, 8, 128, dtype=torch.bfloat16)
+    key = torch.empty(3, 2, 128, dtype=torch.bfloat16)
+    value = torch.empty_like(key)
+    output = torch.empty_like(query)
+    kv_cache = torch.empty(2, 2, 16, 256, dtype=torch.bfloat16)
+    metadata = SimpleNamespace(
+        use_cascade=False,
+        num_actual_tokens=3,
+        block_table=torch.zeros(1, 1, dtype=torch.int32),
+        query_start_loc=torch.tensor([0, 3], dtype=torch.int32),
+        seq_lens=torch.tensor([3], dtype=torch.int32),
+        max_seq_len=3,
+        max_query_len=3,
+        causal=True,
+    )
+    layer = SimpleNamespace(_k_scale=torch.ones(()), _v_scale=torch.ones(()))
+
+    with patch(
+        "vllm.v1.attention.backends.rocm_segmented_attn.segmented_attention"
+    ) as dispatch:
+        result = impl.forward(layer, query, key, value, kv_cache, metadata, output)
+
+    assert result is output
+    dispatch.assert_called_once()
+    assert dispatch.call_args.kwargs["key_cache"].shape == (2, 16, 2, 128)
+
+
 def test_aiter_unified_attention_capture_preserves_query_start_locations():
     from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
         RocmAiterUnifiedAttentionMetadataBuilder,

vllm/config/kernel.py

@@ -278,6 +278,9 @@ class KernelConfig:
     enable_flashinfer_autotune: bool = None  # type: ignore[assignment]
     """If True, run FlashInfer autotuning during kernel warmup."""
 
+    enable_rocm_segmented_attn_autotune: bool = False
+    """If True, autotune ROCm segmented attention during kernel warmup on RDNA GPUs."""
+
     # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
     # to the shared JIT warmup infrastructure.
     # https://github.com/vllm-project/vllm/pull/47451
@@ -406,6 +409,7 @@ class KernelConfig:
             "enable_cutedsl_warmup",
             "enable_jit_warmup",
             "enable_flashinfer_autotune",
+            "enable_rocm_segmented_attn_autotune",
             "ir_op_priority",  # handled separately below
         }
         if self.linear_backend_per_quant is None:

vllm/model_executor/models/qwen3_dflash.py

@@ -34,6 +34,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.multimodal.inputs import NestedTensors
+from vllm.platforms import current_platform
 from vllm.transformers_utils.config import set_default_rope_theta
 from vllm.transformers_utils.repo_utils import get_hf_file_bytes
 from vllm.v1.attention.backend import AttentionType
@@ -269,6 +270,24 @@ class DFlashQwen3Attention(nn.Module):
             sinks=self.attention_sink_bias,
         )
         self.causal = causal
+        if current_platform.is_rocm():
+            from vllm.platforms.rocm import on_gfx1x
+            from vllm.v1.attention.backends.rocm_segmented_attn import (
+                RocmSegmentedAttentionImpl,
+            )
+
+            if on_gfx1x() and isinstance(self.attn.impl, RocmSegmentedAttentionImpl):
+                # Startup segmented tuning needs the draft's actual attention mode.
+                # Causality is normally supplied later through attention metadata.
+                self.attn.segmented_causal = causal
+                speculative_config = get_current_vllm_config().speculative_config
+                if (
+                    speculative_config is not None
+                    and speculative_config.num_speculative_tokens is not None
+                ):
+                    self.attn.segmented_query_limit = (
+                        speculative_config.num_speculative_tokens + 1
+                    )
         self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
         self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
 

vllm/model_executor/warmup/kernel_warmup.py

@@ -220,6 +220,13 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     if process_local_only:
         return
 
+    if current_platform.is_rocm():
+        from vllm.model_executor.warmup.rocm_segmented_attn_autotune_warmup import (
+            rocm_segmented_attn_autotune_warmup,
+        )
+
+        rocm_segmented_attn_autotune_warmup(worker)
+
     flashinfer_sparse_mla_decode_autotune_warmup(worker)
     deepseek_v4_sparse_mla_attention_warmup(worker)
 

vllm/v1/attention/backends/registry.py

@@ -53,6 +53,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
         "vllm.v1.attention.backends.triton_flash_attn.TritonFlashAttentionBackend"
     )
     ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
+    ROCM_SEGMENTED_ATTN = (
+        "vllm.v1.attention.backends.rocm_segmented_attn.RocmSegmentedAttentionBackend"
+    )
     ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
     ROCM_AITER_TRITON_MLA = (
         "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend"

vllm/v1/attention/backends/rocm_segmented_attn.py

@@ -0,0 +1,371 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ROCm attention backend for token-major segmented Triton attention."""
+
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, ClassVar
+
+import torch
+
+import vllm.envs as envs
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.cache import CacheDType
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.v1.attention.backend import AttentionLayer, AttentionType
+from vllm.v1.attention.backends.rocm_attn import (
+    RocmAttentionBackend,
+    RocmAttentionImpl,
+    RocmAttentionMetadata,
+    RocmAttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import get_num_attention_heads_from_layers
+from vllm.v1.attention.ops.segmented_attention import (
+    get_segmented_attention_workspace,
+    segmented_attention,
+    segmented_attention_workspace_size,
+)
+from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
+    triton_reshape_and_cache_flash,
+)
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+if TYPE_CHECKING:
+    from vllm.platforms.interface import DeviceCapability
+
+logger = init_logger(__name__)
+
+
+def is_rdna() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import on_gfx1x
+
+    return on_gfx1x()
+
+
+class RocmSegmentedAttentionMetadataBuilder(RocmAttentionMetadataBuilder):
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ):
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        model_config = vllm_config.model_config
+        self.num_heads_q = get_num_attention_heads_from_layers(
+            vllm_config, layer_names
+        ) or model_config.get_num_attention_heads(vllm_config.parallel_config)
+        self.num_heads_kv = kv_cache_spec.num_kv_heads
+        self.headdim = kv_cache_spec.head_size
+
+
+class RocmSegmentedAttentionBackend(RocmAttentionBackend):
+    """Explicit token-major ROCm backend optimized for segmented prefill."""
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [
+        torch.float16,
+        torch.bfloat16,
+    ]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "float16",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @staticmethod
+    def get_name() -> str:
+        return "ROCM_SEGMENTED_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["RocmSegmentedAttentionImpl"]:
+        return RocmSegmentedAttentionImpl
+
+    @staticmethod
+    def get_builder_cls() -> type["RocmSegmentedAttentionMetadataBuilder"]:
+        return RocmSegmentedAttentionMetadataBuilder
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [64, 128, 256]
+
+    @classmethod
+    def supports_attn_type(cls, attn_type: str) -> bool:
+        return attn_type == AttentionType.DECODER
+
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_sliding_window(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_sink(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_mm_prefix(cls) -> bool:
+        return False
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        device_capability: "DeviceCapability",
+    ) -> str | None:
+        del (
+            head_size,
+            dtype,
+            block_size,
+            use_mla,
+            has_sink,
+            use_sparse,
+            use_mm_prefix,
+            device_capability,
+        )
+        if not is_rdna():
+            return "ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm"
+
+        from vllm.platforms.rocm import on_gfx12x
+
+        if kv_cache_dtype in ("fp8", "fp8_e4m3") and not on_gfx12x():
+            return "FP8 segmented attention requires gfx12"
+        return None
+
+
+class RocmSegmentedAttentionImpl(RocmAttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None = None,
+        attn_type: AttentionType = AttentionType.DECODER,
+        kv_sharing_target_layer_name: int | None = None,
+        sinks: torch.Tensor | None = None,
+    ) -> None:
+        if not is_rdna():
+            raise ValueError("ROCM_SEGMENTED_ATTN requires AMD RDNA GPUs on ROCm")
+
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            sinks,
+        )
+        from vllm.platforms.rocm import on_gfx1x, on_gfx12x
+
+        fp8 = kv_cache_dtype in ("fp8", "fp8_e4m3")
+        arch_supported = on_gfx12x() if fp8 else on_gfx1x()
+        gqa_ratio_supported = (
+            num_kv_heads > 0
+            and num_heads % num_kv_heads == 0
+            and 1 <= self.num_queries_per_kv <= 16
+        )
+        if not (
+            arch_supported
+            and attn_type == AttentionType.DECODER
+            and head_size in (64, 128, 256)
+            and gqa_ratio_supported
+            and kv_cache_dtype in ("auto", "float16", "bfloat16", "fp8", "fp8_e4m3")
+            and alibi_slopes is None
+        ):
+            raise ValueError(
+                "ROCM_SEGMENTED_ATTN requires decoder attention on gfx1x "
+                "(gfx12 for FP8), head size 64, 128 or 256, GQA ratio 1-16, and "
+                "does not support ALiBi."
+            )
+        # Startup tuning and target/draft compilation must share a Triton cache.
+        cache_dir = os.environ.setdefault(
+            "TRITON_CACHE_DIR",
+            str(Path(envs.VLLM_CACHE_ROOT) / "rocm_segmented_attention" / "triton"),
+        )
+        Path(cache_dir).mkdir(parents=True, exist_ok=True)
+        logger.info_once("Using token-major ROCm segmented Triton attention")
+        self._segmented_attention_warmed_up = False
+        self._segmented_attention_config: VllmConfig | None = None
+
+    def process_weights_after_loading(self, act_dtype: torch.dtype):
+        super().process_weights_after_loading(act_dtype)
+        config = get_current_vllm_config()
+        self._segmented_attention_config = config
+
+    def _autotune_enabled(self) -> bool:
+        config = self._segmented_attention_config
+        return (
+            config is not None
+            and config.kernel_config.enable_rocm_segmented_attn_autotune
+        )
+
+    def _get_workspace(self, device):
+        config = self._segmented_attention_config
+        assert config is not None
+        sizes = segmented_attention_workspace_size(
+            config.scheduler_config.max_num_seqs,
+            self.num_heads,
+            self.num_kv_heads,
+            self.head_size,
+            config.model_config.max_model_len,
+            max_tokens=config.scheduler_config.max_num_batched_tokens,
+            fp8=self.kv_cache_dtype in ("fp8", "fp8_e4m3"),
+            autotune=self._autotune_enabled(),
+        )
+        return get_segmented_attention_workspace(device, sizes)
+
+    def _split_kv_cache(
+        self,
+        kv_cache: torch.Tensor,
+        num_kv_heads: int | None = None,
+        head_size: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        num_kv_heads = num_kv_heads or self.num_kv_heads
+        head_size = head_size or self.head_size
+        return (
+            kv_cache[:, 0].unflatten(-1, (num_kv_heads, head_size)),
+            kv_cache[:, 1].unflatten(-1, (num_kv_heads, head_size)),
+        )
+
+    def forward(
+        self,
+        layer: torch.nn.Module,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: RocmAttentionMetadata,
+        output: torch.Tensor,
+        output_scale: torch.Tensor | None = None,
+        output_block_scale: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if output_block_scale is not None:
+            raise NotImplementedError(
+                "fused block_scale output quantization is not supported"
+            )
+        # Profile runs allocate the same fixed scratch used during graph replay.
+        workspace = self._get_workspace(query.device)
+        if attn_metadata is None:
+            return output.fill_(0)
+        assert not attn_metadata.use_cascade
+
+        num_actual_tokens = attn_metadata.num_actual_tokens
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
+        if is_quantized_kv_cache(self.kv_cache_dtype):
+            key_cache = key_cache.view(self.fp8_dtype)
+            value_cache = value_cache.view(self.fp8_dtype)
+            if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0:
+                raise NotImplementedError(
+                    "A non-1.0 q_scale with an FP8 query is not supported"
+                )
+
+        segmented_attention(
+            query=query[:num_actual_tokens],
+            key=key[:num_actual_tokens],
+            value=value[:num_actual_tokens],
+            output=output[:num_actual_tokens],
+            kv_cache_dtype=self.kv_cache_dtype,
+            key_cache=key_cache,
+            value_cache=value_cache,
+            block_table=attn_metadata.block_table,
+            query_start_loc=attn_metadata.query_start_loc,
+            seq_lens=attn_metadata.seq_lens,
+            max_seq_len=attn_metadata.max_seq_len,
+            max_query_len=attn_metadata.max_query_len,
+            k_scale=layer._k_scale,
+            v_scale=layer._v_scale,
+            sm_scale=self.scale,
+            sliding_window=self.sliding_window[0],
+            output_scale=output_scale,
+            sinks=self.sinks,
+            causal=attn_metadata.causal,
+            softcap=self.logits_soft_cap,
+            workspace=workspace,
+            use_tuned_config=self._autotune_enabled(),
+        )
+        return output
+
+    def do_kv_cache_update(
+        self,
+        layer: AttentionLayer,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+    ):
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
+        if is_quantized_kv_cache(self.kv_cache_dtype):
+            key_cache = key_cache.view(self.fp8_dtype)
+            value_cache = value_cache.view(self.fp8_dtype)
+        triton_reshape_and_cache_flash(
+            key,
+            value,
+            key_cache,
+            value_cache,
+            slot_mapping,
+            self.kv_cache_dtype,
+            layer._k_scale,
+            layer._v_scale,
+        )
+
+    def do_rope_and_kv_cache_update(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        positions: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        key_cache, value_cache = self._split_kv_cache(
+            kv_cache,
+            layer.num_kv_heads,  # type: ignore[attr-defined]
+            layer.head_size,  # type: ignore[attr-defined]
+        )
+        is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
+        if is_fp8_kv_cache:
+            key_cache = key_cache.view(self.fp8_dtype)
+            value_cache = value_cache.view(self.fp8_dtype)
+
+        rocm_aiter_ops.triton_rope_and_cache(
+            query,
+            key,
+            value,
+            positions,
+            cos_sin_cache,
+            is_neox,
+            key_cache,
+            value_cache,
+            layer_slot_mapping,
+            layer._k_scale,
+            layer._v_scale,
+            True,
+            is_fp8_kv_cache,
+        )