From 1cb08387214fba82c2705429100c5942fab65f05 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 4 May 2026 18:32:55 -0500 Subject: [PATCH 01/20] [ROCm][CI] Fix MLA prefill scale for DeepSeek GSM8K (#41569) Signed-off-by: Andreas Karatzas --- tests/v1/attention/test_mla_backends.py | 8 ++- .../v1/attention/test_mla_prefill_selector.py | 61 +++++++++++++++++++ .../layers/attention/mla_attention.py | 34 ++++++++++- 3 files changed, 99 insertions(+), 4 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index f91ea85779d..8ab47b61895 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -22,6 +22,7 @@ from vllm.config.vllm import set_current_vllm_config from vllm.model_executor.layers.attention.mla_attention import ( QueryLenSupport, _DecodeConcatQuantFP8, + get_mla_prefill_scale, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -785,7 +786,8 @@ def test_backend_correctness( assert kv_lora_rank + qk_rope_head_dim == head_size, ( f"MLA dimensions don't match: {total_head_size} != {head_size}" ) - scale = 1.0 / (total_head_size**0.5) + decode_scale = 1.0 / (total_head_size**0.5) + prefill_scale = get_mla_prefill_scale(vllm_config.model_config) # 2. Generate data and compute SDPA reference output for MLA all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], [] @@ -902,7 +904,7 @@ def test_backend_correctness( v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2) sdpa_out_i_decode = torch.nn.functional.scaled_dot_product_attention( - q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=decode_scale ) sdpa_out_i_decode = sdpa_out_i_decode.transpose(1, 2).squeeze( 0 @@ -938,7 +940,7 @@ def test_backend_correctness( # Single attention call with custom mask sdpa_out_i_prefill = torch.nn.functional.scaled_dot_product_attention( - q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale + q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=prefill_scale ) sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0) sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2) diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index 068eb43faf4..873cfb18701 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -2,12 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for MLA prefill backend selector.""" +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest import torch from vllm.config import AttentionConfig, ModelConfig, VllmConfig +from vllm.model_executor.layers.attention.mla_attention import get_mla_prefill_scale +from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import ( + yarn_get_mscale, +) from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.mla.prefill.selector import ( @@ -53,6 +58,62 @@ def _make_vllm_config( return mock_vllm_config +class TestMLAPrefillScale: + """Tests for the MLA prefill softmax scale.""" + + def test_uses_qk_head_dim_for_deepseek_v2_style_mla(self): + model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rope_parameters={"rope_type": "default"}, + ) + ) + + assert get_mla_prefill_scale(model_config) == pytest.approx(192**-0.5) + + def test_applies_deepseek_yarn_mscale(self): + model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rope_parameters={ + "rope_type": "yarn", + "factor": 40, + "mscale_all_dim": 0.707, + }, + ) + ) + + mscale = yarn_get_mscale(40, 0.707) + assert get_mla_prefill_scale(model_config) == pytest.approx( + 192**-0.5 * mscale * mscale + ) + + def test_deepseek_v4_style_mla_does_not_apply_yarn_mscale(self): + model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + compress_ratios=[4], + q_lora_rank=1536, + head_dim=128, + qk_rope_head_dim=64, + rope_parameters={ + "rope_type": "yarn", + "factor": 40, + "mscale_all_dim": 0.707, + }, + ) + ) + + assert get_mla_prefill_scale(model_config) == pytest.approx(128**-0.5) + + class TestGetMLAPrefillBackend: """Tests for get_mla_prefill_backend (public API).""" diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 82eecc8cd49..20981f60cd2 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -238,6 +238,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kNvfp4Dynamic, ) +from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import ( + yarn_get_mscale, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer from vllm.utils.math_utils import cdiv, round_down @@ -1327,6 +1330,35 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims: ) +def get_mla_prefill_scale(model_config: ModelConfig) -> float: + hf_text_config = model_config.hf_text_config + mla_dims = get_mla_dims(model_config) + qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim + scale = qk_head_dim**-0.5 + + # Deepseek V4 disables YaRN mscale for attention; Deepseek V2/V3 applies + # the same mscale correction when constructing the MLA attention module. + if hasattr(hf_text_config, "compress_ratios"): + return scale + + rope_parameters = getattr(hf_text_config, "rope_parameters", None) + if rope_parameters is None: + rope_parameters = getattr(hf_text_config, "rope_scaling", None) + + if rope_parameters is None: + return scale + + rope_type = rope_parameters.get("rope_type", rope_parameters.get("type")) + apply_yarn_scaling = rope_parameters.get("apply_yarn_scaling", True) + if rope_type != "default" and apply_yarn_scaling: + mscale_all_dim = rope_parameters.get("mscale_all_dim", False) + scaling_factor = rope_parameters["factor"] + mscale = yarn_get_mscale(float(scaling_factor), float(mscale_all_dim)) + scale *= mscale * mscale + + return scale + + @functools.cache def backend_supports_prefill_query_quantization() -> bool: """Check if the selected MLA prefill backend supports query quantization. @@ -1527,7 +1559,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): prefill_backend_cls = get_mla_prefill_backend(vllm_config) self._prefill_backend = prefill_backend_cls( num_heads=self.num_heads, - scale=self.model_config.get_head_size() ** -0.5, + scale=get_mla_prefill_scale(self.model_config), kv_lora_rank=self.mla_dims.kv_lora_rank, qk_nope_head_dim=self.mla_dims.qk_nope_head_dim, qk_rope_head_dim=self.mla_dims.qk_rope_head_dim, From 577b9623e6f8801698d411f4b04269326f5afbe2 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 4 May 2026 19:37:16 -0400 Subject: [PATCH 02/20] [Bug] Fix status update address for non-MOE model within external dp mode (#40839) Signed-off-by: yewentao256 --- docs/serving/data_parallel_deployment.md | 2 +- tests/test_config.py | 2 -- tests/v1/distributed/test_external_lb_dp.py | 2 +- vllm/config/parallel.py | 6 ++++-- vllm/engine/arg_utils.py | 14 +++++++++++++- 5 files changed, 19 insertions(+), 7 deletions(-) diff --git a/docs/serving/data_parallel_deployment.md b/docs/serving/data_parallel_deployment.md index 7b963b99d56..1f18b92f95b 100644 --- a/docs/serving/data_parallel_deployment.md +++ b/docs/serving/data_parallel_deployment.md @@ -98,7 +98,7 @@ For larger scale deployments especially, it can make sense to handle the orchest In this case, it's more convenient to treat each DP rank like a separate vLLM deployment, with its own endpoint, and have an external router balance HTTP requests between them, making use of appropriate real-time telemetry from each server for routing decisions. -This can already be done trivially for non-MoE models, since each deployed server is fully independent. No data parallel CLI options need to be used for this. +This can already be done trivially for non-MoE models, since each deployed server is fully independent. In that case, launch independent vLLM instances without any `--data-parallel-*` arguments; external DP CLI options are only supported for MoE deployments. We support an equivalent topology for MoE DP+EP which can be configured via the following CLI arguments. diff --git a/tests/test_config.py b/tests/test_config.py index 02e4d1d5d77..57d1e1bc686 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1215,8 +1215,6 @@ def test_scheduler_config_init(): ("facebook/opt-125m", 1, False, False), # Non-MoE model with DP>1 internal LB should need coordinator ("facebook/opt-125m", 2, False, True), - # Non-MoE model with DP>1 external LB should not need coordinator - ("facebook/opt-125m", 2, True, False), # MoE model with DP=1 should not need coordinator ("mistralai/Mixtral-8x7B-Instruct-v0.1", 1, False, False), # MoE model with DP>1 internal LB should need both coordinator diff --git a/tests/v1/distributed/test_external_lb_dp.py b/tests/v1/distributed/test_external_lb_dp.py index cfef8449ebf..06e8e574a05 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -14,7 +14,7 @@ import requests from tests.utils import RemoteOpenAIServer from vllm.platforms import current_platform -MODEL_NAME = "ibm-research/PowerMoE-3b" +MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") # Number of data parallel ranks for external LB testing DP_SIZE = int(os.getenv("DP_SIZE", "2")) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6ba392802e3..95fd8787afe 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -135,8 +135,10 @@ class ParallelConfig: data_parallel_external_lb: bool = False """Whether to use "external" DP LB mode. Applies only to online serving and when data_parallel_size > 0. This is useful for a "one-pod-per-rank" - wide-EP setup in Kubernetes. Set implicitly when --data-parallel-rank - is provided explicitly to vllm serve.""" + wide-EP setup in Kubernetes. Supported only for MoE deployments; non-MoE + models should use independent vLLM instances without --data-parallel-* + arguments. Set implicitly when --data-parallel-rank is provided explicitly + to vllm serve.""" data_parallel_hybrid_lb: bool = False """Whether to use "hybrid" DP LB mode. Applies only to online serving and when data_parallel_size > 0. Enables running an AsyncLLM diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index cd955100333..4c37d5a149c 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -962,7 +962,9 @@ class EngineArgs: "-dpn", type=int, help="Data parallel rank of this instance. " - "When set, enables external load balancer mode.", + "When set, enables external load balancer mode for MoE " + "data-parallel deployments. Unsupported for non-MoE models; " + "launch independent vLLM instances instead.", ) parallel_group.add_argument( "--data-parallel-start-rank", @@ -1793,6 +1795,16 @@ class EngineArgs: data_parallel_external_lb = ( self.data_parallel_external_lb or self.data_parallel_rank is not None ) + if ( + self.data_parallel_size > 1 + and data_parallel_external_lb + and not model_config.is_moe + ): + raise ValueError( + "Non-MoE models do not support external data parallel mode. " + "For external load balancing, launch independent vLLM " + "instances without --data-parallel-* arguments." + ) # Local DP rank = 1, use pure-external LB. if data_parallel_external_lb: assert self.data_parallel_rank is not None, ( From 4f2af1a7c03aae2b3227dd7e69d726104d44a711 Mon Sep 17 00:00:00 2001 From: JartX Date: Tue, 5 May 2026 02:14:01 +0200 Subject: [PATCH 03/20] [Feature] TurboQuant: support hybrid models and uniform quantization (#39931) Signed-off-by: JartX Signed-off-by: Jim Smith Co-authored-by: Jim Smith Co-authored-by: Sandermage Co-authored-by: Claude --- tests/quantization/test_turboquant.py | 86 ++++++++++++++++++- vllm/engine/arg_utils.py | 20 +---- .../layers/quantization/turboquant/config.py | 69 +++++++++++++-- vllm/platforms/interface.py | 36 ++++++++ 4 files changed, 185 insertions(+), 26 deletions(-) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index f074ce119ae..b9567195b3a 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -182,22 +182,100 @@ class TestTurboQuantConfig: # ---- Boundary skip layers ---- + @staticmethod + def _dense_model_config(num_layers): + from types import SimpleNamespace + + return SimpleNamespace( + is_hybrid=False, + hf_text_config=SimpleNamespace(num_hidden_layers=num_layers), + ) + def test_boundary_skip_layers_basic(self): - layers = TurboQuantConfig.get_boundary_skip_layers(32) + mc = self._dense_model_config(32) + layers = TurboQuantConfig.get_boundary_skip_layers(mc) assert layers == ["0", "1", "30", "31"] def test_boundary_skip_layers_zero(self): - assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == [] + mc = self._dense_model_config(32) + assert TurboQuantConfig.get_boundary_skip_layers(mc, 0) == [] def test_boundary_skip_layers_small_model(self): - layers = TurboQuantConfig.get_boundary_skip_layers(4) + mc = self._dense_model_config(4) + layers = TurboQuantConfig.get_boundary_skip_layers(mc) assert layers == ["0", "1", "2", "3"] def test_boundary_skip_layers_cap_at_half(self): - layers = TurboQuantConfig.get_boundary_skip_layers(8, 10) + mc = self._dense_model_config(8) + layers = TurboQuantConfig.get_boundary_skip_layers(mc, 10) assert len(layers) == 8 +class TestHybridAttentionIndices: + """Regression tests for boundary protection on hybrid models. + + Hybrid models (attention + Mamba / linear-attention) identify KV-carrying + layers via layer_types / layers_block_type / attn_type_list. The helper + must return the *global* layer indices of the full-attention layers so + that kv_cache_dtype_skip_layers matches what extract_layer_index(prefix) + reports on the Attention layers at runtime. + """ + + @staticmethod + def _fake_model_config(text_cfg=None, hf_cfg=None): + from types import SimpleNamespace + + return SimpleNamespace( + hf_text_config=text_cfg if text_cfg is not None else SimpleNamespace(), + hf_config=hf_cfg if hf_cfg is not None else SimpleNamespace(), + ) + + def test_layer_types_full_attention(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + cfg = type("C", (), {})() + cfg.layer_types = [ + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + "full_attention", + ] + mc = self._fake_model_config(text_cfg=cfg) + assert _get_full_attention_layer_indices(mc) == [2, 4, 5] + + def test_layers_block_type_jamba(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + cfg = type("C", (), {})() + cfg.layers_block_type = ["mamba", "attention", "mamba", "attention"] + mc = self._fake_model_config(text_cfg=cfg) + assert _get_full_attention_layer_indices(mc) == [1, 3] + + def test_attn_type_list_minimax(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + hf = type("C", (), {})() + hf.attn_type_list = [0, 1, 0, 1, 1] + mc = self._fake_model_config(hf_cfg=hf) + assert _get_full_attention_layer_indices(mc) == [1, 3, 4] + + def test_no_hybrid_hints_returns_empty(self): + from vllm.model_executor.layers.quantization.turboquant.config import ( + _get_full_attention_layer_indices, + ) + + mc = self._fake_model_config() + assert _get_full_attention_layer_indices(mc) == [] + + # ============================================================================ # Centroids tests (CPU-only) # ============================================================================ diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 4c37d5a149c..1b380313921 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1699,29 +1699,15 @@ class EngineArgs: kv_offloading_backend=self.kv_offloading_backend, ) - # TurboQuant: auto-skip first/last 2 layers (boundary protection). - # These layers are most sensitive to quantization error. - # Users can add extra layers via --kv-cache-dtype-skip-layers. if resolved_cache_dtype.startswith("turboquant_"): - if model_config.is_hybrid: - raise NotImplementedError( - "TurboQuant KV cache is not supported for hybrid " - "(attention + Mamba) models. Boundary layer protection " - "requires uniform attention layers." - ) from vllm.model_executor.layers.quantization.turboquant.config import ( TurboQuantConfig, ) - num_layers = model_config.hf_text_config.num_hidden_layers - boundary = TurboQuantConfig.get_boundary_skip_layers(num_layers) + boundary = TurboQuantConfig.get_boundary_skip_layers(model_config) existing = set(cache_config.kv_cache_dtype_skip_layers) - merged = sorted(existing | set(boundary), key=lambda x: int(x)) - cache_config.kv_cache_dtype_skip_layers = merged - logger.info( - "TQ: skipping layers %s for boundary protection (num_layers=%d)", - merged, - num_layers, + cache_config.kv_cache_dtype_skip_layers = sorted( + existing | set(boundary), key=int ) ray_runtime_env = None diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py index f9cfc89c0c1..50beb8d1d9b 100644 --- a/vllm/model_executor/layers/quantization/turboquant/config.py +++ b/vllm/model_executor/layers/quantization/turboquant/config.py @@ -2,8 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """TurboQuant configuration.""" +from __future__ import annotations + +import logging import math from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.config import ModelConfig + +logger = logging.getLogger(__name__) # Named TQ presets: each maps to frozen config parameters. # key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys. @@ -159,12 +168,34 @@ class TurboQuantConfig: return s + (s % 2) # round up to even @staticmethod - def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]: - """Get layer indices to skip TQ compression (boundary protection). + def get_boundary_skip_layers( + model_config: ModelConfig, + n: int = 2, + ) -> list[str]: + """Layer indices to skip TQ compression (boundary protection). - Returns first N and last N layer indices as strings, suitable for - kv_cache_dtype_skip_layers. + For hybrid models (attention + Mamba/linear-attention), boundary + protection is disabled — hybrids typically have only 8-12 + full-attention layers and a hard n=2 on each side would cover + ~40 % of them. The dense GSM8K baselines that motivate n=2 + don't apply to hybrids. + + For dense models, skips first N and last N attention layers. + Empirically required for aggressive presets (k3v4_nc, 3bit_nc) + — without it GSM8K drops ~30 points on Qwen3-4B. """ + if model_config.is_hybrid: + attn_indices = _get_full_attention_layer_indices(model_config) + if not attn_indices: + raise NotImplementedError( + "TurboQuant KV cache requires identifiable " + "full-attention layers, but none were found in " + "the hybrid model config." + ) + logger.info("TQ hybrid: full-attention layers %s", attn_indices) + return [] + + num_layers = model_config.hf_text_config.num_hidden_layers if n <= 0 or num_layers <= 0: return [] n = min(n, num_layers // 2) # don't skip more than half @@ -175,7 +206,7 @@ class TurboQuantConfig: return [str(i) for i in indices] @staticmethod - def from_cache_dtype(cache_dtype: str, head_dim: int) -> "TurboQuantConfig": + def from_cache_dtype(cache_dtype: str, head_dim: int) -> TurboQuantConfig: """Create config from a named preset. Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc. @@ -193,3 +224,31 @@ class TurboQuantConfig: value_quant_bits=preset["value_quant_bits"], norm_correction=preset["norm_correction"], ) + + +def _get_full_attention_layer_indices(model_config: ModelConfig) -> list[int]: + """Global indices of full-attention layers in a hybrid model. + + Covers the conventions used across vLLM: ``layer_types`` (Qwen3.5/Next), + ``layers_block_type`` (Jamba/Zamba2), ``attn_type_list`` (Minimax). + """ + text_cfg = model_config.hf_text_config + hf_cfg = model_config.hf_config + + layer_types = getattr(text_cfg, "layer_types", None) + if layer_types is not None: + return [ + i for i, t in enumerate(layer_types) if t in ("full_attention", "attention") + ] + + layers_block_type = getattr(text_cfg, "layers_block_type", None) + if layers_block_type is not None: + return [ + i for i, t in enumerate(layers_block_type) if t in ("attention", "hybrid") + ] + + attn_type_list = getattr(hf_cfg, "attn_type_list", None) + if attn_type_list is not None: + return [i for i, t in enumerate(attn_type_list) if t == 1] + + return [] diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 2753326755f..80952ced73d 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -545,6 +545,42 @@ class Platform: dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, ).page_size_bytes + elif cache_config.cache_dtype.startswith("turboquant_"): + # TQ has a packed K|V layout; the standard FullAttentionSpec + # formula over-sizes it and trips unify_kv_cache_spec_page_size + # when all attention layers are TQ. With mixed skip+TQ the skip + # layers still use the standard layout — take max so mamba + # padding covers the largest actual page. + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + from vllm.v1.kv_cache_interface import TQFullAttentionSpec + + tq_cfg = TurboQuantConfig.from_cache_dtype( + cache_config.cache_dtype, model_config.get_head_size() + ) + tq_page = TQFullAttentionSpec( + block_size=1, + num_kv_heads=model_config.get_num_kv_heads(parallel_config), + head_size=model_config.get_head_size(), + head_size_v=model_config.get_head_size(), + dtype=kv_cache_dtype, + kv_quant_mode=kv_quant_mode, + tq_slot_size=tq_cfg.slot_size_aligned, + ).page_size_bytes + if cache_config.kv_cache_dtype_skip_layers: + skip_page = FullAttentionSpec( + block_size=1, + num_kv_heads=model_config.get_num_kv_heads(parallel_config), + head_size=model_config.get_head_size(), + dtype=model_config.dtype, + ).page_size_bytes + # lcm, not max: skip_page is often not a multiple of + # tq_page, so max would leave per-layer page sizes + # un-unifiable downstream. + attn_page_size_1_token = lcm(tq_page, skip_page) + else: + attn_page_size_1_token = tq_page else: attn_page_size_1_token = FullAttentionSpec( block_size=1, From e1e4646b06f289475ee57f31e3df817e06351321 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 4 May 2026 17:44:55 -0700 Subject: [PATCH 04/20] [Model Runner V2] Rebuild attn metadata between draft decode steps (#41162) Signed-off-by: Giancarlo Delfin --- vllm/v1/worker/gpu/sample/gumbel.py | 22 +- .../gpu/spec_decode/eagle/speculator.py | 289 +++++++++++------- .../probabilistic_rejection_sampler_utils.py | 6 +- 3 files changed, 197 insertions(+), 120 deletions(-) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 62912491492..a02dd62026a 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -76,6 +76,8 @@ def gumbel_block_argmax( pos_ptr, processed_logits_ptr, processed_logits_stride, + processed_logits_col_ptr, + vocab_size, APPLY_TEMPERATURE: tl.constexpr, ): req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) @@ -88,8 +90,15 @@ def gumbel_block_argmax( if processed_logits_ptr is not None: # Store the temperature-applied logits. + if processed_logits_col_ptr is not None: + col = tl.load(processed_logits_col_ptr) + else: + col = 0 tl.store( - processed_logits_ptr + req_state_idx * processed_logits_stride + block, + processed_logits_ptr + + req_state_idx * processed_logits_stride + + col * vocab_size + + block, logits, mask=mask, ) @@ -121,6 +130,7 @@ def _gumbel_sample_kernel( local_max_stride, processed_logits_ptr, processed_logits_stride, + processed_logits_col_ptr, logits_ptr, logits_stride, expanded_idx_mapping_ptr, @@ -153,6 +163,8 @@ def _gumbel_sample_kernel( pos_ptr, processed_logits_ptr, processed_logits_stride, + processed_logits_col_ptr, + vocab_size, APPLY_TEMPERATURE=APPLY_TEMPERATURE, ) token_id = block_idx * BLOCK_SIZE + idx @@ -167,7 +179,8 @@ def gumbel_sample( seed: torch.Tensor, # [max_num_reqs] pos: torch.Tensor, # [num_tokens] apply_temperature: bool, - processed_logits_out: torch.Tensor | None = None, # [num_reqs, vocab_size] + output_processed_logits: torch.Tensor | None = None, + output_processed_logits_col: torch.Tensor | None = None, ) -> torch.Tensor: num_tokens, vocab_size = logits.shape BLOCK_SIZE = 1024 @@ -179,8 +192,9 @@ def gumbel_sample( local_argmax.stride(0), local_max, local_max.stride(0), - processed_logits_out, - processed_logits_out.stride(0) if processed_logits_out is not None else 0, + output_processed_logits, + output_processed_logits.stride(0) if output_processed_logits is not None else 0, + output_processed_logits_col, logits, logits.stride(0), expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index c6b0aa364f5..efe510f16e2 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -89,9 +89,13 @@ class EagleSpeculator: dtype=torch.int64, device=device, ) + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device) self.last_token_indices = torch.zeros( self.max_num_reqs, dtype=torch.int64, device=device ) + self.arange = torch.arange( + self.max_num_reqs + 1, dtype=torch.int32, device="cpu" + ) self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs( self.draft_model_config @@ -228,9 +232,10 @@ class EagleSpeculator: logits: torch.Tensor, idx_mapping: torch.Tensor, pos: torch.Tensor, - step: int, + draft_step: torch.Tensor, + draft_logits: torch.Tensor | None, ) -> torch.Tensor: - if self.draft_logits is not None: + if draft_logits is not None: # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise # used for draft and target sampling. return gumbel_sample( @@ -240,7 +245,8 @@ class EagleSpeculator: self.seeds, pos + 1, apply_temperature=True, - processed_logits_out=self.draft_logits[:, step], + output_processed_logits=draft_logits, + output_processed_logits_col=draft_step, ) else: return logits.argmax(dim=-1) @@ -274,11 +280,63 @@ class EagleSpeculator: logits, idx_mapping, pos, - step=0, + self.current_draft_step, + self.draft_logits, ) self.hidden_states[:num_reqs] = hidden_states[last_token_indices] self.input_buffers.positions[:num_reqs] = pos + def multi_step_decode( + self, + num_reqs: int, + skip_attn: bool, + batch_desc: BatchExecutionDescriptor, + num_tokens_across_dp: torch.Tensor | None, + ) -> None: + positions = self.input_buffers.positions[:num_reqs] + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] + idx_mapping = self.idx_mapping[:num_reqs] + + for step in range(1, self.num_speculative_steps): + attn_metadata = None + slot_mappings_by_layer = None + if not skip_attn: + # Build attention metadata and slot mappings for each draft + # decode step. It is necessary to rebuild the attention + # metadata even when replaying the FULL graph so that any + # attention metadata builder state is updated. + slot_mappings = self.block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + batch_desc.num_tokens, + ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, self.kv_cache_config + ) + attn_metadata = self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=batch_desc.num_reqs or num_reqs, + num_tokens_padded=batch_desc.num_tokens, + ) + + # Update the current draft step. + self.current_draft_step.fill_(step) + + # Generate draft tokens for the current step. + if batch_desc.cg_mode == CUDAGraphMode.FULL: + assert self.decode_cudagraph_manager is not None + self.decode_cudagraph_manager.run_fullgraph(batch_desc) + else: + self.generate_draft( + num_reqs, + batch_desc.num_tokens, + attn_metadata, + slot_mappings_by_layer, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=batch_desc.cg_mode, + ) + def generate_draft( self, num_reqs: int, @@ -288,59 +346,52 @@ class EagleSpeculator: num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, ) -> None: - pos = self.input_buffers.positions[:num_reqs] - query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1] idx_mapping = self.idx_mapping[:num_reqs] - for step in range(1, self.num_speculative_steps): - # Run the eagle model. - last_hidden_states, hidden_states = self.run_model( - num_tokens_padded, - attn_metadata, - slot_mappings, - num_tokens_across_dp, - cudagraph_runtime_mode, - ) - last_hidden_states = last_hidden_states[:num_reqs] - hidden_states = hidden_states[:num_reqs] - logits = self.model.compute_logits(last_hidden_states) + positions = self.input_buffers.positions[:num_reqs] + # Run the eagle model forward pass. + last_hidden_states, hidden_states = self.run_model( + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + ) + last_hidden_states = last_hidden_states[:num_reqs] - draft_tokens = self._sample_draft( - logits, - idx_mapping, - pos, - step=step, - ) - self.draft_tokens[:num_reqs, step] = draft_tokens + # Sample the draft tokens. + logits = self.model.compute_logits(last_hidden_states) + draft_tokens = self._sample_draft( + logits, + idx_mapping, + positions, + self.current_draft_step, + self.draft_logits, + ) - if step < self.num_speculative_steps - 1: - # Update the inputs for the next step. - update_eagle_inputs( - draft_tokens, - hidden_states, - self.input_buffers, - self.hidden_states, - self.max_model_len, - ) - if attn_metadata is not None: - self.block_tables.compute_slot_mappings( - idx_mapping, query_start_loc, pos, num_tokens_padded - ) + # Update the inputs for the next step. + update_eagle_draft_inputs( + draft_tokens, + self.current_draft_step, + hidden_states, + self.draft_tokens, + self.hidden_states, + self.input_buffers, + num_reqs, + self.max_model_len, + self.num_speculative_steps, + ) def _build_draft_attn_metadata( self, num_reqs: int, num_reqs_padded: int, num_tokens_padded: int, - max_query_len: int, ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None - query_start_loc_cpu = ( - torch.arange(num_reqs_padded + 1, dtype=torch.int32, device="cpu").clamp_( - max=num_reqs - ) - * max_query_len + query_start_loc_cpu = torch.clamp( + self.arange[: num_reqs_padded + 1], max=num_reqs ) block_tables = [ x[:num_reqs_padded] for x in self.block_tables.input_block_tables @@ -354,7 +405,7 @@ class EagleSpeculator: : num_reqs_padded + 1 ], query_start_loc_cpu=query_start_loc_cpu, - max_query_len=max_query_len, + max_query_len=1, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], max_seq_len=self.max_model_len, block_tables=block_tables, @@ -373,7 +424,7 @@ class EagleSpeculator: self.last_token_indices.zero_() # Capture the prefill routine (model forward + compute_logits + - # gumbel_sample). + # sample). # For FULL graphs, the entire routine is recorded as one graph. # For PIECEWISE, only the model's compiled regions are captured # and the rest (compute_logits, gumbel_sample) runs eagerly. @@ -387,10 +438,9 @@ class EagleSpeculator: if self.num_speculative_steps == 1: return - # Capture the decode draft generation loop (model forward + - # compute_logits + gumbel_sample + update_eagle_inputs, for - # each step). For FULL graphs, the entire multi-step loop is - # recorded as one graph. + # Capture the decode draft generation routine (model forward + + # compute_logits + sample + update_eagle_inputs) for a single + # step. assert self.decode_cudagraph_manager is not None self.decode_cudagraph_manager.capture( self.generate_draft, @@ -461,9 +511,10 @@ class EagleSpeculator: # Get the input ids and last token indices for the speculator. prepare_eagle_inputs( + self.last_token_indices, + self.current_draft_step, self.input_buffers, input_batch, - self.last_token_indices, num_sampled, num_rejected, last_sampled, @@ -473,12 +524,18 @@ class EagleSpeculator: # When all requests are decoding (no true prefills), each has # num_speculative_steps + 1 tokens, enabling FULL graph replay. - # Mixed or prefill-only batches fall back to PIECEWISE. + uniform_token_count = get_uniform_token_count( + num_reqs, + # Use the actual number of tokens without padding added by + # the target model during FULL cudagraph. + input_batch.num_tokens, + max_query_len, + ) prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.prefill_cudagraph_manager, num_reqs, num_tokens, - get_uniform_token_count(num_reqs, num_tokens, max_query_len), + uniform_token_count, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, @@ -528,48 +585,21 @@ class EagleSpeculator: need_eager=is_profile, ) - attn_metadata_updated = None - slot_mappings_updated = None - if not (dummy_run and skip_attn_for_dummy_run): - # Build attention metadata and slot mappings for the draft - # decode steps. It is necessary to rebuild the attention - # metadata even when replaying the FULL graph so that any - # attention metadata builder state is updated. - slot_mappings = self.block_tables.compute_slot_mappings( - self.idx_mapping[:num_reqs], - self.input_buffers.query_start_loc[: num_reqs + 1], - self.input_buffers.positions[:num_reqs], - decode_batch_desc.num_tokens, - ) - slot_mappings_updated = build_slot_mappings_by_layer( - slot_mappings, self.kv_cache_config - ) - attn_metadata_updated = self._build_draft_attn_metadata( - num_reqs=num_reqs, - num_reqs_padded=decode_batch_desc.num_reqs or num_reqs, - num_tokens_padded=decode_batch_desc.num_tokens, - max_query_len=1, - ) + # Generate the remaining num_speculative_steps - 1 draft tokens. + self.multi_step_decode( + num_reqs, + dummy_run and skip_attn_for_dummy_run, + decode_batch_desc, + num_tokens_across_dp, + ) - if decode_batch_desc.cg_mode == CUDAGraphMode.FULL: - # Replay the full graph for draft generation. - assert self.decode_cudagraph_manager is not None - self.decode_cudagraph_manager.run_fullgraph(decode_batch_desc) - else: - self.generate_draft( - num_reqs, - decode_batch_desc.num_tokens, - attn_metadata_updated, - slot_mappings_updated, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=decode_batch_desc.cg_mode, - ) return self.draft_tokens[:num_reqs] @triton.jit def _prepare_eagle_inputs_kernel( last_token_indices_ptr, + eagle_current_draft_step_ptr, eagle_input_ids_ptr, eagle_positions_ptr, eagle_query_start_loc_ptr, @@ -630,6 +660,8 @@ def _prepare_eagle_inputs_kernel( # Copy sequence lengths. tl.store(eagle_seq_lens_ptr + req_idx, seq_len) if req_idx == (num_reqs - 1): + # Reset the current draft step to 0. + tl.store(eagle_current_draft_step_ptr, 0) # Pad query_start_loc for CUDA graphs. for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) @@ -648,10 +680,11 @@ def _prepare_eagle_inputs_kernel( def prepare_eagle_inputs( - input_buffers: InputBuffers, - input_batch: InputBatch, # [num_reqs] last_token_indices: torch.Tensor, + current_draft_step: torch.Tensor, + input_buffers: InputBuffers, + input_batch: InputBatch, # [num_reqs] num_sampled: torch.Tensor, # [num_reqs] @@ -665,6 +698,7 @@ def prepare_eagle_inputs( num_reqs = input_batch.num_reqs _prepare_eagle_inputs_kernel[(num_reqs,)]( last_token_indices, + current_draft_step, input_buffers.input_ids, input_buffers.positions, input_buffers.query_start_loc, @@ -685,7 +719,7 @@ def prepare_eagle_inputs( @triton.jit -def _prepare_eagle_docode_kernel( +def _prepare_eagle_decode_kernel( draft_tokens_ptr, draft_tokens_stride, target_seq_lens_ptr, @@ -742,7 +776,7 @@ def prepare_eagle_decode( max_num_reqs: int, ): num_reqs = draft_tokens.shape[0] - _prepare_eagle_docode_kernel[(num_reqs + 1,)]( + _prepare_eagle_decode_kernel[(num_reqs + 1,)]( draft_tokens, draft_tokens.stride(0), target_seq_lens, @@ -758,36 +792,55 @@ def prepare_eagle_decode( @triton.jit -def _update_eagle_inputs_kernel( +def _update_eagle_draft_inputs_kernel( + output_draft_tokens_ptr, + output_draft_tokens_stride, + next_input_hidden_states_ptr, + next_input_hidden_states_stride, input_ids_ptr, positions_ptr, - input_hidden_states_ptr, - input_hidden_states_stride, seq_lens_ptr, - max_model_len, draft_tokens_ptr, - output_hidden_states_ptr, - output_hidden_states_stride, + current_draft_step_ptr, + hidden_states_ptr, + hidden_states_stride, hidden_size, + max_model_len, + num_speculative_steps, BLOCK_SIZE: tl.constexpr, ): req_idx = tl.program_id(0) - # Draft token -> Input ID. + # Write the sampled draft token into self.draft_tokens[req_idx, step]. draft_token = tl.load(draft_tokens_ptr + req_idx) + step = tl.load(current_draft_step_ptr) + tl.store( + output_draft_tokens_ptr + req_idx * output_draft_tokens_stride + step, + draft_token, + ) + + if step >= num_speculative_steps - 1: + # This is the final step. Skip updating draft forward inputs. + return + + # Write the sampled draft token into the input ids tensor for the next + # forward pass. tl.store(input_ids_ptr + req_idx, draft_token) - # Output hidden states -> Input hidden states. + # Copy hidden states into the input hidden states tensor for the next + # forward pass. for i in range(0, hidden_size, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < hidden_size - output_hidden_states = tl.load( - output_hidden_states_ptr + req_idx * output_hidden_states_stride + block, + hidden_states = tl.load( + hidden_states_ptr + req_idx * hidden_states_stride + block, mask=mask, ) tl.store( - input_hidden_states_ptr + req_idx * input_hidden_states_stride + block, - output_hidden_states, + next_input_hidden_states_ptr + + req_idx * next_input_hidden_states_stride + + block, + hidden_states, mask=mask, ) @@ -803,24 +856,32 @@ def _update_eagle_inputs_kernel( tl.store(seq_lens_ptr + req_idx, seq_len) -def update_eagle_inputs( +def update_eagle_draft_inputs( draft_tokens: torch.Tensor, - output_hidden_states: torch.Tensor, - input_buffers: InputBuffers, + current_draft_step: torch.Tensor, hidden_states: torch.Tensor, + output_draft_tokens: torch.Tensor, + next_input_hidden_states: torch.Tensor, + input_buffers: InputBuffers, + num_reqs: int, max_model_len: int, + num_speculative_steps: int, ): - num_reqs, hidden_size = output_hidden_states.shape - _update_eagle_inputs_kernel[(num_reqs,)]( + _, hidden_size = hidden_states.shape + _update_eagle_draft_inputs_kernel[(num_reqs,)]( + output_draft_tokens, + output_draft_tokens.stride(0), + next_input_hidden_states, + next_input_hidden_states.stride(0), input_buffers.input_ids, input_buffers.positions, + input_buffers.seq_lens, + draft_tokens, + current_draft_step, hidden_states, hidden_states.stride(0), - input_buffers.seq_lens, - max_model_len, - draft_tokens, - output_hidden_states, - output_hidden_states.stride(0), hidden_size, + max_model_len, + num_speculative_steps, BLOCK_SIZE=1024, ) diff --git a/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py index 9d86372e624..10b29433efb 100644 --- a/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py @@ -392,8 +392,10 @@ def _resample_kernel( temp_ptr, seed_ptr, pos_ptr, - None, - 0, + None, # processed_logits_ptr + 0, # processed_logits_stride + None, # processed_logits_col_ptr + vocab_size, APPLY_TEMPERATURE=False, ) token_id = block_idx * BLOCK_SIZE + idx From 685bf811d65b58b2f8ef149d7da53dfd0393a912 Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Mon, 4 May 2026 20:07:39 -0500 Subject: [PATCH 05/20] [XPU] enable is_act_and_mul for xpu (#37481) Signed-off-by: Chendi Xue Co-authored-by: Kunshang Ji --- vllm/model_executor/layers/fused_moe/experts/xpu_moe.py | 3 ++- vllm/model_executor/layers/fused_moe/layer.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index e10be4af868..d6bd2b14008 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -62,7 +62,7 @@ class XPUExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_activation(activation: MoEActivation) -> bool: @@ -70,6 +70,7 @@ class XPUExperts(mk.FusedMoEExpertsModular): MoEActivation.SILU, MoEActivation.GELU, MoEActivation.SWIGLUOAI, + MoEActivation.RELU2_NO_MUL, ] @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 3de05cd93d3..456f40bbf7a 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -538,9 +538,11 @@ class FusedMoE(PluggableLayer): # for heuristic purposes, so it must be initialized first. self.quant_method: FusedMoEMethodBase = _get_quant_method() - if not self.moe_config.is_act_and_mul and not current_platform.is_cuda_alike(): + if not self.moe_config.is_act_and_mul and not ( + current_platform.is_cuda_alike() or current_platform.is_xpu() + ): raise NotImplementedError( - "is_act_and_mul=False is supported only for CUDA and ROCm for now" + "is_act_and_mul=False is supported only for CUDA and XPU for now" ) if self.enable_eplb and not self.quant_method.supports_eplb: From 416f9cdede967edbf712727fa8a510c70f18aacb Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 4 May 2026 19:43:25 -0700 Subject: [PATCH 06/20] [Perf][2/n] Eliminate GPU<->CPU syncs in pooling code (#41433) Signed-off-by: Nick Hill --- .../layers/pooler/seqwise/methods.py | 18 +++++---- vllm/model_executor/layers/pooler/special.py | 38 +++++++++++++++++-- .../layers/pooler/tokwise/methods.py | 29 +++++++------- 3 files changed, 58 insertions(+), 27 deletions(-) diff --git a/vllm/model_executor/layers/pooler/seqwise/methods.py b/vllm/model_executor/layers/pooler/seqwise/methods.py index b967ff4ede7..82170b5fbdc 100644 --- a/vllm/model_executor/layers/pooler/seqwise/methods.py +++ b/vllm/model_executor/layers/pooler/seqwise/methods.py @@ -68,21 +68,23 @@ class MeanPool(SequencePoolingMethod): "partial prefill not supported with MEAN pooling" ) - prompt_lens = pooling_cursor.prompt_lens_cpu.to( - hidden_states.device, dtype=torch.int64, non_blocking=True - ) - - num_seqs = prompt_lens.numel() + prompt_lens_cpu = pooling_cursor.prompt_lens_cpu + num_seqs = prompt_lens_cpu.numel() hidden_size = hidden_states.shape[-1] if num_seqs == 0: # early return for empty batch return hidden_states.new_empty((0, hidden_size), dtype=torch.float32) - # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] + # Build segment_ids on CPU so repeat_interleave doesn't need to sync + # GPU->CPU to learn its data-dependent output length, then upload + # non-blocking. eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2] segment_ids = torch.repeat_interleave( - torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long), - prompt_lens, + torch.arange(num_seqs, dtype=torch.long), + prompt_lens_cpu, + ).to(hidden_states.device, non_blocking=True) + prompt_lens = prompt_lens_cpu.to( + hidden_states.device, dtype=torch.int64, non_blocking=True ) segment_sums = torch.zeros( (num_seqs, hidden_size), diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index d06663b5b94..ae5926cd62f 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses from collections.abc import Mapping, Set from itertools import groupby @@ -80,9 +81,11 @@ class DispatchPooler(Pooler): pooling_metadata: PoolingMetadata, ) -> PoolerOutput: poolers_by_task = self.poolers_by_task + cursor = pooling_metadata.pooling_cursor outputs = list[torch.Tensor | None]() offset = 0 + token_offset = 0 for task, group in groupby(pooling_metadata.tasks): if not (pooler := poolers_by_task.get(task)): raise ValueError( @@ -91,10 +94,37 @@ class DispatchPooler(Pooler): ) num_items = len(list(group)) - group_output: PoolerOutput = pooler( - hidden_states, - pooling_metadata[offset : offset + num_items], - ) + group_metadata = pooling_metadata[offset : offset + num_items] + if cursor is None: + group_hidden_states = hidden_states + else: + # Slice out this group's tokens so sub-poolers see only their + # portion of the batch. Token offset is computed from the CPU + # `num_scheduled_tokens_cpu` to avoid a GPU->CPU sync. + group_cursor = group_metadata.pooling_cursor + num_group_tokens = int(group_cursor.num_scheduled_tokens_cpu.sum()) + group_hidden_states = hidden_states[ + token_offset : token_offset + num_group_tokens + ] + if token_offset: + # Shift first/last indices to be relative to the slice + # so seqwise poolers (which index `hidden_states` directly) + # remain correct. + pooling_cursor = dataclasses.replace( + group_cursor, + first_token_indices_gpu=( + group_cursor.first_token_indices_gpu - token_offset + ), + last_token_indices_gpu=( + group_cursor.last_token_indices_gpu - token_offset + ), + ) + group_metadata = dataclasses.replace( + group_metadata, pooling_cursor=pooling_cursor + ) + token_offset += num_group_tokens + + group_output: PoolerOutput = pooler(group_hidden_states, group_metadata) outputs.extend(group_output) offset += num_items diff --git a/vllm/model_executor/layers/pooler/tokwise/methods.py b/vllm/model_executor/layers/pooler/tokwise/methods.py index d3fefb745cf..59b7234661b 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/vllm/model_executor/layers/pooler/tokwise/methods.py @@ -47,17 +47,12 @@ class AllPool(TokenPoolingMethod): pooling_metadata: PoolingMetadata, ) -> list[TokenPoolingMethodOutputItem]: pooling_cursor = pooling_metadata.get_pooling_cursor() - split_sizes = pooling_cursor.num_scheduled_tokens_cpu.tolist() - if split_sizes: - # DispatchPooler passes the full hidden_states tensor. - # slice out the subgroup once, then split it by - # per-request token counts - group_start = int(pooling_cursor.first_token_indices_gpu[0].item()) - group_end = int(pooling_cursor.last_token_indices_gpu[-1].item()) + 1 - hidden_states_group = hidden_states[group_start:group_end] - hidden_states_lst = list(hidden_states_group.split(split_sizes)) - else: - hidden_states_lst = [] + # Use the already-CPU num_scheduled_tokens tensor so `.tolist()` + # doesn't trigger a GPU->CPU sync. torch.split produces the same + # consecutive slices as indexing with first/last per-sequence indices. + hidden_states_lst = list( + torch.split(hidden_states, pooling_cursor.num_scheduled_tokens_cpu.tolist()) + ) if not self.enable_chunked_prefill: return hidden_states_lst @@ -95,12 +90,14 @@ class StepPool(AllPool): pooling_metadata: PoolingMetadata, ) -> list[TokenPoolingMethodOutputItem]: pooled_data_lst = super().forward(hidden_states, pooling_metadata) - prompt_token_ids = pooling_metadata.get_prompt_token_ids() + # Use the CPU copy of prompt_token_ids so the step_tag_id mask can be + # resolved to indices without a d2h sync from boolean indexing. + prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu() pooling_params = pooling_metadata.pooling_params pooled_data = list[torch.Tensor | None]() - for data, token_id, pooling_param in zip( - pooled_data_lst, prompt_token_ids, pooling_params + for data, token_id_cpu, pooling_param in zip( + pooled_data_lst, prompt_token_ids_cpu, pooling_params ): # for unfinished chunked prefill if data is None: @@ -113,7 +110,9 @@ class StepPool(AllPool): data = data[:, returned_token_ids] if step_tag_id is not None: - data = data[token_id == step_tag_id] + idx_cpu = (token_id_cpu == step_tag_id).nonzero(as_tuple=True)[0] + idx = idx_cpu.to(data.device, non_blocking=True) + data = data[idx] pooled_data.append(data) From 1e9500410a21782847ae86561b4de7f3aa69f0bc Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Mon, 4 May 2026 19:50:38 -0700 Subject: [PATCH 07/20] [ROCm][Quantization][2/N] Refactor quark_moe w4a8 w/ oracle (#39136) Signed-off-by: Bowen Bao --- tests/kernels/moe/test_ocp_mx_moe.py | 396 +++++++++++++++++- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 292 +++++++++++++ .../experts/gpt_oss_triton_kernels_moe.py | 123 ------ .../layers/fused_moe/oracle/mxfp4.py | 139 +++++- .../layers/quantization/mxfp4.py | 6 +- .../layers/quantization/quark/quark_moe.py | 266 +++--------- 6 files changed, 875 insertions(+), 347 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index aefc35324d8..8ed7757f655 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -28,6 +28,25 @@ HOPPER_MXFP4_BF16_AVAILABLE = ( and has_flashinfer() ) +# ROCm platform and dependencies +ROCM_AVAILABLE = current_platform.is_rocm() +ROCM_TRITON_KERNELS_AVAILABLE = False +ROCM_AITER_AVAILABLE = False +ROCM_GFX950 = False + +if ROCM_AVAILABLE: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.platforms.rocm import on_gfx950 + from vllm.utils.import_utils import has_triton_kernels + + ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels() + ROCM_GFX950 = on_gfx950() + ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled() + + if ROCM_AITER_AVAILABLE: + from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp + from aiter.ops.triton.quant import dynamic_mxfp4_quant + if TRTLLM_GEN_MXFP4_AVAILABLE: from flashinfer import ( fp4_quantize, @@ -111,6 +130,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase): def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = None): # Note we add an extra bias of 1 to the linear layer + # Uses chunked layout: first half is gate, second half is up x_glu, x_linear = torch.chunk(x, 2, dim=-1) if limit is not None: x_glu = x_glu.clamp(max=limit) @@ -119,6 +139,16 @@ def swiglu(x, alpha: float = 1.702, beta: float = 1.0, limit: float | None = Non return out_glu * (x_linear + beta) +def swigluoai(x, alpha: float = 1.702, limit: float = 7.0): + # OAI swiglu uses interleaved layout: gate/up alternating + # See SwigluOAIAndMul in vllm/model_executor/layers/activation.py + gate, up = x[..., ::2], x[..., 1::2] + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + return (up + 1) * glu + + fp4_lookup_table = [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6] @@ -168,8 +198,20 @@ def reference_moe( beta, limit, act_type, - is_gated, + activation: str = "swiglu", + use_interleaved_layout: bool = False, ): + """ + Reference MoE implementation for accuracy testing. + + Args: + activation: One of "swiglu", "silu", "relu2". Controls the activation + function used after the first MLP. + use_interleaved_layout: If True, uses interleaved gate/up layout + (gate=x[..., ::2], up=x[..., 1::2]) as used by SWIGLUOAI. + If False, uses chunked layout (gate, up = chunk(x, 2)) as used + by standard swiglu/silu. + """ # renormalize routing experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True) expert_weights = torch.nn.functional.softmax(experts.values, dim=1) @@ -179,12 +221,21 @@ def reference_moe( mlp1_weight = w13[expert_indices, ...] mlp1_bias = bias13[expert_indices, ...] t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias - if is_gated: - t = swiglu(t, alpha=alpha, beta=beta, limit=limit) - else: + + # Apply activation + if activation in ("swiglu", "silu"): + if use_interleaved_layout: + # SWIGLUOAI: interleaved gate/up layout + t = swigluoai(t, alpha=alpha, limit=limit) + else: + # Standard swiglu/silu: chunked layout + t = swiglu(t, alpha=alpha, beta=beta, limit=limit) + elif activation == "relu2": # RELU2_NO_MUL: relu(x)^2 t = torch.relu(t) t = t * t + else: + raise ValueError(f"Unknown activation: {activation}") if act_type == "mxfp8": t_quantized, t_scale = mxfp8_quantize( @@ -585,7 +636,8 @@ def test_trtllm_gen_mxfp4_fused_moe( beta, limit, act_type, - is_gated=True, + activation="swiglu", + use_interleaved_layout=False, ) ref_result[start_idx:end_idx].copy_(chunk_result) @@ -722,7 +774,8 @@ def test_flashinfer_cutlass_mxfp4_fused_moe( beta, limit, "bf16", - is_gated=True, + activation="swiglu", + use_interleaved_layout=False, ) from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe @@ -908,7 +961,8 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe( beta, limit, "mxfp8", - is_gated=True, + activation="swiglu", + use_interleaved_layout=False, ) # Prepare inputs for FlashInfer CUTLASS fused MoE @@ -1080,7 +1134,8 @@ def test_trtllm_gen_mxfp8_block_scale_moe( beta=0.0, limit=None, act_type="mxfp8", - is_gated=is_gated, + activation="swiglu" if is_gated else "relu2", + use_interleaved_layout=False, ) # Shuffle weights/scales with the same indexed layout used by TRTLLM kernels. @@ -1150,3 +1205,328 @@ def test_trtllm_gen_mxfp8_block_scale_moe( # Block-scale MXFP8 kernels are approximate; require majority close. check_accuracy(ref, out, atol=0.1, rtol=0.85, percent=0.8) + + +# ----------------------------------------------------------------------------- +# ROCm Oracle-based kernel execution tests +# ----------------------------------------------------------------------------- +# TODO: Further tighten the accuracy threshold. +# - More accurate ref moe to include activation quantization +# - Check aiter kernel accuracy. E.g., quant / dequant details. +ROCM_BACKEND_CONFIGS = { + "TRITON": { + "activation": "SWIGLUOAI", + "rtol": 0.3, + "percent": 0.95, + "requires_aiter": False, + "requires_gfx950": False, + }, + "TRITON_UNFUSED": { + "activation": "SWIGLUOAI", + "rtol": 0.3, + "percent": 0.95, + "requires_aiter": False, + "requires_gfx950": False, + }, + "AITER_MXFP4_BF16": { + "activation": "SILU", + "rtol": 1.0, + "percent": 0.7, + "requires_aiter": True, + "requires_gfx950": True, + }, + "AITER_MXFP4_FP8": { + "activation": "SWIGLUOAI", + "rtol": 0.5, + "percent": 0.9, + "requires_aiter": True, + "requires_gfx950": True, + }, +} + + +@pytest.mark.parametrize("backend_name", list(ROCM_BACKEND_CONFIGS.keys())) +@pytest.mark.parametrize("topk", [4]) +@pytest.mark.parametrize("num_experts", [8]) +@pytest.mark.parametrize("num_tokens,hidden_size,intermediate_size", [(16, 256, 256)]) +@pytest.mark.skipif( + not ROCM_AVAILABLE, + reason="ROCm is required for this test", +) +@torch.inference_mode() +def test_rocm_mxfp4_moe_oracle( + backend_name: str, + topk: int, + num_experts: int, + num_tokens: int, + hidden_size: int, + intermediate_size: int, +): + """ + Test ROCm MXFP4 MoE using oracle functions. + + This test validates that the oracle functions work end-to-end: + - select_mxfp4_moe_backend() selects a valid backend + - convert_to_mxfp4_moe_kernel_format() converts weights without error + - make_mxfp4_moe_quant_config() builds a valid quant config + - make_mxfp4_moe_kernel() creates a kernel that runs without error + - The kernel output is within accuracy tolerance of reference + """ + config = ROCM_BACKEND_CONFIGS[backend_name] + + # Check platform requirements + if not ROCM_TRITON_KERNELS_AVAILABLE: + pytest.skip("triton_kernels required for quantization") + if config["requires_aiter"] and not ROCM_AITER_AVAILABLE: + pytest.skip(f"Backend {backend_name} requires AITER") + if config["requires_gfx950"] and not ROCM_GFX950: + pytest.skip(f"Backend {backend_name} requires GFX950") + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + backend_to_kernel_cls, + convert_to_mxfp4_moe_kernel_format, + make_mxfp4_moe_kernel, + make_mxfp4_moe_quant_config, + ) + from vllm.v1.worker.workspace import init_workspace_manager + + # Initialize workspace manager (needed for modular kernels) + init_workspace_manager(torch.accelerator.current_device_index()) + + # Map string to enum + backend = Mxfp4MoeBackend[backend_name] + + # Get experts class from oracle + experts_cls_list = backend_to_kernel_cls(backend) + if experts_cls_list is None or len(experts_cls_list) == 0: + pytest.skip(f"Backend {backend_name} not available") + + # Use first experts class + experts_cls = experts_cls_list[0] + + torch.manual_seed(42) + dtype = torch.bfloat16 + device = "cuda:0" + + # Create MoE config with Renormalize routing (required by monolithic kernels) + from vllm.model_executor.layers.fused_moe import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + RoutingMethodType, + ) + + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=hidden_size, + intermediate_size_per_partition=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation[config["activation"]], + in_dtype=dtype, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + ) + + # Create float weights in checkpoint format: + # w13: [num_experts, 2*intermediate_size, hidden_size] + # w2: [num_experts, hidden_size, intermediate_size] + w13_float = torch.randn( + num_experts, 2 * intermediate_size, hidden_size, dtype=dtype, device=device + ) + w2_float = torch.randn( + num_experts, hidden_size, intermediate_size, dtype=dtype, device=device + ) + + # dynamic_mxfp4_quant expects 2D input, so reshape 3D weights + # w13: [E, 2*I, H] -> [E*2*I, H] -> quantize -> [E, 2*I, H//2] + # w2: [E, H, I] -> [E*H, I] -> quantize -> [E, H, I//2] + w13_2d = w13_float.reshape(-1, hidden_size) + w13_quant_2d, w13_scale_2d = dynamic_mxfp4_quant(w13_2d) + w13_quant = w13_quant_2d.reshape(num_experts, 2 * intermediate_size, -1) + w13_scale = w13_scale_2d.reshape(num_experts, 2 * intermediate_size, -1) + + w2_2d = w2_float.reshape(-1, intermediate_size) + w2_quant_2d, w2_scale_2d = dynamic_mxfp4_quant(w2_2d) + w2_quant = w2_quant_2d.reshape(num_experts, hidden_size, -1) + w2_scale = w2_scale_2d.reshape(num_experts, hidden_size, -1) + + w13_bias = torch.randn( + num_experts, 2 * intermediate_size, dtype=dtype, device=device + ) + w2_bias = torch.randn(num_experts, hidden_size, dtype=dtype, device=device) + + # Create static input scales for W4A8 backend (AITER_MXFP4_FP8) + w13_input_scale: torch.Tensor | None = None + w2_input_scale: torch.Tensor | None = None + if backend_name == "AITER_MXFP4_FP8": + # Static FP8 scales: one scale per expert + w13_input_scale = torch.ones(num_experts, dtype=torch.float32, device=device) + w2_input_scale = torch.ones(num_experts, dtype=torch.float32, device=device) + + # Create mock layer for oracle functions + class MockLayer: + w13_weight: torch.Tensor + w2_weight: torch.Tensor + w13_weight_scale: torch.Tensor + w2_weight_scale: torch.Tensor + w13_input_scale: torch.Tensor | None + w2_input_scale: torch.Tensor | None + + layer = MockLayer() + layer.w13_weight = w13_quant + layer.w2_weight = w2_quant + layer.w13_weight_scale = w13_scale + layer.w2_weight_scale = w2_scale + layer.w13_input_scale = w13_input_scale + layer.w2_input_scale = w2_input_scale + + # Convert weights using oracle + w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = ( + convert_to_mxfp4_moe_kernel_format( + mxfp4_backend=backend, + layer=layer, # type: ignore[arg-type] + w13_weight=w13_quant, + w2_weight=w2_quant, + w13_weight_scale=w13_scale, + w2_weight_scale=w2_scale, + w13_bias=w13_bias, + w2_bias=w2_bias, + ) + ) + + # Build quant config using oracle + quant_config = make_mxfp4_moe_quant_config( + mxfp4_backend=backend, + w1_scale=w13_scale_conv, + w2_scale=w2_scale_conv, + w1_bias=w13_bias_conv, + w2_bias=w2_bias_conv, + a1_scale=w13_input_scale, + a2_scale=w2_input_scale, + ) + + # Select activation based on backend + activation_name = str(config["activation"]) + activation = MoEActivation[activation_name] + + # Build kernel using oracle + assert quant_config is not None, "Failed to create quant config" + with set_current_vllm_config(VllmConfig()): + kernel = make_mxfp4_moe_kernel( + moe_quant_config=quant_config, + moe_config=moe_config, + mxfp4_backend=backend, + experts_cls=experts_cls, + routing_tables=None, + shared_experts=None, + ) + + # Create inputs + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + router_logits = torch.randn( + num_tokens, num_experts, dtype=torch.float32, device=device + ) + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) + + # Run kernel - use appropriate method based on impl type + if kernel.is_monolithic: + # Monolithic impl uses router_logits + out = kernel.apply_monolithic( + hidden_states=x, + w1=w13_conv, + w2=w2_conv, + router_logits=router_logits, + activation=activation, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + else: + # Modular impl uses topk_weights and topk_ids + out = kernel.apply( + hidden_states=x, + w1=w13_conv, + w2=w2_conv, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + + # Verify output is valid (no NaN/Inf) and has expected shape + assert out.shape == (num_tokens, hidden_size), f"Unexpected shape: {out.shape}" + assert not torch.any(torch.isnan(out)), "Output contains NaN" + assert not torch.any(torch.isinf(out)), "Output contains Inf" + + # Verify output has reasonable magnitude (not all zeros) + assert out.abs().max() > 0.01, "Output is effectively zero" + + # Dequantize weights for reference computation + w13_dq = upcast_from_mxfp( + w13_quant.view(torch.uint8), w13_scale, torch.bfloat16, axis=-1 + ) + w2_dq = upcast_from_mxfp( + w2_quant.view(torch.uint8), w2_scale, torch.bfloat16, axis=-1 + ) + + # Determine activation type and layout + # SWIGLUOAI uses interleaved layout (gate/up alternating) + # SILU uses chunked layout (first half gate, second half up) + use_interleaved = activation == MoEActivation.SWIGLUOAI + if activation in [MoEActivation.SWIGLUOAI, MoEActivation.SILU]: + act_name = "swiglu" + else: + act_name = "relu2" + + ref = reference_moe( + router_logits, + topk, + num_experts, + x.to(torch.float32), + w13_dq.to(torch.float32), + w13_bias.to(torch.float32), + w2_dq.to(torch.float32), + w2_bias.to(torch.float32), + alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0, + beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0, + limit=7.0 if activation == MoEActivation.SWIGLUOAI else None, + act_type="bf16", + activation=act_name, + use_interleaved_layout=use_interleaved, + ) + + # Compute and print accuracy statistics + diff = (ref.float() - out.float()).abs() + rel_diff = diff / (ref.float().abs() + 1e-6) + + print(f"\n[{backend_name}] Accuracy statistics:") + print( + f" Reference: min={ref.min():.4f}, max={ref.max():.4f}, mean={ref.mean():.4f}" + ) + print( + f" Output: min={out.min():.4f}, max={out.max():.4f}, mean={out.mean():.4f}" + ) + print( + f" Abs diff: min={diff.min():.4f}, max={diff.max():.4f}, " + f"mean={diff.mean():.4f}" + ) + print( + f" Rel diff: min={rel_diff.min():.4f}, max={rel_diff.max():.4f}, " + f"mean={rel_diff.mean():.4f}" + ) + + # Check what percentage of values are within various tolerances + for rtol in [0.1, 0.5, 1.0, 2.0]: + within_tol = (diff <= rtol * out.float().abs()).float().mean() + print(f" Within rtol={rtol}: {within_tol * 100:.1f}%") + + # Check accuracy using per-backend thresholds + check_accuracy(ref, out, atol=0.1, rtol=config["rtol"], percent=config["percent"]) diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py new file mode 100644 index 00000000000..3906a7e057c --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +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.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8StaticTensorSym, + kMxfp4Static, +) + +__all__ = [ + "AiterW4A8ExpertsMonolithic", + "aiter_triton_kernel_w4a8_moe_forward", +] + + +def aiter_triton_kernel_w4a8_moe_forward( + hidden_states: torch.Tensor, + w1, # Tensor or triton_kernels.Tensor + w2, # Tensor or triton_kernels.Tensor + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + quant_config: FusedMoEQuantConfig | None = None, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, +): + assert ( + quant_config is not None + and quant_config.use_mxfp4_w4a8 + and rocm_aiter_ops.is_enabled() + ) + from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) + return triton_kernel_fused_mxfp4_w4a8_experts( + None, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + activation=activation.value, + quant_config=quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + unpadded_N_w1=unpadded_N_w1, + unpadded_K_w1=unpadded_K_w1, + unpadded_N_w2=unpadded_N_w2, + unpadded_K_w2=unpadded_K_w2, + ) + + +def triton_kernel_fused_mxfp4_w4a8_experts( + output_tensor: torch.Tensor, + hidden_states: torch.Tensor, + w1, # Tensor or triton_kernels.Tensor + w2, # Tensor or triton_kernels.Tensor + routing_data, # RoutingData + gather_indx, # GatherIndx + scatter_indx, # ScatterIndx + activation: str = "silu", + quant_config: FusedMoEQuantConfig | None = None, + swiglu_alpha: float = 1.702, + swiglu_limit: float = 7.0, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + a1q_scale: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, +) -> torch.Tensor: + assert quant_config is not None + # type check, uint8 means mxfp4 + assert hidden_states.dtype == torch.bfloat16 + assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32 + assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32 + + # Shape check: weights are padded (e.g. hidden_size padded for + # GFX950 swizzle). + assert hidden_states.shape[-1] == w1.shape[-2] + assert w2.shape[-1] == w1.shape[1] + + E, _, N = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + gammas = routing_data.gate_scal if routing_data else None + + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + + assert quant_config.w1_precision is not None, ( + "w1_precision in quant config can't be None" + ) + assert quant_config.w2_precision is not None, ( + "w2_precision in quant config can't be None" + ) + + hidden_states = downcast_to_static_fp8( + hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale + ) + + intermediate_cache1 = moe_gemm_a8w4( + hidden_states, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + quant_config.w1_precision.flex_ctx.lhs_data.scale, + quant_config.w2_precision.flex_ctx.lhs_data.scale, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale="CDNA4_SCALE", + out_dtype=torch.float8_e4m3fn, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + + intermediate_cache3 = moe_gemm_a8w4( + intermediate_cache1, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + quant_config.w2_precision.flex_ctx.lhs_data.scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale="CDNA4_SCALE", + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + + return intermediate_cache3 + + +class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + """ + Monolithic MXFP4 W4A8 expert using AITER triton kernels. + + This backend uses: + - aiter.ops.triton.moe_routing.routing for routing + - aiter.ops.triton.moe_op_gemm_a8w4.moe_gemm_a8w4 for computation + + Weight format: MXFP4 weights with GFX950 swizzle + Activation: Static FP8 quantization + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self.topk = moe_config.experts_per_token + self.renormalize = moe_config.routing_method in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + # Requires AITER and GFX950 + if not rocm_aiter_ops.is_enabled(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # W4A8: MXFP4 weights with static FP8 activations + SUPPORTED_W_A = [ + (kMxfp4Static, kFp8StaticTensorSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + # Only SILU activation (swiglu) is supported + return activation == MoEActivation.SWIGLUOAI + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False # Expert parallelism not yet supported + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert self.moe_config.intermediate_size_per_partition_unpadded is not None + assert self.moe_config.hidden_dim_unpadded is not None + return aiter_triton_kernel_w4a8_moe_forward( + hidden_states=hidden_states, + w1=w1, + w2=w2, + gating_output=router_logits, + topk=self.topk, + renormalize=self.renormalize, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=self.quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2, + unpadded_K_w1=self.moe_config.hidden_dim_unpadded, + unpadded_N_w2=self.moe_config.hidden_dim_unpadded, + unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index ac317ac7762..e10514debd0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -5,7 +5,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops -from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -286,35 +285,6 @@ def triton_kernel_moe_forward( unpadded_N_w2=None, unpadded_K_w2=None, ) -> torch.Tensor: - if ( - quant_config is not None - and quant_config.use_mxfp4_w4a8 - and rocm_aiter_ops.is_enabled() - ): - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing - - routing_data, gather_idx, scatter_idx = aiter_routing( - gating_output, topk, sm_first=not renormalize - ) - return triton_kernel_fused_mxfp4_w4a8_experts( - None, - hidden_states, - w1, - w2, - routing_data, - gather_idx, - scatter_idx, - activation=activation.value, - quant_config=quant_config, - apply_router_weight_on_input=apply_router_weight_on_input, - global_num_experts=global_num_experts, - expert_map=expert_map, - unpadded_N_w1=unpadded_N_w1, - unpadded_K_w1=unpadded_K_w1, - unpadded_N_w2=unpadded_N_w2, - unpadded_K_w2=unpadded_K_w2, - ) - from triton_kernels.topk import topk as topk_fn sm_first = not renormalize @@ -471,99 +441,6 @@ def triton_kernel_fused_experts( return output_tensor -# This is a triton implementation of the fused_experts function -def triton_kernel_fused_mxfp4_w4a8_experts( - output_tensor: torch.Tensor, - hidden_states: torch.Tensor, - w1, # Tensor or triton_kernels.Tensor - w2, # Tensor or triton_kernels.Tensor - routing_data, # RoutingData - gather_indx, # GatherIndx - scatter_indx, # ScatterIndx - activation: str = "silu", - quant_config: FusedMoEQuantConfig | None = None, - swiglu_alpha: float = 1.702, - swiglu_limit: float = 7.0, - apply_router_weight_on_input: bool = False, - global_num_experts: int = -1, - expert_map: torch.Tensor | None = None, - a1q_scale: torch.Tensor | None = None, - unpadded_N_w1=None, - unpadded_K_w1=None, - unpadded_N_w2=None, - unpadded_K_w2=None, -) -> torch.Tensor: - assert quant_config is not None - # type check, uint8 means mxfp4 - assert hidden_states.dtype == torch.bfloat16 - assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32 - assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32 - - # Shape check: weights are padded (e.g. hidden_size padded for - # GFX950 swizzle). - assert hidden_states.shape[-1] == w1.shape[-2] - assert w2.shape[-1] == w1.shape[1] - - E, _, N = w1.shape - - if global_num_experts == -1: - global_num_experts = E - - gammas = routing_data.gate_scal if routing_data else None - - from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 - from aiter.ops.triton.quant_moe import downcast_to_static_fp8 - - assert quant_config.w1_precision is not None, ( - "w1_precision in quant config can't be None" - ) - assert quant_config.w2_precision is not None, ( - "w2_precision in quant config can't be None" - ) - - hidden_states = downcast_to_static_fp8( - hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale - ) - - intermediate_cache1 = moe_gemm_a8w4( - hidden_states, - w1.storage.data, - None, - quant_config.w1_precision.weight_scale.storage.data, - quant_config.w1_precision.flex_ctx.lhs_data.scale, - quant_config.w2_precision.flex_ctx.lhs_data.scale, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - gammas=gammas if apply_router_weight_on_input else None, - swizzle_mx_scale="CDNA4_SCALE", - out_dtype=torch.float8_e4m3fn, - apply_swiglu=True, - alpha=swiglu_alpha, - limit=swiglu_limit, - unpadded_N=unpadded_N_w1, - unpadded_K=unpadded_K_w1, - ) - - intermediate_cache3 = moe_gemm_a8w4( - intermediate_cache1, - w2.storage.data, - None, - quant_config.w2_precision.weight_scale.storage.data, - quant_config.w2_precision.flex_ctx.lhs_data.scale, - None, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - gammas=None if apply_router_weight_on_input else gammas, - swizzle_mx_scale="CDNA4_SCALE", - unpadded_N=unpadded_N_w2, - unpadded_K=unpadded_K_w2, - ) - - return intermediate_cache3 - - def make_routing_data( topk_ids: torch.Tensor, topk_weights: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3f2aca27716..437da8e6438 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, FusedMoEQuantDesc, mxfp4_mxfp8_moe_quant_config, + mxfp4_w4a8_moe_quant_config, mxfp4_w4a16_moe_quant_config, ocp_mx_moe_quant_config, ) @@ -26,9 +27,11 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_m from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, + kFp8StaticTensorSym, kMxfp4Static, kMxfp8Dynamic, ) +from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d from vllm.platforms import current_platform from vllm.utils.import_utils import has_triton_kernels from vllm.utils.math_utils import round_up @@ -59,8 +62,9 @@ class Mxfp4MoeBackend(Enum): # Marlin BATCHED_MARLIN = "BATCHED_MARLIN" MARLIN = "MARLIN" - # ROCm AITER - AITER = "AITER" + # ROCm AITER backends + AITER_MXFP4_BF16 = "AITER_MXFP4_BF16" # W4A16: CK kernel + AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel # Triton TRITON = "TRITON" TRITON_UNFUSED = "TRITON_UNFUSED" @@ -72,6 +76,13 @@ class Mxfp4MoeBackend(Enum): HUMMING = "HUMMING" +# AITER backends group +AITER_BACKENDS = ( + Mxfp4MoeBackend.AITER_MXFP4_BF16, + Mxfp4MoeBackend.AITER_MXFP4_FP8, +) + + # Backends that share the same TRTLLM weight format TRTLLM_BACKENDS = ( Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, @@ -159,13 +170,20 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] - elif backend == Mxfp4MoeBackend.AITER: + elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( AiterExperts, ) return [AiterExperts] + elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( + AiterW4A8ExpertsMonolithic, + ) + + return [AiterW4A8ExpertsMonolithic] + elif backend == Mxfp4MoeBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMXFp4 @@ -194,7 +212,8 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend: "triton_unfused": Mxfp4MoeBackend.TRITON_UNFUSED, "humming": Mxfp4MoeBackend.HUMMING, "marlin": Mxfp4MoeBackend.MARLIN, - "aiter": Mxfp4MoeBackend.AITER, + "aiter": Mxfp4MoeBackend.AITER_MXFP4_BF16, + "aiter_mxfp4_fp8": Mxfp4MoeBackend.AITER_MXFP4_FP8, "xpu": Mxfp4MoeBackend.XPU, "emulation": Mxfp4MoeBackend.EMULATION, } @@ -213,7 +232,8 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: """ _AVAILABLE_BACKENDS = [ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, - Mxfp4MoeBackend.AITER, + Mxfp4MoeBackend.AITER_MXFP4_BF16, + Mxfp4MoeBackend.AITER_MXFP4_FP8, Mxfp4MoeBackend.TRITON, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, # TRITON_UNFUSED has bug with MTP support @@ -254,16 +274,28 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, ): return kMxfp8Dynamic - return None + if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: + return kFp8StaticTensorSym + return None # BF16 activation -def select_gpt_oss_mxfp4_moe_backend( +def select_mxfp4_moe_backend( config: FusedMoEConfig, + activation_key: QuantKey | None = None, ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: """ Select the primary MXFP4 MoE backend. + + Args: + config: MoE configuration + activation_key: Optional activation quantization key. If provided, + overrides the default activation key for backend selection. + Use kFp8StaticTensorSym for W4A8 scheme. + Note: Shape-specific fallbacks may still occur at runtime. """ + # If activation_key is explicitly provided (e.g., W4A8), use it + requested_activation_key = activation_key device_capability = current_platform.get_device_capability() triton_kernels_supported = ( has_triton_kernels() @@ -332,11 +364,17 @@ def select_gpt_oss_mxfp4_moe_backend( and requested_backend == Mxfp4MoeBackend.MARLIN ): requested_backend = Mxfp4MoeBackend.BATCHED_MARLIN + # Use requested_activation_key if provided, otherwise use backend default + act_key = ( + requested_activation_key + if requested_activation_key is not None + else _backend_activation_key(requested_backend) + ) return _return_or_raise( requested_backend, config, kMxfp4Static, - _backend_activation_key(requested_backend), + act_key, activation_format, ) @@ -408,10 +446,15 @@ def select_gpt_oss_mxfp4_moe_backend( ) for backend in AVAILABLE_BACKENDS: - activation_key = _backend_activation_key(backend) + # Use requested_activation_key if provided, otherwise use backend default + act_key = ( + requested_activation_key + if requested_activation_key is not None + else _backend_activation_key(backend) + ) for k_cls in backend_to_kernel_cls(backend): supported, reason = k_cls.is_supported_config( - k_cls, config, kMxfp4Static, activation_key, activation_format + k_cls, config, kMxfp4Static, act_key, activation_format ) if supported: logger.info_once(_make_log_backend(backend)) @@ -438,7 +481,7 @@ def select_gpt_oss_mxfp4_moe_backend( return Mxfp4MoeBackend.NONE, None -def select_mxfp4_moe_backend( +def select_deepseek_v4_mxfp4_moe_backend( config: FusedMoEConfig, ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: """ @@ -836,7 +879,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER: + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: from vllm._aiter_ops import rocm_aiter_ops if w13_bias is not None: @@ -898,6 +941,63 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( w2_bias, ) + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: + # W4A8: MXFP4 weights + static FP8 activations (triton kernel) + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig + from triton_kernels.numerics import InFlexData + + if w13_bias is not None: + w13_bias = w13_bias.to(torch.float32) + if w2_bias is not None: + w2_bias = w2_bias.to(torch.float32) + + # Process static FP8 input scales (reduce to scalar, warn if not uniform) + w13_input_scale = layer.w13_input_scale + w2_input_scale = layer.w2_input_scale + if w13_input_scale is None or w2_input_scale is None: + raise ValueError( + "W4A8 (AITER_MXFP4_FP8) requires static input scales, but found " + "w13_input_scale or w2_input_scale is None." + ) + if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale): + logger.warning_once( + "Found input_scales that are not equal for " + "fp8 MoE layer. Using the maximum across experts " + "for each layer." + ) + w13_input_scale = w13_input_scale.max().to(torch.float32) + w2_input_scale = w2_input_scale.max().to(torch.float32) + + # Swizzle weights for GFX950 + w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(w13_weight, w13_weight_scale) + w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(w2_weight, w2_weight_scale) + + # Create InFlexData for activation scales + lhs_data13 = InFlexData(scale=w13_input_scale) + lhs_data2 = InFlexData(scale=w2_input_scale) + + # Create PrecisionConfig with both weight and activation info + w13_precision_config = PrecisionConfig( + weight_scale=w13_scale, + flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13), + ) + w2_precision_config = PrecisionConfig( + weight_scale=w2_scale, + flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), + ) + + del layer.w13_weight + del layer.w2_weight + + return ( + w13_weight, + w2_weight, + w13_precision_config, + w2_precision_config, + w13_bias, + w2_bias, + ) + elif mxfp4_backend in TRITON_BACKENDS: from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig @@ -1220,6 +1320,8 @@ def make_mxfp4_moe_quant_config( swiglu_limit: float | None = None, w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, + a1_scale: torch.Tensor | None = None, + a2_scale: torch.Tensor | None = None, layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig | None: """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" @@ -1262,6 +1364,17 @@ def make_mxfp4_moe_quant_config( gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, ) + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: + # W4A8: MXFP4 weights + static FP8 activations + return mxfp4_w4a8_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + w1_bias=w1_bias, + w2_bias=w2_bias, + block_shape=None, + ) elif mxfp4_backend in ( Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, @@ -1269,7 +1382,7 @@ def make_mxfp4_moe_quant_config( Mxfp4MoeBackend.TRITON_UNFUSED, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, - Mxfp4MoeBackend.AITER, + Mxfp4MoeBackend.AITER_MXFP4_BF16, ): return mxfp4_w4a16_moe_quant_config( w1_bias=w1_bias, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 2be77f2b8b8..d6fef0b3d3d 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -24,7 +24,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, mxfp4_round_up_hidden_size_and_intermediate_size, - select_gpt_oss_mxfp4_moe_backend, + select_deepseek_v4_mxfp4_moe_backend, select_mxfp4_moe_backend, ) from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod @@ -140,7 +140,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): def __init__(self, moe: FusedMoEConfig): super().__init__(moe) self.weight_dtype = "gpt_oss_mxfp4" - self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe) + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) self.max_capture_size = ( get_current_vllm_config().compilation_config.max_cudagraph_capture_size @@ -468,7 +468,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): def __init__(self, moe: FusedMoEConfig): super().__init__(moe) self.weight_dtype = "mxfp4" - self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) + self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe) self.max_capture_size = ( get_current_vllm_config().compilation_config.max_cudagraph_capture_size diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 1eeca142343..a14bfbc9c19 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -35,19 +35,19 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, mxfp4_round_up_hidden_size_and_intermediate_size, - select_gpt_oss_mxfp4_moe_backend, + select_mxfp4_moe_backend, ) from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( prepare_fp8_moe_layer_for_marlin, ) -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( - _swizzle_mxfp4, -) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, OCP_MX_Scheme, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + kFp8StaticTensorSym, +) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( all_close_1d, normalize_e4m3fn_to_e4m3fnuz, @@ -62,7 +62,6 @@ logger = init_logger(__name__) __all__ = [ "QuarkMoEMethod", "QuarkOCP_MX_MoEMethod", - "QuarkOCP_MX_MoEMethod_OSS", ] @@ -94,22 +93,9 @@ class QuarkMoEMethod(FusedMoEMethodBase): elif quant_config._is_fp8_w8a8(weight_config, input_config): return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config) elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config): - emulate = not current_platform.supports_mx() or not ( - rocm_aiter_ops.is_fused_moe_enabled() - ) - if ( - input_config is not None - and input_config.get("dtype") == "fp8_e4m3" - and not input_config.get("is_dynamic") - and not emulate - ): - return QuarkOCP_MX_MoEMethod_OSS( - weight_config, input_config, module.moe_config - ) - else: - return QuarkOCP_MX_MoEMethod( - weight_config, input_config, module.moe_config - ) + # All OCP MX schemes (W4A16, W4A8, etc.) handled by QuarkOCP_MX_MoEMethod + # Backend selection happens inside via oracle + return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config) elif quant_config._is_static_tensor_w8a8( weight_config, input_config ) or quant_config._is_dynamic_per_token_w8a8(weight_config, input_config): @@ -993,7 +979,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): self.experts_cls: type[mk.FusedMoEExperts] | None = None self.moe_kernel: mk.FusedMoEKernel | None = None - # Used for triton kernel precision configs + # Used for triton kernel precision configs (W4A8, TRITON backends) self.w13_precision_config = None self.w2_precision_config = None @@ -1002,6 +988,17 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): else: self.static_input_scales = False + # Select backend based on OCP MX scheme + if self.ocp_mx_scheme == "w_mxfp4": + # W4A16: weight-only MXFP4 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) + elif self.ocp_mx_scheme == "w_mxfp4_a_fp8" and self.static_input_scales: + # W4A8: MXFP4 weights + static FP8 activations + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend( + moe, activation_key=kFp8StaticTensorSym + ) + + # Validation for unsupported schemes if any( self.ocp_mx_scheme.endswith(a_scheme) for a_scheme in ["a_mxfp4", "a_mxfp6_e3m2", "a_mxfp6_e2m3"] @@ -1026,7 +1023,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): ) # TODO: Remove once all OCP MX schemes use the kernel abstraction - _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4", "w_mxfp4_a_mxfp4") + _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4", "w_mxfp4_a_mxfp4", "w_mxfp4_a_fp8") self.emulate = ( not current_platform.supports_mx() or self.ocp_mx_scheme not in _AITER_NATIVE_OCP_MX_SCHEMES @@ -1034,9 +1031,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe ) - if self.ocp_mx_scheme == "w_mxfp4": - self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe) - if self.emulate: # We use the same code path between MXFP4/MXFP6 emulation. self.mxfp4_backend = Mxfp4MoeBackend.EMULATION @@ -1046,7 +1040,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): if self.mxfp4_backend != Mxfp4MoeBackend.NONE: self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] - if self.emulate: + # Log backend selection + if self.mxfp4_backend != Mxfp4MoeBackend.NONE: + logger.info_once( + f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}" + ) + elif self.emulate: logger.warning_once( f"The current mode (supports_mx={current_platform.supports_mx()}, " f"use_rocm_aiter_moe={self.use_rocm_aiter_moe}, " @@ -1056,10 +1055,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): "QDQ (quantize and dequantize) will be used, with the linear " "layers computed in high precision." ) - else: - logger.warning_once( - "The current mode supports native MoE MXFP4 computation" - ) def maybe_roundup_sizes( self, @@ -1204,6 +1199,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): layer.w2_input_scale = None def process_weights_after_loading(self, layer): + # For MXFP4 schemes with native backend, use oracle + if self.mxfp4_backend != Mxfp4MoeBackend.NONE: + self._setup_kernel(layer) + return + if self.static_input_scales and self.input_dtype == "fp8": # firstly, process activations if fp8 static input if layer.w13_input_scale is None or layer.w2_input_scale is None: @@ -1252,14 +1252,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): w2_input_scale, requires_grad=False ) - # For w_mxfp4, use oracle functions - if self.emulate or ( - self.ocp_mx_scheme == "w_mxfp4" - and self.mxfp4_backend != Mxfp4MoeBackend.NONE - ): - self._setup_kernel_via_oracle(layer) - return - # TODO(bowenbao): gradually migrate to oracles. # Existing AITER path for w_mxfp4_a_mxfp4 and other schemes from aiter.utility.fp4_utils import e8m0_shuffle @@ -1298,46 +1290,48 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) torch.accelerator.empty_cache() - def _setup_kernel_via_oracle(self, layer: FusedMoE): - """Setup kernel using oracle functions for w_mxfp4 scheme.""" - w13 = layer.w13_weight - w2 = layer.w2_weight - w13_scale = layer.w13_weight_scale - w2_scale = layer.w2_weight_scale + def _setup_kernel(self, layer: FusedMoE): + """Setup kernel using oracle functions for MXFP4 schemes (W4A16, W4A8).""" w13_bias = getattr(layer, "w13_bias", None) w2_bias = getattr(layer, "w2_bias", None) - # Convert weights to kernel format + # Convert weights to kernel format (handles all backend-specific logic) w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend=self.mxfp4_backend, layer=layer, - w13_weight=w13, - w2_weight=w2, - w13_weight_scale=w13_scale, - w2_weight_scale=w2_scale, + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, w13_bias=w13_bias, w2_bias=w2_bias, ) ) - # For TRITON backends, weights are wrapped tensors from triton_kernels - # that don't support .detach(). Manually assign parameters. - if self.mxfp4_backend not in TRITON_BACKENDS: - replace_parameter(layer, "w13_weight", w13) - replace_parameter(layer, "w2_weight", w2) - replace_parameter(layer, "w13_weight_scale", w13_scale) - replace_parameter(layer, "w2_weight_scale", w2_scale) - else: + # Handle weight/scale assignment based on backend type + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_FP8, + ): + # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor + # Store on layer for apply(), scales are PrecisionConfig layer.w13_weight = w13 layer.w2_weight = w2 self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale + else: + # Standard backends: replace parameters + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) if w13_bias is not None and w2_bias is not None: replace_parameter(layer, "w13_bias", w13_bias) replace_parameter(layer, "w2_bias", w2_bias) + torch.accelerator.empty_cache() + # Build quant config and kernel self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config is not None and self.experts_cls is not None: @@ -1353,22 +1347,26 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: - # For w_mxfp4 with oracle backend, use oracle function - if self.ocp_mx_scheme == "w_mxfp4" and self.mxfp4_backend not in ( - Mxfp4MoeBackend.NONE, - Mxfp4MoeBackend.EMULATION, - ): - w1_scale = layer.w13_weight_scale - w2_scale = layer.w2_weight_scale - if self.mxfp4_backend in TRITON_BACKENDS: + # For oracle-based backends (W4A16, W4A8), use make_mxfp4_moe_quant_config + if self.mxfp4_backend not in (Mxfp4MoeBackend.NONE, Mxfp4MoeBackend.EMULATION): + # Determine scale source based on backend type + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_FP8, + ): w1_scale = self.w13_precision_config w2_scale = self.w2_precision_config + else: + w1_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale + return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=w1_scale, w2_scale=w2_scale, w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), + a1_scale=getattr(layer, "w13_input_scale", None), + a2_scale=getattr(layer, "w2_input_scale", None), ) # Emulation and other schemes @@ -1421,7 +1419,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - # For oracle kernel or emulation kernel + # For oracle-based kernels (W4A16, W4A8) or emulation kernel if self.moe_kernel is not None: return self.moe_kernel.apply( hidden_states=x, @@ -1473,135 +1471,3 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, ) - - -class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): - def __init__( - self, - weight_config: dict[str, Any], - input_config: dict[str, Any], - moe: FusedMoEConfig, - ): - super().__init__(weight_config, input_config, moe) - - def process_weights_after_loading(self, layer): - from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig - - w13_bias = layer.w13_bias.to(torch.float32) - w2_bias = layer.w2_bias.to(torch.float32) - - layer.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False) - layer.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) - - # FIXME warp need to be adjusted based on batch size - # only apply to batched mode - if self.moe.use_ep: - num_warps = 4 if self.moe.max_num_tokens <= 512 else 8 - else: - num_warps = 8 - - w13_weight, w13_flex, w13_scale = _swizzle_mxfp4( - layer.w13_weight, layer.w13_weight_scale, num_warps - ) - w2_weight, w2_flex, w2_scale = _swizzle_mxfp4( - layer.w2_weight, layer.w2_weight_scale, num_warps - ) - - self.w13_weight_triton_tensor = w13_weight - self.w2_weight_triton_tensor = w2_weight - - # need to delete the original weights to save memory on single GPU - del layer.w13_weight - del layer.w2_weight - layer.w13_weight = None - layer.w2_weight = None - torch.accelerator.empty_cache() - - if self.static_input_scales: - if layer.w13_input_scale is None or layer.w2_input_scale is None: - raise ValueError( - "QuantConfig has static quantization, but found " - "activation scales are None." - ) - if not all_close_1d(layer.w13_input_scale) or not all_close_1d( - layer.w2_input_scale - ): - logger.warning_once( - "Found input_scales that are not equal for " - "fp8 MoE layer. Using the maximum across experts " - "for each layer." - ) - - layer.w13_input_scale = torch.nn.Parameter( - layer.w13_input_scale.max().to(torch.float32), requires_grad=False - ) - layer.w2_input_scale = torch.nn.Parameter( - layer.w2_input_scale.max().to(torch.float32), requires_grad=False - ) - - from triton_kernels.numerics import InFlexData - - lhs_data13 = InFlexData(scale=layer.w13_input_scale) - lhs_data2 = InFlexData(scale=layer.w2_input_scale) - - self.w13_precision_config = PrecisionConfig( - weight_scale=w13_scale, - flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13), - ) - - self.w2_precision_config = PrecisionConfig( - weight_scale=w2_scale, - flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), - ) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return mxfp4_w4a8_moe_quant_config( - w1_scale=self.w13_precision_config, - w2_scale=self.w2_precision_config, - a1_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, - w1_bias=layer.w13_bias, - w2_bias=layer.w2_bias, - block_shape=None, - ) - - @property - def is_monolithic(self) -> bool: - return True - - def apply_monolithic( - self, - layer: FusedMoE, - x: torch.Tensor, - router_logits: torch.Tensor, - input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: - if layer.enable_eplb: - raise NotImplementedError( - f"EPLB not supported for {self.__class__.__name__} yet." - ) - - from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 - triton_kernel_moe_forward, - ) - - assert self.moe.hidden_dim_unpadded is not None - assert self.moe.intermediate_size_per_partition_unpadded is not None - return triton_kernel_moe_forward( - hidden_states=x, - w1=self.w13_weight_triton_tensor, - w2=self.w2_weight_triton_tensor, - gating_output=router_logits, - topk=layer.top_k, - renormalize=layer.renormalize, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - quant_config=self.moe_quant_config, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - unpadded_N_w1=self.moe.intermediate_size_per_partition_unpadded * 2, - unpadded_K_w1=self.moe.hidden_dim_unpadded, - unpadded_N_w2=self.moe.hidden_dim_unpadded, - unpadded_K_w2=self.moe.intermediate_size_per_partition_unpadded, - ) From 420b0a5c95187809b2701323f1112472b2f3b707 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar <61960177+Akashcodes732@users.noreply.github.com> Date: Tue, 5 May 2026 09:21:09 +0530 Subject: [PATCH 08/20] [Hardware][Power]Add Power VSX Attention Backend and fix l2 Cache Crash (#40451) Signed-off-by: Akash Kaothalkar Signed-off-by: Akash Kaothalkar Signed-off-by: Akash kaothalkar Co-authored-by: Akash Kaothalkar Co-authored-by: Akash Kaothalkar Co-authored-by: Li, Jiang --- csrc/cpu/cpu_attn.cpp | 4 + csrc/cpu/cpu_attn_impl.hpp | 5 +- csrc/cpu/cpu_attn_vec.hpp | 4 +- csrc/cpu/cpu_attn_vsx.hpp | 359 +++++++++++++++++++++++++ csrc/cpu/cpu_types_vsx.hpp | 4 + csrc/cpu/generate_cpu_attn_dispatch.py | 15 +- csrc/cpu/utils.hpp | 2 +- vllm/v1/attention/backends/cpu_attn.py | 15 +- 8 files changed, 399 insertions(+), 9 deletions(-) create mode 100644 csrc/cpu/cpu_attn_vsx.hpp diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 18afe4b7925..4750dd78838 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata( isa = cpu_attention::ISA::NEON; } else if (isa_hint == "vxe") { isa = cpu_attention::ISA::VXE; + } else if (isa_hint == "vsx") { + isa = cpu_attention::ISA::VSX; } else { TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint); } @@ -129,6 +131,8 @@ void cpu_attn_reshape_and_cache( return cpu_attention::ISA::NEON; } else if (isa == "vxe") { return cpu_attention::ISA::VXE; + } else if (isa == "vsx") { + return cpu_attention::ISA::VSX; } else { TORCH_CHECK(false, "Invalid ISA type: " + isa); } diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index f5b473bd262..b9987fb26c1 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -12,7 +12,7 @@ #include "cpu/utils.hpp" namespace cpu_attention { -enum class ISA { AMX, VEC, VEC16, NEON, VXE }; +enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX }; // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly. enum class Fp8KVCacheDataType { @@ -164,6 +164,9 @@ struct AttentionMetadata { case ISA::VXE: ss << "VXE, "; break; + case ISA::VSX: + ss << "VSX, "; + break; } ss << "workitem_group_num: " << workitem_group_num << ", reduction_item_num: " << reduction_item_num diff --git a/csrc/cpu/cpu_attn_vec.hpp b/csrc/cpu/cpu_attn_vec.hpp index 61cae12d67d..c3983e0578a 100644 --- a/csrc/cpu/cpu_attn_vec.hpp +++ b/csrc/cpu/cpu_attn_vec.hpp @@ -27,8 +27,8 @@ FORCE_INLINE std::pair load_b_pair_vec( return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)}; } else { using load_vec_t = typename VecTypeTrait::vec_t; - return {vec_op::FP32Vec16(load_vec_t(ptr)), - vec_op::FP32Vec16(load_vec_t(ptr + 16))}; + return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)), + vec_op::FP32Vec16(load_vec_t(ptr + 16))); } } diff --git a/csrc/cpu/cpu_attn_vsx.hpp b/csrc/cpu/cpu_attn_vsx.hpp new file mode 100644 index 00000000000..c7e1502bcb0 --- /dev/null +++ b/csrc/cpu/cpu_attn_vsx.hpp @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#ifndef CPU_ATTN_VSX_HPP +#define CPU_ATTN_VSX_HPP + +#include "cpu_attn_impl.hpp" +#include +#include + +namespace cpu_attention { + +namespace { + +// ppc64le Vector = 16 bytes (128 bits) +#define BLOCK_SIZE_ALIGNMENT 32 +#define HEAD_SIZE_ALIGNMENT 32 +#define MAX_Q_HEAD_NUM_PER_ITER 16 + +template +FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0, + __vector float& b1); + +// [1] Float Specialization +template <> +FORCE_INLINE void load_row8_B_as_f32(const float* p, __vector float& b0, + __vector float& b1) { + b0 = vec_xl(0, const_cast(p)); + b1 = vec_xl(0, const_cast(p + 4)); +} + +// [2] BFloat16 Specialization (Little Endian ppc64le) +// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32. +// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)] +// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed. +// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]} +// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair: +// uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs +// uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits +// Cast to float32 gives exactly (bf16_bits << 16) per element. +template <> +FORCE_INLINE void load_row8_B_as_f32(const c10::BFloat16* p, + __vector float& b0, + __vector float& b1) { + __vector unsigned short raw = vec_xl( + 0, reinterpret_cast(const_cast(p))); + __vector unsigned short zeros = vec_splat_u16(0); + + // LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32 + b0 = (__vector float)vec_mergeh(zeros, raw); + b1 = (__vector float)vec_mergel(zeros, raw); +} + +// Note: c10::Half (FP16) is not supported on PowerPC architecture + +template +FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4( + const float* __restrict A, // [M x K] + const kv_cache_t* __restrict B, // [K x 8] + float* __restrict C, // [M x 8] + int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) { + static_assert(1 <= M && M <= 8, "M must be in [1,8]"); + +#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7) +#define IF_M(i) if constexpr (M > (i)) + + // 1. Define A pointers +#define DECL_A(i) const float* a##i = A + (i) * lda; + ROWS_APPLY(DECL_A) +#undef DECL_A + + // 2. Define Accumulators (2 vectors covers 8 columns) +#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1; + ROWS_APPLY(DECL_ACC) +#undef DECL_ACC + + // 3. Initialize Accumulators (Load C or Zero) +#define INIT_ACC(i) \ + IF_M(i) { \ + if (accumulate) { \ + acc##i##_0 = vec_xl(0, const_cast(C + (i) * ldc + 0)); \ + acc##i##_1 = vec_xl(0, const_cast(C + (i) * ldc + 4)); \ + } else { \ + acc##i##_0 = vec_splats(0.0f); \ + acc##i##_1 = vec_splats(0.0f); \ + } \ + } + ROWS_APPLY(INIT_ACC) +#undef INIT_ACC + + int32_t k = 0; + + for (; k + 3 < K; k += 4) { + // Load 4 values of A for each Row M: A[k...k+3] +#define LOAD_A4(i) \ + __vector float a##i##v; \ + IF_M(i) a##i##v = vec_xl(0, const_cast(a##i + k)); + ROWS_APPLY(LOAD_A4) +#undef LOAD_A4 + + // FMA for specific lane L of A + // ppc64le: vec_madd(b, vec_splat(a, lane), acc) +#define FMAS_LANE(i, aiv, L) \ + IF_M(i) { \ + __vector float a_broad = vec_splat(aiv, L); \ + acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \ + acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \ + } + + // Unroll K=0..3 + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 0) * ldb, b0, b1); +#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0) + ROWS_APPLY(STEP_K0) +#undef STEP_K0 + } + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 1) * ldb, b0, b1); +#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1) + ROWS_APPLY(STEP_K1) +#undef STEP_K1 + } + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 2) * ldb, b0, b1); +#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2) + ROWS_APPLY(STEP_K2) +#undef STEP_K2 + } + { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)(k + 3) * ldb, b0, b1); +#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3) + ROWS_APPLY(STEP_K3) +#undef STEP_K3 + } +#undef FMAS_LANE + } + + for (; k < K; ++k) { + __vector float b0, b1; + load_row8_B_as_f32(B + (int64_t)k * ldb, b0, b1); +#define TAIL_ROW(i) \ + IF_M(i) { \ + __vector float ai = vec_splats(*(a##i + k)); \ + acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \ + acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \ + } + ROWS_APPLY(TAIL_ROW) +#undef TAIL_ROW + } + +#define STORE_ROW(i) \ + IF_M(i) { \ + vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \ + vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \ + } + ROWS_APPLY(STORE_ROW) +#undef STORE_ROW + +#undef ROWS_APPLY +#undef IF_M +} + +template +FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A, + const kv_cache_t* __restrict B, + float* __restrict C, int32_t M, + int32_t K, int64_t lda, + int64_t ldb, int64_t ldc, + bool accumulate) { + static_assert(N % 8 == 0, "N must be a multiple of 8"); + for (int32_t m = 0; m < M;) { + int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1; + const float* Ab = A + m * lda; + float* Cb = C + m * ldc; + + for (int32_t n = 0; n < N; n += 8) { + const kv_cache_t* Bn = B + n; + float* Cn = Cb + n; + switch (mb) { + case 8: + gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + case 4: + gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + case 2: + gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + default: + gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + } + } + m += mb; + } +} + +template +class TileGemmPPC64 { + public: + template + FORCE_INLINE static void gemm(const int32_t m_size, + float* __restrict__ a_tile, + kv_cache_t* __restrict__ b_tile, + float* __restrict__ c_tile, const int64_t lda, + const int64_t ldb, const int64_t ldc, + const int32_t block_size, + const int32_t dynamic_k_size, + const bool accum_c) { + if constexpr (phase == AttentionGemmPhase::QK) { + gemm_macro_ppc64le_Mx8_Ku4( + a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c); + } else { + gemm_macro_ppc64le_Mx8_Ku4( + a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc, + accum_c); + } + } +}; + +} // namespace + +template +class AttentionImpl { + public: + using query_t = scalar_t; + using q_buffer_t = float; + using kv_cache_t = scalar_t; + using logits_buffer_t = float; + using partial_output_buffer_t = float; + using prob_buffer_t = float; + + constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT; + constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT; + constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER; + constexpr static int64_t HeadDim = head_dim; + constexpr static ISA ISAType = ISA::VSX; + constexpr static bool scale_on_logits = + false; // Scale is applied to Q during copy + + public: + AttentionImpl() {} + + template