[ROCm][Perf] HY4 preview perf boost (#58509)

Signed-off-by: Qiang Li <qiang.li2@amd.com>

.buildkite/test_areas/models_basic.yaml

@@ -102,9 +102,11 @@ steps:
       - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
       - tests/models/test_dsv41_mono_decode.py
       - tests/models/test_deepseek_v4_vl_stream_weights.py
+      - tests/models/hy_v4/test_triton_ihc.py
       commands:
       - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_deepseek_v4_vl_stream_weights.py
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
+      - pytest -v -s models/hy_v4/test_triton_ihc.py
 
 - label: ":amd: (MI355) MiniMax M3 MonoKernel"
   key: minimax-m3-monokernel-mi355

benchmarks/kernels/benchmark_hy_v4_ihc.py

@@ -1,24 +1,33 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark the HY V4 Triton iHC pre/post kernels against eager PyTorch."""
+"""Benchmark HY V4 Triton iHC kernels on CUDA or ROCm against eager PyTorch."""
 
 import os
 import subprocess
 from functools import partial
-from importlib.metadata import version
+from importlib.metadata import PackageNotFoundError, version
 from statistics import median
 
 import torch
 import torch.nn.functional as F
 
 import vllm
-from vllm.models.hy_v4.nvidia.triton_ihc import (
-    triton_ihc_post,
-    triton_ihc_pre,
-)
+from vllm.platforms import current_platform
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+if current_platform.is_rocm():
+    from vllm.models.hy_v4.amd.triton_ihc import (
+        triton_ihc_post,
+        triton_ihc_post_pre_rms_norm,
+        triton_ihc_pre,
+    )
+else:
+    from vllm.models.hy_v4.nvidia.triton_ihc import (
+        triton_ihc_post,
+        triton_ihc_pre,
+    )
+
 
 def eager_pre(
     x: torch.Tensor,
@@ -51,22 +60,90 @@ def eager_post(
     )
 
 
-def _timer(method: str):
-    if method == "cupti":
-        from flashinfer.testing import bench_gpu_time_with_cupti
+def eager_post_pre_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    attn_post: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    norm_weight: torch.Tensor,
+    magnitude: float,
+    hc_eps: float,
+    norm_eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    y = eager_post(x, residual, attn_post)
+    output, post = eager_pre(y, weight, scale, base, magnitude, hc_eps, norm_eps)
+    output_float = output.float()
+    output = (
+        output_float
+        * torch.rsqrt(output_float.square().mean(-1, keepdim=True) + norm_eps)
+        * norm_weight
+    ).to(output.dtype)
+    return output, post, y
 
+
+def _timer(method: str):
+    """Return the explicitly requested GPU timer or fail if unavailable."""
+    if method == "torch_event":
+        return _torch_event_timer
+
+    if current_platform.is_rocm():
+        raise RuntimeError(
+            f"timing method {method!r} is CUDA-only; use --method torch_event on ROCm"
+        )
+
+    if method == "cupti":
+        try:
+            from flashinfer.testing import bench_gpu_time_with_cupti
+        except ImportError as error:
+            raise RuntimeError(
+                "timing method 'cupti' requires flashinfer-python with "
+                "bench_gpu_time_with_cupti"
+            ) from error
         return partial(
             bench_gpu_time_with_cupti,
             use_cuda_graph=True,
             cold_l2_cache=True,
         )
     if method == "cudagraph":
-        from flashinfer.testing import bench_gpu_time_with_cudagraph
-
+        try:
+            from flashinfer.testing import bench_gpu_time_with_cudagraph
+        except ImportError as error:
+            raise RuntimeError(
+                "timing method 'cudagraph' requires flashinfer-python with "
+                "bench_gpu_time_with_cudagraph"
+            ) from error
         return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True)
     raise ValueError(f"unknown timing method: {method}")
 
 
+def _torch_event_timer(fn):
+    for _ in range(10):
+        fn()
+    torch.accelerator.synchronize()
+    start = torch.Event(enable_timing=True)
+    end = torch.Event(enable_timing=True)
+    start.record()
+    for _ in range(20):
+        fn()
+    end.record()
+    end.synchronize()
+    return [start.elapsed_time(end) / 20]
+
+
+def _git_metadata(arguments: list[str], *, empty: str) -> str:
+    result = subprocess.run(
+        ["git", *arguments],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        return "<unknown>"
+    return result.stdout.strip() or empty
+
+
 @torch.inference_mode()
 def run_benchmark(
     token_counts: list[int],
@@ -85,30 +162,39 @@ def run_benchmark(
     )
     scale = torch.randn(2, device=device, dtype=torch.float32) * 0.01
     base = torch.randn(2 * hc_mult, device=device, dtype=torch.float32)
+    norm_weight = torch.randn(hidden_size, device=device, dtype=torch.float32)
     timer = _timer(method)
-
     properties = torch.cuda.get_device_properties(device)
-    git_branch = subprocess.run(
-        ["git", "branch", "--show-current"],
-        check=False,
-        capture_output=True,
-        text=True,
-    ).stdout.strip()
-    git_commit = subprocess.run(
-        ["git", "rev-parse", "--short", "HEAD"],
-        check=False,
-        capture_output=True,
-        text=True,
-    ).stdout.strip()
-    print(f"device: {properties.name}")
+    git_branch = _git_metadata(["branch", "--show-current"], empty="<detached>")
+    git_commit = _git_metadata(["rev-parse", "--short", "HEAD"], empty="<unknown>")
+    device_name = (
+        current_platform.get_device_name()
+        if current_platform.is_rocm()
+        else properties.name
+    )
+    platform_description = (
+        f"HIP: {torch.version.hip}"
+        if current_platform.is_rocm()
+        else f"CUDA: {torch.version.cuda}"
+    )
+    print(f"device: {device_name}")
     print(f"branch: {git_branch}; commit: {git_commit}")
     print(
-        f"vllm: {vllm.__version__}; torch: {torch.__version__}; "
-        f"CUDA: {torch.version.cuda}"
+        f"vllm: {vllm.__version__}; torch: {torch.__version__}; {platform_description}"
     )
+    if method == "torch_event":
+        flashinfer_version = "not used"
+        event_backend = "HIP" if current_platform.is_rocm() else "CUDA"
+        time_description = f"method: PyTorch {event_backend} events; cache: warm"
+    else:
+        try:
+            flashinfer_version = version("flashinfer-python")
+        except PackageNotFoundError:
+            flashinfer_version = "unavailable"
+        time_description = f"method: {method}; cache: cold L2"
     print(
-        f"triton: {triton.__version__}; flashinfer: {version('flashinfer-python')}; "
-        f"dtype: {dtype}; method: {method}; cache: cold L2"
+        f"triton: {triton.__version__}; flashinfer: {flashinfer_version}; "
+        f"dtype: {dtype}; {time_description}"
     )
     print(
         "env: VLLM_ENABLE_HPC_OPS="
@@ -145,8 +231,7 @@ def run_benchmark(
             atol=0,
             rtol=0,
         )
-
-        benchmarks = (
+        benchmarks = [
             (
                 "pre",
                 partial(eager_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
@@ -164,7 +249,73 @@ def run_benchmark(
                 partial(triton_ihc_post, block_output, residual, post),
                 block_output.nbytes + residual.nbytes + post.nbytes + residual.nbytes,
             ),
-        )
+        ]
+        if current_platform.is_rocm():
+            fused_expected = eager_post_pre_rms_norm(
+                block_output,
+                residual,
+                post,
+                weight,
+                scale,
+                base,
+                norm_weight,
+                2.0,
+                1e-6,
+                1e-5,
+            )
+            fused_actual = triton_ihc_post_pre_rms_norm(
+                block_output,
+                residual,
+                post,
+                weight,
+                scale,
+                base,
+                norm_weight,
+                2.0,
+                1e-6,
+                1e-5,
+            )
+            for actual, expected in zip(fused_actual, fused_expected):
+                torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+            benchmarks.append(
+                (
+                    "fused",
+                    partial(
+                        eager_post_pre_rms_norm,
+                        block_output,
+                        residual,
+                        post,
+                        weight,
+                        scale,
+                        base,
+                        norm_weight,
+                        2.0,
+                        1e-6,
+                        1e-5,
+                    ),
+                    partial(
+                        triton_ihc_post_pre_rms_norm,
+                        block_output,
+                        residual,
+                        post,
+                        weight,
+                        scale,
+                        base,
+                        norm_weight,
+                        2.0,
+                        1e-6,
+                        1e-5,
+                    ),
+                    block_output.nbytes
+                    + residual.nbytes
+                    + post.nbytes
+                    + weight.nbytes
+                    + scale.nbytes
+                    + base.nbytes
+                    + norm_weight.nbytes
+                    + sum(tensor.nbytes for tensor in fused_expected),
+                )
+            )
         for op_name, eager_fn, triton_fn, logical_bytes in benchmarks:
             eager_us = median(timer(eager_fn)) * 1e3
             triton_us = median(timer(triton_fn)) * 1e3
@@ -189,7 +340,15 @@ if __name__ == "__main__":
     )
     parser.add_argument("--hidden-size", type=int, choices=[4096, 6144], default=6144)
     parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16")
-    parser.add_argument("--method", choices=["cupti", "cudagraph"], default="cupti")
+    parser.add_argument(
+        "--method",
+        choices=["torch_event", "cupti", "cudagraph"],
+        default="torch_event" if current_platform.is_rocm() else "cupti",
+        help=(
+            "timing backend; cupti/cudagraph require flashinfer-python on CUDA, "
+            "while torch_event uses warm-cache PyTorch GPU events"
+        ),
+    )
     args = parser.parse_args()
     run_benchmark(
         args.token_counts,

tests/kernels/test_minimax_m3_amd_ops.py

@@ -526,6 +526,7 @@ def _capture_expert_mask(expert_mask, *, global_num_experts):
     from unittest import mock
 
     from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
     from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import (
         AiterMxfp8Experts,
     )
@@ -556,7 +557,7 @@ def _capture_expert_mask(expert_mask, *, global_num_experts):
             w2=w2,
             topk_weights=tw,
             topk_ids=ti,
-            activation=None,
+            activation=MoEActivation.SILU,
             global_num_experts=global_num_experts,
             expert_map=expert_mask,
             a1q_scale=None,

tests/models/hy_v4/test_triton_ihc.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.models.hy_v4.amd.triton_ihc import (
+    triton_ihc_post,
+    triton_ihc_post_pre_rms_norm,
+    triton_ihc_pre,
+    triton_ihc_supported,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.utils.torch_utils import set_random_seed
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm() or not HAS_TRITON or not torch.cuda.is_available(),
+    reason="Requires ROCm, Triton, and a GPU",
+)
+
+_HC_MULT = 4
+_HIDDEN_SIZE = 6144
+_DTYPE = torch.bfloat16
+_TOKEN_COUNTS = [1, 3, 8, 17, 32, 33, 127, 128]
+
+
+@pytest.fixture(autouse=True)
+def _reset_random_seed() -> None:
+    set_random_seed(0)
+
+
+def _pre_reference(x, weight, scale, base, magnitude, hc_eps, norm_eps):
+    flat = x.float().flatten(1)
+    reciprocal_rms = torch.rsqrt(flat.square().mean(-1, keepdim=True) + norm_eps)
+    mixes = flat @ weight.float().transpose(0, 1) * reciprocal_rms
+    pre = torch.sigmoid(mixes[:, :_HC_MULT] * scale[0] + base[:_HC_MULT]) + hc_eps
+    post = (
+        magnitude
+        * torch.sigmoid(mixes[:, _HC_MULT:] * scale[1] + base[_HC_MULT : 2 * _HC_MULT])
+        + hc_eps
+    )
+    return (pre.unsqueeze(-1) * x.float()).sum(1).to(x.dtype), post
+
+
+def _rms_norm_reference(x, weight, eps):
+    x_float = x.float()
+    return (
+        x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + eps) * weight
+    ).to(x.dtype)
+
+
+def _post_pre_rms_norm_reference(
+    x,
+    residual,
+    attn_post,
+    weight,
+    scale,
+    base,
+    norm_weight,
+    magnitude,
+    hc_eps,
+    norm_eps,
+):
+    y = attn_post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()
+    y = y.to(x.dtype)
+    reduced, mlp_post = _pre_reference(
+        y, weight, scale, base, magnitude, hc_eps, norm_eps
+    )
+    return _rms_norm_reference(reduced, norm_weight, norm_eps), mlp_post, y
+
+
+def _make_fused_inputs(
+    num_tokens: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+]:
+    device = torch.device("cuda")
+    x = torch.randn((num_tokens, hidden_size), device=device, dtype=dtype)
+    residual = torch.randn(
+        (num_tokens, _HC_MULT, hidden_size), device=device, dtype=dtype
+    )
+    attn_post = torch.randn((num_tokens, _HC_MULT), device=device, dtype=torch.float32)
+    weight = torch.randn(
+        (2 * _HC_MULT, _HC_MULT * hidden_size),
+        device=device,
+        dtype=torch.float32,
+    )
+    scale = torch.randn(2, device=device, dtype=torch.float32)
+    base = torch.randn(2 * _HC_MULT, device=device, dtype=torch.float32)
+    norm_weight = torch.randn(hidden_size, device=device, dtype=torch.float32)
+    return x, residual, attn_post, weight, scale, base, norm_weight
+
+
+@pytest.mark.parametrize("num_tokens", _TOKEN_COUNTS)
+def test_rocm_triton_ihc_pre_matches_eager(num_tokens):
+    device = torch.device("cuda")
+    x = torch.randn((num_tokens, _HC_MULT, _HIDDEN_SIZE), device=device, dtype=_DTYPE)
+    weight = torch.randn(
+        (2 * _HC_MULT, _HC_MULT * _HIDDEN_SIZE),
+        device=device,
+        dtype=torch.float32,
+    )
+    scale = torch.randn(2, device=device, dtype=torch.float32)
+    base = torch.randn(2 * _HC_MULT, device=device, dtype=torch.float32)
+    args = (2.0, 1e-6, 1e-5)
+
+    actual_output, actual_post = triton_ihc_pre(x, weight, scale, base, *args)
+    expected_output, expected_post = _pre_reference(x, weight, scale, base, *args)
+    torch.testing.assert_close(actual_output, expected_output, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(actual_post, expected_post, rtol=2e-3, atol=2e-3)
+
+
+@pytest.mark.parametrize("num_tokens", _TOKEN_COUNTS)
+def test_rocm_triton_ihc_post_matches_eager(num_tokens):
+    device = torch.device("cuda")
+    x = torch.randn((num_tokens, _HIDDEN_SIZE), device=device, dtype=_DTYPE)
+    residual = torch.randn(
+        (num_tokens, _HC_MULT, _HIDDEN_SIZE), device=device, dtype=_DTYPE
+    )
+    post = torch.randn((num_tokens, _HC_MULT), device=device, dtype=torch.float32)
+
+    actual = triton_ihc_post(x, residual, post)
+    expected = (
+        post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()
+    ).to(_DTYPE)
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.parametrize("num_tokens", _TOKEN_COUNTS)
+def test_rocm_triton_ihc_post_pre_rms_norm_matches_eager(num_tokens):
+    device = torch.device("cuda")
+    x = torch.randn((num_tokens, _HIDDEN_SIZE), device=device, dtype=_DTYPE)
+    residual = torch.randn(
+        (num_tokens, _HC_MULT, _HIDDEN_SIZE), device=device, dtype=_DTYPE
+    )
+    attn_post = torch.randn((num_tokens, _HC_MULT), device=device, dtype=torch.float32)
+    weight = torch.randn(
+        (2 * _HC_MULT, _HC_MULT * _HIDDEN_SIZE),
+        device=device,
+        dtype=torch.float32,
+    )
+    scale = torch.randn(2, device=device, dtype=torch.float32)
+    base = torch.randn(2 * _HC_MULT, device=device, dtype=torch.float32)
+    norm_weight = torch.randn(_HIDDEN_SIZE, device=device, dtype=torch.float32)
+    args = (2.0, 1e-6, 1e-5)
+
+    actual_reduced, actual_post, actual_residual = triton_ihc_post_pre_rms_norm(
+        x,
+        residual,
+        attn_post,
+        weight,
+        scale,
+        base,
+        norm_weight,
+        *args,
+    )
+    expected_reduced, expected_post, expected_residual = _post_pre_rms_norm_reference(
+        x, residual, attn_post, weight, scale, base, norm_weight, *args
+    )
+    torch.testing.assert_close(actual_reduced, expected_reduced, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(actual_post, expected_post, rtol=2e-3, atol=2e-3)
+    torch.testing.assert_close(actual_residual, expected_residual, rtol=2e-2, atol=2e-2)
+
+
+def test_rocm_triton_ihc_dispatch():
+    x = torch.empty((1, _HC_MULT, _HIDDEN_SIZE), device="cuda", dtype=_DTYPE)
+    assert triton_ihc_supported(x)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("hidden_size", [4096, 6144])
+def test_rocm_triton_ihc_fused_production_shapes(dtype, hidden_size):
+    inputs = _make_fused_inputs(2, hidden_size, dtype)
+    args = (2.0, 1e-6, 1e-5)
+
+    actual = triton_ihc_post_pre_rms_norm(*inputs, *args)
+    expected = _post_pre_rms_norm_reference(*inputs, *args)
+    for actual_tensor, expected_tensor in zip(actual, expected):
+        torch.testing.assert_close(actual_tensor, expected_tensor, rtol=2e-2, atol=2e-2)
+
+
+def test_rocm_triton_ihc_empty_and_noncontiguous_inputs():
+    device = torch.device("cuda")
+    weight = torch.randn(
+        (2 * _HC_MULT, _HC_MULT * 4096), device=device, dtype=torch.float32
+    )
+    scale = torch.randn(2, device=device)
+    base = torch.randn(2 * _HC_MULT, device=device)
+    empty = torch.empty((0, _HC_MULT, 4096), device=device, dtype=_DTYPE)
+    output, post = triton_ihc_pre(empty, weight, scale, base, 2.0, 1e-6, 1e-5)
+    assert output.shape == (0, 4096)
+    assert post.shape == (0, _HC_MULT)
+
+    x = torch.randn((2, _HC_MULT, 8192), device=device, dtype=_DTYPE)[..., ::2]
+    noncontiguous_weight = torch.randn((_HC_MULT * 4096, 2 * _HC_MULT), device=device).T
+    assert not x.is_contiguous()
+    assert not noncontiguous_weight.is_contiguous()
+    actual = triton_ihc_pre(x, noncontiguous_weight, scale, base, 2.0, 1e-6, 1e-5)
+    expected = _pre_reference(x, noncontiguous_weight, scale, base, 2.0, 1e-6, 1e-5)
+    torch.testing.assert_close(actual[0], expected[0], rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(actual[1], expected[1], rtol=2e-3, atol=2e-3)
+
+
+def test_rocm_triton_ihc_post_rejects_invalid_gate_shape():
+    x = torch.empty((2, 4096), device="cuda", dtype=_DTYPE)
+    residual = torch.empty((2, _HC_MULT, 4096), device="cuda", dtype=_DTYPE)
+    post = torch.empty((1, _HC_MULT), device="cuda", dtype=torch.float32)
+    with pytest.raises(AssertionError):
+        triton_ihc_post(x, residual, post)
+
+
+def test_rocm_triton_ihc_fused_compiles_and_captures():
+    inputs = _make_fused_inputs(2, 4096, _DTYPE)
+    args = (*inputs, 2.0, 1e-6, 1e-5)
+    compiled = torch.compile(triton_ihc_post_pre_rms_norm, fullgraph=True)
+    eager = triton_ihc_post_pre_rms_norm(*args)
+    compiled_output = compiled(*args)
+    for actual, expected in zip(compiled_output, eager):
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+
+    for _ in range(2):
+        triton_ihc_post_pre_rms_norm(*args)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = triton_ihc_post_pre_rms_norm(*args)
+    graph.replay()
+    torch.accelerator.synchronize()
+    for actual, expected in zip(captured, eager):
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)

tests/models/quantization/test_mxfp8.py

@@ -92,7 +92,7 @@ def test_mxfp8_logprobs(
     reason="AITER MXFP8 MoE backend is ROCm-only.",
 )
 @pytest.mark.quant_model
-def test_mxfp8_aiter_requires_swigluoai_activation(
+def test_mxfp8_aiter_accepts_silu_activation(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -127,8 +127,12 @@ def test_mxfp8_aiter_requires_swigluoai_activation(
         moe_backend="aiter",
     )
 
-    with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"):
-        select_mxfp8_moe_backend(config)
+    from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
+
+    backend, experts_cls = select_mxfp8_moe_backend(config)
+
+    assert backend == Fp8MoeBackend.AITER_MXFP8
+    assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts
 
 
 @pytest.mark.skipif(

vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py

@@ -45,7 +45,14 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
 
     @staticmethod
     def _supports_current_device() -> bool:
-        return current_platform.supports_mx() and rocm_aiter_ops.is_fused_moe_enabled()
+        # ROCm ``supports_mx`` deliberately excludes gfx942, so gfx942 may
+        # import this module during backend discovery but never selects or
+        # executes the AITER MXFP8 experts path.
+        return (
+            current_platform.supports_mx()
+            and rocm_aiter_ops.is_fused_moe_enabled()
+            and rocm_aiter_ops.fused_moe_supports_gate_mode()
+        )
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config) -> bool:
@@ -62,20 +69,23 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
         is_supported, reason = super().is_supported_config(
             cls, moe_config, weight_key, activation_key, activation_format
         )
-        if (
-            is_supported
-            and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE
+        if is_supported and moe_config.activation not in (
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
         ):
             return False, (
-                "kernel hardcodes SwiGLU-OAI activation and requires "
-                f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; "
-                f"got activation={moe_config.activation.value}"
+                "kernel supports only SiLU or SwiGLU-OAI activation; got "
+                f"activation={moe_config.activation.value}"
+            )
+        if (
+            is_supported
+            and moe_config.activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
+            and (
+                moe_config.swiglu_alpha is None
+                or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA)
+                or moe_config.swiglu_beta is None
+                or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA)
             )
-        if is_supported and (
-            moe_config.swiglu_alpha is None
-            or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA)
-            or moe_config.swiglu_beta is None
-            or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA)
         ):
             return False, (
                 "kernel hardcodes SwiGLU-OAI with "
@@ -93,7 +103,7 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
         w2: torch.Tensor,
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
-        activation,
+        activation: MoEActivation,
         global_num_experts: int,
         expert_map: torch.Tensor | None,
         a1q_scale: torch.Tensor | None,
@@ -110,6 +120,16 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
 
         limit = self.quant_config.gemm1_clamp_limit
         swiglu_limit = 0.0 if limit is None else float(limit)
+        if activation == MoEActivation.SILU:
+            activation_method = ActivationType.Silu.value
+        elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
+            activation_method = ActivationType.Swiglu.value
+        else:
+            activation_name = getattr(activation, "value", activation)
+            raise ValueError(
+                "AITER MXFP8 MoE supports only SiLU or packed SwiGLU-OAI; "
+                f"got activation={activation_name!r}"
+            )
 
         # RoutedExperts.expert_map hands AITER experts the precomputed 0/1
         # expert_mask (with trailing sentinel) instead of the vLLM expert_map.
@@ -126,13 +146,16 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
             topk_weights.to(torch.float32),
             topk_ids.to(torch.int32),
             expert_mask=expert_mask,
-            activation_method=ActivationType.Swiglu.value,
+            activation_method=activation_method,
             quant_method=QuantType.per_1x32.value,
             doweight_stage1=apply_router_weight_on_input,
             w1_scale=self.w1_scale_val,
             w2_scale=self.w2_scale_val,
             a1_scale=None,
             a2_scale=None,
+            # MXFP8 w1 and its scales are preshuffled into AITER's GUGU layout
+            # during model loading. Declare that layout explicitly so AITER can
+            # dispatch the gfx950-optimized FlyDSL kernel instead of CK.
             gate_mode=GateMode.INTERLEAVE.value,
             swiglu_limit=swiglu_limit,
             output_dtype=output.dtype,

vllm/models/hy_v4/amd/hc.py

@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""iHC (independent Hyper-Connections) layers for HY V4 (ROCm).
+
+iHC replaces the single residual stream of a standard transformer with
+``hc_mult`` parallel residual channels. Each decoder sub-block reduces the
+channels to one hidden state (``HYV4HCPreLayer``), runs the sub-block, then
+scatters the result back over the channels (``HYV4HCPostLayer``). The final
+``HYV4HCHeadLayer`` merges the channels before the model's output norm.
+
+NOTE: Each of the three steps has an optional single-kernel HPC replacement
+(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). Pre and post fall back to
+in-tree Triton kernels on ROCm when HPC is unavailable, then to the eager path.
+The AMD Triton implementation also provides a fused cross-layer post+pre path.
+"""
+
+import torch
+from torch import nn
+from transformers.configuration_utils import PreTrainedConfig
+
+from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre
+from vllm.model_executor.layers.linear import ReplicatedLinear
+
+from .triton_ihc import (
+    triton_ihc_post,
+    triton_ihc_pre,
+    triton_ihc_supported,
+)
+
+
+class HYV4HCPreLayer(nn.Module):
+    """iHC pre-processing layer (2D-activation adaptation).
+
+    Steps:
+        1. RMS-normalize the flattened ``[num_tokens, hc * d]`` input.
+        2. Project to the pre/post gating logits.
+        3. Turn the logits into sigmoid gates.
+        4. Reduce over the channel dim with the pre gates.
+    """
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        hidden_dim: int,
+        hc_mult: int = 4,
+        magnitude: float = 2.0,
+        init_std: float = 6e-3,
+        base_noise_std: float = 0.0,
+        hc_eps: float = 1e-6,
+        layernorm_epsilon: float = 1e-5,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_dim = hidden_dim
+        self.hc_mult = hc_mult
+        self.magnitude = magnitude
+        self.hc_eps = hc_eps
+        self.layernorm_epsilon = layernorm_epsilon
+        hc_in_dim = hc_mult * hidden_dim
+        mix_hc = 2 * hc_mult  # pre + post only, no comb
+
+        self.hc_fn = ReplicatedLinear(
+            input_size=hc_in_dim,
+            output_size=mix_hc,
+            params_dtype=torch.float32,
+            bias=False,
+            prefix=f"{prefix}.hc_fn",
+        )
+        self.hc_scale = nn.Parameter(torch.empty(2, dtype=torch.float32))
+        self.hc_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
+
+        self.reset_parameters(init_std, base_noise_std)
+
+        # Optional single-kernel HPC replacement for the whole forward below.
+        # ``norm_owner`` is left unset: the RMSNorm that follows the pre block
+        # lives in the decoder layer and is not folded in yet.
+        self.hpc_op: HpcIHCPre | None = None
+        if HpcIHCPre.support(hc_mult, hidden_dim):
+            self.hpc_op = HpcIHCPre(
+                hc_mult=hc_mult,
+                hidden_size=hidden_dim,
+                magnitude=magnitude,
+                hc_eps=hc_eps,
+                norm_eps=layernorm_epsilon,
+                fallback_op=self,
+            )
+
+    def reset_parameters(self, init_std: float, base_noise_std: float = 0.0) -> None:
+        """Initialize the gate scale and per-channel gate bias."""
+        del init_std  # hc_fn is initialized by ReplicatedLinear
+        nn.init.constant_(self.hc_scale, 0.01)
+        with torch.no_grad():
+            self.hc_base[: self.hc_mult].fill_(
+                -torch.log(torch.tensor(self.hc_mult - 1.0, dtype=self.hc_base.dtype))
+            )
+            self.hc_base[self.hc_mult : 2 * self.hc_mult].fill_(0.0)
+            if base_noise_std > 0.0:
+                self.hc_base.add_(torch.randn_like(self.hc_base) * base_noise_std)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Reduce the iHC channels and emit the post gates.
+
+        Args:
+            x: Input of shape ``[num_tokens, hc, d]``.
+
+        Returns:
+            A tuple of the pre-gated reduction ``[num_tokens, d]`` and the post
+            gates ``[num_tokens, hc]`` consumed by `HYV4HCPostLayer`.
+
+        """
+        if self.hpc_op is not None:
+            return self.hpc_op(x)
+        if triton_ihc_supported(x):
+            return triton_ihc_pre(
+                x,
+                self.hc_fn.weight,
+                self.hc_scale,
+                self.hc_base,
+                self.magnitude,
+                self.hc_eps,
+                self.layernorm_epsilon,
+            )
+
+        shape = x.size()  # [num_tokens, hc, d]
+        hc = self.hc_mult
+        hc_eps = self.hc_eps
+
+        x_flat = x.flatten(1).float()  # [num_tokens, hc*d]
+
+        rsqrt = torch.rsqrt(
+            x_flat.square().mean(-1, keepdim=True) + self.layernorm_epsilon
+        )
+        mixes = self.hc_fn(x_flat)[0] * rsqrt  # [num_tokens, 2*hc]
+
+        pre_raw = mixes[..., :hc]
+        post_raw = mixes[..., hc : 2 * hc]
+
+        pre = (
+            torch.sigmoid(
+                pre_raw * self.hc_scale[0].float() + self.hc_base[:hc].float()
+            )
+            + hc_eps
+        )
+        post = (
+            self.magnitude
+            * torch.sigmoid(
+                post_raw * self.hc_scale[1].float() + self.hc_base[hc : 2 * hc].float()
+            )
+            + hc_eps
+        )
+
+        y = torch.sum(pre.unsqueeze(-1) * x.reshape(shape), dim=1)  # [num_tokens, d]
+        return y.to(x.dtype), post
+
+
+class HYV4HCPostLayer(nn.Module):
+    """iHC post-processing layer (2D-activation adaptation).
+
+    Applies post-gating to the sub-block output and adds the multi-channel
+    residual (no comb mixing)::
+
+        y[n, i, d] = post[n, i] * x[n, d] + residual[n, i, d]
+    """
+
+    def __init__(self, config: PreTrainedConfig):
+        super().__init__()
+        self.config = config
+
+        # Optional single-kernel HPC replacement for the whole forward below.
+        # Only constructed under enable_ihc (see HYV4HCLayer), so hc_mult is
+        # present; getattr keeps this robust if that ever changes.
+        self.hpc_op: HpcIHCPost | None = None
+        hc_mult = getattr(config, "hc_mult", 0)
+        if HpcIHCPost.support(hc_mult, config.hidden_size):
+            self.hpc_op = HpcIHCPost(hc_mult=hc_mult, hidden_size=config.hidden_size)
+
+    def forward(
+        self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor
+    ) -> torch.Tensor:
+        """Scatter the sub-block output back onto the iHC channels.
+
+        Args:
+            x: Attention/MLP output of shape ``[num_tokens, d]``.
+            residual: Multi-channel residual ``[num_tokens, hc, d]``.
+            post: Post gates ``[num_tokens, hc]`` from `HYV4HCPreLayer`.
+
+        Returns:
+            The updated residual channels ``[num_tokens, hc, d]``.
+
+        """
+        if self.hpc_op is not None:
+            return self.hpc_op(x, residual, post)
+        if triton_ihc_supported(x):
+            return triton_ihc_post(x, residual, post)
+
+        dtype = x.dtype
+        x = x.float()
+        residual = residual.float()
+        post = post.float()
+
+        post_gated = post.unsqueeze(-1) * x.unsqueeze(-2)  # [num_tokens, hc, d]
+        y = post_gated + residual
+        return y.to(dtype)
+
+
+class HYV4HCHeadLayer(nn.Module):
+    """iHC head layer (2D-activation adaptation).
+
+    Merges the iHC channels back into a single hidden state before the final
+    layer norm, using an RMS-normed projection plus sigmoid-gated reduction.
+    """
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        hidden_size: int,
+        hc_mult: int = 4,
+        hc_eps: float = 1e-6,
+        init_std: float = 6e-3,
+        base_noise_std: float = 0.0,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = hidden_size
+        self.hc_mult = hc_mult
+        self.hc_eps = hc_eps
+        self.hc_head_fn = ReplicatedLinear(
+            input_size=hc_mult * hidden_size,
+            output_size=hc_mult,
+            params_dtype=torch.float32,
+            bias=False,
+            prefix=f"{prefix}.hc_head_fn",
+        )
+        self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32))
+        self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32))
+
+        self.reset_parameters(init_std, base_noise_std)
+
+        # Optional single-kernel HPC replacement for the whole forward below.
+        self.hpc_op: HpcIHCHead | None = None
+        if HpcIHCHead.support(hc_mult, hidden_size):
+            self.hpc_op = HpcIHCHead(
+                hc_mult=hc_mult,
+                hidden_size=hidden_size,
+                hc_eps=hc_eps,
+                norm_eps=config.rms_norm_eps,
+                fallback_op=self,
+            )
+
+    def reset_parameters(
+        self, init_std: float = 6e-3, base_noise_std: float = 0.0
+    ) -> None:
+        """Initialize the head gate scale and per-channel gate bias."""
+        del init_std  # hc_head_fn is initialized by ReplicatedLinear
+        nn.init.constant_(self.hc_head_scale, 0.01)
+        with torch.no_grad():
+            self.hc_head_base.fill_(
+                -torch.log(
+                    torch.tensor(self.hc_mult - 1.0, dtype=self.hc_head_base.dtype)
+                )
+            )
+            if base_noise_std > 0.0:
+                self.hc_head_base.add_(
+                    torch.randn_like(self.hc_head_base) * base_noise_std
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Merge the iHC channels into a single hidden state.
+
+        Args:
+            x: Input of shape ``[num_tokens, hc, d]``.
+
+        Returns:
+            The merged hidden state ``[num_tokens, d]``.
+
+        """
+        if self.hpc_op is not None:
+            return self.hpc_op(x)
+
+        shape, x_dtype = x.size(), x.dtype
+
+        x = x.flatten(1).float()  # [num_tokens, hc*d]
+        rsqrt = torch.rsqrt(
+            x.square().mean(-1, keepdim=True) + self.config.rms_norm_eps
+        )
+        mixes = self.hc_head_fn(x)[0] * rsqrt  # [num_tokens, hc]
+        pre = (
+            torch.sigmoid(
+                mixes * self.hc_head_scale.float() + self.hc_head_base.float()
+            )
+            + self.hc_eps
+        )
+        y = torch.sum(pre.unsqueeze(-1) * x.reshape(shape), dim=1)  # [num_tokens, d]
+        return y.to(x_dtype)
+
+
+class HYV4HCLayer(nn.Module):
+    """Wrapper owning one iHC boundary (pre + post) of a decoder sub-block."""
+
+    def __init__(
+        self,
+        config: PreTrainedConfig,
+        layer_idx: int,
+        init_std: float = 6e-3,
+        base_noise_std: float = 0.0,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.enable_ihc = getattr(config, "enable_ihc", False)
+        if self.enable_ihc:
+            self.hc_pre = HYV4HCPreLayer(
+                config,
+                config.hidden_size,
+                config.hc_mult,
+                config.hc_magnitude,
+                init_std,
+                base_noise_std,
+                config.hc_eps,
+                config.rms_norm_eps,
+                prefix=f"{prefix}.hc_pre",
+            )
+            self.hc_post = HYV4HCPostLayer(config)
+
+    def prepare_input(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Normalize the sub-block input to 3D when iHC is enabled."""
+        if not self.enable_ihc:
+            return hidden_states
+        return self._prepare_input_to_3d(hidden_states)
+
+    def _prepare_input_to_3d(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Reshape the iHC input to ``[num_tokens, hc, h]``.
+
+        Accepted inputs are ``[num_tokens, hc, h]`` (no-op), ``[num_tokens, h]``
+        (broadcast over the channels) and ``[num_tokens, hc * h]`` (reshape).
+        """
+        if hidden_states.dim() == 3:
+            return hidden_states
+        if hidden_states.dim() != 2:
+            raise RuntimeError(
+                f"HC expects a 2D/3D tensor, got shape={tuple(hidden_states.shape)}"
+            )
+
+        d0, d1 = hidden_states.shape
+        h = self.config.hidden_size
+        hc = self.config.hc_mult
+
+        if d1 == h:
+            return hidden_states.unsqueeze(1).repeat(1, hc, 1)
+
+        expected = hc * h
+        if d1 == expected:
+            return hidden_states.reshape(d0, hc, h)
+
+        raise RuntimeError(
+            f"HC expects last dim to be hidden_size ({h})"
+            f"or hc_mult*hidden_size ({expected}), got {d1}. "
+            f"hc_mult={hc}, hidden_size={h}."
+        )
+
+    def pre(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
+        """Reduce the iHC channels and produce the post gates.
+
+        Returns:
+            A tuple of the reduced hidden states ``[num_tokens, d]``, the post
+            gates ``[num_tokens, hc]`` (``None`` when iHC is disabled) and the
+            residual (the untouched input).
+
+        """
+        if not self.enable_ihc:
+            return hidden_states, None, hidden_states
+        reduced, post_gates = self.hc_pre(hidden_states)
+        return reduced, post_gates, hidden_states
+
+    def post(
+        self,
+        output_with_bias: torch.Tensor,
+        residual: torch.Tensor,
+        post_gates: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Apply post-gating and add the residual."""
+        if not self.enable_ihc:
+            return output_with_bias + residual
+        assert post_gates is not None
+        return self.hc_post(output_with_bias, residual, post_gates)

vllm/models/hy_v4/amd/model.py

@@ -67,8 +67,13 @@ from ..nvidia.attention import (
     compute_skip_topk_layers,
     is_skip_topk_indexer_weight,
 )
-from ..nvidia.hc import HYV4HCHeadLayer, HYV4HCLayer
 from ..nvidia.moe import HYV4FeedForward, HYV4MoEFused
+from .hc import (
+    HYV4HCHeadLayer,
+    HYV4HCLayer,
+    HYV4HCPreLayer,
+)
+from .triton_ihc import triton_ihc_post_pre_rms_norm, triton_ihc_supported
 
 logger = init_logger(__name__)
 
@@ -192,15 +197,35 @@ class HYV4DecoderLayer(nn.Module):
         hidden_states = self.hc_attn_layer.prepare_input(hidden_states)
         hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states)
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(
+        attention_output = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
         )
-        hidden_states = self.hc_attn_layer.post(hidden_states, residual, post_gates)
-
-        hidden_states = self.hc_mlp_layer.prepare_input(hidden_states)
-        hidden_states, post_gates, residual = self.hc_mlp_layer.pre(hidden_states)
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        can_fuse_ihc_post_pre_rms_norm = (
+            triton_ihc_supported(attention_output)
+            and self.hc_attn_layer.hc_post.hpc_op is None
+            and self.hc_mlp_layer.hc_pre.hpc_op is None
+        )
+        if can_fuse_ihc_post_pre_rms_norm:
+            hidden_states, post_gates, residual = triton_ihc_post_pre_rms_norm(
+                attention_output,
+                residual,
+                post_gates,
+                self.hc_mlp_layer.hc_pre.hc_fn.weight,
+                self.hc_mlp_layer.hc_pre.hc_scale,
+                self.hc_mlp_layer.hc_pre.hc_base,
+                self.post_attention_layernorm.weight,
+                self.hc_mlp_layer.hc_pre.magnitude,
+                self.hc_mlp_layer.hc_pre.hc_eps,
+                self.hc_mlp_layer.hc_pre.layernorm_epsilon,
+            )
+        else:
+            hidden_states = self.hc_attn_layer.post(
+                attention_output, residual, post_gates
+            )
+            hidden_states = self.hc_mlp_layer.prepare_input(hidden_states)
+            hidden_states, post_gates, residual = self.hc_mlp_layer.pre(hidden_states)
+            hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = self.hc_mlp_layer.post(hidden_states, residual, post_gates)
 
@@ -598,6 +623,15 @@ class HYV4Model(nn.Module):
 
         return loaded_params
 
+    @torch.no_grad()
+    def process_weights_after_loading(self) -> None:
+        """Make static iHC weights contiguous for the ROCm Triton kernels."""
+        for layer in self.modules():
+            if isinstance(layer, HYV4HCPreLayer):
+                layer.hc_fn.weight.data = layer.hc_fn.weight.data.contiguous()
+            elif isinstance(layer, HYV4HCHeadLayer):
+                layer.hc_head_fn.weight.data = layer.hc_head_fn.weight.data.contiguous()
+
 
 class HYV4ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
     hf_to_vllm_mapper = WeightsMapper(
@@ -694,5 +728,8 @@ class HYV4ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper | drop_mapper)
 
+    def process_weights_after_loading(self) -> None:
+        self.model.process_weights_after_loading()
+
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         return self.model.get_expert_mapping()

vllm/models/hy_v4/amd/triton_ihc.py

@@ -0,0 +1,550 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""ROCm Triton iHC kernels for HY V4.
+
+This module is intentionally separate from the CUDA implementation.  The
+kernels use fp32 reductions and fixed launch shapes that are suitable for
+HIP/Triton graph capture.
+"""
+
+import torch
+
+from vllm import envs
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+_BLOCK_K = 1024
+_BLOCK_D = 1024
+_PRE_STAGE1_WARPS = 8
+_PRE_STAGE2_WARPS = 4
+_POST_WARPS = 4
+
+
+def triton_ihc_supported(x: torch.Tensor) -> bool:
+    return (
+        HAS_TRITON
+        and current_platform.is_rocm()
+        and x.is_cuda
+        and x.dtype in (torch.float16, torch.bfloat16)
+        and not envs.VLLM_BATCH_INVARIANT
+    )
+
+
+@triton.jit
+def _ihc_pre_stage1(
+    x_ptr,
+    weight_ptr,
+    partial_ptr,
+    K_TOTAL: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    HC_POW2: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    PARTIAL_STRIDE: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    split_idx = tl.program_id(1)
+    hc_idx = tl.arange(0, HC_POW2)
+    hc_mask = hc_idx < HC_MULT
+    k_offsets = split_idx * BLOCK_K + tl.arange(0, BLOCK_K)
+    k_mask = k_offsets < K_TOTAL
+    x = tl.load(x_ptr + token_idx * K_TOTAL + k_offsets, mask=k_mask, other=0.0).to(
+        tl.float32
+    )
+    weight_offsets = hc_idx[:, None] * K_TOTAL + k_offsets[None, :]
+    weight_mask = hc_mask[:, None] & k_mask[None, :]
+    pre_mix = tl.sum(
+        tl.load(
+            weight_ptr + weight_offsets,
+            mask=weight_mask,
+            other=0.0,
+        )
+        * x[None, :],
+        axis=1,
+    )
+    post_mix = tl.sum(
+        tl.load(
+            weight_ptr + HC_MULT * K_TOTAL + weight_offsets,
+            mask=weight_mask,
+            other=0.0,
+        )
+        * x[None, :],
+        axis=1,
+    )
+    partial = partial_ptr + (token_idx * NUM_SPLITS + split_idx) * PARTIAL_STRIDE
+    tl.store(partial, tl.sum(x * x, axis=0))
+    tl.store(partial + 1 + hc_idx, pre_mix, mask=hc_mask)
+    tl.store(partial + 1 + HC_POW2 + hc_idx, post_mix, mask=hc_mask)
+
+
+@triton.jit
+def _ihc_pre_stage2(
+    x_ptr,
+    partial_ptr,
+    scale_ptr,
+    base_ptr,
+    output_ptr,
+    post_ptr,
+    HIDDEN_SIZE: tl.constexpr,
+    K_TOTAL: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    HC_POW2: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
+    PARTIAL_STRIDE: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    MAGNITUDE: tl.constexpr,
+    NORM_EPS: tl.constexpr,
+    HC_EPS: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    hidden_block_idx = tl.program_id(1)
+    hc_idx = tl.arange(0, HC_POW2)
+    hc_mask = hc_idx < HC_MULT
+    partial_row = partial_ptr + token_idx * NUM_SPLITS * PARTIAL_STRIDE
+    sum_squares = tl.zeros((), dtype=tl.float32)
+    pre_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
+    post_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
+    for split_idx in tl.static_range(NUM_SPLITS):
+        partial = partial_row + split_idx * PARTIAL_STRIDE
+        sum_squares += tl.load(partial)
+        pre_mix += tl.load(partial + 1 + hc_idx, mask=hc_mask, other=0.0)
+        post_mix += tl.load(partial + 1 + HC_POW2 + hc_idx, mask=hc_mask, other=0.0)
+
+    x_row = x_ptr + token_idx * K_TOTAL
+    reciprocal_rms = tl.rsqrt(sum_squares / K_TOTAL + NORM_EPS)
+    pre_scale = tl.load(scale_ptr)
+    post_scale = tl.load(scale_ptr + 1)
+    pre_base = tl.load(base_ptr + hc_idx, mask=hc_mask, other=0.0)
+    post_base = tl.load(base_ptr + HC_MULT + hc_idx, mask=hc_mask, other=0.0)
+    pre = tl.sigmoid(pre_mix * reciprocal_rms * pre_scale + pre_base) + HC_EPS
+    post = (
+        MAGNITUDE * tl.sigmoid(post_mix * reciprocal_rms * post_scale + post_base)
+        + HC_EPS
+    )
+    if hidden_block_idx == 0:
+        tl.store(post_ptr + token_idx * HC_MULT + hc_idx, post, mask=hc_mask)
+
+    hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
+    hidden_mask = hidden_offsets < HIDDEN_SIZE
+    output = tl.zeros((BLOCK_D,), dtype=tl.float32)
+    for channel_idx in tl.static_range(HC_MULT):
+        channel = tl.load(
+            x_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
+            mask=hidden_mask,
+            other=0.0,
+        ).to(tl.float32)
+        gate = tl.sum(tl.where(hc_idx == channel_idx, pre, 0.0), axis=0)
+        output += gate * channel
+    tl.store(
+        output_ptr + token_idx * HIDDEN_SIZE + hidden_offsets,
+        output.to(output_ptr.dtype.element_ty),
+        mask=hidden_mask,
+    )
+
+
+@triton.jit
+def _ihc_post_kernel(
+    x_ptr,
+    residual_ptr,
+    post_ptr,
+    output_ptr,
+    HIDDEN_SIZE: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    HC_POW2: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    hidden_block_idx = tl.program_id(1)
+    hc_idx = tl.arange(0, HC_POW2)
+    hc_mask = hc_idx < HC_MULT
+    post = tl.load(post_ptr + token_idx * HC_MULT + hc_idx, mask=hc_mask, other=0.0).to(
+        tl.float32
+    )
+    offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
+    mask = offsets < HIDDEN_SIZE
+    x = tl.load(x_ptr + token_idx * HIDDEN_SIZE + offsets, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    residual_row = residual_ptr + token_idx * HC_MULT * HIDDEN_SIZE
+    output_row = output_ptr + token_idx * HC_MULT * HIDDEN_SIZE
+    for channel_idx in tl.static_range(HC_MULT):
+        residual = tl.load(
+            residual_row + channel_idx * HIDDEN_SIZE + offsets,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        gate = tl.sum(tl.where(hc_idx == channel_idx, post, 0.0), axis=0)
+        tl.store(
+            output_row + channel_idx * HIDDEN_SIZE + offsets,
+            (gate * x + residual).to(output_ptr.dtype.element_ty),
+            mask=mask,
+        )
+
+
+def triton_ihc_pre(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    magnitude: float,
+    hc_eps: float,
+    norm_eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert x.dim() == 3 and x.is_cuda
+    assert x.dtype in (torch.float16, torch.bfloat16)
+    assert weight.dtype == torch.float32 and scale.dtype == torch.float32
+    assert base.dtype == torch.float32
+    assert weight.is_cuda and scale.is_cuda and base.is_cuda
+    assert all(tensor.device == x.device for tensor in (weight, scale, base))
+    assert weight.ndim == 2
+    # Static iHC weights are made contiguous once by
+    # HYV4Model.process_weights_after_loading. Keep this guard for reloads or
+    # alternate loaders that may provide a non-contiguous parameter.
+    x = x.contiguous()
+    scale, base = scale.contiguous(), base.contiguous()
+    num_tokens, hc_mult, hidden_size = x.shape
+    k_total = hc_mult * hidden_size
+    hc_pow2 = triton.next_power_of_2(hc_mult)
+    weight = weight.contiguous()
+    assert weight.shape == (2 * hc_mult, k_total)
+    assert scale.shape == (2,) and base.shape == (2 * hc_mult,)
+    if num_tokens == 0:
+        return (
+            torch.empty((0, hidden_size), dtype=x.dtype, device=x.device),
+            torch.empty((0, hc_mult), dtype=torch.float32, device=x.device),
+        )
+    output = torch.empty((num_tokens, hidden_size), dtype=x.dtype, device=x.device)
+    post = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=x.device)
+    num_splits = triton.cdiv(k_total, _BLOCK_K)
+    partial_stride = 1 + 2 * hc_pow2
+    partial = torch.empty(
+        (num_tokens, num_splits, partial_stride),
+        dtype=torch.float32,
+        device=x.device,
+    )
+    _ihc_pre_stage1[(num_tokens, num_splits)](
+        x,
+        weight,
+        partial,
+        K_TOTAL=k_total,
+        HC_MULT=hc_mult,
+        HC_POW2=hc_pow2,
+        NUM_SPLITS=num_splits,
+        BLOCK_K=_BLOCK_K,
+        PARTIAL_STRIDE=partial_stride,
+        num_warps=_PRE_STAGE1_WARPS,
+        enable_fp_fusion=False,
+    )
+    _ihc_pre_stage2[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))](
+        x,
+        partial,
+        scale,
+        base,
+        output,
+        post,
+        HIDDEN_SIZE=hidden_size,
+        K_TOTAL=k_total,
+        HC_MULT=hc_mult,
+        HC_POW2=hc_pow2,
+        NUM_SPLITS=num_splits,
+        PARTIAL_STRIDE=partial_stride,
+        BLOCK_D=_BLOCK_D,
+        MAGNITUDE=magnitude,
+        NORM_EPS=norm_eps,
+        HC_EPS=hc_eps,
+        num_warps=_PRE_STAGE2_WARPS,
+        enable_fp_fusion=False,
+    )
+    return output, post
+
+
+@triton.jit
+def _ihc_post_pre_rms_norm_stage1(
+    x_ptr,
+    residual_ptr,
+    attn_post_ptr,
+    weight_ptr,
+    partial_ptr,
+    y_ptr,
+    K_TOTAL: tl.constexpr,
+    HIDDEN_SIZE: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    HC_POW2: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    PARTIAL_STRIDE: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    split_idx = tl.program_id(1)
+    hc_idx = tl.arange(0, HC_POW2)
+    hc_mask = hc_idx < HC_MULT
+    offsets = split_idx * BLOCK_K + tl.arange(0, BLOCK_K)
+    k_mask = offsets < K_TOTAL
+    channel_idx = offsets // HIDDEN_SIZE
+    hidden_idx = offsets % HIDDEN_SIZE
+    gates = tl.load(
+        attn_post_ptr + token_idx * HC_MULT + channel_idx,
+        mask=k_mask,
+        other=0.0,
+    ).to(tl.float32)
+    x = tl.load(
+        x_ptr + token_idx * HIDDEN_SIZE + hidden_idx,
+        mask=k_mask,
+        other=0.0,
+    ).to(tl.float32)
+    residual = tl.load(
+        residual_ptr + token_idx * K_TOTAL + offsets,
+        mask=k_mask,
+        other=0.0,
+    ).to(tl.float32)
+    y = gates * x + residual
+    y = y.to(y_ptr.dtype.element_ty)
+    tl.store(y_ptr + token_idx * K_TOTAL + offsets, y, mask=k_mask)
+    y = y.to(tl.float32)
+
+    weight_offsets = hc_idx[:, None] * K_TOTAL + offsets[None, :]
+    weight_mask = hc_mask[:, None] & k_mask[None, :]
+    pre_mix = tl.sum(
+        tl.load(weight_ptr + weight_offsets, mask=weight_mask, other=0.0) * y[None, :],
+        axis=1,
+    )
+    post_mix = tl.sum(
+        tl.load(
+            weight_ptr + HC_MULT * K_TOTAL + weight_offsets,
+            mask=weight_mask,
+            other=0.0,
+        )
+        * y[None, :],
+        axis=1,
+    )
+    partial = partial_ptr + (token_idx * NUM_SPLITS + split_idx) * PARTIAL_STRIDE
+    tl.store(partial, tl.sum(y * y, axis=0))
+    tl.store(partial + 1 + hc_idx, pre_mix, mask=hc_mask)
+    tl.store(partial + 1 + HC_POW2 + hc_idx, post_mix, mask=hc_mask)
+
+
+@triton.jit
+def _ihc_post_pre_rms_norm_stage2(
+    y_ptr,
+    partial_ptr,
+    scale_ptr,
+    base_ptr,
+    norm_weight_ptr,
+    output_ptr,
+    post_ptr,
+    HIDDEN_SIZE: tl.constexpr,
+    K_TOTAL: tl.constexpr,
+    HC_MULT: tl.constexpr,
+    HC_POW2: tl.constexpr,
+    NUM_SPLITS: tl.constexpr,
+    PARTIAL_STRIDE: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    NUM_HIDDEN_BLOCKS: tl.constexpr,
+    MAGNITUDE: tl.constexpr,
+    NORM_EPS: tl.constexpr,
+    HC_EPS: tl.constexpr,
+):
+    token_idx = tl.program_id(0).to(tl.int64)
+    hc_idx = tl.arange(0, HC_POW2)
+    hc_mask = hc_idx < HC_MULT
+    partial_row = partial_ptr + token_idx * NUM_SPLITS * PARTIAL_STRIDE
+    sum_squares = tl.zeros((), dtype=tl.float32)
+    pre_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
+    post_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
+    for split_idx in tl.static_range(NUM_SPLITS):
+        partial = partial_row + split_idx * PARTIAL_STRIDE
+        sum_squares += tl.load(partial)
+        pre_mix += tl.load(partial + 1 + hc_idx, mask=hc_mask, other=0.0)
+        post_mix += tl.load(partial + 1 + HC_POW2 + hc_idx, mask=hc_mask, other=0.0)
+
+    reciprocal_rms = tl.rsqrt(sum_squares / K_TOTAL + NORM_EPS)
+    pre_scale = tl.load(scale_ptr)
+    post_scale = tl.load(scale_ptr + 1)
+    pre_base = tl.load(base_ptr + hc_idx, mask=hc_mask, other=0.0)
+    post_base = tl.load(base_ptr + HC_MULT + hc_idx, mask=hc_mask, other=0.0)
+    pre = tl.sigmoid(pre_mix * reciprocal_rms * pre_scale + pre_base) + HC_EPS
+    post = (
+        MAGNITUDE * tl.sigmoid(post_mix * reciprocal_rms * post_scale + post_base)
+        + HC_EPS
+    )
+    tl.store(post_ptr + token_idx * HC_MULT + hc_idx, post, mask=hc_mask)
+
+    y_row = y_ptr + token_idx * K_TOTAL
+    output_sum_squares = tl.zeros((), dtype=tl.float32)
+    for hidden_block_idx in tl.static_range(NUM_HIDDEN_BLOCKS):
+        offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
+        hidden_mask = offsets < HIDDEN_SIZE
+        output = tl.zeros((BLOCK_D,), dtype=tl.float32)
+        for channel_idx in tl.static_range(HC_MULT):
+            channel = tl.load(
+                y_row + channel_idx * HIDDEN_SIZE + offsets,
+                mask=hidden_mask,
+                other=0.0,
+            ).to(tl.float32)
+            gate = tl.sum(tl.where(hc_idx == channel_idx, pre, 0.0), axis=0)
+            output += gate * channel
+        output = output.to(output_ptr.dtype.element_ty).to(tl.float32)
+        output_sum_squares += tl.sum(
+            tl.where(hidden_mask, output * output, 0.0), axis=0
+        )
+
+    reciprocal_rms = tl.rsqrt(output_sum_squares / HIDDEN_SIZE + NORM_EPS)
+    for hidden_block_idx in tl.static_range(NUM_HIDDEN_BLOCKS):
+        offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
+        hidden_mask = offsets < HIDDEN_SIZE
+        output = tl.zeros((BLOCK_D,), dtype=tl.float32)
+        for channel_idx in tl.static_range(HC_MULT):
+            channel = tl.load(
+                y_row + channel_idx * HIDDEN_SIZE + offsets,
+                mask=hidden_mask,
+                other=0.0,
+            ).to(tl.float32)
+            gate = tl.sum(tl.where(hc_idx == channel_idx, pre, 0.0), axis=0)
+            output += gate * channel
+        output = output.to(output_ptr.dtype.element_ty).to(tl.float32)
+        norm_weight = tl.load(
+            norm_weight_ptr + offsets, mask=hidden_mask, other=0.0
+        ).to(tl.float32)
+        tl.store(
+            output_ptr + token_idx * HIDDEN_SIZE + offsets,
+            (output * reciprocal_rms * norm_weight).to(output_ptr.dtype.element_ty),
+            mask=hidden_mask,
+        )
+
+
+def triton_ihc_post_pre_rms_norm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    attn_post: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    norm_weight: torch.Tensor,
+    magnitude: float,
+    hc_eps: float,
+    norm_eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    assert x.dim() == 2 and residual.dim() == 3 and attn_post.dim() == 2
+    assert x.is_cuda and residual.is_cuda and attn_post.is_cuda
+    assert weight.is_cuda and scale.is_cuda and base.is_cuda
+    assert norm_weight.is_cuda
+    assert all(
+        tensor.device == x.device
+        for tensor in (residual, attn_post, weight, scale, base, norm_weight)
+    )
+    assert x.dtype in (torch.float16, torch.bfloat16)
+    assert residual.dtype == x.dtype and attn_post.dtype == torch.float32
+    assert weight.dtype == torch.float32 and scale.dtype == torch.float32
+    assert base.dtype == torch.float32
+    x, residual, attn_post = (
+        x.contiguous(),
+        residual.contiguous(),
+        attn_post.contiguous(),
+    )
+    weight, scale, base = (
+        weight.contiguous(),
+        scale.contiguous(),
+        base.contiguous(),
+    )
+    # Normal serving makes this loaded parameter contiguous in the post-load
+    # hook; keep this guard for direct callers and alternate loading paths.
+    norm_weight = norm_weight.contiguous()
+    num_tokens, hidden_size = x.shape
+    hc_mult = residual.shape[1]
+    k_total = hc_mult * hidden_size
+    assert residual.shape == (num_tokens, hc_mult, hidden_size)
+    assert attn_post.shape == (num_tokens, hc_mult)
+    assert weight.shape == (2 * hc_mult, k_total)
+    assert scale.shape == (2,) and base.shape == (2 * hc_mult,)
+    assert norm_weight.shape == (hidden_size,)
+    if num_tokens == 0:
+        return (
+            torch.empty((0, hidden_size), dtype=x.dtype, device=x.device),
+            torch.empty((0, hc_mult), dtype=torch.float32, device=x.device),
+            torch.empty_like(residual),
+        )
+
+    output = torch.empty((num_tokens, hidden_size), dtype=x.dtype, device=x.device)
+    post = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=x.device)
+    y = torch.empty_like(residual)
+    hc_pow2 = triton.next_power_of_2(hc_mult)
+    num_splits = triton.cdiv(k_total, _BLOCK_K)
+    num_hidden_blocks = triton.cdiv(hidden_size, _BLOCK_D)
+    partial_stride = 1 + 2 * hc_pow2
+    partial = torch.empty(
+        (num_tokens, num_splits, partial_stride),
+        dtype=torch.float32,
+        device=x.device,
+    )
+    _ihc_post_pre_rms_norm_stage1[(num_tokens, num_splits)](
+        x,
+        residual,
+        attn_post,
+        weight,
+        partial,
+        y,
+        K_TOTAL=k_total,
+        HIDDEN_SIZE=hidden_size,
+        HC_MULT=hc_mult,
+        HC_POW2=hc_pow2,
+        NUM_SPLITS=num_splits,
+        BLOCK_K=_BLOCK_K,
+        PARTIAL_STRIDE=partial_stride,
+        num_warps=_PRE_STAGE1_WARPS,
+        enable_fp_fusion=False,
+    )
+    _ihc_post_pre_rms_norm_stage2[(num_tokens,)](
+        y,
+        partial,
+        scale,
+        base,
+        norm_weight,
+        output,
+        post,
+        HIDDEN_SIZE=hidden_size,
+        K_TOTAL=k_total,
+        HC_MULT=hc_mult,
+        HC_POW2=hc_pow2,
+        NUM_SPLITS=num_splits,
+        PARTIAL_STRIDE=partial_stride,
+        BLOCK_D=_BLOCK_D,
+        NUM_HIDDEN_BLOCKS=num_hidden_blocks,
+        MAGNITUDE=magnitude,
+        NORM_EPS=norm_eps,
+        HC_EPS=hc_eps,
+        num_warps=_PRE_STAGE2_WARPS,
+        enable_fp_fusion=False,
+    )
+    return output, post, y
+
+
+def triton_ihc_post(
+    x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor
+) -> torch.Tensor:
+    assert x.dim() == 2 and residual.dim() == 3 and post.dim() == 2
+    assert x.is_cuda and residual.is_cuda and post.is_cuda
+    assert residual.device == x.device and post.device == x.device
+    assert x.dtype in (torch.float16, torch.bfloat16)
+    assert residual.dtype == x.dtype
+    assert post.dtype == torch.float32
+    num_tokens, hidden_size = x.shape
+    hc_mult = residual.shape[1]
+    assert residual.shape == (num_tokens, hc_mult, hidden_size)
+    assert post.shape == (num_tokens, hc_mult)
+    x, residual, post = x.contiguous(), residual.contiguous(), post.contiguous()
+    if num_tokens == 0:
+        return torch.empty_like(residual)
+    output = torch.empty_like(residual)
+    _ihc_post_kernel[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))](
+        x,
+        residual,
+        post,
+        output,
+        HIDDEN_SIZE=hidden_size,
+        HC_MULT=hc_mult,
+        HC_POW2=triton.next_power_of_2(hc_mult),
+        BLOCK_D=_BLOCK_D,
+        num_warps=_POST_WARPS,
+        enable_fp_fusion=False,
+    )
+    return output

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
 import importlib
+import inspect
 import math
 from collections.abc import Callable
 from importlib.util import find_spec
@@ -701,9 +702,35 @@ def paged_mqa_logits_module():
     if paged_mqa_logits_module_path is not None:
         try:
             module = importlib.import_module(paged_mqa_logits_module_path)
-            return module
-        except ImportError:
+        except ImportError as error:
+            logger.warning_once(
+                "AITER paged-MQA logits module %s could not be imported: %s",
+                paged_mqa_logits_module_path,
+                str(error),
+            )
             return None
+        if _ON_GFX942 or _ON_GFX950:
+            function_name = "deepgemm_fp8_paged_mqa_logits"
+            required_parameters = {
+                "Preshuffle",
+                "KVBlockSize",
+                "ChunkK",
+                "WavePerEU",
+            }
+        else:
+            function_name = "deepgemm_fp8_paged_mqa_logits_stage1"
+            required_parameters = {"ChunkQ"}
+        function = getattr(module, function_name, None)
+        if not callable(function) or not _accepts_parameters(
+            function, required_parameters
+        ):
+            logger.warning_once(
+                "AITER paged-MQA logits module %s lacks a compatible %s API.",
+                paged_mqa_logits_module_path,
+                function_name,
+            )
+            return None
+        return module
     return None
 
 
@@ -715,6 +742,17 @@ def _per_sequence_context_lens(context_lens: torch.Tensor) -> torch.Tensor:
     return context_lens
 
 
+def _accepts_parameters(function: Callable, required: set[str]) -> bool:
+    try:
+        parameters = inspect.signature(function).parameters
+    except (TypeError, ValueError):
+        return False
+    return required.issubset(parameters) or any(
+        parameter.kind == inspect.Parameter.VAR_KEYWORD
+        for parameter in parameters.values()
+    )
+
+
 def rocm_fp8_paged_mqa_logits(
     q_fp8: torch.Tensor,
     kv_cache_fp8: torch.Tensor,
@@ -775,6 +813,10 @@ def rocm_fp8_paged_mqa_logits(
 
     if aiter_paged_mqa_logits_module is not None:
         if _ON_GFX942 or _ON_GFX950:
+            logger.info_once(
+                "Using AITER paged-MQA logits from %s.",
+                aiter_paged_mqa_logits_module.__name__,
+            )
             deepgemm_fp8_paged_mqa_logits = (
                 aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
             )
@@ -816,6 +858,11 @@ def rocm_fp8_paged_mqa_logits(
         )
         return out_qk.sum(dim=0)
     else:
+        logger.warning_once(
+            "AITER paged-MQA logits is unavailable; using the PyTorch "
+            "fallback. Set VLLM_ROCM_USE_AITER=1 and ensure "
+            "aiter.ops.triton.attention.pa_mqa_logits is importable."
+        )
         return fp8_paged_mqa_logits_torch(
             q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
         )