diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 24bfbd012ec..365a814dc59 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.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 diff --git a/benchmarks/kernels/benchmark_hy_v4_ihc.py b/benchmarks/kernels/benchmark_hy_v4_ihc.py index 96220a66cbf..f843a6a8507 100644 --- a/benchmarks/kernels/benchmark_hy_v4_ihc.py +++ b/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 "" + 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="") + git_commit = _git_metadata(["rev-parse", "--short", "HEAD"], empty="") + 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, diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 1c4f3cd2909..868c2fcac7b 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/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, diff --git a/tests/models/hy_v4/test_triton_ihc.py b/tests/models/hy_v4/test_triton_ihc.py new file mode 100644 index 00000000000..77bfc3dae8a --- /dev/null +++ b/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) diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py index 06b2b3cccd2..e6f709e2dd8 100644 --- a/tests/models/quantization/test_mxfp8.py +++ b/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( diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index 1efb058ef07..d5c299e6dfc 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/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, diff --git a/vllm/models/hy_v4/amd/hc.py b/vllm/models/hy_v4/amd/hc.py new file mode 100644 index 00000000000..3f2768c2744 --- /dev/null +++ b/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) diff --git a/vllm/models/hy_v4/amd/model.py b/vllm/models/hy_v4/amd/model.py index 50925b060c6..ceb06113680 100644 --- a/vllm/models/hy_v4/amd/model.py +++ b/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() diff --git a/vllm/models/hy_v4/amd/triton_ihc.py b/vllm/models/hy_v4/amd/triton_ihc.py new file mode 100644 index 00000000000..ece62d47f08 --- /dev/null +++ b/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 diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 77303beae89..2b882dc298c 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/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 )