From 51295793a2eed0eefc7505cb9a7d5f96effd7773 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 1 May 2026 13:02:03 -0400 Subject: [PATCH 0001/1083] [Model Runner V2] Add `logprob_token_ids` support (#40559) Signed-off-by: yewentao256 Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- vllm/sampling_params.py | 25 +++++ vllm/v1/core/sched/scheduler.py | 2 +- vllm/v1/engine/logprobs.py | 2 +- vllm/v1/worker/gpu/sample/logprob.py | 142 +++++++++++++++++++++++++-- vllm/v1/worker/gpu/sample/sampler.py | 22 ++++- 5 files changed, 179 insertions(+), 14 deletions(-) diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 77fa6402180e..88b1b0b8e8e9 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -25,6 +25,10 @@ _SAMPLING_EPS = 1e-5 _MAX_TEMP = 1e-2 +MAX_LOGPROB_TOKEN_IDS = 128 +"""Upper bound on `SamplingParams.logprob_token_ids` list length. Must match +the per-request row width allocated by the sampler's `LogprobTokenIdsState`.""" + class SamplingType(IntEnum): GREEDY = 0 @@ -628,6 +632,16 @@ def bad_words_token_ids(self) -> list[list[int]] | None: # For internal use only. Backward compatibility not guaranteed return self._bad_words_token_ids + @property + def num_logprobs(self) -> int | None: + """Number of sample logprobs to return per output token, or `None` if + no sample logprobs were requested. Takes `logprob_token_ids` into + account: when `logprobs` is unset but `logprob_token_ids` is set, + returns `len(logprob_token_ids)`.""" + if self.logprobs is not None: + return self.logprobs + return len(self.logprob_token_ids) if self.logprob_token_ids else None + def clone(self) -> "SamplingParams": """If skip_clone is True, uses shallow copy instead of deep copy.""" if self.skip_clone: @@ -666,6 +680,17 @@ def _validate_logprobs(self, model_config: ModelConfig) -> None: value=num_logprobs, ) + # Validate logprob_token_ids. + if self.logprob_token_ids is not None: + n = len(self.logprob_token_ids) + if n > MAX_LOGPROB_TOKEN_IDS: + raise VLLMValidationError( + f"Requested logprob_token_ids of length {n}, " + f"which is greater than max allowed: {MAX_LOGPROB_TOKEN_IDS}", + parameter="logprob_token_ids", + value=n, + ) + # Validate prompt logprobs. if num_prompt_logprobs := self.prompt_logprobs: if num_prompt_logprobs == -1: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 7820b858bdf2..032767cdf3b0 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1435,7 +1435,7 @@ def update_from_output( # Extract sample logprobs if needed. if ( request.sampling_params is not None - and request.sampling_params.logprobs is not None + and request.sampling_params.num_logprobs is not None and logprobs ): new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) diff --git a/vllm/v1/engine/logprobs.py b/vllm/v1/engine/logprobs.py index 9ada6eda48ce..74a45ab1e4d4 100644 --- a/vllm/v1/engine/logprobs.py +++ b/vllm/v1/engine/logprobs.py @@ -47,7 +47,7 @@ def from_new_request( ) -> "LogprobsProcessor": sampling_params = request.sampling_params assert sampling_params is not None - num_logprobs = sampling_params.logprobs + num_logprobs = sampling_params.num_logprobs num_prompt_logprobs = sampling_params.prompt_logprobs return cls( tokenizer=tokenizer, diff --git a/vllm/v1/worker/gpu/sample/logprob.py b/vllm/v1/worker/gpu/sample/logprob.py index 4317cad9ce7f..7530337fcd12 100644 --- a/vllm/v1/worker/gpu/sample/logprob.py +++ b/vllm/v1/worker/gpu/sample/logprob.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import torch +from vllm.sampling_params import MAX_LOGPROB_TOKEN_IDS, SamplingParams from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsTensors +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor @triton.jit @@ -75,6 +78,9 @@ def _ranks_kernel( def compute_token_logprobs( logits: torch.Tensor, token_ids: torch.Tensor ) -> torch.Tensor: + # NOTE(woosuk): To save GPU memory, we do not materialize the full + # [batch_size, vocab_size] logprobs tensor. The kernel computes + # max + logsumexp per row and only emits logprobs at `token_ids`. batch_size, vocab_size = logits.shape token_ids = token_ids.to(torch.int64) num_logprobs = token_ids.shape[1] @@ -97,18 +103,52 @@ def compute_topk_logprobs( num_logprobs: int, sampled_token_ids: torch.Tensor, cu_num_logits: list[int] | None = None, + logprob_token_ids_state: "LogprobTokenIdsState | None" = None, + expanded_idx_mapping: torch.Tensor | None = None, + max_per_req_token_ids: int = 0, ) -> LogprobsTensors: assert num_logprobs >= 0 batch_size, vocab_size = logits.shape - logprob_token_ids = sampled_token_ids.unsqueeze(-1) - if num_logprobs > 0: - topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices - logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1) - - # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full - # logprobs tensor. Instead, we only compute and return the logprobs of - # the topk + 1 tokens. - logprobs = compute_token_logprobs(logits, logprob_token_ids) + + if max_per_req_token_ids == 0: + # Fast path: no request asked for custom logprob_token_ids. + logprob_token_ids = sampled_token_ids.unsqueeze(-1) + if num_logprobs > 0: + topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices + logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1) + logprobs = compute_token_logprobs(logits, logprob_token_ids) + else: + # Some requests specified logprob_token_ids. Build the [batch_size, + # 1 + max_cols] token_ids matrix and validity mask on the GPU via a + # single triton kernel, overriding the topk columns with per-request + # tokens where applicable. + assert logprob_token_ids_state is not None + assert expanded_idx_mapping is not None + topk_indices = None + if num_logprobs > 0: + topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices + + num_cols = max(num_logprobs, max_per_req_token_ids) + logprob_token_ids = sampled_token_ids.new_zeros((batch_size, 1 + num_cols)) + valid_mask = torch.zeros_like(logprob_token_ids, dtype=torch.bool) + _fill_logprob_token_ids_kernel[(batch_size,)]( + logprob_token_ids, + logprob_token_ids.stride(0), + valid_mask, + valid_mask.stride(0), + sampled_token_ids, + topk_indices if topk_indices is not None else logprob_token_ids, + topk_indices.stride(0) if topk_indices is not None else 0, + expanded_idx_mapping, + logprob_token_ids_state.num_token_ids.gpu, + logprob_token_ids_state.token_ids.gpu, + logprob_token_ids_state.token_ids.gpu.stride(0), + NUM_TOPK=num_logprobs, + PADDED_COLS=triton.next_power_of_2(num_cols), + ) + logprobs = compute_token_logprobs(logits, logprob_token_ids) + logprobs = logprobs.masked_fill(~valid_mask, float("-inf")) + token_ranks = torch.empty(batch_size, dtype=torch.int64, device=logits.device) _ranks_kernel[(batch_size,)]( token_ranks, @@ -124,3 +164,87 @@ def compute_topk_logprobs( selected_token_ranks=token_ranks, cu_num_generated_tokens=cu_num_logits, ) + + +@triton.jit +def _fill_logprob_token_ids_kernel( + # [batch_size, 1 + num_cols] + out_token_ids_ptr, + out_token_ids_stride, + # [batch_size, 1 + num_cols] + out_valid_mask_ptr, + out_valid_mask_stride, + sampled_token_ids_ptr, # [batch_size] + topk_indices_ptr, # [batch_size, NUM_TOPK] (unused when NUM_TOPK == 0) + topk_indices_stride, + expanded_idx_mapping_ptr, # [batch_size] -> req_state_idx + num_per_req_token_ids_ptr, # [max_num_reqs] + per_req_token_ids_ptr, # [max_num_reqs, MAX_LOGPROB_TOKEN_IDS] + per_req_token_ids_stride, + NUM_TOPK: tl.constexpr, + PADDED_COLS: tl.constexpr, +): + batch_idx = tl.program_id(0) + + # Column 0: always the sampled token, always valid. + sampled = tl.load(sampled_token_ids_ptr + batch_idx) + tl.store(out_token_ids_ptr + batch_idx * out_token_ids_stride, sampled) + tl.store(out_valid_mask_ptr + batch_idx * out_valid_mask_stride, 1) + + req_state_idx = tl.load(expanded_idx_mapping_ptr + batch_idx) + num_custom = tl.load(num_per_req_token_ids_ptr + req_state_idx) + + col = tl.arange(0, PADDED_COLS) + tid_base = out_token_ids_ptr + batch_idx * out_token_ids_stride + 1 + mask_base = out_valid_mask_ptr + batch_idx * out_valid_mask_stride + 1 + + if num_custom > 0: + # Override topk with per-request custom tokens. + src = per_req_token_ids_ptr + req_state_idx * per_req_token_ids_stride + valid = col < num_custom + # per_req_token_ids is int32; output is int64. + tokens = tl.load(src + col, mask=valid, other=0).to(tl.int64) + else: + # Fill with topk indices (no-op when NUM_TOPK == 0). + src = topk_indices_ptr + batch_idx * topk_indices_stride + valid = col < NUM_TOPK + tokens = tl.load(src + col, mask=valid, other=0) + + tl.store(tid_base + col, tokens, mask=valid) + tl.store(mask_base + col, tl.full([PADDED_COLS], 1, tl.int1), mask=valid) + + +class LogprobTokenIdsState: + """Per-request override of which token ids' logprobs to return. + + See `SamplingParams.logprob_token_ids`. + """ + + def __init__(self, max_num_reqs: int, device: torch.device): + self.max_num_reqs = max_num_reqs + self.num_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + self.token_ids = StagedWriteTensor( + (max_num_reqs, MAX_LOGPROB_TOKEN_IDS), + dtype=torch.int32, + device=device, + ) + + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: + token_ids = sampling_params.logprob_token_ids + if not token_ids: + self.num_token_ids.np[req_idx] = 0 + return + n = len(token_ids) + if n > MAX_LOGPROB_TOKEN_IDS: + raise ValueError( + f"Too many logprob_token_ids: {n}. The max is {MAX_LOGPROB_TOKEN_IDS}." + ) + self.num_token_ids.np[req_idx] = n + self.token_ids.stage_write(req_idx, 0, token_ids) + + def apply_staged_writes(self) -> None: + self.num_token_ids.copy_to_uva() + self.token_ids.apply_write() + + def max_num_token_ids(self, idx_mapping_np: np.ndarray) -> int: + return int(self.num_token_ids.np[idx_mapping_np].max(initial=0)) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 6f73ca87ac67..5d91d5b2f097 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -12,7 +12,10 @@ from vllm.v1.worker.gpu.sample.bad_words import BadWordsState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState -from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs +from vllm.v1.worker.gpu.sample.logprob import ( + LogprobTokenIdsState, + compute_topk_logprobs, +) from vllm.v1.worker.gpu.sample.output import SamplerOutput from vllm.v1.worker.gpu.sample.penalties import PenaltiesState from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates @@ -38,6 +41,7 @@ def __init__( self.penalties_state = PenaltiesState(req_states) self.logit_bias_state = LogitBiasState(max_num_reqs, device) self.bad_words_state = BadWordsState(req_states) + self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) self.num_speculative_tokens = num_speculative_tokens def add_request( @@ -47,12 +51,14 @@ def add_request( self.penalties_state.add_request(req_idx, sampling_params) self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params) self.bad_words_state.add_request(req_idx, sampling_params) + self.logprob_token_ids_state.add_request(req_idx, sampling_params) def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() self.penalties_state.apply_staged_writes() self.logit_bias_state.apply_staged_writes() self.bad_words_state.apply_staged_writes() + self.logprob_token_ids_state.apply_staged_writes() def __call__( self, @@ -79,13 +85,23 @@ def __call__( ) max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) - if max_num_logprobs != NO_LOGPROBS: + max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids( + idx_mapping_np + ) + if max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0: if self.logprobs_mode == "processed_logprobs": logits = processed_logits expanded_logits = logits.shape[0] != idx_mapping_np.shape[0] cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None + num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0 logprobs_tensors = compute_topk_logprobs( - logits, max_num_logprobs, sampled, cu_num_logits + logits, + num_logprobs, + sampled, + cu_num_logits, + logprob_token_ids_state=self.logprob_token_ids_state, + expanded_idx_mapping=input_batch.expanded_idx_mapping, + max_per_req_token_ids=max_per_req_token_ids, ) else: logprobs_tensors = None From f3fef123504db07b3ac83ad4ef677915b53e8386 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Fri, 1 May 2026 13:36:20 -0400 Subject: [PATCH 0002/1083] [Attention] Abstract the MLA prefill backends and eliminate cuDNN (#32623) Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Co-authored-by: Michael Goin Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/design/attention_backends.md | 20 +- tests/engine/test_arg_utils.py | 3 - tests/v1/attention/test_mla_backends.py | 1 + .../v1/attention/test_mla_prefill_selector.py | 304 ++++++++ .../generate_attention_backend_docs.py | 341 ++++++--- vllm/config/attention.py | 59 +- .../layers/attention/mla_attention.py | 653 ++---------------- vllm/platforms/interface.py | 3 + .../backends/mla/prefill/__init__.py | 11 + .../v1/attention/backends/mla/prefill/base.py | 125 ++++ .../backends/mla/prefill/flash_attn.py | 180 +++++ .../backends/mla/prefill/flashinfer.py | 211 ++++++ .../backends/mla/prefill/registry.py | 53 ++ .../backends/mla/prefill/selector.py | 183 +++++ .../backends/mla/prefill/trtllm_ragged.py | 178 +++++ vllm/v1/attention/backends/mla/triton_mla.py | 12 - 16 files changed, 1629 insertions(+), 708 deletions(-) create mode 100644 tests/v1/attention/test_mla_prefill_selector.py create mode 100644 vllm/v1/attention/backends/mla/prefill/__init__.py create mode 100644 vllm/v1/attention/backends/mla/prefill/base.py create mode 100644 vllm/v1/attention/backends/mla/prefill/flash_attn.py create mode 100644 vllm/v1/attention/backends/mla/prefill/flashinfer.py create mode 100644 vllm/v1/attention/backends/mla/prefill/registry.py create mode 100644 vllm/v1/attention/backends/mla/prefill/selector.py create mode 100644 vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index bc24a8a5a1f3..dc4b5402cab6 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -192,21 +192,25 @@ MLA uses separate backends for prefill and decode phases. ### Prefill Backends -The prefill backend is selected at runtime based on hardware and -configuration. +To explicitly select a prefill backend, use +`-ac.mla_prefill_backend=` (e.g., `FLASH_ATTN`, `FLASHINFER`). +Otherwise, the prefill backend is selected automatically at runtime based on +hardware and configuration. -| Backend | Description | Compute Cap. | Enable | Disable | Notes | -| ------- | ----------- | ------------ | ------ | ------- | ----- | -| TRT-LLM Ragged‡ | TensorRT-LLM ragged attention | 10.x | Default on SM100 | `-ac.use_trtllm_ragged_deepseek_prefill=0` | DeepSeek R1 dims only | -| FlashInfer | FlashInfer CUTLASS backend | 10.x | `-ac.disable_flashinfer_prefill=0` | `-ac.disable_flashinfer_prefill=1` | DeepSeek R1 dims only | -| cuDNN | cuDNN-based attention | 10.x | `-ac.use_cudnn_prefill=1` | `-ac.use_cudnn_prefill=0` | | -| FlashAttention | FlashAttention varlen (FA2/FA3) | Any | Default fallback | Use other backends | FA3 on SM90, FA2 otherwise | +| Backend | Description | Dtypes | Compute Cap. | Notes | +| ------- | ----------- | ------ | ------------ | ----- | +| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise | +| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only | +| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only | > **‡** TRT-LLM Ragged is the default on Blackwell (SM100). > On other GPUs, FlashAttention is used as the default. ### Decode Backends +MLA decode backends are selected using the standard +`-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`). + | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ | | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index bf3b400d9d7e..1ab4949c4003 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -333,8 +333,6 @@ def test_attention_config(): "true", "--attention-config.flash_attn_max_num_splits_for_cuda_graph", "16", - "--attention-config.use_cudnn_prefill", - "true", "--attention-config.use_trtllm_ragged_deepseek_prefill", "true", "--attention-config.use_trtllm_attention", @@ -352,7 +350,6 @@ def test_attention_config(): assert engine_args.attention_config.flash_attn_version == 3 assert engine_args.attention_config.use_prefill_decode_attention is True assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16 - assert engine_args.attention_config.use_cudnn_prefill is True assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is True assert engine_args.attention_config.use_trtllm_attention is True assert engine_args.attention_config.disable_flashinfer_prefill is True diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index e65d1d604029..f91ea85779d5 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -672,6 +672,7 @@ def run_attention_backend( def test_backend_correctness( default_vllm_config, dist_init, + workspace_init, batch_spec_name: str, model: str, tensor_parallel_size: int, diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py new file mode 100644 index 000000000000..068eb43faf40 --- /dev/null +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MLA prefill backend selector.""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.config import AttentionConfig, ModelConfig, VllmConfig +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 ( + MLAPrefillSelectorConfig, + _auto_select_mla_prefill_backend, + get_mla_prefill_backend, + is_deepseek_r1_mla_compatible, +) + + +@pytest.fixture(autouse=True) +def clear_cache(): + """Clear lru cache to ensure each test case runs without caching.""" + _auto_select_mla_prefill_backend.cache_clear() + + +def _make_mock_model_config( + qk_nope_head_dim: int = 128, + qk_rope_head_dim: int = 64, + v_head_dim: int = 128, + dtype: torch.dtype = torch.bfloat16, +) -> ModelConfig: + mock_config = MagicMock(spec=ModelConfig) + mock_config.dtype = dtype + mock_config.hf_text_config = MagicMock() + mock_config.hf_text_config.qk_nope_head_dim = qk_nope_head_dim + mock_config.hf_text_config.qk_rope_head_dim = qk_rope_head_dim + mock_config.hf_text_config.v_head_dim = v_head_dim + return mock_config + + +def _make_vllm_config( + model_config: ModelConfig | None = None, + mla_prefill_backend: MLAPrefillBackendEnum | None = None, +) -> VllmConfig: + if model_config is None: + model_config = _make_mock_model_config() + + attention_config = AttentionConfig(mla_prefill_backend=mla_prefill_backend) + mock_vllm_config = MagicMock(spec=VllmConfig) + mock_vllm_config.model_config = model_config + mock_vllm_config.attention_config = attention_config + return mock_vllm_config + + +class TestGetMLAPrefillBackend: + """Tests for get_mla_prefill_backend (public API).""" + + def test_no_device_capability_returns_flash_attn(self): + vllm_config = _make_vllm_config() + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_device_capability.return_value = None + + backend = get_mla_prefill_backend(vllm_config) + assert backend.get_name() == "FLASH_ATTN" + + def test_explicit_flash_attn_selection(self): + try: + flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() + except ImportError: + pytest.skip("FLASH_ATTN backend not available") + return + + vllm_config = _make_vllm_config( + mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN, + ) + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_device_capability.return_value = DeviceCapability( + major=9, minor=0 + ) + + with patch.object( + flash_attn_cls, + "validate_configuration", + return_value=[], + ): + backend = get_mla_prefill_backend(vllm_config) + assert backend.get_name() == "FLASH_ATTN" + + def test_explicit_backend_invalid_raises_error(self): + vllm_config = _make_vllm_config( + mla_prefill_backend=MLAPrefillBackendEnum.FLASHINFER, + ) + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_device_capability.return_value = DeviceCapability( + major=9, minor=0 + ) + + with pytest.raises(ValueError, match="is not valid"): + get_mla_prefill_backend(vllm_config) + + def test_explicit_backend_import_error_raises(self): + vllm_config = _make_vllm_config( + mla_prefill_backend=MLAPrefillBackendEnum.TRTLLM_RAGGED, + ) + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_device_capability.return_value = DeviceCapability( + major=10, minor=0 + ) + + with ( + patch.object( + MLAPrefillBackendEnum.TRTLLM_RAGGED, + "get_class", + side_effect=ImportError("trtllm not installed"), + ), + pytest.raises(ValueError, match="is not valid"), + ): + get_mla_prefill_backend(vllm_config) + + def test_auto_selection_on_hopper(self): + try: + flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() + except ImportError: + pytest.skip("FLASH_ATTN backend not available") + return + + vllm_config = _make_vllm_config() + + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_device_capability.return_value = DeviceCapability( + major=9, minor=0 + ) + + with patch.object( + flash_attn_cls, + "validate_configuration", + return_value=[], + ): + backend = get_mla_prefill_backend(vllm_config) + assert backend.get_name() == "FLASH_ATTN" + + +class TestAutoSelectMLAPrefillBackend: + """Tests for fallback and error paths in auto-selection.""" + + def test_blackwell_falls_back_to_trtllm(self): + vllm_config = _make_vllm_config() + capability = DeviceCapability(major=10, minor=0) + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config), + ) + + try: + trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class() + except ImportError: + pytest.skip("TRTLLM_RAGGED backend not available") + return + + with ( + patch.object( + MLAPrefillBackendEnum.FLASH_ATTN, + "get_class", + side_effect=ImportError("FLASH_ATTN not available"), + ), + patch.object(trtllm_cls, "validate_configuration", return_value=[]), + ): + backend = _auto_select_mla_prefill_backend( + capability, + selector_config, + ) + assert backend.get_name() == "TRTLLM_RAGGED" + + def test_all_fail_raises_error(self): + vllm_config = _make_vllm_config() + capability = DeviceCapability(major=10, minor=0) + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config), + ) + + def mock_get_class(backend_enum): # noqa: ARG001 + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls + + with patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class): + _auto_select_mla_prefill_backend.cache_clear() + with pytest.raises(ValueError, match="No valid MLA"): + _auto_select_mla_prefill_backend( + capability, + selector_config, + ) + + +class TestBackendValidation: + """Tests for backend validation logic.""" + + def test_r1_dimension_requirement(self): + try: + from vllm.v1.attention.backends.mla.prefill.flashinfer import ( + FlashInferPrefillBackend, + ) + except ImportError: + pytest.skip("FlashInfer prefill backend not available") + return + + assert FlashInferPrefillBackend.requires_r1_mla_dimensions is True + + vllm_config = _make_vllm_config( + model_config=_make_mock_model_config( + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) + ) + capability = DeviceCapability(major=10, minor=0) + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config), + ) + + with patch.object(FlashInferPrefillBackend, "is_available", return_value=True): + invalid_reasons = FlashInferPrefillBackend.validate_configuration( + capability, + selector_config, + ) + assert len(invalid_reasons) == 0 + + vllm_config_invalid = _make_vllm_config( + model_config=_make_mock_model_config( + qk_nope_head_dim=64, + qk_rope_head_dim=64, + v_head_dim=128, + ) + ) + selector_config_invalid = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config_invalid), + ) + + with patch.object(FlashInferPrefillBackend, "is_available", return_value=True): + invalid_reasons = FlashInferPrefillBackend.validate_configuration( + capability, + selector_config_invalid, + ) + assert len(invalid_reasons) == 1 + assert "DeepSeek R1 MLA dimensions" in invalid_reasons[0] + + +class TestMLAPrefillBackendParsing: + """Tests for string-based mla_prefill_backend parsing from CLI args.""" + + def test_valid_string_parses_to_enum(self): + config = AttentionConfig( + mla_prefill_backend="FLASH_ATTN", # type: ignore[arg-type] + ) + assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN + + def test_invalid_string_raises_error(self): + with pytest.raises(ValueError, match="Unknown MLA prefill backend"): + AttentionConfig( + mla_prefill_backend="NONEXISTENT", # type: ignore[arg-type] + ) + + +class TestDeprecatedFlagMigration: + """Tests for _migrate_deprecated_mla_prefill_flags in AttentionConfig.""" + + def test_no_deprecated_flags_leaves_backend_none(self): + config = AttentionConfig() + assert config.mla_prefill_backend is None + + def test_use_trtllm_ragged_migrates_to_trtllm_ragged(self): + config = AttentionConfig(use_trtllm_ragged_deepseek_prefill=True) + assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED + + def test_disable_flashinfer_prefill_migrates_to_flash_attn(self): + config = AttentionConfig(disable_flashinfer_prefill=True) + assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN + + def test_explicit_backend_ignores_deprecated_flags(self): + config = AttentionConfig( + mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN, + use_cudnn_prefill=True, + ) + assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN + + def test_cudnn_raises_error(self): + match = "cuDNN MLA prefill backend has been removed" + with pytest.raises(ValueError, match=match): + AttentionConfig(use_cudnn_prefill=True) + + def test_trtllm_takes_priority_over_disable_flashinfer(self): + config = AttentionConfig( + use_trtllm_ragged_deepseek_prefill=True, + disable_flashinfer_prefill=True, + ) + assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index d131c9cc01fb..73ef8b915821 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -30,7 +30,6 @@ RELEVANT_PATTERNS = [ "vllm/v1/attention/backends/*.py", "vllm/v1/attention/backends/**/*.py", - "vllm/v1/attention/backends/fa_utils.py", "vllm/model_executor/layers/attention/mla_attention.py", "vllm/platforms/cuda.py", "tools/pre_commit/generate_attention_backend_docs.py", @@ -68,6 +67,11 @@ def is_relevant_file(filepath: str) -> bool: return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS) +MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" +MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" +MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" + + # --------------------------------------------------------------------------- # AST utility helpers # --------------------------------------------------------------------------- @@ -293,6 +297,242 @@ def get_file_from_class_path(class_path: str) -> Path | None: return py_file if py_file.exists() else None +def parse_mla_prefill_registry() -> dict[str, str]: + """Parse MLAPrefillBackendEnum from the prefill registry. + + Returns: + A dict mapping backend names to their class paths. + """ + if not MLA_PREFILL_REGISTRY_FILE.exists(): + return {} + + try: + tree = ast.parse(MLA_PREFILL_REGISTRY_FILE.read_text()) + except Exception: + return {} + + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == "MLAPrefillBackendEnum": + return _extract_enum_values(node) + return {} + + +def parse_mla_prefill_priorities() -> dict[str, list[str]]: + """Parse MLA prefill backend priorities from selector.py. + + Returns: + A dict with keys like 'blackwell' and 'default' containing + lists of backend enum names in priority order. + """ + if not MLA_PREFILL_SELECTOR_FILE.exists(): + return {} + + try: + tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text()) + except Exception: + return {} + + priorities: dict[str, list[str]] = {} + + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef): + continue + if node.name != "_get_mla_prefill_backend_priorities": + continue + + # Look for if statements checking device_capability.major + for stmt in ast.walk(node): + if not isinstance(stmt, ast.If): + continue + + # Check if it's a capability.major == 10 check (Blackwell) + is_blackwell = ( + isinstance(stmt.test, ast.Compare) + and isinstance(stmt.test.left, ast.Attribute) + and stmt.test.left.attr == "major" + and stmt.test.comparators + and isinstance(stmt.test.comparators[0], ast.Constant) + and stmt.test.comparators[0].value == 10 + ) + + # Extract backends from return statements + for body_stmt in stmt.body: + if isinstance(body_stmt, ast.Return) and isinstance( + body_stmt.value, ast.List + ): + backends = [] + for elt in body_stmt.value.elts: + if isinstance(elt, ast.Attribute): + backends.append(elt.attr) + if is_blackwell: + priorities["blackwell"] = backends + else: + priorities["default"] = backends + + # Extract from else branch + for else_stmt in stmt.orelse: + if isinstance(else_stmt, ast.Return) and isinstance( + else_stmt.value, ast.List + ): + backends = [] + for elt in else_stmt.value.elts: + if isinstance(elt, ast.Attribute): + backends.append(elt.attr) + priorities["default"] = backends + + return priorities + + +def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None: + """Parse a single MLA prefill backend file to extract its properties. + + Args: + class_path: The fully qualified class path. + + Returns: + A dict with backend properties, or None if parsing fails. + """ + file_path = get_file_from_class_path(class_path) + if file_path is None: + return None + + try: + tree = ast.parse(file_path.read_text()) + except Exception: + return None + + class_name = class_path.rsplit(".", 1)[1] + class_node = find_class_in_ast(tree, class_name) + if class_node is None: + return None + + info: dict[str, Any] = { + "compute_capability": "Any", + "requires_r1_dims": False, + "dtypes": "fp16, bf16", # Default from base class + } + + # Parse class variables + for item in class_node.body: + if isinstance(item, ast.Assign): + for target in item.targets: + if ( + isinstance(target, ast.Name) + and target.id == "requires_r1_mla_dimensions" + and isinstance(item.value, ast.Constant) + ): + info["requires_r1_dims"] = item.value.value + + # Parse supported_dtypes class variable + if ( + isinstance(item, ast.AnnAssign) + and isinstance(item.target, ast.Name) + and item.target.id == "supported_dtypes" + and isinstance(item.value, ast.List) + ): + dtype_map = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"} + dtypes = [] + for elt in item.value.elts: + if isinstance(elt, ast.Attribute): + dtypes.append(dtype_map.get(elt.attr, elt.attr)) + if dtypes: + info["dtypes"] = ", ".join(dtypes) + + # Parse get_name static method + get_name_method = find_method(class_node, "get_name") + if get_name_method: + for n in ast.walk(get_name_method): + if isinstance(n, ast.Return) and isinstance(n.value, ast.Constant): + info["name"] = n.value.value + + # Parse supports_compute_capability classmethod + cc_method = find_method(class_node, "supports_compute_capability") + if cc_method: + for n in ast.walk(cc_method): + # Look for capability.major == 10 style checks + if ( + isinstance(n, ast.Compare) + and isinstance(n.left, ast.Attribute) + and n.left.attr == "major" + and n.comparators + and isinstance(n.comparators[0], ast.Constant) + ): + major = n.comparators[0].value + info["compute_capability"] = f"{major}.x" + + return info + + +def parse_mla_prefill_backends() -> list[dict[str, Any]]: + """Parse MLA prefill backend options from the prefill registry. + + MLA uses different backends for prefill vs decode. The decode backends are + registered in the main registry, but prefill backends have their own + registry at vllm/v1/attention/backends/mla/prefill/registry.py. + + Returns a list of prefill backend info dicts with their requirements. + """ + registry = parse_mla_prefill_registry() + priorities = parse_mla_prefill_priorities() + + if not registry: + return [] + + # Get the priority order (Blackwell order shows all backends) + priority_order = priorities.get("blackwell", list(registry.keys())) + + prefill_backends: list[dict[str, Any]] = [] + + # Backend-specific metadata that can't be easily parsed from code + backend_metadata = { + "TRTLLM_RAGGED": { + "description": "TensorRT-LLM ragged attention", + }, + "FLASHINFER": { + "description": "FlashInfer CUTLASS backend", + }, + "FLASH_ATTN": { + "description": "FlashAttention varlen (FA2/FA3/FA4)", + }, + } + + for backend_name in priority_order: + if backend_name not in registry: + continue + + class_path = registry[backend_name] + backend_info = parse_mla_prefill_backend_file(class_path) + if backend_info is None: + continue + + metadata = backend_metadata.get(backend_name, {}) + display_name = backend_info.get("name", backend_name) + + # Add marker for default Blackwell backend + marker = "" + if backend_name == priority_order[0] and priorities.get("blackwell"): + marker = "‡" + + notes = "" + if backend_info.get("requires_r1_dims"): + notes = "DeepSeek R1 dims only" + elif backend_name == "FLASH_ATTN": + notes = "FA4 on SM100+, FA3 on SM90, FA2 otherwise" + + prefill_backends.append( + { + "name": display_name, + "marker": marker, + "description": metadata.get("description", ""), + "dtypes": backend_info.get("dtypes", "fp16, bf16"), + "compute_capability": backend_info.get("compute_capability", "Any"), + "notes": notes, + } + ) + + return prefill_backends + + # --------------------------------------------------------------------------- # Backend feature extraction from AST # --------------------------------------------------------------------------- @@ -807,86 +1047,6 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]: } -def parse_mla_prefill_backends() -> list[dict[str, Any]]: - """Parse MLA prefill backend options from mla_attention.py. - - MLA uses different backends for prefill vs decode. The decode backends are - registered in the registry, but prefill backends are selected at runtime - based on conditions in MLACommonImpl.__init__. - - Returns a list of prefill backend info dicts with their requirements. - """ - if not MLA_ATTENTION_FILE.exists(): - return [] - - try: - tree = ast.parse(MLA_ATTENTION_FILE.read_text()) - except Exception: - return [] - - # Find compute capability requirements by parsing use_* functions - trtllm_cc = _find_cc_in_function(tree, "use_trtllm_ragged_deepseek_prefill") - flashinfer_cc = _find_cc_in_function(tree, "use_flashinfer_prefill") - cudnn_cc = _find_cc_in_function(tree, "use_cudnn_prefill") - - # Build prefill backend list based on what we found - # Order matches the priority in MLACommonImpl.__init__ - prefill_backends: list[dict[str, Any]] = [] - - # TRT-LLM Ragged (highest priority if available) - if trtllm_cc: - prefill_backends.append( - { - "name": "TRT-LLM Ragged‡", - "description": "TensorRT-LLM ragged attention", - "compute_capability": trtllm_cc, - "enable": "Default on SM100", - "disable": "`-ac.use_trtllm_ragged_deepseek_prefill=0`", - "notes": "DeepSeek R1 dims only", - } - ) - - # FlashInfer prefill - if flashinfer_cc: - prefill_backends.append( - { - "name": "FlashInfer", - "description": "FlashInfer CUTLASS backend", - "compute_capability": flashinfer_cc, - "enable": "`-ac.disable_flashinfer_prefill=0`", - "disable": "`-ac.disable_flashinfer_prefill=1`", - "notes": "DeepSeek R1 dims only", - } - ) - - # cuDNN prefill - if cudnn_cc: - prefill_backends.append( - { - "name": "cuDNN", - "description": "cuDNN-based attention", - "compute_capability": cudnn_cc, - "enable": "`-ac.use_cudnn_prefill=1`", - "disable": "`-ac.use_cudnn_prefill=0`", - "notes": "", - } - ) - - # FlashAttention is always available as fallback - prefill_backends.append( - { - "name": "FlashAttention", - "description": "FlashAttention varlen (FA2/FA3)", - "compute_capability": "Any", - "enable": "Default fallback", - "disable": "Use other backends", - "notes": "FA3 on SM90, FA2 otherwise", - } - ) - - return prefill_backends - - # --------------------------------------------------------------------------- # Backend variant expansion (FA2/FA3/FA4, FlashInfer native/TRTLLM) # --------------------------------------------------------------------------- @@ -1415,20 +1575,22 @@ def generate_mla_section( "", "### Prefill Backends", "", - "The prefill backend is selected at runtime based on hardware and", - "configuration.", + "To explicitly select a prefill backend, use", + "`-ac.mla_prefill_backend=` (e.g., `FLASH_ATTN`, `FLASHINFER`).", + "Otherwise, the prefill backend is selected automatically at runtime based on", + "hardware and configuration.", "", - "| Backend | Description | Compute Cap. | Enable | Disable | Notes |", - "| ------- | ----------- | ------------ | ------ | ------- | ----- |", + "| Backend | Description | Dtypes | Compute Cap. | Notes |", + "| ------- | ----------- | ------ | ------------ | ----- |", ] for backend in prefill_backends: - row = "| {} | {} | {} | {} | {} | {} |".format( + row = "| `{}`{} | {} | {} | {} | {} |".format( backend["name"], + backend.get("marker", ""), backend["description"], + backend.get("dtypes", "fp16, bf16"), backend["compute_capability"], - backend["enable"], - backend["disable"], backend.get("notes", ""), ) lines.append(row.replace(" ", " ")) @@ -1441,6 +1603,9 @@ def generate_mla_section( "", "### Decode Backends", "", + "MLA decode backends are selected using the standard", + "`-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`).", + "", ] ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 18973f5d66d8..b5dc7a5bf602 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -6,8 +6,12 @@ from pydantic import field_validator from vllm.config.utils import config +from vllm.logger import init_logger +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum +logger = init_logger(__name__) + @config class AttentionConfig: @@ -33,7 +37,7 @@ class AttentionConfig: and buffers can be pre-allocated to avoid inflating the memory estimate.""" use_cudnn_prefill: bool = False - """Whether to use cudnn prefill.""" + """Deprecated: cuDNN prefill backend has been removed.""" use_trtllm_ragged_deepseek_prefill: bool = False """Whether to use TRTLLM ragged deepseek prefill.""" @@ -42,12 +46,18 @@ class AttentionConfig: """If set to True/False, use or don't use the TRTLLM attention backend in flashinfer. If None, auto-detect the attention backend in flashinfer.""" - disable_flashinfer_prefill: bool = True + disable_flashinfer_prefill: bool | None = None """Whether to disable flashinfer prefill.""" disable_flashinfer_q_quantization: bool = False """If set, when using fp8 kv, do not quantize Q to fp8.""" + mla_prefill_backend: MLAPrefillBackendEnum | None = None + """MLA prefill backend to use. If None, will be selected automatically. + Valid options: FLASH_ATTN (FA3/FA4), FLASHINFER, TRTLLM_RAGGED. + This option supersedes use_trtllm_ragged_deepseek_prefill + and disable_flashinfer_prefill which are deprecated.""" + use_prefill_query_quantization: bool = False """If set, quantize query for attention in prefill.""" @@ -84,3 +94,48 @@ def validate_backend_before(cls, value: Any) -> Any: return None return AttentionBackendEnum[value.upper()] return value + + @field_validator("mla_prefill_backend", mode="before") + @classmethod + def validate_mla_prefill_backend_before(cls, value: Any) -> Any: + """Enable parsing of the `mla_prefill_backend` enum type from string.""" + if isinstance(value, str): + return MLAPrefillBackendEnum[value.upper()] + return value + + def __post_init__(self) -> None: + self._migrate_deprecated_mla_prefill_flags() + + def _migrate_deprecated_mla_prefill_flags(self) -> None: + """Migrate deprecated MLA prefill flags to mla_prefill_backend.""" + # If the new option is already set, it takes precedence + if self.mla_prefill_backend is not None: + return + + # Check for deprecated flags and migrate them. + # Only the first flag encountered sets the backend. + if self.use_cudnn_prefill: + raise ValueError( + "The cuDNN MLA prefill backend has been removed. " + "Use --attention-config.mla_prefill_backend=FLASH_ATTN or " + "FLASHINFER or TRTLLM_RAGGED instead." + ) + + if self.use_trtllm_ragged_deepseek_prefill: + if self.mla_prefill_backend is None: + self.mla_prefill_backend = MLAPrefillBackendEnum.TRTLLM_RAGGED + logger.warning_once( + "use_trtllm_ragged_deepseek_prefill is deprecated and " + "will be removed in v0.22. Use " + "--attention-config.mla_prefill_backend=TRTLLM_RAGGED " + "instead." + ) + + if self.disable_flashinfer_prefill: + if self.mla_prefill_backend is None: + self.mla_prefill_backend = MLAPrefillBackendEnum.FLASH_ATTN + logger.warning_once( + "disable_flashinfer_prefill is deprecated and will be removed " + "in v0.22. Use --attention-config.mla_prefill_backend=" + "FLASH_ATTN instead." + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4afe2319570e..82eecc8cd49b 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -189,12 +189,9 @@ import functools from abc import abstractmethod -from dataclasses import dataclass, field +from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast - -if TYPE_CHECKING: - from flashinfer import BatchPrefillWithRaggedKVCacheWrapper +from typing import ClassVar, Generic, TypeVar, cast import torch import torch.nn as nn @@ -242,7 +239,7 @@ kNvfp4Dynamic, ) from vllm.platforms import current_platform -from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory +from vllm.utils.flashinfer import has_flashinfer from vllm.utils.math_utils import cdiv, round_down from vllm.utils.torch_utils import ( LayerNameType, @@ -262,11 +259,9 @@ MLAAttentionImpl, SparseMLAAttentionImpl, ) -from vllm.v1.attention.backends.fa_utils import get_flash_attn_version +from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackend from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, - get_per_layer_parameters, - infer_global_hyperparameters, split_decodes_and_prefills, ) from vllm.v1.attention.ops.common import cp_lse_ag_out_rs @@ -1123,33 +1118,6 @@ class QueryLenSupport(Enum): VARLEN = "varlen" -try: - from vllm.vllm_flash_attn import ( # type: ignore[attr-defined] - flash_attn_varlen_func, - ) - - is_vllm_fa = True -except ImportError: - is_vllm_fa = False - flash_attn_varlen_func = None # type: ignore[assignment] - # On ROCm, vllm_flash_attn is not available, try upstream flash_attn instead. - # On CUDA, vllm_flash_attn should always be available (built with vLLM), - # so we don't attempt the fallback there. - if current_platform.is_rocm(): - try: - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] - except ImportError: - logger.debug( - "flash_attn not available on ROCm; " - "MLA models using TRITON_MLA will require flash_attn. " - "AITER_MLA backends use aiter kernels instead." - ) - elif current_platform.is_xpu(): - from vllm._xpu_ops import xpu_ops - - flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[no-redef,attr-defined,assignment] - - def dynamic_per_batched_tensor_quant( x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn ): @@ -1161,9 +1129,6 @@ def dynamic_per_batched_tensor_quant( return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() -logger = init_logger(__name__) - - @CustomOp.register( "mla_decode_concat_quant_fp8", dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0}, @@ -1197,9 +1162,6 @@ def forward( forward_hip = _make_forward(QuantFP8.forward_hip) # type: ignore[arg-type] -CUDNN_WORKSPACE_SIZE = 12800 - - class MLACommonBackend(AttentionBackend): @staticmethod def get_name() -> str: @@ -1268,26 +1230,9 @@ class ChunkedContextMetadata: query_start_loc: torch.Tensor max_query_len: int chunked_context: ChunkedContextMetadata | None = None - query_seq_lens: torch.Tensor | None = None - workspace_buffer: torch.Tensor | None = None q_data_type: torch.dtype | None = None output_dtype: torch.dtype | None = None - - -@dataclass -class FlashInferPrefillMetadata(MLACommonPrefillMetadata): - prefill_main: "BatchPrefillWithRaggedKVCacheWrapper | None" = None - prefill_chunks: "list[BatchPrefillWithRaggedKVCacheWrapper]" = field( - default_factory=list - ) - - -@dataclass -class CudnnPrefillMetadata(MLACommonPrefillMetadata): - class ChunkedContextMetadata(MLACommonPrefillMetadata.ChunkedContextMetadata): - seq_lens: torch.Tensor - - cudnn_workspace: torch.Tensor | None = None + prefill_backend: MLAPrefillBackend | None = None @dataclass @@ -1333,13 +1278,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]): # The dimension of the attention heads head_dim: int | None = None + prefill: MLACommonPrefillMetadata | None = None decode: D | None = None - prefill: ( - MLACommonPrefillMetadata - | FlashInferPrefillMetadata - | CudnnPrefillMetadata - | None - ) = None def __post_init__(self): if self.head_dim is not None and not MLACommonBackend.supports_head_size( @@ -1352,64 +1292,6 @@ def __post_init__(self): A = TypeVar("A", bound=AttentionMetadata) -def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool: - # Check if model has DeepSeek R1 compatible MLA dimensions: - # qk_nope_head_dim = 128, qk_rope_head_dim = 64, v_head_dim = 128 - # which results in query/key head dim = 192. - if vllm_config.model_config is None: - return False - hf_text_config = vllm_config.model_config.hf_text_config - qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) - qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 1) - v_head_dim = getattr(hf_text_config, "v_head_dim", 1) - return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128 - - -@functools.cache -def use_flashinfer_prefill() -> bool: - from vllm.config import get_current_vllm_config - - vllm_config = get_current_vllm_config() - if not ( - not vllm_config.attention_config.disable_flashinfer_prefill - and has_flashinfer() - and not vllm_config.attention_config.use_cudnn_prefill - and current_platform.is_device_capability_family(100) - ): - return False - - return is_deepseek_r1_mla_compatible(vllm_config) - - -@functools.cache -def use_cudnn_prefill() -> bool: - from vllm.config import get_current_vllm_config - - vllm_config = get_current_vllm_config() - return ( - has_flashinfer() - and vllm_config.attention_config.use_cudnn_prefill - and current_platform.is_device_capability_family(100) - and has_nvidia_artifactory() - ) - - -@functools.cache -def use_trtllm_ragged_deepseek_prefill() -> bool: - """Check if TRT-LLM ragged DeepSeek prefill should be used.""" - from vllm.config import get_current_vllm_config - - vllm_config = get_current_vllm_config() - if not ( - has_flashinfer() - and vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill - and current_platform.is_device_capability_family(100) - ): - return False - - return is_deepseek_r1_mla_compatible(vllm_config) - - @dataclass class MLADims: q_lora_rank: int | None @@ -1447,15 +1329,14 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims: @functools.cache def backend_supports_prefill_query_quantization() -> bool: - """Check if the selected MLA backend supports prefill query quantization. + """Check if the selected MLA prefill backend supports query quantization. Currently supported backends: - - FlashInfer prefill - - TRT-LLM ragged DeepSeek prefill + - FlashInfer + - TRT-LLM Ragged Not supported: - - cuDNN Prefill - - FlashAttention + - FlashAttention (FA3/FA4) - Non-GB200 devices (FP8 prefill requires device capability 100) """ # FP8 prefill query quantization requires GB200 (device capability 100) @@ -1463,7 +1344,15 @@ def backend_supports_prefill_query_quantization() -> bool: if not current_platform.is_device_capability_family(100): return False - return use_flashinfer_prefill() or use_trtllm_ragged_deepseek_prefill() + from vllm.config import get_current_vllm_config + from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend + + vllm_config = get_current_vllm_config() + backend_cls = get_mla_prefill_backend(vllm_config) + return backend_cls.get_name() in ( + "FLASHINFER", + "TRTLLM_RAGGED", + ) class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): @@ -1574,7 +1463,6 @@ def __init__( metadata_cls if metadata_cls is not None else MLACommonMetadata ) self.kv_cache_spec = kv_cache_spec - scheduler_config = vllm_config.scheduler_config self.model_config = vllm_config.model_config parallel_config = vllm_config.parallel_config self.compilation_config = vllm_config.compilation_config @@ -1634,139 +1522,32 @@ def __init__( device=device, ) - self._use_cudnn_prefill = use_cudnn_prefill() - self._use_fi_prefill = use_flashinfer_prefill() - self._use_trtllm_ragged_prefill = use_trtllm_ragged_deepseek_prefill() - self.prefill_metadata_cls = ( - FlashInferPrefillMetadata - if self._use_fi_prefill - else CudnnPrefillMetadata - if self._use_cudnn_prefill - else MLACommonPrefillMetadata - ) - - if self._use_fi_prefill: - self._workspace_buffer = torch.empty( - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, - dtype=torch.uint8, - device=device, - ) - - self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None - self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = [] - - self._global_hyperparameters = infer_global_hyperparameters( - get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract] - ) - - if self._use_trtllm_ragged_prefill: - self._workspace_buffer = torch.empty( - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, - dtype=torch.uint8, - device=device, - ) + from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend - if self._use_cudnn_prefill: - self.cudnn_workspace = torch.empty( - CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs, - dtype=torch.int8, - device=device, - ) + 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, + 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, + v_head_dim=self.mla_dims.v_head_dim, + vllm_config=vllm_config, + device=device, + layer_names=layer_names, + ) supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY self._init_reorder_batch_threshold( self.reorder_batch_threshold, supports_spec_decode, supports_dcp_with_varlen ) - # Validate consistency between query_len_support and reorder_batch_threshold if self.query_len_support == QueryLenSupport.SINGLE_ONLY: assert self.reorder_batch_threshold == 1, ( f"reorder_batch_threshold must be 1 when query_len_support is " f"SINGLE_ONLY, got {self.reorder_batch_threshold}" ) - def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata): - qo_indptr = prefill.query_start_loc - - has_context = False - if prefill.chunked_context is not None: - chunked_context = prefill.chunked_context - has_context = True - - if self._fi_prefill_main is None: - from flashinfer import BatchPrefillWithRaggedKVCacheWrapper - - self._fi_prefill_main = BatchPrefillWithRaggedKVCacheWrapper( - self._workspace_buffer, "NHD", backend="cutlass" - ) - - if has_context: - num_chunks = chunked_context.cu_seq_lens.shape[0] - # Allocate more prefill chunk wrappers if needed - if len(self._fi_prefill_chunks) < num_chunks: - from flashinfer import BatchPrefillWithRaggedKVCacheWrapper - - for _ in range(len(self._fi_prefill_chunks), num_chunks): - self._fi_prefill_chunks.append( - BatchPrefillWithRaggedKVCacheWrapper( - self._workspace_buffer, "NHD", backend="cutlass" - ) - ) - assert num_chunks <= len(self._fi_prefill_chunks) - - # In MLA, the non-latent num_qo_heads == num_kv_heads - num_qo_heads = self.num_heads - num_kv_heads = num_qo_heads - - # Sanity: Verify that num_kv_heads == 1 since it is latent space - assert self.kv_cache_spec.num_kv_heads == 1 - - # Get non-latent head_dim_qk and head_dim_vo - head_dim_qk = self.mla_dims.qk_nope_head_dim + self.mla_dims.qk_rope_head_dim - head_dim_vo = self.mla_dims.v_head_dim - - # For main run, qo_indptr == kv_indptr - kv_indptr = qo_indptr.clone() - - # Prepare main prefill - self._fi_prefill_main.plan( - qo_indptr=qo_indptr, - kv_indptr=kv_indptr, - num_qo_heads=num_qo_heads, - num_kv_heads=num_kv_heads, - head_dim_qk=head_dim_qk, - head_dim_vo=head_dim_vo, - causal=True, # This is main run - sm_scale=self._global_hyperparameters.sm_scale, - window_left=self._global_hyperparameters.window_left, - logits_soft_cap=self._global_hyperparameters.logits_soft_cap, - q_data_type=self.q_data_type, - o_data_type=prefill.output_dtype, - ) - - # Prepare context prefills - if has_context: - for i in range(num_chunks): - kv_indptr_chunk = chunked_context.cu_seq_lens[i] - - self._fi_prefill_chunks[i].plan( - qo_indptr=qo_indptr, - kv_indptr=kv_indptr_chunk, - num_qo_heads=num_qo_heads, - num_kv_heads=num_kv_heads, - head_dim_qk=head_dim_qk, - head_dim_vo=head_dim_vo, - causal=False, # This is context run - sm_scale=self._global_hyperparameters.sm_scale, - window_left=self._global_hyperparameters.window_left, - logits_soft_cap=self._global_hyperparameters.logits_soft_cap, - q_data_type=self.q_data_type, - o_data_type=prefill.output_dtype, - ) - - prefill.prefill_main = self._fi_prefill_main - prefill.prefill_chunks = self._fi_prefill_chunks - def _build_decode( self, block_table_tensor: torch.Tensor, @@ -1972,18 +1753,14 @@ def build( dtype=torch.int32, ) - chunked_context_metadata_cls = ( - CudnnPrefillMetadata.ChunkedContextMetadata - if self._use_cudnn_prefill - else MLACommonPrefillMetadata.ChunkedContextMetadata - ) prefill_tokens_with_context = None if num_prefills_with_context_cpu > 0: prefill_tokens_with_context = prefill_query_start_loc_cpu[ num_prefills_with_context_cpu ].item() + _ChunkedMetadata = MLACommonPrefillMetadata.ChunkedContextMetadata if self.dcp_world_size > 1: - chunked_context_metadata = chunked_context_metadata_cls( + chunked_context_metadata = _ChunkedMetadata( cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True), starts=local_chunk_starts.to(device, non_blocking=True), seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(), @@ -2004,7 +1781,7 @@ def build( prefill_tokens_with_context=prefill_tokens_with_context, ) else: - chunked_context_metadata = chunked_context_metadata_cls( + chunked_context_metadata = _ChunkedMetadata( cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True), starts=chunk_starts.to(device, non_blocking=True), seq_tot=chunk_seq_lens.sum(dim=1).tolist(), @@ -2018,35 +1795,22 @@ def build( prefill_tokens_with_context=prefill_tokens_with_context, ) - if self._use_cudnn_prefill: - chunked_context_metadata.seq_lens = chunk_seq_lens - assert ( max(chunked_context_metadata.max_seq_lens) <= self.chunked_prefill_workspace_size ) - prefill_metadata = self.prefill_metadata_cls( + prefill_metadata = MLACommonPrefillMetadata( block_table=block_table_tensor[reqs_start:, ...], query_start_loc=prefill_query_start_loc, max_query_len=max_query_len, chunked_context=chunked_context_metadata, output_dtype=self.model_config.dtype, q_data_type=self.q_data_type, + prefill_backend=self._prefill_backend, ) - if self._use_cudnn_prefill: - assert isinstance(prefill_metadata, CudnnPrefillMetadata) - prefill_metadata.query_seq_lens = ( - prefill_query_start_loc[1:] - prefill_query_start_loc[:-1] - ) - prefill_metadata.cudnn_workspace = self.cudnn_workspace - - if self._use_trtllm_ragged_prefill: - prefill_metadata.query_seq_lens = ( - prefill_query_start_loc[1:] - prefill_query_start_loc[:-1] - ) - prefill_metadata.workspace_buffer = self._workspace_buffer + self._prefill_backend.prepare_metadata(prefill_metadata) decode_metadata = None if num_decodes > 0: @@ -2091,10 +1855,6 @@ def build( decode=decode_metadata, ) - if self._use_fi_prefill and num_prefills > 0: - assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata) - self._build_fi_prefill_wrappers(attn_metadata.prefill) - return attn_metadata # type: ignore[return-value] @@ -2240,308 +2000,12 @@ def __init__( and (self.qk_rope_head_dim == 64) ) - if use_trtllm_ragged_deepseek_prefill(): - logger.info_once("Using TRT-LLM ragged DeepSeek prefill for MLA") - self._run_prefill_context_chunk = ( - self._run_prefill_context_chunk_trtllm_ragged - ) - self._run_prefill_new_tokens = self._run_prefill_new_tokens_trtllm_ragged - self._pad_v = False - elif use_flashinfer_prefill(): - logger.info_once("Using FlashInfer prefill for MLA") - self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi - self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi - self._pad_v = False - elif use_cudnn_prefill(): - logger.info_once("Using CUDNN prefill for MLA") - self._run_prefill_context_chunk = self._run_prefill_context_chunk_cudnn - self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn - self._pad_v = False - else: # Use FlashAttention - if flash_attn_varlen_func is None: - raise RuntimeError( - "MLA attention requires FlashAttention but it is not " - "available. Please install flash_attn or use " - "--attention-backend ROCM_AITER_MLA." - ) - logger.info_once("Using FlashAttention prefill for MLA") - self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa - self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa - - # Handle the differences between the flash_attn_varlen from - # flash_attn and the one from vllm_flash_attn. The former is used on - # RoCM and the latter has an additional parameter to control - # FA2 vs FA3 - self.flash_attn_varlen_func = flash_attn_varlen_func - self.vllm_flash_attn_version = get_flash_attn_version( - head_size=self.qk_head_dim - ) - if self.vllm_flash_attn_version is not None: - self.flash_attn_varlen_func = functools.partial( - flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version - ) - - # For MLA the v head dim is smaller than qk head dim so we pad out - # v with 0s to match the qk head dim for attention backends that do - # not support different headdims. - # FA3 on Hopper (SM90) and FA4 natively handle diff headdims. - device_capability = current_platform.get_device_capability() - self._pad_v = self.vllm_flash_attn_version is None or not ( - ( - self.vllm_flash_attn_version == 3 - and device_capability is not None - and device_capability[0] == 9 - ) - or self.vllm_flash_attn_version == 4 - ) - self.dcp_world_size: int = -1 self.cp_kv_cache_interleave_size: int = ( get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size ) - def _flash_attn_varlen_diff_headdims( - self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs - ): - maybe_padded_v = v - if self._pad_v: - maybe_padded_v = torch.nn.functional.pad( - v, [0, q.shape[-1] - v.shape[-1]], value=0 - ) - - if is_vllm_fa: - kwargs["return_softmax_lse"] = return_softmax_lse - else: - # ROCm leverages the upstream flash_attn, which takes a parameter - # called "return_attn_probs" instead of return_softmax_lse - kwargs["return_attn_probs"] = return_softmax_lse - if envs.VLLM_BATCH_INVARIANT: - kwargs["num_splits"] = 1 - - attn_out = self.flash_attn_varlen_func( - q=q, - k=k, - v=maybe_padded_v, - softmax_scale=softmax_scale, - **kwargs, - ) - - # Unpack the output if there is multiple results - lse = None - if isinstance(attn_out, tuple): - attn_out, lse = attn_out[0], attn_out[1] - - # Remain consistent with old `flash_attn_varlen_func` where there - # is only one output tensor if `return_softmax_lse` is False. - if return_softmax_lse: - return attn_out, lse - return attn_out - - def _run_prefill_new_tokens_fa( - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse - ): - return self._flash_attn_varlen_diff_headdims( - q=q, - k=k, - v=v, - cu_seqlens_q=prefill.query_start_loc, - cu_seqlens_k=prefill.query_start_loc, - max_seqlen_q=prefill.max_query_len, - max_seqlen_k=prefill.max_query_len, - softmax_scale=self.scale, - causal=True, - return_softmax_lse=return_softmax_lse, - ) - - def _run_prefill_new_tokens_fi( - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse - ): - assert isinstance(prefill, FlashInferPrefillMetadata) - assert prefill.prefill_main is not None - - ret = prefill.prefill_main.run( - q=q, - k=k, - v=v, - return_lse=return_softmax_lse, - ) - - if isinstance(ret, tuple): - return ret[0], ret[1].transpose(0, 1).contiguous() - return ret - - def _run_prefill_new_tokens_cudnn( - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse - ): - assert isinstance(prefill, CudnnPrefillMetadata) - assert prefill.query_seq_lens is not None - from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache - - output, lse = cudnn_batch_prefill_with_kv_cache( - q=q, - k_cache=k, - v_cache=v, - scale=self.scale, - workspace_buffer=prefill.cudnn_workspace, - max_token_per_sequence=prefill.max_query_len, - max_sequence_kv=prefill.max_query_len, - actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1), - actual_seq_lens_kv=prefill.query_seq_lens.view(-1, 1, 1, 1), - causal=True, - # Do not support False for now - return_lse=True, - # Indicates actual_seq_lens are on GPU or CPU. - is_cuda_graph_compatible=True, - ) - if return_softmax_lse: - return output, lse - return output - - def _run_prefill_context_chunk_fa( - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v - ): - assert prefill.chunked_context is not None - return self._flash_attn_varlen_diff_headdims( - q=q, - k=k, - v=v, - cu_seqlens_q=prefill.query_start_loc, - cu_seqlens_k=prefill.chunked_context.cu_seq_lens[chunk_idx], - max_seqlen_q=prefill.max_query_len, - max_seqlen_k=prefill.chunked_context.max_seq_lens[chunk_idx], - softmax_scale=self.scale, - causal=False, # Context is unmasked - return_softmax_lse=True, - ) - - def _run_prefill_context_chunk_fi( - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v - ): - assert isinstance(prefill, FlashInferPrefillMetadata) - - attn_out, lse = prefill.prefill_chunks[chunk_idx].run( - q=q, - k=k, - v=v, - return_lse=True, - ) - - # Convert from (q_len, num_heads) to (num_heads, q_len) - return attn_out, lse.transpose(0, 1).contiguous() - - def _run_prefill_context_chunk_cudnn( - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v - ): - assert isinstance(prefill, CudnnPrefillMetadata) - assert prefill.chunked_context is not None - assert prefill.chunked_context.seq_lens[chunk_idx] is not None - assert prefill.query_seq_lens is not None - from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache - - return cudnn_batch_prefill_with_kv_cache( - q=q, - k_cache=k, - v_cache=v, - scale=self.scale, - workspace_buffer=prefill.cudnn_workspace, - max_token_per_sequence=prefill.max_query_len, - max_sequence_kv=prefill.chunked_context.max_seq_lens[chunk_idx], - actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1), - actual_seq_lens_kv=prefill.chunked_context.seq_lens[chunk_idx].view( - -1, 1, 1, 1 - ), - causal=False, - return_lse=True, - # Indicates actual_seq_lens are on GPU or CPU. - is_cuda_graph_compatible=True, - ) - - def _run_prefill_new_tokens_trtllm_ragged( - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse - ): - """TRT-LLM ragged attention for new tokens (causal).""" - from flashinfer.prefill import trtllm_ragged_attention_deepseek - - assert prefill.query_seq_lens is not None - assert prefill.workspace_buffer is not None - # allocate BF16 / FP16 output tensor for TRT-LLM ragged attention - out = torch.empty( - q.shape[0], - q.shape[1], - v.shape[2], - device=q.device, - dtype=prefill.output_dtype, - ) - - ret = trtllm_ragged_attention_deepseek( - query=q, - key=k, - value=v, - workspace_buffer=prefill.workspace_buffer, - seq_lens=prefill.query_seq_lens, - max_q_len=prefill.max_query_len, - max_kv_len=prefill.max_query_len, - bmm1_scale=self.scale, - bmm2_scale=1.0, - o_sf_scale=1.0, - batch_size=prefill.query_seq_lens.shape[0], - window_left=-1, - cum_seq_lens_q=prefill.query_start_loc, - cum_seq_lens_kv=prefill.query_start_loc, - enable_pdl=False, - is_causal=True, - return_lse=return_softmax_lse, - out=out, - ) - - if isinstance(ret, tuple): - # Convert from (q_len, num_heads) to (num_heads, q_len) - return ret[0], ret[1].transpose(0, 1).contiguous() - return ret - - def _run_prefill_context_chunk_trtllm_ragged( - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v - ): - """TRT-LLM ragged attention for context chunks (non-causal).""" - from flashinfer.prefill import trtllm_ragged_attention_deepseek - - assert prefill.chunked_context is not None - assert prefill.chunked_context.seq_lens[chunk_idx] is not None - assert prefill.workspace_buffer is not None - - out = torch.empty( - q.shape[0], - q.shape[1], - v.shape[2], - device=q.device, - dtype=prefill.output_dtype, - ) - - attn_out, lse = trtllm_ragged_attention_deepseek( - query=q, - key=k, - value=v, - workspace_buffer=prefill.workspace_buffer, - seq_lens=prefill.chunked_context.seq_lens[chunk_idx], - max_q_len=prefill.max_query_len, - max_kv_len=prefill.chunked_context.max_seq_lens[chunk_idx], - bmm1_scale=self.scale, - bmm2_scale=1.0, - o_sf_scale=1.0, - batch_size=prefill.chunked_context.seq_lens[chunk_idx].shape[0], - window_left=-1, - cum_seq_lens_q=prefill.query_start_loc, - cum_seq_lens_kv=prefill.chunked_context.cu_seq_lens[chunk_idx], - enable_pdl=False, - is_causal=False, - return_lse=True, - out=out, - ) - - # Convert from (q_len, num_heads) to (num_heads, q_len) - return attn_out, lse.transpose(0, 1).contiguous() - def _concat_k_nope_k_pe( self, k_nope: torch.Tensor, k_pe: torch.Tensor ) -> torch.Tensor: @@ -2582,6 +2046,7 @@ def _compute_prefill_context( ): assert attn_metadata.prefill is not None prefill_metadata = attn_metadata.prefill + assert prefill_metadata.prefill_backend is not None assert prefill_metadata.chunked_context is not None use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype() @@ -2649,12 +2114,13 @@ def _compute_prefill_context( k = self._concat_k_nope_k_pe(k_nope, k_pe) - attn_output, attn_softmax_lse = self._run_prefill_context_chunk( - prefill=prefill_metadata, - chunk_idx=i, - q=q, - k=k, - v=v, + attn_output, attn_softmax_lse = ( + prefill_metadata.prefill_backend.run_prefill_context_chunk( + chunk_idx=i, + q=q, + k=k, + v=v, + ) ) if output is None: @@ -2687,6 +2153,7 @@ def _context_parallel_compute_prefill_context( assert k_scale is None, "DCP not support scaled kvcache now." assert attn_metadata.prefill is not None prefill_metadata = attn_metadata.prefill + assert prefill_metadata.prefill_backend is not None assert prefill_metadata.chunked_context is not None assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None assert prefill_metadata.chunked_context.local_context_lens_allranks is not None @@ -2753,12 +2220,13 @@ def _context_parallel_compute_prefill_context( k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) k = self._concat_k_nope_k_pe(k_nope, k_pe) - attn_output, attn_softmax_lse = self._run_prefill_context_chunk( - prefill=prefill_metadata, - chunk_idx=i, - q=q, - k=k, - v=v, + attn_output, attn_softmax_lse = ( + prefill_metadata.prefill_backend.run_prefill_context_chunk( + chunk_idx=i, + q=q, + k=k, + v=v, + ) ) if output is None: @@ -2790,11 +2258,11 @@ def forward_mha( k_scale: torch.Tensor, output: torch.Tensor, ) -> None: - # TODO (zyongye): Prefill function here assert attn_metadata.prefill is not None assert self.dcp_world_size != -1 prefill_metadata = attn_metadata.prefill + assert prefill_metadata.prefill_backend is not None use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype() # Convert q to FP8 if FP8 prefill attention is enabled @@ -2813,8 +2281,7 @@ def forward_mha( k = k.to(prefill_metadata.q_data_type) v = v.to(prefill_metadata.q_data_type) - output_prefill = self._run_prefill_new_tokens( - prefill=prefill_metadata, + output_prefill = prefill_metadata.prefill_backend.run_prefill_new_tokens( q=q, k=k, v=v, @@ -2839,11 +2306,6 @@ def forward_mha( q, kv_c_and_k_pe_cache, attn_metadata, k_scale ) - # unpad if necessary - if self._pad_v: - context_output = context_output[..., : v.shape[-1]] - suffix_output = suffix_output[..., : v.shape[-1]] - output = output.view(-1, self.num_heads, self.v_head_dim) merge_attn_states( output=output, @@ -2854,7 +2316,8 @@ def forward_mha( prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context, ) else: - output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2) + assert isinstance(output_prefill, torch.Tensor) + output_prefill = output_prefill.flatten(start_dim=-2) output.copy_(output_prefill) @abstractmethod diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index c0d52620c086..2753326755fb 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -86,6 +86,9 @@ def __gt__(self, other: Any) -> bool: return NotImplemented return (self.major, self.minor) > (other.major, other.minor) + def __hash__(self) -> int: + return hash((self.major, self.minor)) + def as_version_str(self) -> str: return f"{self.major}.{self.minor}" diff --git a/vllm/v1/attention/backends/mla/prefill/__init__.py b/vllm/v1/attention/backends/mla/prefill/__init__.py new file mode 100644 index 000000000000..ae5b7ae82598 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum +from vllm.v1.attention.backends.mla.prefill.selector import get_mla_prefill_backend + +__all__ = [ + "MLAPrefillBackend", + "MLAPrefillBackendEnum", + "get_mla_prefill_backend", +] diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py new file mode 100644 index 000000000000..9c850a0b1d99 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Abstract base class for MLA prefill backends.""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, ClassVar + +import torch + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonPrefillMetadata, + ) + from vllm.platforms.interface import DeviceCapability + from vllm.v1.attention.backends.mla.prefill.selector import ( + MLAPrefillSelectorConfig, + ) + + +class MLAPrefillBackend(ABC): + """Abstract base class for MLA prefill backends.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + requires_r1_mla_dimensions: ClassVar[bool] = False + + @staticmethod + @abstractmethod + def get_name() -> str: + raise NotImplementedError + + @classmethod + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + return True + + @classmethod + def supports_dtype(cls, dtype: torch.dtype) -> bool: + return dtype in cls.supported_dtypes + + @classmethod + def is_available(cls) -> bool: + return True + + @classmethod + def validate_configuration( + cls, + device_capability: "DeviceCapability", + selector_config: "MLAPrefillSelectorConfig", + ) -> list[str]: + invalid_reasons: list[str] = [] + + if not cls.supports_compute_capability(device_capability): + invalid_reasons.append( + f"compute capability {device_capability.major}." + f"{device_capability.minor} not supported" + ) + + if not cls.supports_dtype(selector_config.dtype): + invalid_reasons.append(f"dtype {selector_config.dtype} not supported") + + if not cls.is_available(): + invalid_reasons.append("required dependencies not available") + + if cls.requires_r1_mla_dimensions and not selector_config.is_r1_compatible: + invalid_reasons.append( + "model does not have DeepSeek R1 MLA dimensions " + "(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128)" + ) + + return invalid_reasons + + def __init__( + self, + num_heads: int, + scale: float, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + vllm_config: "VllmConfig", + device: torch.device, + layer_names: list[str] | None = None, + ) -> None: + self.num_heads = num_heads + self.scale = scale + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.vllm_config = vllm_config + self.device = device + self.layer_names = layer_names + + def prepare_metadata( # noqa: B027 + self, + prefill_metadata: "MLACommonPrefillMetadata", + ) -> None: + """Prepare backend-specific metadata before the forward pass. + + Called by the metadata builder after constructing the prefill metadata. + """ + self._prefill_metadata = prefill_metadata + + @abstractmethod + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + + @abstractmethod + def run_prefill_context_chunk( + self, + chunk_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py new file mode 100644 index 000000000000..42d77b12d0de --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashAttention backend for MLA prefill.""" + +import functools +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm.platforms import current_platform +from vllm.v1.attention.backends.fa_utils import ( + get_flash_attn_version, + is_flash_attn_varlen_func_available, +) +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +if is_flash_attn_varlen_func_available(): + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func +else: + flash_attn_varlen_func = None # type: ignore[assignment] + + +class FlashAttnPrefillBackend(MLAPrefillBackend): + """FlashAttention backend for MLA prefill.""" + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN" + + @classmethod + def is_available(cls) -> bool: + return is_flash_attn_varlen_func_available() + + def __init__( + self, + num_heads: int, + scale: float, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + vllm_config: "VllmConfig", + device: torch.device, + layer_names: list[str] | None = None, + ) -> None: + super().__init__( + num_heads=num_heads, + scale=scale, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + vllm_config=vllm_config, + device=device, + layer_names=layer_names, + ) + + # Handle the differences between the flash_attn_varlen from + # flash_attn and the one from vllm_flash_attn + assert flash_attn_varlen_func is not None, ( + "FlashAttnPrefillBackend requires flash_attn_varlen_func. " + "Ensure FlashAttnPrefillBackend.is_available() is checked first." + ) + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.flash_attn_varlen_func = flash_attn_varlen_func + self.vllm_flash_attn_version = get_flash_attn_version(head_size=qk_head_dim) + if self.vllm_flash_attn_version is not None: + self.flash_attn_varlen_func = functools.partial( + flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version + ) + + # Determine if we need to pad V + # For MLA the v head dim is smaller than qk head dim so we pad out + # v with 0s to match the qk head dim for attention backends that do + # not support different headdims. + # FA3 on Hopper (SM90) and FA4 natively handle diff headdims. + device_capability = current_platform.get_device_capability() + self.requires_v_padding = self.vllm_flash_attn_version is None or not ( + ( + self.vllm_flash_attn_version == 3 + and device_capability is not None + and device_capability[0] == 9 + ) + or self.vllm_flash_attn_version == 4 + ) + + # Track whether we're using vllm's FA or upstream (for ROCm) + self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu() + + def _flash_attn_varlen_diff_headdims( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool = False, + softmax_scale: float | None = None, + **kwargs, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + maybe_padded_v = v + if self.requires_v_padding: + maybe_padded_v = torch.nn.functional.pad( + v, [0, q.shape[-1] - v.shape[-1]], value=0 + ) + + if self._is_vllm_fa: + kwargs["return_softmax_lse"] = return_softmax_lse + else: + # ROCm leverages the upstream flash_attn, which takes a parameter + # called "return_attn_probs" instead of return_softmax_lse + kwargs["return_attn_probs"] = return_softmax_lse + if envs.VLLM_BATCH_INVARIANT: + kwargs["num_splits"] = 1 + + attn_out = self.flash_attn_varlen_func( + q=q, + k=k, + v=maybe_padded_v, + softmax_scale=softmax_scale, + **kwargs, + ) + + # Unpack the output if there are multiple results + lse = None + if isinstance(attn_out, tuple): + attn_out, lse = attn_out[0], attn_out[1] + + # Unpad output back to v_head_dim if we padded V + if self.requires_v_padding: + attn_out = attn_out[..., : v.shape[-1]] + + # Remain consistent with old `flash_attn_varlen_func` where there + # is only one output tensor if `return_softmax_lse` is False. + if return_softmax_lse: + return attn_out, lse + return attn_out + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return self._flash_attn_varlen_diff_headdims( + q=q, + k=k, + v=v, + cu_seqlens_q=self._prefill_metadata.query_start_loc, + cu_seqlens_k=self._prefill_metadata.query_start_loc, + max_seqlen_q=self._prefill_metadata.max_query_len, + max_seqlen_k=self._prefill_metadata.max_query_len, + softmax_scale=self.scale, + causal=True, + return_softmax_lse=return_softmax_lse, + ) + + def run_prefill_context_chunk( + self, + chunk_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert self._prefill_metadata.chunked_context is not None + return self._flash_attn_varlen_diff_headdims( + q=q, + k=k, + v=v, + cu_seqlens_q=self._prefill_metadata.query_start_loc, + cu_seqlens_k=self._prefill_metadata.chunked_context.cu_seq_lens[chunk_idx], + max_seqlen_q=self._prefill_metadata.max_query_len, + max_seqlen_k=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx], + softmax_scale=self.scale, + causal=False, # Context is unmasked + return_softmax_lse=True, + ) diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py new file mode 100644 index 000000000000..a1107df304f3 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer backend for MLA prefill.""" + +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +from vllm.v1.attention.backends.utils import ( + get_per_layer_parameters, + infer_global_hyperparameters, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonPrefillMetadata, + ) + from vllm.platforms.interface import DeviceCapability + +try: + from flashinfer import BatchPrefillWithRaggedKVCacheWrapper +except ImportError: + BatchPrefillWithRaggedKVCacheWrapper = object # type: ignore[misc,assignment] + +_DEFAULT_NUM_CHUNKS = 32 + + +class FlashInferPrefillBackend(MLAPrefillBackend): + """FlashInfer backend for MLA prefill.""" + + requires_r1_mla_dimensions = True + + @staticmethod + def get_name() -> str: + return "FLASHINFER" + + @classmethod + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + return device_capability.major == 10 + + @classmethod + def is_available(cls) -> bool: + try: + from flashinfer import ( + BatchPrefillWithRaggedKVCacheWrapper, # noqa: F401 + ) + + return True + except ImportError: + return False + + def __init__( + self, + num_heads: int, + scale: float, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + vllm_config: "VllmConfig", + device: torch.device, + layer_names: list[str] | None = None, + ) -> None: + super().__init__( + num_heads=num_heads, + scale=scale, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + vllm_config=vllm_config, + device=device, + layer_names=layer_names, + ) + + self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None + self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = [] + if layer_names is None: + raise ValueError( + "FlashInferPrefillBackend requires layer_names to " + "initialize global hyperparameters." + ) + + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonImpl, + ) + + self._global_hyperparameters = infer_global_hyperparameters( + get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract] + ) + + def _ensure_chunks( + self, + num_chunks: int, + workspace_buffer: torch.Tensor, + ) -> None: + if len(self._prefill_chunks) < num_chunks: + for _ in range(len(self._prefill_chunks), num_chunks): + self._prefill_chunks.append( + BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, "NHD", backend="cutlass" + ) + ) + + def prepare_metadata( + self, + prefill_metadata: "MLACommonPrefillMetadata", + ) -> None: + qo_indptr = prefill_metadata.query_start_loc + has_context = prefill_metadata.chunked_context is not None + (workspace_buffer,) = current_workspace_manager().get_simultaneous( + ((envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8), + ) + + if self._prefill_main is None: + self._prefill_main = BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, "NHD", backend="cutlass" + ) + self._ensure_chunks(_DEFAULT_NUM_CHUNKS, workspace_buffer) + + if has_context: + chunked_context = prefill_metadata.chunked_context + assert chunked_context is not None + num_chunks = chunked_context.cu_seq_lens.shape[0] + self._ensure_chunks(num_chunks, workspace_buffer) + + num_qo_heads = self.num_heads + num_kv_heads = num_qo_heads + + head_dim_qk = self.qk_nope_head_dim + self.qk_rope_head_dim + head_dim_vo = self.v_head_dim + kv_indptr = qo_indptr.clone() + + assert self._prefill_main is not None + self._prefill_main.plan( + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim_qk=head_dim_qk, + head_dim_vo=head_dim_vo, + causal=True, + sm_scale=self._global_hyperparameters.sm_scale, + window_left=self._global_hyperparameters.window_left, + logits_soft_cap=self._global_hyperparameters.logits_soft_cap, + q_data_type=prefill_metadata.q_data_type, + o_data_type=prefill_metadata.output_dtype, + ) + + if has_context: + chunked_context = prefill_metadata.chunked_context + assert chunked_context is not None + for i in range(num_chunks): + kv_indptr_chunk = chunked_context.cu_seq_lens[i] + + self._prefill_chunks[i].plan( + qo_indptr=qo_indptr, + kv_indptr=kv_indptr_chunk, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim_qk=head_dim_qk, + head_dim_vo=head_dim_vo, + causal=False, + sm_scale=self._global_hyperparameters.sm_scale, + window_left=self._global_hyperparameters.window_left, + logits_soft_cap=self._global_hyperparameters.logits_soft_cap, + q_data_type=prefill_metadata.q_data_type, + o_data_type=prefill_metadata.output_dtype, + ) + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert self._prefill_main is not None + + ret = self._prefill_main.run( + q=q, + k=k, + v=v, + return_lse=return_softmax_lse, + ) + + if isinstance(ret, tuple): + # Convert from (q_len, num_heads) to (num_heads, q_len) + return ret[0], ret[1].transpose(0, 1).contiguous() + return ret + + def run_prefill_context_chunk( + self, + chunk_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + attn_out, lse = self._prefill_chunks[chunk_idx].run( + q=q, + k=k, + v=v, + return_lse=True, + ) + + # Convert from (q_len, num_heads) to (num_heads, q_len) + return attn_out, lse.transpose(0, 1).contiguous() diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py new file mode 100644 index 000000000000..3a3242f60365 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Registry for MLA prefill backends. + +This module provides an enumeration of all available MLA prefill backends +and utilities for loading them. +""" + +from enum import Enum, EnumMeta +from typing import TYPE_CHECKING + +from vllm.utils.import_utils import resolve_obj_by_qualname + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + + +class _MLAPrefillBackendEnumMeta(EnumMeta): + """Metaclass for MLAPrefillBackendEnum to provide better error messages.""" + + def __getitem__(cls, name: str): + try: + return super().__getitem__(name) + except KeyError: + members = cls.__members__.keys() + valid_backends = ", ".join(members) + raise ValueError( + f"Unknown MLA prefill backend: '{name}'. " + f"Valid options are: {valid_backends}" + ) from None + + +class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): + """Enumeration of all supported MLA prefill backends.""" + + FLASH_ATTN = ( + "vllm.v1.attention.backends.mla.prefill.flash_attn.FlashAttnPrefillBackend" + ) + FLASHINFER = ( + "vllm.v1.attention.backends.mla.prefill.flashinfer.FlashInferPrefillBackend" + ) + TRTLLM_RAGGED = ( + "vllm.v1.attention.backends.mla.prefill.trtllm_ragged." + "TrtllmRaggedPrefillBackend" + ) + + def get_path(self) -> str: + """Get the fully qualified class path for this backend.""" + return self.value + + def get_class(self) -> "type[MLAPrefillBackend]": + """Lazy load and return the backend class.""" + return resolve_obj_by_qualname(self.get_path()) diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py new file mode 100644 index 000000000000..fdb8be6d65d9 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Selector for MLA prefill backends. + +This module provides functions for selecting the appropriate MLA prefill +backend based on device capabilities and configuration. +""" + +from functools import cache +from typing import TYPE_CHECKING, NamedTuple + +import torch + +from vllm.logger import init_logger +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + +logger = init_logger(__name__) + + +class MLAPrefillSelectorConfig(NamedTuple): + """Hashable configuration for MLA prefill backend selection. + + This is analogous to AttentionSelectorConfig and contains model-specific + configuration needed to select an MLA prefill backend, extracted from + VllmConfig into a hashable form for caching. + """ + + dtype: torch.dtype + is_r1_compatible: bool + + +def is_deepseek_r1_mla_compatible(vllm_config: "VllmConfig") -> bool: + """Check if model has DeepSeek R1 compatible MLA dimensions. + + DeepSeek R1 MLA dimensions are: + - qk_nope_head_dim = 128 + - qk_rope_head_dim = 64 + - v_head_dim = 128 + """ + if vllm_config.model_config is None: + return False + hf_text_config = vllm_config.model_config.hf_text_config + qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) + qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 1) + v_head_dim = getattr(hf_text_config, "v_head_dim", 1) + return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128 + + +def _get_mla_prefill_backend_priorities( + device_capability: DeviceCapability, +) -> list[MLAPrefillBackendEnum]: + """Get MLA prefill backend priorities based on device capability. + + Args: + device_capability: The device's compute capability. + + Returns: + List of backends in priority order (highest priority first). + """ + if device_capability.major == 10: # Blackwell + return [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + ] + else: # Hopper (SM90) and older + return [ + MLAPrefillBackendEnum.FLASH_ATTN, + ] + + +def get_mla_prefill_backend( + vllm_config: "VllmConfig", +) -> "type[MLAPrefillBackend]": + """Select the MLA prefill backend based on configuration and device. + + This function first checks for explicit user preferences via + mla_prefill_backend in AttentionConfig, then falls back to automatic + priority-based selection. + + Args: + vllm_config: The vLLM configuration. + + Returns: + The selected prefill backend class. + """ + from vllm.platforms import current_platform + + device_capability = current_platform.get_device_capability() + if device_capability is None: + logger.info_once( + "Device capability not available, using FlashAttention MLA prefill backend." + ) + return MLAPrefillBackendEnum.FLASH_ATTN.get_class() + + attention_config = vllm_config.attention_config + + selector_config = MLAPrefillSelectorConfig( + dtype=vllm_config.model_config.dtype, + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config), + ) + + if attention_config.mla_prefill_backend is not None: + selected_backend = attention_config.mla_prefill_backend + backend_cls: type[MLAPrefillBackend] | None = None + try: + backend_cls = selected_backend.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if invalid_reasons: + raise ValueError( + f"Selected MLA prefill backend {selected_backend.name} " + f"is not valid for this configuration. " + f"Reason: {invalid_reasons}" + ) + assert backend_cls is not None + logger.info("Using %s MLA prefill backend.", selected_backend.name) + return backend_cls + + return _auto_select_mla_prefill_backend( + device_capability, + selector_config, + ) + + +@cache +def _auto_select_mla_prefill_backend( + device_capability: DeviceCapability, + selector_config: MLAPrefillSelectorConfig, +) -> "type[MLAPrefillBackend]": + """Auto-select the best available MLA prefill backend. + + Args: + device_capability: The device's compute capability. + selector_config: Hashable configuration for backend selection. + + Returns: + The selected prefill backend class. + """ + priorities = _get_mla_prefill_backend_priorities(device_capability) + all_invalid_reasons: dict[str, list[str]] = {} + + for backend_enum in priorities: + backend_cls: type[MLAPrefillBackend] | None = None + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + assert backend_cls is not None + logger.info_once("Using %s MLA prefill backend.", backend_enum.name) + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + config_str = repr(selector_config) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + config_str, + reasons_str, + ) + + raise ValueError( + f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}." + ) diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py new file mode 100644 index 000000000000..2ffe244ebe0d --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TRT-LLM Ragged backend for MLA prefill.""" + +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonPrefillMetadata, + ) + from vllm.platforms.interface import DeviceCapability + + +class TrtllmRaggedPrefillBackend(MLAPrefillBackend): + """TRT-LLM Ragged backend for MLA prefill.""" + + requires_r1_mla_dimensions = True + + @staticmethod + def get_name() -> str: + return "TRTLLM_RAGGED" + + @classmethod + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + return device_capability.major == 10 + + @classmethod + def is_available(cls) -> bool: + try: + from flashinfer.prefill import ( + trtllm_ragged_attention_deepseek, # noqa: F401 + ) + + return True + except ImportError: + return False + + def __init__( + self, + num_heads: int, + scale: float, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + vllm_config: "VllmConfig", + device: torch.device, + layer_names: list[str] | None = None, + ) -> None: + super().__init__( + num_heads=num_heads, + scale=scale, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + vllm_config=vllm_config, + device=device, + layer_names=layer_names, + ) + + def _get_workspace_buffer(self) -> torch.Tensor: + (workspace_buffer,) = current_workspace_manager().get_simultaneous( + ( + (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), + torch.uint8, + ), + ) + return workspace_buffer + + def prepare_metadata( + self, + prefill_metadata: "MLACommonPrefillMetadata", + ) -> None: + super().prepare_metadata(prefill_metadata) + self._query_seq_lens = ( + prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1] + ) + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + from flashinfer.prefill import trtllm_ragged_attention_deepseek + + workspace_buffer = self._get_workspace_buffer() + out = torch.empty( + q.shape[0], + q.shape[1], + v.shape[2], + device=q.device, + dtype=self._prefill_metadata.output_dtype, + ) + + ret = trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=workspace_buffer, + seq_lens=self._query_seq_lens, + max_q_len=self._prefill_metadata.max_query_len, + max_kv_len=self._prefill_metadata.max_query_len, + bmm1_scale=self.scale, + bmm2_scale=1.0, + o_sf_scale=1.0, + batch_size=self._query_seq_lens.shape[0], + window_left=-1, + cum_seq_lens_q=self._prefill_metadata.query_start_loc, + cum_seq_lens_kv=self._prefill_metadata.query_start_loc, + enable_pdl=False, + is_causal=True, + return_lse=return_softmax_lse, + out=out, + ) + + if isinstance(ret, tuple): + # Convert from (q_len, num_heads) to (num_heads, q_len) + return ret[0], ret[1].transpose(0, 1).contiguous() + return ret + + def run_prefill_context_chunk( + self, + chunk_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + from flashinfer.prefill import trtllm_ragged_attention_deepseek + + assert self._prefill_metadata.chunked_context is not None + assert self._prefill_metadata.chunked_context.seq_lens[chunk_idx] is not None + workspace_buffer = self._get_workspace_buffer() + + out = torch.empty( + q.shape[0], + q.shape[1], + v.shape[2], + device=q.device, + dtype=self._prefill_metadata.output_dtype, + ) + + attn_out, lse = trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=workspace_buffer, + seq_lens=self._prefill_metadata.chunked_context.seq_lens[chunk_idx], + max_q_len=self._prefill_metadata.max_query_len, + max_kv_len=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx], + bmm1_scale=self.scale, + bmm2_scale=1.0, + o_sf_scale=1.0, + batch_size=self._prefill_metadata.chunked_context.seq_lens[chunk_idx].shape[ + 0 + ], + window_left=-1, + cum_seq_lens_q=self._prefill_metadata.query_start_loc, + cum_seq_lens_kv=self._prefill_metadata.chunked_context.cu_seq_lens[ + chunk_idx + ], + enable_pdl=False, + is_causal=False, + return_lse=True, + out=out, + ) + + # Convert from (q_len, num_heads) to (num_heads, q_len) + return attn_out, lse.transpose(0, 1).contiguous() diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 7aa8a646f415..c45a631008e5 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -123,18 +123,6 @@ def __init__( self._sm_count = current_platform.num_compute_units() - def _flash_attn_varlen_diff_headdims( - self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs - ): - return super()._flash_attn_varlen_diff_headdims( - q, - k, - v, - return_softmax_lse=return_softmax_lse, - softmax_scale=softmax_scale, - **kwargs, - ) - def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], From a9484dac7b734096ed26db4902454da7e497d2c3 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sat, 2 May 2026 03:01:17 +0800 Subject: [PATCH 0003/1083] [Perf] Intergrate Tile Kernels `head_compute_mix_kernel` for Deepseek-V4 (#41255) Signed-off-by: Isotr0py Co-authored-by: Roger Wang --- vllm/model_executor/layers/mhc.py | 134 ++++++++++++++++++++++ vllm/model_executor/models/deepseek_v4.py | 28 +++-- 2 files changed, 153 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 1521a6b601bf..f5c1f06844b0 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -448,3 +448,137 @@ def _mhc_post_fake( mutates_args=[], fake_impl=_mhc_post_fake, ) + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10, + }, +) +def hc_head_fuse_tilelang( + residual, + fn, + hc_scale, + hc_base, + out, + hidden_size: int, + rms_eps: float, + hc_eps: float, + hc_mult: int = 4, + n_thr: int = 128, + h_blk: int = 1024, +): + """Two-pass fused kernel for hc_head. + + Pass 1: accumulate per-token squared sum and hc_mult dot-products + (projections onto fn rows) using cross-thread reducers. + Pass 2: apply sigmoid-gated weighted sum of residual channels to output. + + Avoids materialising mixes / rsqrt / pre tensors to global memory. + """ + num_tokens = T.dynamic("num_tokens") + hc_dim = hc_mult * hidden_size + h_block = math.gcd(h_blk, hidden_size) + n_h = hidden_size // h_block + + residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type] + hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type] + hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type] + out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] + + with T.Kernel(num_tokens, threads=n_thr) as i: + T.pdl_sync() + + # ------------------------------------------------------------------ + # Pass 1 – for each residual channel m_c and h_block: + # • accumulate squared sum (for RMS norm denominator) + # • accumulate hc_mult dot-products with fn rows + # ------------------------------------------------------------------ + sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all") + mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all") + T.fill(sqrsum_r, 0.0) + T.fill(mixes_r, 0.0) + + for m_c in T.serial(hc_mult): + for i_h in T.serial(n_h): + x_local = T.alloc_fragment(h_block, T.float32) + T.copy(residual[i, m_c, i_h * h_block], x_local) + + for k in T.Parallel(h_block): + sqrsum_r[0] += x_local[k] * x_local[k] + + for m_m in T.unroll(hc_mult): + fn_local = T.alloc_fragment(h_block, T.float32) + T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local) + for k in T.Parallel(h_block): + mixes_r[m_m] += x_local[k] * fn_local[k] + + T.finalize_reducer(sqrsum_r) + T.finalize_reducer(mixes_r) + + # ------------------------------------------------------------------ + # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps + # ------------------------------------------------------------------ + pre_mix_shared = T.alloc_shared(hc_mult, T.float32) + rsqrt_val = T.alloc_fragment(1, T.float32) + rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps) + for m in T.Parallel(hc_mult): + pre_mix_shared[m] = ( + T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps + ) + + # ------------------------------------------------------------------ + # Pass 2 – apply_mix: pipelined weighted sum over residual channels + # ------------------------------------------------------------------ + for i0_h in T.Pipelined(n_h, num_stages=2): + xs = T.alloc_shared((hc_mult, h_block), T.bfloat16) + xl = T.alloc_fragment((hc_mult, h_block), T.float32) + T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True) + T.copy(xs, xl) + + ol = T.alloc_fragment(h_block, T.float32) + T.clear(ol) + for i_hc in T.serial(hc_mult): + pre = pre_mix_shared[i_hc] + for i1_h in T.Parallel(h_block): + ol[i1_h] += pre * xl[i_hc, i1_h] + + T.copy(ol, out[i, i0_h * h_block], disable_tma=True) + + T.pdl_trigger() + + +def _hc_head_fused_kernel( + hs_flat: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + out: torch.Tensor, + hidden_size: int, + rms_eps: float, + hc_eps: float, + hc_mult: int, +) -> None: + """Fill pre-allocated `out` (T, H) in-place with the hc_head result.""" + if hs_flat.shape[0] > 0: + hc_head_fuse_tilelang( + hs_flat, + fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_eps, + hc_eps, + hc_mult, + ) + + +direct_register_custom_op( + op_name="hc_head_fused_kernel", + op_func=_hc_head_fused_kernel, + mutates_args=["out"], +) diff --git a/vllm/model_executor/models/deepseek_v4.py b/vllm/model_executor/models/deepseek_v4.py index 5521a9764a9c..0b762d50fe72 100644 --- a/vllm/model_executor/models/deepseek_v4.py +++ b/vllm/model_executor/models/deepseek_v4.py @@ -7,7 +7,6 @@ import regex as re import torch import torch.nn as nn -import torch.nn.functional as F from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig, get_current_vllm_config @@ -1456,14 +1455,25 @@ def hc_head( rms_norm_eps: float, hc_eps: float, ) -> torch.Tensor: - x = hidden_states - shape, dtype = x.size(), x.dtype - x = x.flatten(1).float() - rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_norm_eps) - mixes = F.linear(x, hc_fn) * rsqrt - pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps - y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1) - return y.to(dtype) + hc_mult, hidden_size = hidden_states.shape[-2:] + outer_shape = hidden_states.shape[:-2] + hs_flat = hidden_states.view(-1, hc_mult, hidden_size) + num_tokens = hs_flat.shape[0] + out = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device + ) + torch.ops.vllm.hc_head_fused_kernel( + hs_flat, + hc_fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_norm_eps, + hc_eps, + hc_mult, + ) + return out.view(*outer_shape, hidden_size) def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: From bcf5cac9fb956788f649d1f5297b74c886a9d6d3 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 1 May 2026 15:23:17 -0400 Subject: [PATCH 0004/1083] [DSV4] Add knob to enable pre-attn gemm (#41443) Signed-off-by: Yongye Zhu --- vllm/envs.py | 12 ++++++++++++ vllm/model_executor/layers/deepseek_v4_attention.py | 3 +++ vllm/utils/multi_stream_utils.py | 12 +++++++++--- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 4191cd6a9743..24ec92c3d755 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -245,6 +245,7 @@ VLLM_DEBUG_WORKSPACE: bool = False VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 + VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 4096 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" VLLM_USE_V2_MODEL_RUNNER: bool = False VLLM_LOG_MODEL_INSPECTION: bool = False @@ -1678,6 +1679,17 @@ def _get_or_set_default() -> str: "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int( int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256)) ), + # Token-count cutoff for multi-stream overlap of the attention input + # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer + # weights / kv-score projections in DeepSeek-V4). At or below this many + # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs + # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device + # and the cross-stream sync becomes pure overhead. Set to 0 to disable + # the multi-stream path entirely. Empirical crossover on B300 (148 SMs) + # is ~4096; B200 (132 SMs) is expected ~3072. + "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int( + os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "4096") + ), # Format for saving torch.compile cache artifacts # - "binary": saves as binary file # Safe for multiple vllm serve processes accessing the same torch compile cache. diff --git a/vllm/model_executor/layers/deepseek_v4_attention.py b/vllm/model_executor/layers/deepseek_v4_attention.py index 74b494dc4851..847c3eee55a8 100644 --- a/vllm/model_executor/layers/deepseek_v4_attention.py +++ b/vllm/model_executor/layers/deepseek_v4_attention.py @@ -13,6 +13,7 @@ import torch.nn.functional as F from transformers import DeepseekV2Config, DeepseekV3Config +import vllm.envs as envs from vllm.model_executor.layers.linear import ( ReplicatedLinear, ) @@ -385,6 +386,8 @@ def fused_wqa_wkv() -> torch.Tensor: self.ln_events[0], self.ln_events[1:4], self.aux_stream_list[:3], + enable=hidden_states.shape[0] + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, ) return qr_kv, kv_score, indexer_kv_score, indexer_weights diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index c00f08f93329..2203221c5a14 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -64,6 +64,7 @@ def execute_in_parallel( start_event: torch.cuda.Event, done_events: list[torch.cuda.Event], aux_streams: list[torch.cuda.Stream] | None = None, + enable: bool = False, ) -> tuple[Any, list[Any]]: """Run default_fn on the current stream and aux_fns concurrently on aux_streams. @@ -74,8 +75,9 @@ def execute_in_parallel( start_event fans out from the current stream to every launched aux stream; done_events[i] is recorded after aux_fns[i] so the current stream joins - before returning. When aux_streams is None, all aux_fns run sequentially - on the current stream. + before returning. Falls back to sequential execution on the current stream + when aux_streams is None or enable is False; in that case default_fn runs + first, then aux_fns in order. Args: default_fn: Callable for the default (current) stream. @@ -86,13 +88,17 @@ def execute_in_parallel( corresponding aux_fn. Length must match aux_fns. aux_streams: Per-aux CUDA streams. Length must match aux_fns. Multi-stream is disabled when None. + enable: Opt-in switch for the multi-stream path. Defaults to False, + so callers that pass aux_streams must also pass enable=True + (typically gated by an env var) to actually overlap. When False, + execution falls back to sequential on the current stream. Returns: Tuple of (default_result, aux_results) where aux_results[i] is the result of aux_fns[i] (or None when skipped). """ aux_results: list[Any] - if aux_streams is None: + if aux_streams is None or not enable: default_result = default_fn() aux_results = [fn() if fn is not None else None for fn in aux_fns] return default_result, aux_results From edd60ac93a3247c7ef1bf1e2a3e9c0e95bc83bf6 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 1 May 2026 17:42:52 -0400 Subject: [PATCH 0005/1083] [Bugfix] Fix persistent_topk inter-CTA init race on RadixRowState (#41444) Signed-off-by: Yongye Zhu --- csrc/persistent_topk.cuh | 25 ++++++------------------- csrc/topk.cu | 23 +++++++++++++++++++++++ 2 files changed, 29 insertions(+), 19 deletions(-) diff --git a/csrc/persistent_topk.cuh b/csrc/persistent_topk.cuh index d6162d52998b..8b9d10ff83dd 100644 --- a/csrc/persistent_topk.cuh +++ b/csrc/persistent_topk.cuh @@ -887,27 +887,14 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) uint32_t* shared_ordered = reinterpret_cast(smem_raw + kFixedSmemLarge); - // RadixRowState for multi-CTA cooperative radix + // RadixRowState for multi-CTA cooperative radix. + // Zero-initialization is done host-side via cudaMemsetAsync in topk.cu + // before launch — that gives a stream-ordered happens-before edge for all + // CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA + // __syncthreads) did not provide and which manifested as a race against + // CTA-1+'s first red_release on arrival_counter. RadixRowState* state = ¶ms.row_states[group_id]; - // -- Initialize RadixRowState (only needed if large rows exist) -- - if (params.max_seq_len > RADIX_THRESHOLD) { - if (cta_in_group == 0) { - for (uint32_t buf = 0; buf < 3; buf++) { - for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { - state->histogram[buf][i] = 0; - } - } - if (tx == 0) { - state->remaining_k = 0; - state->prefix = 0; - state->arrival_counter = 0; - state->output_counter = 0; - } - } - __syncthreads(); - } - int barrier_phase = 0; const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups; diff --git a/csrc/topk.cu b/csrc/topk.cu index b0f612ba6e4b..68352629ef02 100644 --- a/csrc/topk.cu +++ b/csrc/topk.cu @@ -153,6 +153,29 @@ void launch_persistent_topk(const torch::Tensor& logits, TORCH_CHECK(workspace.size(0) >= static_cast(state_bytes), "workspace too small, need ", state_bytes, " bytes"); + // Zero the per-group RadixRowState region before launch — only when the + // radix path will actually run (max_seq_len > RADIX_THRESHOLD). The + // RadixRowState fields (arrival_counter, histograms) are only touched by + // radix_topk; the decode/medium paths inside the persistent kernel + // operate purely in shared memory and never read these globals, so a + // stale workspace is harmless for them. + // + // Why we need the memset (when needs_cooperative is true): + // 1. arrival_counter accumulates within a launch and is never reset, + // so a prior call leaves it at a large positive value. Without this + // reset, the very first wait_ge in the next call sees counter >> + // target and returns instantly, breaking the barrier. + // 2. The previous in-kernel init only ran in CTA-0 with intra-CTA + // __syncthreads(), so it had no happens-before edge to CTA-1+'s + // first red_release. cudaMemsetAsync is stream-ordered: the zero + // is globally visible before any CTA runs. + if (needs_cooperative) { + cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr(), 0, + state_bytes, stream); + TORCH_CHECK(mz_err == cudaSuccess, + "row_states memset failed: ", cudaGetErrorString(mz_err)); + } + P::PersistentTopKParams params; params.input = logits.data_ptr(); params.output = output.data_ptr(); From 0c99629ede51524f00b88cb758c895fd76a5f6f9 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 1 May 2026 17:45:03 -0400 Subject: [PATCH 0006/1083] [Build] Make bundled DeepGEMM wheel portable across Python versions (#41476) Signed-off-by: mgoin --- cmake/external_projects/deepgemm.cmake | 27 ++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 0d7ea43fb7d0..b821b90ec8e9 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -59,11 +59,26 @@ if(DEEPGEMM_ARCHS) # Build the _C pybind11 extension from DeepGEMM's C++ source. # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime. # - Python_add_library(_deep_gemm_C MODULE WITH_SOABI - "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") + # Free-threaded Python doesn't yet support the stable ABI, so skip USE_SABI + # there. (The other vLLM extensions get this guard for free via + # define_extension_target; this target uses raw Python_add_library.) + run_python(IS_FREETHREADED_PYTHON + "import sysconfig; print(1 if sysconfig.get_config_var(\"Py_GIL_DISABLED\") else 0)" + "Failed to determine whether interpreter is free-threaded") + if (NOT IS_FREETHREADED_PYTHON) + Python_add_library(_deep_gemm_C MODULE WITH_SOABI USE_SABI 3 + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") + else() + Python_add_library(_deep_gemm_C MODULE WITH_SOABI + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") + endif() # The pybind11 module name must be _C to match DeepGEMM's Python imports. - set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C") + # Place the build artifact in a subdir so it doesn't collide with vLLM's own + # `_C.abi3.so` in the build tree (the install destination still differs). + set_target_properties(_deep_gemm_C PROPERTIES + OUTPUT_NAME "_C" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm") target_compile_definitions(_deep_gemm_C PRIVATE "-DTORCH_EXTENSION_NAME=_C") @@ -75,11 +90,15 @@ if(DEEPGEMM_ARCHS) "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include" "${deepgemm_SOURCE_DIR}/third-party/fmt/include") + # Keep Stable ABI for the module, but *not* for CUDA/C++ files. + # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles. target_compile_options(_deep_gemm_C PRIVATE $<$:-std=c++17> $<$:-O3> $<$:-Wno-psabi> - $<$:-Wno-deprecated-declarations>) + $<$:-Wno-deprecated-declarations> + $<$:-UPy_LIMITED_API> + $<$:-UPy_LIMITED_API>) # torch_python is required because DeepGEMM uses pybind11 type casters # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which From 5737770c6c346d918fdfb13e9378f9514f616186 Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Sat, 2 May 2026 00:01:37 +0100 Subject: [PATCH 0007/1083] Re-enable allreduce rms fusion for DP / PP (#41458) Signed-off-by: Andy Lo --- vllm/config/vllm.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index bb3ea81bce52..8d2c2608e56b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -138,12 +138,6 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool: current_platform.is_device_capability_family(100) or current_platform.is_device_capability(90) ) - # tp-dp combination broken: - # https://github.com/vllm-project/vllm/issues/34458 - and cfg.parallel_config.data_parallel_size == 1 - # tp-pp combination broken: - # https://github.com/vllm-project/vllm/issues/35426 - and cfg.parallel_config.pipeline_parallel_size == 1 ) From c408fdd663afb34ab82a10b26f553bec9e8052d9 Mon Sep 17 00:00:00 2001 From: FredericOdermatt <50372080+FredericOdermatt@users.noreply.github.com> Date: Sat, 2 May 2026 05:06:54 +0200 Subject: [PATCH 0008/1083] [Fix] Sync gemma4 chat template from hf (#39570) Signed-off-by: Frederic Odermatt --- examples/tool_chat_template_gemma4.jinja | 111 ++++++++++++++--------- 1 file changed, 67 insertions(+), 44 deletions(-) diff --git a/examples/tool_chat_template_gemma4.jinja b/examples/tool_chat_template_gemma4.jinja index 15c5238ac332..f62ca843a405 100644 --- a/examples/tool_chat_template_gemma4.jinja +++ b/examples/tool_chat_template_gemma4.jinja @@ -1,9 +1,9 @@ -{%- macro format_parameters(properties, required) -%} +{%- macro format_parameters(properties, required, filter_keys=false) -%} {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} {%- set ns = namespace(found_first=false) -%} {%- for key, value in properties | dictsort -%} {%- set add_comma = false -%} - {%- if key not in standard_keys -%} + {%- if not filter_keys or key not in standard_keys -%} {%- if ns.found_first %},{% endif -%} {%- set ns.found_first = true -%} {{ key }}:{ @@ -11,34 +11,15 @@ description:<|"|>{{ value['description'] }}<|"|> {%- set add_comma = true -%} {%- endif -%} - {%- if value['nullable'] %} - {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} - nullable:true - {%- endif -%} {%- if value['type'] | upper == 'STRING' -%} {%- if value['enum'] -%} {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} enum:{{ format_argument(value['enum']) }} {%- endif -%} - {%- elif value['type'] | upper == 'OBJECT' -%} - ,properties:{ - {%- if value['properties'] is defined and value['properties'] is mapping -%} - {{- format_parameters(value['properties'], value['required'] | default([])) -}} - {%- elif value is mapping -%} - {{- format_parameters(value, value['required'] | default([])) -}} - {%- endif -%} - } - {%- if value['required'] -%} - ,required:[ - {%- for item in value['required'] | default([]) -%} - <|"|>{{- item -}}<|"|> - {%- if not loop.last %},{% endif -%} - {%- endfor -%} - ] - {%- endif -%} {%- elif value['type'] | upper == 'ARRAY' -%} {%- if value['items'] is mapping and value['items'] -%} - ,items:{ + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ {%- set ns_items = namespace(found_first=false) -%} {%- for item_key, item_value in value['items'] | dictsort -%} {%- if item_value is not none -%} @@ -71,6 +52,32 @@ } {%- endif -%} {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} type:<|"|>{{ value['type'] | upper }}<|"|>} {%- endif -%} @@ -167,20 +174,25 @@ {%- set ns = namespace(prev_message_type=None) -%} {%- set loop_messages = messages -%} -{{ bos_token }} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} {%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%} {{- '<|turn>system\n' -}} - + {#- Inject Thinking token at the very top of the FIRST system turn -#} {%- if enable_thinking is defined and enable_thinking -%} - {{- '<|think|>' -}} + {{- '<|think|>\n' -}} {%- set ns.prev_message_type = 'think' -%} {%- endif -%} - {%- if messages[0]['role'] in ['system', 'developer'] -%} - {{- messages[0]['content'] | trim -}} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} {%- set loop_messages = messages[1:] -%} {%- endif -%} - {%- if tools -%} {%- for tool in tools %} {{- '<|tool>' -}} @@ -189,10 +201,10 @@ {%- endfor %} {%- set ns.prev_message_type = 'tool' -%} {%- endif -%} - {{- '\n' -}} {%- endif %} +{#- Pre-scan: find last user message index for reasoning guard -#} {%- set ns_turn = namespace(last_user_idx=-1) -%} {%- for i in range(loop_messages | length) -%} {%- if loop_messages[i]['role'] == 'user' -%} @@ -200,12 +212,12 @@ {%- endif -%} {%- endfor -%} +{#- Loop through messages -#} {%- for message in loop_messages -%} {%- if message['role'] != 'tool' -%} {%- set ns.prev_message_type = None -%} {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} - {#- OpenAI may emit multiple assistant messages in one tool loop (user → asst → tool → asst → tool). - Only the first of those should open <|turn>model; later ones continue the same model turn. -#} + {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#} {%- set prev_nt = namespace(role=None, found=false) -%} {%- if loop.index0 > 0 -%} {%- for j in range(loop.index0 - 1, -1, -1) -%} @@ -222,8 +234,10 @@ {{- '<|turn>' + role + '\n' }} {%- endif -%} - {%- if message.get('reasoning') and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%} - {{- '<|channel>thought\n' + message['reasoning'] + '\n'}} + {#- Render reasoning/reasoning_content as thinking channel -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} {%- endif -%} {%- if message['tool_calls'] -%} @@ -247,14 +261,14 @@ {%- set ns_tr_out = namespace(flag=false) -%} {%- if message.get('tool_responses') -%} - {#- Legacy: tool_responses embedded on the assistant message -#} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} {%- for tool_response in message['tool_responses'] -%} {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}} {%- set ns_tr_out.flag = true -%} {%- set ns.prev_message_type = 'tool_response' -%} {%- endfor -%} {%- elif message.get('tool_calls') -%} - {#- OpenAI Chat Completions: consecutive following messages with role "tool" (no break/continue; range scan) -#} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} {%- set ns_tool_scan = namespace(stopped=false) -%} {%- for k in range(loop.index0 + 1, loop_messages | length) -%} {%- if ns_tool_scan.stopped -%} @@ -262,12 +276,14 @@ {%- set ns_tool_scan.stopped = true -%} {%- else -%} {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%} {%- for tc in message['tool_calls'] -%} {%- if tc.get('id') == follow.get('tool_call_id') -%} {%- set ns_tname.name = tc['function']['name'] -%} {%- endif -%} {%- endfor -%} + {#- Handle content as string or content-parts array -#} {%- set tool_body = follow.get('content') -%} {%- if tool_body is string -%} {{- format_tool_response_block(ns_tname.name, tool_body) -}} @@ -288,6 +304,7 @@ {%- endfor -%} {%- endif -%} + {%- set captured_content -%} {%- if message['content'] is string -%} {%- if role == 'model' -%} {{- strip_thinking(message['content']) -}} @@ -303,29 +320,35 @@ {{- item['text'] | trim -}} {%- endif -%} {%- elif item['type'] == 'image' -%} - {{- '\n\n<|image|>\n\n' -}} + {{- '<|image|>' -}} {%- set ns.prev_message_type = 'image' -%} {%- elif item['type'] == 'audio' -%} {{- '<|audio|>' -}} {%- set ns.prev_message_type = 'audio' -%} {%- elif item['type'] == 'video' -%} - {{- '\n\n<|video|>\n\n' -}} + {{- '<|video|>' -}} {%- set ns.prev_message_type = 'video' -%} {%- endif -%} {%- endfor -%} {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} - {%- if not (ns_tr_out.flag and not message.get('content')) -%} + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif not (ns_tr_out.flag and not has_content) -%} {{- '\n' -}} {%- endif -%} {%- endif -%} {%- endfor -%} {%- if add_generation_prompt -%} - {%- if ns.prev_message_type != 'tool_response' -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} {{- '<|turn>model\n' -}} + {%- if not enable_thinking | default(false) -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} {%- endif -%} - {%- if not enable_thinking | default(false) -%} - {{- '<|channel>thought\n' -}} - {%- endif -%} -{%- endif -%} +{%- endif -%} \ No newline at end of file From 964a4bc2a57aca2a42d04538b27cab4d333d0f5d Mon Sep 17 00:00:00 2001 From: John Calderon <81483067+johncalesp@users.noreply.github.com> Date: Fri, 1 May 2026 23:10:14 -0400 Subject: [PATCH 0009/1083] [MM][CG] Support ViT CG for Qwen2.5-VL (#40830) Signed-off-by: John Calderon --- docs/design/cuda_graphs_multimodal.md | 2 + .../multimodal/vision_language_offline.py | 1 + .../multimodal/generation/test_qwen2_5_vl.py | 95 ++++ .../generation/test_vit_cudagraph.py | 13 +- vllm/model_executor/models/qwen2_5_vl.py | 450 +++++++++++++++++- 5 files changed, 539 insertions(+), 22 deletions(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index e32010232ef0..f44ef359df38 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -86,9 +86,11 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | Architecture | Models | CG for Image | CG for Video | | ------------ | ------ | ------------ | ------------ | | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | +| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | !!! note Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. + For Qwen2.5-VL only FA2 and FA3 has been tested. ## Configuration diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 87d42c036ec1..794f20dd0a52 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -2466,6 +2466,7 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData: MODELS_SUPPORT_VIT_CUDA_GRAPH = [ "qwen3_vl", "qwen3_vl_moe", + "qwen2_5_vl", ] diff --git a/tests/models/multimodal/generation/test_qwen2_5_vl.py b/tests/models/multimodal/generation/test_qwen2_5_vl.py index 3ba665710af4..791bb3b3088f 100644 --- a/tests/models/multimodal/generation/test_qwen2_5_vl.py +++ b/tests/models/multimodal/generation/test_qwen2_5_vl.py @@ -3,6 +3,7 @@ import pytest +from vllm.assets.image import ImageAsset from vllm.multimodal.video import sample_frames_from_video from ....conftest import VIDEO_ASSETS @@ -11,6 +12,7 @@ target_dtype = "bfloat16" VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>" +IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" def qwen2_5_vl_chat_template(*query): @@ -28,6 +30,25 @@ def qwen2_5_vl_chat_template(*query): ) +WINDOW_ATTN_IMAGE_PROMPT = qwen2_5_vl_chat_template( + IMAGE_PLACEHOLDER, + "Describe the image.", +) + + +def _window_attention_regression_image(): + # image from regression issue: https://github.com/vllm-project/vllm/issues/15122 + image = ImageAsset("hato").pil_image + return image.resize((image.width // 2, image.height // 2)) + + +def _encoder_cudagraph_config(*, max_vision_items: int) -> dict: + return { + "cudagraph_mm_encoder": True, + "encoder_cudagraph_max_vision_items_per_batch": max_vision_items, + } + + @pytest.mark.core_model @pytest.mark.parametrize("model", models) @pytest.mark.parametrize("video_pruning_rate", [0.0, 0.75]) @@ -146,3 +167,77 @@ def test_qwen2_5_vl_evs_batched_videos( # Ensure the output is a string assert isinstance(output_text, str) + + +@pytest.mark.core_model +@pytest.mark.parametrize("model", models) +@pytest.mark.parametrize("dtype", [target_dtype]) +@pytest.mark.parametrize("max_tokens", [128]) +@pytest.mark.parametrize("use_bytecode_hook", [True, False]) +def test_qwen2_5_vl_window_attention_image( + vllm_runner, + model, + dtype: str, + max_tokens: int, + use_bytecode_hook: bool, + monkeypatch, +) -> None: + """Regression test for Qwen2.5 window-attention image path.""" + monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0") + + prompt = [WINDOW_ATTN_IMAGE_PROMPT] + images = [[_window_attention_regression_image()]] + + with vllm_runner( + model, + runner="generate", + max_model_len=4096, + dtype=dtype, + limit_mm_per_prompt={"image": 1}, + compilation_config=_encoder_cudagraph_config(max_vision_items=1), + ) as vllm_model: + outputs = vllm_model.generate_greedy(prompt, max_tokens, images=images) + + assert len(outputs) == 1 + output_ids, output_text = outputs[0] + assert len(output_ids) > 0 + assert len(output_text) > 0 + assert isinstance(output_text, str) + + +@pytest.mark.core_model +@pytest.mark.parametrize("model", models) +@pytest.mark.parametrize("dtype", [target_dtype]) +@pytest.mark.parametrize("max_tokens", [128]) +@pytest.mark.parametrize("use_bytecode_hook", [True, False]) +def test_qwen2_5_vl_window_attention_image_batch( + vllm_runner, + model, + dtype: str, + max_tokens: int, + use_bytecode_hook: bool, + monkeypatch, +) -> None: + """Regression test window-attention with a small image batch.""" + monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0") + + image = _window_attention_regression_image() + prompts = [WINDOW_ATTN_IMAGE_PROMPT, WINDOW_ATTN_IMAGE_PROMPT] + images = [[image], [image]] + + with vllm_runner( + model, + runner="generate", + max_model_len=4096, + max_num_seqs=2, + dtype=dtype, + limit_mm_per_prompt={"image": 1}, + compilation_config=_encoder_cudagraph_config(max_vision_items=2), + ) as vllm_model: + outputs = vllm_model.generate_greedy(prompts, max_tokens, images=images) + + assert len(outputs) == 2 + for output_ids, output_text in outputs: + assert len(output_ids) > 0 + assert len(output_text) > 0 + assert isinstance(output_text, str) diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 7adea0771b6d..fb7bdfc8625d 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -54,7 +54,18 @@ def qwen_vl_chat_template(content: str) -> str: needs_video_metadata=True, marks=[pytest.mark.core_model], ), - # TODO: Add more models below. + "qwen2_5_vl": VitCudagraphTestConfig( + model="Qwen/Qwen2.5-VL-3B-Instruct", + image_prompt=qwen_vl_chat_template( + "<|vision_start|><|image_pad|><|vision_end|>What is in this image?" + ), + video_prompt=qwen_vl_chat_template( + "<|vision_start|><|video_pad|><|vision_end|>" + "Describe this video in one sentence." + ), + needs_video_metadata=False, + marks=[pytest.mark.core_model], + ), } diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index c11684b4b89b..54334c91bfa6 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -85,11 +85,13 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers from .interfaces import ( MultiModalEmbeddings, SupportsEagle, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMRoPE, SupportsMultiModal, @@ -771,22 +773,54 @@ def invert_permutation(perm: torch.Tensor) -> torch.Tensor: inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype) return inv - def forward( + def prepare_encoder_metadata( self, - x: torch.Tensor, grid_thw: list[list[int]], - ) -> torch.Tensor: + *, + max_batch_size: int | None = None, + max_frames_per_batch: int | None = None, + max_window_seqs_per_batch: int | None = None, + max_seqlen_override: int | None = None, + max_seqlen_window_override: int | None = None, + device: torch.device | None = None, + ) -> dict[str, torch.Tensor]: + """Compute encoder metadata from grid_thw. + + Shared by the eager forward path, CUDA graph capture, and + CUDA graph replay to avoid duplicated implementation. + + Args: + grid_thw: Grid configurations as list of [t, h, w]. + max_batch_size: If set, pad cu_seqlens to this size + (needed for CUDA graph capture/replay). + max_frames_per_batch: If set, overrides max_batch_size for + cu_seqlens padding. For video inputs each item contributes + T attention sequences (frames); this sizes the buffer to + the total frame budget so video replays never overflow. + max_window_seqs_per_batch: If set, pad cu_window_seqlens to this + number of window sequences. This keeps cu_window_seqlens shape + stable across capture/replay for CUDA graph safety. + max_seqlen_override: If set, use this value for max_seqlen + instead of computing from cu_seqlens (needed for CUDA + graph capture to cover worst-case replay scenarios). + max_seqlen_window_override: If set, use this value for + window-attention max_seqlen instead of computing from + cu_window_seqlens (needed for CUDA graph capture to + cover worst-case replay scenarios). + device: Device to place tensors on. Defaults to self.device. + """ + + if device is None: + device = self.device + metadata: dict[str, torch.Tensor] = {} + # patchify - seq_len, _ = x.size() rotary_pos_emb_cos = [] rotary_pos_emb_sin = [] window_index: list = [] cu_window_seqlens: list = [torch.tensor([0], dtype=torch.int32)] cu_seqlens: list = [] - hidden_states = x.to(device=self.device, dtype=self.dtype) - hidden_states = self.patch_embed(hidden_states) - window_index_id = 0 cu_window_seqlens_last = 0 for t, h, w in grid_thw: @@ -825,23 +859,99 @@ def forward( cu_seqlens = torch.cumsum(cu_seqlens, dim=0, dtype=torch.int32) cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0) + # Pad cu_seqlens to the required number of sequences. + # For videos each item contributes T frames = T attention sequences, + # so the total can exceed max_batch_size. max_frames_per_batch + # overrides the pad target when set. + pad_to = ( + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size + ) + if pad_to is not None: + num_seqs = len(cu_seqlens) - 1 + if num_seqs < pad_to: + cu_seqlens = torch.cat( + ( + cu_seqlens, + torch.full( + (pad_to - num_seqs,), + cu_seqlens[-1], + dtype=cu_seqlens.dtype, + device=cu_seqlens.device, + ), + ) + ) + + # Pad cu_window_seqlens to a stable number of window sequences. + # Like cu_seqlens, we repeat the last cumulative offset so padded + # entries represent empty sequences. + if max_window_seqs_per_batch is not None: + num_window_seqs = len(cu_window_seqlens) - 1 + if num_window_seqs < max_window_seqs_per_batch: + cu_window_seqlens = torch.cat( + ( + cu_window_seqlens, + torch.full( + (max_window_seqs_per_batch - num_window_seqs,), + cu_window_seqlens[-1], + dtype=cu_window_seqlens.dtype, + device=cu_window_seqlens.device, + ), + ) + ) + # transformers # pre-compute seqlens for window/full attn to reduce cuMemcpy operations - max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens) - max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens) + if max_seqlen_override is None: + max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens) + else: + max_seqlen_full = torch.tensor(max_seqlen_override, dtype=torch.int32) + if max_seqlen_window_override is None: + max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens) + else: + max_seqlen_window = torch.tensor( + max_seqlen_window_override, dtype=torch.int32 + ) - cu_seqlens = cu_seqlens.to(device=self.device, non_blocking=True) - cu_window_seqlens = cu_window_seqlens.to(device=self.device, non_blocking=True) - rotary_pos_emb_cos = rotary_pos_emb_cos.to( - device=self.device, non_blocking=True - ) - rotary_pos_emb_sin = rotary_pos_emb_sin.to( - device=self.device, non_blocking=True - ) - window_index = window_index.to(device=hidden_states.device, non_blocking=True) - reverse_indices = reverse_indices.to( - device=hidden_states.device, non_blocking=True - ) + cu_seqlens = cu_seqlens.to(device=device, non_blocking=True) + cu_window_seqlens = cu_window_seqlens.to(device=device, non_blocking=True) + rotary_pos_emb_cos = rotary_pos_emb_cos.to(device=device, non_blocking=True) + rotary_pos_emb_sin = rotary_pos_emb_sin.to(device=device, non_blocking=True) + window_index = window_index.to(device=device, non_blocking=True) + reverse_indices = reverse_indices.to(device=device, non_blocking=True) + + metadata["rotary_pos_emb_cos"] = rotary_pos_emb_cos + metadata["rotary_pos_emb_sin"] = rotary_pos_emb_sin + metadata["window_index"] = window_index + metadata["reverse_indices"] = reverse_indices + metadata["cu_seqlens"] = cu_seqlens + metadata["cu_window_seqlens"] = cu_window_seqlens + metadata["max_seqlen_full"] = max_seqlen_full + metadata["max_seqlen_window"] = max_seqlen_window + + return metadata + + def forward( + self, + x: torch.Tensor, + grid_thw: list[list[int]], + *, + encoder_metadata: dict[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + hidden_states = x.to(device=self.device, dtype=self.dtype) + hidden_states = self.patch_embed(hidden_states) + + seq_len = hidden_states.shape[0] + if encoder_metadata is None: + encoder_metadata = self.prepare_encoder_metadata(grid_thw) + + rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"] + rotary_pos_emb_sin = encoder_metadata["rotary_pos_emb_sin"] + window_index = encoder_metadata["window_index"] + reverse_indices = encoder_metadata["reverse_indices"] + cu_seqlens = encoder_metadata["cu_seqlens"] + cu_window_seqlens = encoder_metadata["cu_window_seqlens"] + max_seqlen_full = encoder_metadata["max_seqlen_full"] + max_seqlen_window = encoder_metadata["max_seqlen_window"] hidden_states = hidden_states.reshape( seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 @@ -1003,6 +1113,7 @@ def get_replacement_qwen2vl(item_idx: int, modality: str): class Qwen2_5_VLForConditionalGeneration( nn.Module, SupportsMultiModal, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsPP, SupportsQuant, @@ -1124,6 +1235,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.config = config + self.model_config = vllm_config.model_config self.vllm_config = vllm_config self.multimodal_config = multimodal_config self.video_pruning_rate = multimodal_config.video_pruning_rate @@ -1447,6 +1559,302 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: multimodal_embeddings += tuple(video_embeddings) return multimodal_embeddings + # -- SupportsEncoderCudaGraph protocol methods -- + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphConfig, + ) + + # NOTE: With EVS pruning enabled, multimodal embeddings are post-processed + # (append positions for image and prune+append positions for video) in + # embed_multimodal(). The encoder CUDA graph path bypasses that postprocess + # hook, so disable CUDA graph for all modalities to avoid inconsistent + # embedding formats between eager and cudagraph paths. + modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"] + + return EncoderCudaGraphConfig( + modalities=modalities, + input_key_by_modality={ + "image": "pixel_values", + "video": "pixel_values_videos", + }, + buffer_keys=[ + "rotary_pos_emb_cos", + "rotary_pos_emb_sin", + "window_index", + "reverse_indices", + "cu_seqlens", + "cu_window_seqlens", + "max_seqlen_full", + "max_seqlen_window", + ], + out_hidden_size=self.visual.out_hidden_size, + ) + + def get_input_modality( + self, + mm_kwargs: dict[str, Any], + ) -> str: + if "image_grid_thw" in mm_kwargs: + return "image" + return "video" + + def get_max_frames_per_video(self) -> int: + mm_registry = MULTIMODAL_REGISTRY + info = mm_registry.get_processing_info(self.model_config) + max_frames_per_video = info.get_num_frames_with_most_features( + seq_len=self.model_config.max_model_len, + mm_counts={"video": self.multimodal_config.get_limit_per_prompt("video")}, + ) + return max_frames_per_video + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + # Min: estimated smallest possible encoder input. + # 224x224 image → 16x16 patches (patch_size=14) + # spatial_merge_size=2 → 8x8 = 64 tokens + min_budget = 64 + # Max: capped by max_num_batched_tokens + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def _get_pixel_values_by_modality( + self, + mm_kwargs: dict[str, Any], + ) -> torch.Tensor: + if self.get_input_modality(mm_kwargs) == "image": + pixel_values = mm_kwargs["pixel_values"] + else: + pixel_values = mm_kwargs["pixel_values_videos"] + return pixel_values + + def _get_grid_thw_by_modality( + self, + mm_kwargs: dict[str, Any], + ) -> list[tuple[int, int, int]]: + grid_thw_key = f"{self.get_input_modality(mm_kwargs)}_grid_thw" + grid_thw = mm_kwargs[grid_thw_key] + if not isinstance(grid_thw, list): + grid_thw = grid_thw.tolist() + return grid_thw + + def get_encoder_cudagraph_num_items( + self, + mm_kwargs: dict[str, Any], + ) -> int: + return len(self._get_grid_thw_by_modality(mm_kwargs)) + + def get_encoder_cudagraph_per_item_output_tokens( + self, + mm_kwargs: dict[str, Any], + ) -> list[int]: + m = self.visual.spatial_merge_size + grid_thw = self._get_grid_thw_by_modality(mm_kwargs) + return [t * (h // m) * (w // m) for t, h, w in grid_thw] + + def get_encoder_cudagraph_per_item_input_sizes( + self, + mm_kwargs: dict[str, Any], + ) -> list[int]: + grid_thw = self._get_grid_thw_by_modality(mm_kwargs) + return [t * h * w for t, h, w in grid_thw] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + grid_thw = self._get_grid_thw_by_modality(mm_kwargs) + pixel_values = self._get_pixel_values_by_modality(mm_kwargs) + + if len(indices) == 0: + if self.get_input_modality(mm_kwargs) == "image": + return { + "pixel_values": pixel_values[:0], + "image_grid_thw": [], + } + else: + return { + "pixel_values_videos": pixel_values[:0], + "video_grid_thw": [], + } + + # Compute cumulative patch offsets for slicing pixel_values + patches_per_item = [t * h * w for t, h, w in grid_thw] + cum_patches = [0] + for p in patches_per_item: + cum_patches.append(cum_patches[-1] + p) + + selected_pv = torch.cat( + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices] + ) + selected_grid = [grid_thw[i] for i in indices] + + if self.get_input_modality(mm_kwargs) == "image": + return { + "pixel_values": selected_pv, + "image_grid_thw": selected_grid, + } + else: + return { + "pixel_values_videos": selected_pv, + "video_grid_thw": selected_grid, + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + spatial_merge_size = self.visual.spatial_merge_size + max_window_seqs_per_batch = min( + self.vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + # Use ceil here (not floor) so total captured capacity is never smaller + # than token_budget when token_budget is not divisible by max_batch_size + # (e.g., 324 budget with max_batch_size=8). Floor under-allocates + # input_buffer and can fail replay copy for valid single-item batches. + per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size + + frames_per_item = max_frames_per_batch // max_batch_size + if frames_per_item > 1: + # Build the capture grid using a video-format layout so that + # cu_seqlens is sized for video replays from the start. + # cu_seqlens has one entry per attention sequence (one per frame), + # so using T > 1 per item makes the buffer large enough without + # relying solely on padding. + # Ceiling ensures frames_per_item * tokens_per_frame >= per_mm_item_output + # so the pixel_values buffer covers any valid single-item replay. + tokens_per_frame = ( + per_mm_item_output + frames_per_item - 1 + ) // frames_per_item + # Video-format grid_config (T=frames_per_item). + grid_config = [ + [ + frames_per_item, + spatial_merge_size, + tokens_per_frame * spatial_merge_size, + ] + for _ in range(max_batch_size) + ] + else: + # Image-format grid_config (T=1). + grid_config = [ + [1, spatial_merge_size, per_mm_item_output * spatial_merge_size] + for _ in range(max_batch_size) + ] + + # Create dummy pixel_values + patch_embed = self.visual.patch_embed + in_channels = patch_embed.proj.in_channels + patch_size = patch_embed.patch_size + temporal_patch_size = patch_embed.temporal_patch_size + total_patches = sum(t * h * w for t, h, w in grid_config) + flattened_patch_size = ( + in_channels * temporal_patch_size * patch_size * patch_size + ) + dummy_pixel_values = torch.randn( + total_patches, flattened_patch_size, device=device, dtype=dtype + ) + + # Override max_seqlen with a safe upper bound for capture. + # max_seqlen.item() gets baked into the CUDA graph (not replayed), + # so the capture value must cover any replay scenario. + # Worst case: 1 item consuming the full budget -> + # seq_len = token_budget * spatial_merge_size^2. + # For window-attention, each local window is bounded by fixed geometry: + # (window_size / patch_size / spatial_merge_size)^2 windows in merged + # token space, multiplied by spatial_merge_size^2 to map back to the + # unmerged sequence length used by attention kernels. + vit_merger_window_size = ( + self.visual.window_size + // self.visual.spatial_merge_size + // self.visual.patch_size + ) + max_seqlen_window_override = vit_merger_window_size**2 * (spatial_merge_size**2) + buffers = self.visual.prepare_encoder_metadata( + grid_config, + max_batch_size=max_batch_size, + max_frames_per_batch=max_frames_per_batch, + max_window_seqs_per_batch=max_window_seqs_per_batch, + max_seqlen_override=token_budget * (spatial_merge_size**2), + max_seqlen_window_override=max_seqlen_window_override, + device=device, + ) + + # Just use image-modality dummy input_buffer for capturing, since it's also + # compatible for video inputs (has the same shape: [num_patches, C*T*P*P]). + mm_kwargs = { + "pixel_values": dummy_pixel_values, + "image_grid_thw": grid_config, + } + + return EncoderCudaGraphCaptureInputs( + mm_kwargs=mm_kwargs, + buffers=buffers, + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + ): + modality = self.get_input_modality(mm_kwargs) + grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs) + + if modality == "image": + buffers = self.visual.prepare_encoder_metadata( + grid_thw_list, + max_batch_size=max_batch_size, + max_window_seqs_per_batch=min( + self.vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ), + ) + else: + buffers = self.visual.prepare_encoder_metadata( + grid_thw_list, + max_frames_per_batch=max_frames_per_batch, + max_window_seqs_per_batch=min( + self.vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ), + ) + + return EncoderCudaGraphReplayBuffers(buffers=buffers) + + def encoder_cudagraph_forward( + self, + mm_kwargs: dict[str, Any], + buffers: dict[str, torch.Tensor], + ) -> torch.Tensor: + pixel_values = self._get_pixel_values_by_modality(mm_kwargs) + grid_thw = self._get_grid_thw_by_modality(mm_kwargs) + return self.visual(pixel_values, grid_thw, encoder_metadata=buffers) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + ) -> torch.Tensor: + pixel_values = self._get_pixel_values_by_modality(mm_kwargs) + grid_thw = self._get_grid_thw_by_modality(mm_kwargs) + return self.visual(pixel_values, grid_thw) + def forward( self, input_ids: torch.Tensor | None, From 3e49479c4b766a601804f0c6f5f1c9a3def5ad0c Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 1 May 2026 23:19:07 -0400 Subject: [PATCH 0010/1083] Limit concurrency on `test_transcription_api_correctness.py` (#41478) Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- .../openai/correctness/test_transcription_api_correctness.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py index f17f6f5f90c2..a3df30fb02b2 100644 --- a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py @@ -26,6 +26,9 @@ from ....models.registry import HF_EXAMPLE_MODELS from ....utils import RemoteOpenAIServer +# Tuned to prevent OOM on 18GB GPUs in transcription correctness tests. +MAX_SEQS_FOR_TRANSCRIPTION_TEST = 32 + def to_bytes(y, sr): buffer = io.BytesIO() @@ -184,6 +187,7 @@ def test_wer_correctness( server_args = [ "--enforce-eager", f"--tokenizer_mode={model_info.tokenizer_mode}", + f"--max_num_seqs={MAX_SEQS_FOR_TRANSCRIPTION_TEST}", ] if model_info.trust_remote_code: server_args.append("--trust-remote-code") From d58c42e19cb792e24eb335b75164356a4f71bff0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luka=20Govedi=C4=8D?= Date: Fri, 1 May 2026 23:41:15 -0400 Subject: [PATCH 0011/1083] [vLLM IR] 2/N fused_add_rms_norm and maybe_inplace overload (#36823) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Luka Govedič Signed-off-by: Luka Govedič --- docs/design/debug_vllm_compile.md | 32 +- docs/design/vllm_ir.md | 615 ++++++++++++++++++ docs/mkdocs/hooks/generate_argparse.py | 28 +- tests/compile/backend.py | 17 +- .../distributed/test_sequence_parallelism.py | 29 +- tests/compile/passes/ir/test_clone_cleanup.py | 412 ++++++++++++ .../ir/test_inplace_functionalization.py | 465 +++++++++++++ .../compile/passes/test_functionalization.py | 2 +- tests/compile/passes/test_fuse_act_padding.py | 2 +- tests/compile/passes/test_fusion.py | 19 +- tests/ir/test_inplace_op.py | 91 +++ tests/ir/test_op.py | 43 +- tests/kernels/ir/test_layernorm.py | 206 +++++- .../model_executor/test_enabled_custom_ops.py | 29 +- tests/model_executor/test_oink_integration.py | 129 ++-- tests/test_config.py | 40 +- .../test_rms_norm_batch_invariant.py | 8 +- vllm/_aiter_ops.py | 100 +-- vllm/_oink_ops.py | 96 --- vllm/compilation/backends.py | 19 + .../passes/fusion/allreduce_rms_fusion.py | 38 +- .../passes/fusion/matcher_utils.py | 67 -- .../passes/fusion/rms_quant_fusion.py | 38 +- .../passes/fusion/rocm_aiter_fusion.py | 46 +- .../passes/fusion/sequence_parallelism.py | 20 +- vllm/compilation/passes/inductor_pass.py | 3 + .../passes/ir/clone_elimination.py | 126 ++++ .../passes/ir/inplace_functionalization.py | 101 +++ vllm/compilation/passes/ir/lowering_pass.py | 44 +- vllm/compilation/passes/ir/utils.py | 40 ++ vllm/compilation/passes/pass_manager.py | 8 +- vllm/config/kernel.py | 3 + vllm/config/vllm.py | 7 +- vllm/envs.py | 5 +- vllm/ir/op.py | 156 ++++- vllm/ir/ops/__init__.py | 4 +- vllm/ir/ops/layernorm.py | 38 +- vllm/kernels/aiter_ops.py | 69 ++ vllm/kernels/oink_ops.py | 60 +- vllm/kernels/vllm_c.py | 30 + vllm/kernels/xpu_ops.py | 28 + vllm/model_executor/layers/layernorm.py | 281 +------- vllm/platforms/cuda.py | 4 +- vllm/platforms/rocm.py | 22 +- 44 files changed, 2837 insertions(+), 783 deletions(-) create mode 100644 docs/design/vllm_ir.md create mode 100644 tests/compile/passes/ir/test_clone_cleanup.py create mode 100644 tests/compile/passes/ir/test_inplace_functionalization.py create mode 100644 tests/ir/test_inplace_op.py delete mode 100644 vllm/_oink_ops.py create mode 100644 vllm/compilation/passes/ir/clone_elimination.py create mode 100644 vllm/compilation/passes/ir/inplace_functionalization.py create mode 100644 vllm/compilation/passes/ir/utils.py diff --git a/docs/design/debug_vllm_compile.md b/docs/design/debug_vllm_compile.md index fbee9f4c3e3e..7edda6fa6476 100644 --- a/docs/design/debug_vllm_compile.md +++ b/docs/design/debug_vllm_compile.md @@ -5,12 +5,14 @@ TL;DR: - use tlparse to acquire torch.compile logs. Include these logs in bug reports and/or support asks. - The vLLM-torch.compile integration is multiple pieces. vLLM exposes flags to turn off each piece: -| Online Flag | Offline Flag | Result | -| ----------- | ------------ | ------ | -| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs | -| -cc.mode=0 | mode=CompilationMode.NONE | Turn off torch.compile only | -| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only | -| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor | +| Online Flag | Offline Flag | Result | +|--------------------------------|--------------------------------------------------------------------------------|------------------------------------------------------| +| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs | +| -cc.mode=0 | compilation_config=CompilationConfig(mode=CompilationMode.NONE) | Turn off torch.compile only | +| -cc.mode=1 | compilation_config=CompilationConfig(mode=CompilationMode.STOCK_TORCH_COMPILE) | Turn off vLLM-compile modifications to torch.compile | +| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only | +| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor | +| -cc.ir_enable_torch_wrap=False | compilation_config=CompilationConfig(ir_enable_torch_wrap=False) | Turn off vLLM IR wrapping | ## vLLM-torch.compile overview @@ -22,7 +24,7 @@ Most notably, vLLM-compile is NOT torch.compile, it is a custom compiler built u - Given a model, we do a full graph capture via TorchDynamo that is dynamic on the batch size (number of tokens) - vLLM then optionally splits and/or specializes this graph and then uses TorchInductor to compile each graph into a compiled artifact. -This step may use vLLM custom Inductor passes to further optimize the graph. +This step may use vLLM custom Inductor passes to further optimize the graph. This includes vLLM IR lowering to remove dispatch overhead. - The compiled artifact is saved to vLLM's compile cache so that it can be loaded in the future. - vLLM applies CUDAGraphs to reduce CPU overheads. @@ -34,6 +36,7 @@ For more details on the design, please see the following resources: - [Introduction to vLLM-torch.compile blogpost](https://blog.vllm.ai/2025/08/20/torch-compile.html) - [vLLM-torch.compile integration design](./torch_compile.md) +- [vLLM IR design](./vllm_ir.md) - [vLLM Office Hours #26](https://www.youtube.com/live/xLyxc7hxCJc?si=Xulo9pe53C6ywf0V&t=561) - [Talk at PyTorch Conference 2025](https://youtu.be/1wV1ESbGrVQ?si=s1GqymUfwiwOrDTg&t=725) @@ -117,6 +120,21 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode LLM(model, compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE)) ``` +vLLM IR makes heavy use of the compilation pipeline, from functionalization, custom fusions, and lowering. +To turn that off and capture eager-mode dispatching behavior of vLLM IR, run with `ir_enable_torch_wrap=False`. +IR torch wrap is only enabled by default when using `mode=VLLM_COMPILE` and `backend="inductor"` (default). + +```sh +# Online +vllm serve -cc.ir_enable_torch_wrap=False +``` + +```py +# Offline +from vllm.config.compilation import CompilationConfig +LLM(model, compilation_config=CompilationConfig(ir_enable_torch_wrap=False)) +``` + ## Debugging TorchDynamo vLLM requires model code be capturable into a full graph via TorchDynamo (torch.compile's frontend). diff --git a/docs/design/vllm_ir.md b/docs/design/vllm_ir.md new file mode 100644 index 000000000000..82628f3762fe --- /dev/null +++ b/docs/design/vllm_ir.md @@ -0,0 +1,615 @@ +# vLLM IR: Functional Intermediate Representation + +## Motivation + +vLLM IR is a **functional intermediate representation (IR)** that fills the gap between +low-level `torch` ops and vLLM layers like `RMSNorm` and quantization operators, +By separating operator **semantics** from the **implementation** and **dispatching**, +vLLM IR simplifies both compilation and kernel registration & dispatching simultaneously. +It operates as a **dialect** in the torch FX representation, allowing full interoperability +with “regular” torch ops & custom torch ops/kernels, as well as a piecewise migration from +the previous `CustomOp` approach. + +Key design principles: + +- **Eager-compile consistency**: identical behavior (barring minor numerics) in eager and compiled modes +- **Simple, transparent, yet powerful kernel selection**: good visibility and control allowing easy debugging +- **Convention over configuration**: near-zero boilerplate required to register ops and implementations +- **Extensibility**: ops and implementations can be registered anywhere, in-tree or out-of-tree +- **Interoperability**: fully compatible with “regular” torch ops & custom torch ops/kernels, +reducing developer friction and allowing piecewise migration + +The clean semantics/implementation separation enables a unified and extensible dispatching mechanism, +allowing multiple kernels per-platform and powerful kernel selection. The separation also facilitates +cleaner testing and benchmarking, removing much of the boilerplate standard for legacy approaches. + +By delaying kernel selection until late in the compilation process, the compiler can operate on +a higher-level representation, which has the following main benefits: + +- Pattern matching in fusion/transformation passes only requires a single, simple pattern per op +- OOT compiler backends can lower from the higher-level representation (in-progress) +- The compiler can autotune over available implementations (future feature) + +## Quick Overview + +### Declaring an IR Operation + +IR operations are declared using the `@register_op` decorator with a native PyTorch implementation that defines the op's semantics: + +```python +# vllm/ir/ops/layernorm.py +from torch import Tensor +from vllm.ir import register_op + +@register_op +def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor: + """Weighted root-mean-square layer normalization""" + orig_dtype = x.dtype + x = x.to(torch.float32) + x_var = x if variance_size is None else x[..., :variance_size] + variance = x_var.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + epsilon) + x = x.to(orig_dtype) + if weight is not None: + x = x * weight + return x +``` + +The native implementation serves three purposes: + +1. **Semantic definition**: Specifies the exact semantics of the operation, including shapes and strides +2. **Default implementation**: Used when no other (better) implementation is available +3. **Reference for testing**: Other implementations must match these semantics + +### Registering Implementations + +Kernel implementations are registered using the `register_impl` decorator on the IR op object: + +```python +# vllm/kernels/vllm_c.py +from vllm import ir + +rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None + +@ir.ops.rms_norm.register_impl("vllm_c", supports_args=rms_norm_no_var, supported=current_platform.is_cuda_alike()) +def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor: + output = torch.empty_like(x) + torch.ops._C.rms_norm(output, x, weight, epsilon) + return output +``` + +Implementations can specify: + +- `supported`: Static boolean indicating if this implementation is available +- `supports_args`: Function checking if the implementation supports specific arguments +- `inplace`: Whether this implementation reuses input memory for outputs + +### Using IR Operations in Models + +IR operations are imported and called directly in model code: + +```python +# vllm/model_executor/layers/layernorm.py +from vllm import ir + +class RMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, x: Tensor, residual: Tensor | None = None): + if residual is None: + return ir.ops.rms_norm(x, self.weight, self.variance_epsilon) + + # Use maybe_inplace overload to allow implementation to reuse input memory for outputs + # (using x or residual after this call is undefined behavior) + return ir.ops.fused_add_rms_norm.maybe_inplace( + x, residual, self.weight, self.variance_epsilon + ) +``` + +### Configuring Kernel Selection + +Kernel selection is controlled via priority lists in the configuration. +Priority lists specify the order in which implementations are considered, +with the first supported implementation being selected. +This includes the static support check (`supported=...`) and +the dynamic arg support check (`supports_args=...`). + +#### Command Line Configuration + +Use `--ir-op-priority.=,,...`: + +```bash +# CUDA: Use vllm_c implementation for rms_norm +vllm serve meta-llama/Llama-3.2-1B \ + --ir-op-priority.rms_norm=vllm_c + +# ROCm: Try aiter first, fall back to vllm_c, then native +vllm serve meta-llama/Llama-3.2-1B \ + --ir-op-priority.rms_norm=aiter,vllm_c,native + +# Configure multiple operations +vllm serve meta-llama/Llama-3.2-1B \ + --ir-op-priority.rms_norm=vllm_c \ + --ir-op-priority.fused_add_rms_norm=vllm_c +``` + +#### Python Configuration + +```python +from vllm import LLM +from vllm.config import VllmConfig, KernelConfig + +llm = LLM( + model="meta-llama/Llama-3.2-1B", + vllm_config=VllmConfig( + kernel_config=KernelConfig( + ir_op_priority={ + "rms_norm": ["vllm_c", "native"], + "fused_add_rms_norm": ["vllm_c", "native"], + } + ) + ) +) +``` + +#### Platform Defaults + +Each platform provides default priority lists that are automatically applied: + +```python +# CUDA/XPU/ROCm platform defaults (when compiling with Inductor) +{ + "rms_norm": ["native"], # Native torch is default + "fused_add_rms_norm": ["native"], +} + +# CUDA platform defaults (eager or Dynamo-only) +{ + "rms_norm": ["vllm_c", "native"], + "fused_add_rms_norm": ["vllm_c", "native"], +} + +# ROCm platform defaults (future - currently same as CUDA) +{ + "rms_norm": ["aiter", "vllm_c", "native"], + "fused_add_rms_norm": ["aiter", "vllm_c", "native"], +} + +# XPU platform defaults (eager or Dynamo-only) +{ + "rms_norm": ["xpu_kernels", "native"], + "fused_add_rms_norm": ["xpu_kernels", "native"], +} +``` + +User-specified priorities are prepended to platform defaults, +so you only need to specify the out-of-order implementations, +other implementations are appended automatically. + +## Compilation Pipeline + +vLLM IR heavily customizes the `torch.compile`-based compilation process to allow custom compile +passes to operate on high-level IR while still producing efficient low-level code at the end. +The compilation pipeline consists of several stages: + +### 1. Dynamo Tracing + +When `torch.compile` traces the model's forward pass, vLLM IR operations appear as custom operations +in the `vllm_ir` torch library. These operations are opaque to Dynamo, meaning they appear directly +in the FX graph without decomposition: + +```python +# Python code (epsilon=1e-5) +x1 = ir.ops.rms_norm(x, weight, epsilon) +x2, residual_out = ir.ops.fused_add_rms_norm.maybe_inplace(x1, residual, weight, epsilon) + +# FX graph after Dynamo tracing +x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None +out = torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace(x1, residual, weight, 1e-5); x1 = residual = None +x2 = out[0] +residual_out = out[1] +``` + +### 2. AOTAutograd and Functionalization + +AOTAutograd functionalizes the graph, converting any mutating operations to functional equivalents. +For vLLM IR operations with `maybe_inplace` overloads, we perform this manually before AOTAutograd, +converting them to the functional `default` overload using the pre-grad custom pass hook. + +```python +# After functionalization +x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None +out = torch.ops.vllm_ir.fused_add_rms_norm.default(x1, residual, weight, 1e-5); x1 = residual = None +x2 = out[0] +residual_out = out[1] +``` + +The pass also tracks which inputs were "donated" (passed to `maybe_inplace`), +storing this information in vLLM's `PassContext` for later use in clone elimination. + +### 3. IR Fusion and Transformation Passes + +After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations. +These passes can perform fusion, distribute operations for sequence parallelism, and other transformations: + +```python +# Example: Sequence Parallelism (see SequenceParallelismPass) +# Before SP pass + +all_reduce = torch.ops.vllm.all_reduce(x, "tp:0") +rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) + +# after SP pass +reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0") +rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5) +all_gather = torch.ops.vllm.all_gather(x, "tp:0") +``` + +Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations, +handle different kernel implementations separately, or deal with functionalization of custom kernels. + +### 4. IR Lowering + +The lowering pass (`VllmIRLoweringPass`) replaces each vLLM IR operation with its selected implementation. +The implementation is chosen based on the priority list and support predicates, +using the **fake tensors** in the graph's metadata in place of op arguments: + +```python +# Implementation selection, same in eager dispatch and compile lowering +def dispatch(*args) -> IrOpImpl: + for provider in priority_list: # e.g., ["vllm_c", "native"] + impl = ir_op.impls[provider] + if not impl.supported: + continue + if impl.supports_args and not impl.supports_args(*args): + continue + return impl + +# make_fx uses torch.fx.symbolic_trace +impl_graph = make_fx(selected_impl.impl_fn) +# Replace IR op node with impl_graph's nodes +match.replace_by_example(selected_impl.impl_fn, node.args) +``` + +For example, lowering `rms_norm` with the `vllm_c` implementation: + +```python +# Before lowering (IR op) +rms_norm = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5) + +# After lowering (vllm_c implementation traced) +# Note: Lowering does not currently functionalize, this will likely change in the future. +empty = torch.ops.aten.empty.memory_format(x.shape, ...) +rms_norm = torch.ops._C.rms_norm(empty, x, weight, 1e-5) +``` + +When lowering an implementation that mutates inputs (`inplace=True`), +the lowering pass inserts clones to preserve functional semantics: + +```python +# vllm_c implementation for fused_add_rms_norm mutates its first two arguments +# Lowered with clones for safety +clone_default = torch.ops.aten.clone.default(x) +clone_default_1 = torch.ops.aten.clone.default(residual) +fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(clone_default, clone_default_1, weight, 1e-5) +``` + +### 5. Clone Cleanup + +After lowering, the clone elimination pass (`UnsafeCloneEliminationPass`) removes unnecessary clones introduced during lowering. +This pass is essential for achieving zero-copy behavior when using in-place kernels with `maybe_inplace`. +The pass removes a clone if: + +- the cloned input is created in the graph and not used again in the graph +- the cloned input is a graph parameter, marked as donated + +```python +# After cleanup (donated inputs, no subsequent uses) +fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(x, residual, weight, 1e-5) +``` + +The combination of inplace functionalization (tracking donated inputs) and clone cleanup enables the compiler to safely +use in-place kernels without adding redundant copies or increasing the memory usage. + +### 6. Inductor Optimization and Codegen + +After IR lowering and cleanup, the graph contains only standard PyTorch operations and platform-specific custom ops. +Inductor then performs its standard codegen: + +- **Inductor lowering and pointwise fusion**: Fusing element-wise operations, reductions, etc. +- **Memory planning**: Determining buffer allocation and reuse +- **Kernel generation**: Generating Triton or C++ code for fused operations +- **Autotuning**: Selecting the best kernel configurations + +### Pipeline Summary + +```text +Model Forward Pass + ↓ +[Dynamo Tracing] → FX Graph with vllm_ir.* ops + ↓ +[Pre-grad: Inplace Functionalization] → maybe_inplace → default, track donated inputs + ↓ +[AOTAutograd] → Functionalization + ↓ +[Post-grad: IR Fusion Passes] → Fuse high-level IR ops (e.g., rms_norm + quant) + ↓ +[Post-grad: IR Lowering] → vllm_ir.* ops → impl ops (with clones if needed) + ↓ +[Post-grad: Clone Cleanup] → Remove unnecessary clones using donated input info + ↓ +[Inductor] → Pattern matching, fusion, memory planning, codegen + ↓ +Compiled Code +``` + +## Core vLLM IR Concepts + +### Operation Declaration + +Operations are declared with the `@register_op` decorator, which creates an `IrOp` object: + +```python +@register_op( + name=None, # Operation name (defaults to function name) + activations=None, # List of activation parameters (defaults to params starting with 'x') + allow_inplace=False, # Whether to create a maybe_inplace overload +) +def op_name(...): + ... +``` + +**Parameters:** + +- `activations`: List of parameter names considered "activations" (typically consumed by `maybe_inplace`). Defaults to parameters starting with `x`. +- `allow_inplace`: Creates a `maybe_inplace` overload for memory-efficient execution (see below). + +### The `maybe_inplace` Overload + +The `maybe_inplace` overload is a critical feature for memory efficiency in LLM inference. +It signals that the caller doesn't need to preserve the activation inputs after the operation, +allowing in-place implementations to reuse input memory for outputs. + +#### Semantics and Usage + +```python +# Standard usage: inputs are preserved +out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon) +# x and residual are unchanged, out and res_out are new tensors + +# maybe_inplace: inputs may be modified +out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon) +# x and residual may be modified (undefined behavior to use them after this) +# out and res_out may alias x and residual +``` + +Using an activation input after passing it to `maybe_inplace` is **undefined behavior**: + +```python +# WRONG: Using x after donating it +out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon) +result = out + x # ERROR: x was donated! +``` + +If you need to preserve an input, either use the default overload or clone manually: + +```python +# Option 1: Use default overload +out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon) +result = out + x # OK: x is preserved + +# Option 2: Clone before maybe_inplace +out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x.clone(), residual, weight, epsilon) +result = out + x # OK: x is preserved, clone was donated +``` + +#### Compilation Behavior + +During compilation, the inplace functionalization pass validates that donated inputs are +not used again and converts `maybe_inplace` to the functional `default` overload: + +```python +# Inplace functionalization pass (pre-grad) +for node in graph.nodes: + if node.target == torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace: + # Check that activation inputs aren't used after this node + for activation_arg in activation_inputs: + for user in activation_arg.users: + if user appears after node: + raise ValueError(f"Input {activation_arg} donated but used again") + + # Convert to default overload + node.target = torch.ops.vllm_ir.fused_add_rms_norm.default + + # Track donated graph inputs for later clone elimination + for i, arg in enumerate(node.args): + if arg.op == "placeholder" and i in activation_indices: + pass_context.donated_input_ids.add(node_to_idx[arg]) +``` + +The donated input information is then used by the clone cleanup pass to eliminate +unnecessary copies when in-place kernels are lowered. + +#### Eager Mode Behavior + +In eager mode (without `torch.compile`), `maybe_inplace` enables **maximally memory-efficient** +execution by allowing the IR operation to dispatch directly to in-place implementations: + +```python +# Eager dispatch logic for maybe_inplace +impl: IrOpImpl = ir_op.dispatch(*args) +return impl.impl_fn(*args) + +# Eager dispatch logic for default: +impl: IrOpImpl = ir_op.dispatch(*args) +if impl.inplace: + args = [ + arg.clone() if i in ir_op.activations else arg + for i, arg in enumerate(args) + ] +return impl.impl_fn(*args) +``` + +The combination of `maybe_inplace` in model code and in-place kernel implementations provides optimal memory efficiency +in both eager and compiled modes, with identical semantics in both cases. + +#### Memory Savings Example + +Consider a transformer layer with residual connections: + +```python +# Without maybe_inplace (2 allocations per layer) +hidden_states = self.attention(input) +normed, residual = ir.ops.fused_add_rms_norm(hidden_states, input, weight, eps) +# Memory: input (preserved), hidden_states (preserved), normed (new), residual (new) + +# With maybe_inplace (0 allocations per layer when using in-place kernel) +hidden_states = self.attention(input) +normed, residual = ir.ops.fused_add_rms_norm.maybe_inplace(hidden_states, input, weight, eps) +# Memory: normed (reuses hidden_states), residual (reuses input) +``` + +### Implementation Registration + +Implementations are registered using the `register_impl` method: + +```python +@ir.ops.op_name.register_impl( + provider="provider_name", # Unique identifier (e.g., "vllm_c", "aiter", "triton") + supported=True, # Static availability check + supports_args=None, # Dynamic argument support check +) +def impl_fn(...): + ... +``` + +**Provider naming conventions:** + +- `native`: Reserved for the native torch implementation (declared with `@register_op`) +- `vllm_c`: C++/CUDA kernels via `torch.ops._C` +- `aiter`: AMD AITER library +- `xpu_kernels`: SYCL/SYCLTLA kernels implemented in `vllm-xpu-kernels` +- `triton_*`: Triton kernels +- Platform/library names for other implementations + +**Support checking:** + +- `supported`: Static boolean, checked once at import time (e.g., `HAS_TRITON`, `is_cuda_alike()`) +- `supports_args`: Function `(*args, **kwargs) -> bool` checking argument compatibility + - Called with **fake tensors** during compilation for zero-cost checking + - Called with **real tensors** during eager mode dispatch + - Should NOT check batch sizes or add guards based on values + +Example support predicate: + +```python +def aiter_rms_norm_supports(x, weight, epsilon, variance_size=None): + # Check dtype (OK: doesn't depend on batch size) + if x.dtype not in [torch.float16, torch.bfloat16]: + return False + # Check optional parameter (OK: static check) + if variance_size is not None: + return False + return True + +@ir.ops.rms_norm.register_impl("aiter", supports_args=aiter_rms_norm_supports) +def rms_norm(...): + ... +``` + +Batch-invariant kernels are automatically selected when `VLLM_BATCH_INVARIANT=1` is set. + +### Eager Mode vs Compile Mode + +vLLM IR operations behave identically in eager and compile modes: + +**Eager mode:** + +- Direct dispatch to implementation based on priority list +- Support checked with real tensor arguments +- Minimal overhead (can be optimized further if needed) + +**Compile mode:** + +- IR ops appear in FX graph as `torch.ops.vllm_ir.*` custom ops +- Lowering selects implementation using fake tensors +- Full integration with Inductor optimizations + +This consistency enables: + +- Prototyping in eager mode with confidence +- Debugging by disabling compilation +- Gradual migration from eager to compiled execution + +## Other Topics + +### Out-of-Tree Implementations + +External platforms can register implementations without modifying vLLM: + +```python +# In external package +from vllm import ir + +@ir.ops.rms_norm.register_impl("my_platform", supported=is_my_platform()) +def rms_norm(x, weight, epsilon, variance_size=None): + return my_platform.rms_norm(x, weight, epsilon) +``` + +Then configure priority to use your implementation: + +```python +class MyPlatform(Platform): + def get_default_ir_op_priority(self): + return IrOpPriorityConfig(rms_norm=['my_platform', 'native']) + +# Users can still override priority in the same way +llm = LLM(ir_op_priority=IrOpPriorityConfig(rms_norm=['custom_oot_kernel'])) +``` + +### Debugging and Observability + +!!! note + Please let us know how observability can be improved for your use-case! + +Enable debug logging to see kernel selection: + +```bash +VLLM_LOGGING_LEVEL=DEBUG vllm serve ... +``` + +This logs: + +- Which implementations are selected for each operation +- Why implementations were rejected (unsupported, args not supported) +- Compilation cache hits/misses +- IR lowering statistics + +Check selected implementations in compiled graphs: + +```python +# After compilation, inspect the lowering pass +lowering_pass = backend.lowering_pass +print(lowering_pass.selected_impls) +# Output: {'rms_norm': {'node_123': 'vllm_c', 'node_456': 'vllm_c'}} +``` + +## Migration from CustomOp + +vLLM IR is designed to coexist with and gradually replace `CustomOp`: + +1. **Op declaration**: Convert `CustomOp` class `PluggableLayer` and move `forward_native` to `@register_op` function +2. **Implementation registration**: Use `@ir.ops.op_name.register_impl` instead of overriding methods +3. **Layer usage**: Replace `self.op(...)` with `ir.ops.op_name(...)` +4. **Configuration**: Migrate `--compilation-config.custom-ops` to `--ir-op-priority` + +The migration can be done incrementally, one operation at a time. + +## See Also + +- [torch.compile Integration](torch_compile.md) - General compilation infrastructure +- [Fusions](fusions.md) - Custom fusion and transformation passes in vLLM +- [Custom Operations](custom_op.md) - Legacy custom op system diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 3266b80e5dc0..2c19dc1763f6 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -7,7 +7,7 @@ import textwrap import traceback from argparse import SUPPRESS, Action, HelpFormatter -from collections.abc import Iterable +from collections.abc import Callable, Iterable from importlib.machinery import ModuleSpec from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -48,6 +48,7 @@ def decorator(cls): mock_if_no_torch("vllm._C", MagicMock()) +mock_if_no_torch("vllm._C_stable_libtorch", MagicMock()) mock_if_no_torch( "vllm.model_executor.custom_op", MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer), @@ -67,6 +68,31 @@ def decorator(cls): mock_if_no_torch("torch.nn", MagicMock(Parameter=object)) +# Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__ +# We need to return the corresponding number of inputs, as IR infra will assert it +def get_outputs(native_fn: Callable) -> str: + """ + Extract output schema from function's return type annotation, + e.g. 'Tensor' or 'Tensor, Tensor'. + """ + import typing + + return_type = typing.get_type_hints(native_fn)["return"] + origin = typing.get_origin(return_type) + arg_name = lambda a: a.__name__ if hasattr(a, "__name__") else str(a) + if origin is tuple: + args = typing.get_args(return_type) + return ", ".join(arg_name(arg) for arg in args) + else: + return f"{arg_name(return_type)}" + + +mock_if_no_torch( + "torch.library", + MagicMock(infer_schema=lambda fn, **k: f"(Tensor x) -> {get_outputs(fn)}"), +) + + class PydanticMagicMock(MagicMock): """`MagicMock` that's able to generate pydantic-core schemas.""" diff --git a/tests/compile/backend.py b/tests/compile/backend.py index d61c128a59b6..87f98946a8ad 100644 --- a/tests/compile/backend.py +++ b/tests/compile/backend.py @@ -12,10 +12,17 @@ from torch.fx._utils import lazy_format_graph_code from vllm.compilation.passes.fx_utils import find_op_nodes -from vllm.compilation.passes.inductor_pass import InductorPass +from vllm.compilation.passes.inductor_pass import ( + InductorPass, + pass_context, +) +from vllm.compilation.passes.ir.inplace_functionalization import ( + VllmIRInplaceFunctionalizationPass, +) from vllm.compilation.passes.pass_manager import with_pattern_match_debug from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.utils import Range from vllm.logger import init_logger logger = init_logger("vllm.tests.compile.backend") @@ -53,11 +60,17 @@ def __init__(self, *passes: InductorPass | Callable[[fx.Graph], None]): self.custom_passes = list(passes) vllm_config = get_current_vllm_config() compile_config = vllm_config.compilation_config + self.range = Range(1, vllm_config.scheduler_config.max_num_batched_tokens) # Deepcopy to allow multiple TestBackend instances to use the same VllmConfig self.inductor_config = deepcopy(compile_config.inductor_compile_config) self.inductor_config["force_disable_caches"] = True self.inductor_config["post_grad_custom_post_pass"] = self.post_pass + # Add VllmIRInplaceFunctionalizationPass as pre-grad pass by default + self.inductor_config["pre_grad_custom_pass"] = ( + VllmIRInplaceFunctionalizationPass(vllm_config) + ) + if debug_dump_path := vllm_config.compile_debug_dump_path(): logger.debug("Dumping depyf output to %s", debug_dump_path) self.debug_ctx = depyf.prepare_debug(debug_dump_path.as_posix()) @@ -68,7 +81,7 @@ def __call__(self, graph: fx.GraphModule, example_inputs): self.graph_pre_compile = deepcopy(graph) from torch._inductor.compile_fx import compile_fx - with self.debug_ctx: + with self.debug_ctx, pass_context(self.range): return compile_fx( graph, example_inputs, config_patches=self.inductor_config ) diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py index 1f1eeb8b4789..c40d75f6754a 100644 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ b/tests/compile/passes/distributed/test_sequence_parallelism.py @@ -88,14 +88,10 @@ def ops_in_model_after(self): ] def ops_in_model(self): - return ( - [torch.ops.vllm_ir.rms_norm] - + [ - torch.ops._C.fused_add_rms_norm.default, - ] - if RMSNorm.enabled() - else [] - ) + return [ + torch.ops.vllm_ir.rms_norm, + torch.ops.vllm_ir.fused_add_rms_norm, + ] class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module): @@ -152,16 +148,17 @@ def ops_in_model_before(self): def ops_in_model(self): if self.vllm_config.compilation_config.pass_config.fuse_norm_quant: return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default] - elif RMSNorm.enabled(): - return [ - torch.ops._C.fused_add_rms_norm.default, - ] - elif any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers): + else: + quant_ops = ( + [torch.ops._C.static_scaled_fp8_quant.default] + if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers) + else [torch.ops.aten.reciprocal] + ) return [ - torch.ops._C.static_scaled_fp8_quant.default, + torch.ops.vllm_ir.rms_norm, + torch.ops.vllm_ir.fused_add_rms_norm, + *quant_ops, ] - else: - return [] @multi_gpu_test(num_gpus=2) diff --git a/tests/compile/passes/ir/test_clone_cleanup.py b/tests/compile/passes/ir/test_clone_cleanup.py new file mode 100644 index 000000000000..9fedb5fc9177 --- /dev/null +++ b/tests/compile/passes/ir/test_clone_cleanup.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Comprehensive tests for UnsafeCloneEliminationPass. + +This test suite exercises all possible valid FX graph patterns involving clones: +1. Clone with no users (dead code) +2. Clone with read-only users +3. Clone with mutation users +4. Clone of graph input +5. Clone with original used after mutation +6. Clone chains +""" + +import pytest +import torch +from torch import fx +from torch.fx.experimental.proxy_tensor import make_fx + +from vllm.compilation.passes.fx_utils import find_op_nodes +from vllm.compilation.passes.inductor_pass import get_pass_context, pass_context +from vllm.compilation.passes.ir.clone_elimination import ( + UnsafeCloneEliminationPass, + user_writes_to_node, +) +from vllm.config import VllmConfig +from vllm.config.utils import Range + + +def count_clones(graph: fx.Graph) -> int: + """Count clone nodes in a graph.""" + return len(list(find_op_nodes(torch.ops.aten.clone.default, graph))) + + +@pytest.fixture(scope="function") +def clone_cleanup_pass(): + return UnsafeCloneEliminationPass(VllmConfig()) + + +@pytest.fixture(autouse=True) +def setup_pass_context(): + """Set up pass context for each test.""" + with pass_context(compile_range=Range(1, 8192)): + yield + + +class TestCloneCleanup: + """Test UnsafeCloneEliminationPass behavior on various graph patterns.""" + + def test_remove_clone_readonly_users(self, clone_cleanup_pass): + """Clone with only read-only users should be removed.""" + + def f(x: torch.Tensor) -> torch.Tensor: + x_clone = x.clone() + return x_clone + 1 + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 1 + + expected = graph_module(inp) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + actual = graph_module(inp) + + assert count_clones(graph_module.graph) == 0 + torch.testing.assert_close(actual, expected) + + def test_keep_clone_with_mutation_and_original_used_after(self, clone_cleanup_pass): + """Clone must be kept if it's mutated AND original is used after mutation.""" + + def f(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x = x.relu() # not a graph param + x_clone = x.clone() + x_clone.add_(1) + return x, x_clone + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 1 + + expected = graph_module(inp) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + actual = graph_module(inp) + + # Clone should be KEPT because original is used after mutation + assert count_clones(graph_module.graph) == 1 + torch.testing.assert_close(actual[0], expected[0]) + torch.testing.assert_close(actual[1], expected[1]) + + def test_remove_clone_with_mutation_no_original_use(self, clone_cleanup_pass): + """Clone can be removed if it's mutated but original is not used after.""" + + def f(x: torch.Tensor) -> torch.Tensor: + x = x.relu() # not a graph param + x_clone = x.clone() + x_clone.add_(1) + return x_clone + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 1 + + expected = graph_module(inp) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + actual = graph_module(inp) + + assert count_clones(graph_module.graph) == 0 + torch.testing.assert_close(actual, expected) + + def test_clone_chain(self, clone_cleanup_pass): + """Test handling of clone chains: x -> clone1 -> clone2.""" + + def f(x: torch.Tensor) -> torch.Tensor: + x = x.relu() # not a graph param + x1 = x.clone() + x2 = x1.clone() + return x2 + 1 + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 2 + + expected = graph_module(inp) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + actual = graph_module(inp) + + # Both clones should be removed + assert count_clones(graph_module.graph) == 0 + torch.testing.assert_close(actual, expected) + + def test_multiple_clones_of_same_input(self, clone_cleanup_pass): + """Test multiple independent clones of the same input.""" + + def f(x: torch.Tensor) -> torch.Tensor: + x1 = x.clone() + x2 = x.clone() + return x1 + x2 + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 2 + + expected = graph_module(inp) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + actual = graph_module(inp) + + # Both clones should be removed (only readonly uses) + assert count_clones(graph_module.graph) == 0 + torch.testing.assert_close(actual, expected) + + def test_no_clones_in_graph(self, clone_cleanup_pass): + """Test pass behavior when graph has no clones.""" + + def f(x: torch.Tensor) -> torch.Tensor: + return x + 1 + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 0 + + expected = graph_module(inp) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + actual = graph_module(inp) + + assert count_clones(graph_module.graph) == 0 + torch.testing.assert_close(actual, expected) + + def test_multiple_passes(self, clone_cleanup_pass): + """Test running the pass multiple times (should be idempotent).""" + + def f(x: torch.Tensor) -> torch.Tensor: + x1 = x.clone() + return x1 + 1 + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 1 + + expected = graph_module(inp) + + clone_cleanup_pass(graph_module.graph) + assert count_clones(graph_module.graph) == 0 + graph_module.recompile() + actual = graph_module(inp) + torch.testing.assert_close(actual, expected) + + clone_cleanup_pass(graph_module.graph) + assert count_clones(graph_module.graph) == 0 + graph_module.recompile() + actual = graph_module(inp) + torch.testing.assert_close(actual, expected) + + def test_output_node_no_write(self): + """Output nodes never write to their inputs.""" + + def f(x: torch.Tensor) -> torch.Tensor: + return x + + graph_module = make_fx(f)(torch.randn(2, 3)) + x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0] + output_node = [n for n in graph_module.graph.nodes if n.op == "output"][0] + + assert not user_writes_to_node(output_node, x_node) + + def test_readonly_op_no_write(self): + """Readonly operations don't write to inputs.""" + + def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3)) + placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"] + add_node = [ + n + for n in graph_module.graph.nodes + if n.op == "call_function" and n.target == torch.ops.aten.add.Tensor + ][0] + + assert not user_writes_to_node(add_node, placeholders[0]) + assert not user_writes_to_node(add_node, placeholders[1]) + + def test_inplace_op_writes(self): + """Inplace operations write to first argument.""" + + def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + x.add_(y) + return x + + graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3)) + placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"] + add_node = [ + n + for n in graph_module.graph.nodes + if n.op == "call_function" and "add_" in str(n.target) + ][0] + + # add_ writes to first arg but not second + assert user_writes_to_node(add_node, placeholders[0]) + assert not user_writes_to_node(add_node, placeholders[1]) + + def test_copy_writes(self): + """copy_ operation writes to first argument.""" + + def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + x.copy_(y) + return x + + graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3)) + placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"] + copy_node = [ + n + for n in graph_module.graph.nodes + if n.op == "call_function" and "copy_" in str(n.target) + ][0] + + assert user_writes_to_node(copy_node, placeholders[0]) + assert not user_writes_to_node(copy_node, placeholders[1]) + + def test_auto_functionalized_not_a_write(self): + """auto_functionalized ops are follow-up uses, not writes.""" + from torch._higher_order_ops.auto_functionalize import auto_functionalized + + def f(x: torch.Tensor) -> torch.Tensor: + return x + + graph_module = make_fx(f)(torch.randn(2, 3)) + x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0] + + # Create an auto_functionalized node in the graph + with graph_module.graph.inserting_before(None): + af_node = graph_module.graph.call_function( + auto_functionalized, kwargs={"input": x_node} + ) + + # auto_functionalized should not be treated as a write + assert not user_writes_to_node(af_node, x_node) + + def test_higher_order_op_conservatively_writes(self): + """Other higher-order operators are conservatively treated as writes.""" + from torch._ops import HigherOrderOperator + + def f(x: torch.Tensor) -> torch.Tensor: + return x + + graph_module = make_fx(f)(torch.randn(2, 3)) + x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0] + + # Create a concrete higher-order operator subclass + class MockHigherOrderOp(HigherOrderOperator): + def __call__(self, *args, **kwargs): + return args[0] if args else None + + mock_hoo = MockHigherOrderOp("mock_higher_order_op") + + with graph_module.graph.inserting_before(None): + hoo_node = graph_module.graph.call_function(mock_hoo, args=(x_node,)) + + # Should be conservative and assume it could write + assert user_writes_to_node(hoo_node, x_node) + + +class TestCloneCleanupWithDonatedInputs: + """Test UnsafeCloneEliminationPass with donated input tracking via PassContext.""" + + @pytest.fixture(autouse=True) + def setup_pass_context(self): + """Set up pass context for each test.""" + with pass_context(compile_range=Range(1, 8192)): + yield + + def test_donated_input_clone_removed(self, clone_cleanup_pass): + """Clone of donated input should be removed.""" + + def f(x: torch.Tensor) -> torch.Tensor: + x_clone = x.clone() + x_clone.add_(1) + return x_clone + + inp = torch.randn(2, 3) + graph_module = make_fx(f)(inp) + assert count_clones(graph_module.graph) == 1 + + # Mark first parameter as donated + get_pass_context().donated_input_ids = {0} + + expected = graph_module(inp.clone()) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + + # Clone should be removed since input is donated + assert count_clones(graph_module.graph) == 0 + + # Input can be mutated (donated) + inp_copy = inp.clone() + actual = graph_module(inp_copy) + torch.testing.assert_close(actual, expected) + + def test_non_donated_input_clone_kept(self, clone_cleanup_pass): + """Clone of non-donated input with mutation should be kept.""" + + def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x_clone = x.clone() + x_clone.add_(1) + return x, x_clone + + inp_x = torch.randn(2, 3) + inp_y = torch.randn(2, 3) + graph_module = make_fx(f)(inp_x, inp_y) + assert count_clones(graph_module.graph) == 1 + + # No donated inputs + get_pass_context().donated_input_ids = set() + + expected = graph_module(inp_x.clone(), inp_y.clone()) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + + # Clone should be kept since input is not donated and original is used + assert count_clones(graph_module.graph) == 1 + + # Verify inputs are not mutated + inp_x_before = inp_x.clone() + inp_y_before = inp_y.clone() + actual = graph_module(inp_x, inp_y) + torch.testing.assert_close( + inp_x, inp_x_before, msg="Input x should not be mutated" + ) + torch.testing.assert_close( + inp_y, inp_y_before, msg="Input y should not be mutated" + ) + torch.testing.assert_close(actual[0], expected[0]) + torch.testing.assert_close(actual[1], expected[1]) + + def test_mixed_donated_inputs(self, clone_cleanup_pass): + """Test with some inputs donated and some not.""" + + def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x_clone = x.clone() + x_clone.add_(1) + y_clone = y.clone() + y_clone.add_(2) + return x_clone, y_clone + + inp_x = torch.randn(2, 3) + inp_y = torch.randn(2, 3) + graph_module = make_fx(f)(inp_x, inp_y) + assert count_clones(graph_module.graph) == 2 + + # Only x is donated + get_pass_context().donated_input_ids = {0} + + expected = graph_module(inp_x.clone(), inp_y.clone()) + clone_cleanup_pass(graph_module.graph) + graph_module.recompile() + + # x_clone removed (x is donated), y_clone kept (y is not donated) + assert count_clones(graph_module.graph) == 1 + + # Verify y is not mutated (x can be mutated since it's donated) + inp_y_before = inp_y.clone() + actual = graph_module(inp_x.clone(), inp_y) + torch.testing.assert_close( + inp_y, inp_y_before, msg="Input y should not be mutated" + ) + torch.testing.assert_close(actual[0], expected[0]) + torch.testing.assert_close(actual[1], expected[1]) diff --git a/tests/compile/passes/ir/test_inplace_functionalization.py b/tests/compile/passes/ir/test_inplace_functionalization.py new file mode 100644 index 000000000000..1e8d5662162f --- /dev/null +++ b/tests/compile/passes/ir/test_inplace_functionalization.py @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for IR inplace functionalization pass integration. + +This test suite verifies that the inplace functionalization pass, lowering pass, +and clone cleanup pass work together correctly with donated buffer tracking. +""" + +from collections.abc import Callable + +import pytest +import torch +import torch._dynamo.exc +from torch import nn + +import vllm.kernels # noqa: F401 to register kernels +from vllm.compilation.passes.inductor_pass import InductorPass, get_pass_context +from vllm.compilation.passes.ir.clone_elimination import ( + UnsafeCloneEliminationPass, +) +from vllm.compilation.passes.ir.inplace_functionalization import ( + VllmIRInplaceFunctionalizationPass, +) +from vllm.compilation.passes.ir.lowering_pass import VllmIRLoweringPass +from vllm.config import VllmConfig +from vllm.ir import ops +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +from ...backend import TestBackend + + +class StoreDonationInfoPass(InductorPass): + def __init__(self): + self.donated_input_ids_sets: list[set[int]] = [] + + def __call__(self, *args, **kwargs): + ctx = get_pass_context() + self.donated_input_ids_sets += [ctx.donated_input_ids] + + +class MaybeInplaceModel(nn.Module): + """Model using only maybe_inplace variants.""" + + def __init__(self, hidden_size=16): + super().__init__() + self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + + def forward( + self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor + ): + # First maybe_inplace - x & residual1 are donated + x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace( + x, residual1, self.weight1, 1e-5 + ) + # Second maybe_inplace - residual2 is donated + x_normed2, residual_out2 = ops.fused_add_rms_norm.maybe_inplace( + x_normed1, residual2, self.weight2, 1e-5 + ) + return x_normed2, residual_out1, residual_out2 + + +class FunctionalModel(nn.Module): + """Model using only functional (default) variants.""" + + def __init__(self, hidden_size=16): + super().__init__() + self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + + def forward( + self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor + ): + # First functional - no donation + x_normed1, residual_out1 = ops.fused_add_rms_norm( + x, residual1, self.weight1, 1e-5 + ) + # Second functional - no donation + x_normed2, residual_out2 = ops.fused_add_rms_norm( + x_normed1, residual2, self.weight2, 1e-5 + ) + return x_normed2, residual_out1, residual_out2 + + +class MixedModel(nn.Module): + """Model mixing maybe_inplace and functional variants.""" + + def __init__(self, hidden_size=16): + super().__init__() + self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + + def forward( + self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor + ): + # First maybe_inplace - x & residual1 are donated + x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace( + x, residual1, self.weight1, 1e-5 + ) + # Second functional - no donation, x_normed1 must be preserved as it's returned + x_normed2, residual_out2 = ops.fused_add_rms_norm( + x_normed1, residual2, self.weight2, 1e-5 + ) + # Return both to prevent x_normed1 from being optimized away + return x_normed1, x_normed2, residual_out1, residual_out2 + + +class ModelWithTritonAfterMaybeInplace(nn.Module): + """ + Model using maybe_inplace followed by a Triton kernel. + Test clone elimination can handle Triton in the graph + """ + + def __init__(self, hidden_size=16): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + + @triton.jit + def _triton_add_kernel( + x_ptr, + y_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = x + 0.1 + tl.store(y_ptr + offsets, y, mask=mask) + + def triton_add(x: torch.Tensor) -> torch.Tensor: + """Simple Triton add kernel.""" + y = torch.empty_like(x) + n_elements = x.numel() + grid = (triton.cdiv(n_elements, 256),) + _triton_add_kernel[grid](x, y, n_elements, BLOCK_SIZE=256) + return y + + self.triton_add = triton_add + + def forward(self, x: torch.Tensor, residual: torch.Tensor, residual2: torch.Tensor): + x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace( + x, residual, self.weight, 1e-5 + ) + + x_processed = self.triton_add(x_normed) + + # x_processed does not need to be cloned, residual2 does + x_normed2, residual_out2 = ops.fused_add_rms_norm( + x_processed, residual2, self.weight, 1e-5 + ) + return x_normed2, residual_out2 + + +skipif_no_triton = pytest.mark.skipif(not HAS_TRITON, reason="Requires Triton") + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Only test on cuda and rocm platform", +) +@pytest.mark.parametrize( + "model_class,expected_functionalized,expected_donated,expected_clones", + [ + # 2 inplace calls, all activations donated, all clones eliminated + (MaybeInplaceModel, 2, 3, 0), + # No inplace calls, no donations, 3 clones (one eliminated) + (FunctionalModel, 0, 0, 3), + # One inplace call, two donated activations, 2 clones + (MixedModel, 1, 2, 2), + # One inplace call, two donated, 1 clone remaining + pytest.param(ModelWithTritonAfterMaybeInplace, 1, 2, 1, marks=skipif_no_triton), + ], +) +def test_inplace_functionalization( + default_vllm_config: VllmConfig, + model_class, + expected_functionalized: int, + expected_clones: int, + expected_donated: int, +): + """Test inplace functionalization, lowering, and clone cleanup.""" + torch.set_default_device(current_platform.device_type) + + # Use vllm_c so inplace path is triggered + default_vllm_config.kernel_config.ir_op_priority.fused_add_rms_norm = [ + "vllm_c", + "native", + ] + + # Create passes in order they run during compilation + functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config) + lowering_pass = VllmIRLoweringPass(default_vllm_config) + donated_info_pass = StoreDonationInfoPass() + cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config) + + # Set up backend with pre-grad pass + backend = TestBackend(lowering_pass, donated_info_pass, cleanup_pass) + backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass + + model = model_class() + x = torch.randn(8, 16, dtype=torch.bfloat16) + residual1 = torch.randn(8, 16, dtype=torch.bfloat16) + residual2 = torch.randn(8, 16, dtype=torch.bfloat16) + + with default_vllm_config.kernel_config.ir_op_priority.set_priority(): + # Reference output without optimization + ref_output = model(x.clone(), residual1.clone(), residual2.clone()) + + # Compile with inplace optimization + compiled_model = torch.compile(model, backend=backend, fullgraph=True) + output = compiled_model(x.clone(), residual1.clone(), residual2.clone()) + + # Verify correctness (relaxed tolerance for bfloat16) + for i in range(len(ref_output)): + torch.testing.assert_close(output[i], ref_output[i], rtol=1e-2, atol=1e-2) + + # Verify expected number of ops were functionalized + func_ops = functionalization_pass.functionalized_ops + assert len(func_ops) == int(bool(expected_functionalized)) + if expected_functionalized > 0: + assert "fused_add_rms_norm" in func_ops + assert func_ops["fused_add_rms_norm"] == expected_functionalized + + # Verify lowering happened (2 ops in all cases) + assert "fused_add_rms_norm" in lowering_pass.selected_impls + assert len(lowering_pass.selected_impls["fused_add_rms_norm"]) == 2 + assert all( + provider == "vllm_c" + for node, provider in lowering_pass.selected_impls["fused_add_rms_norm"].items() + ), lowering_pass.selected_impls + + # Verify correct number of donated IDs + assert len(donated_info_pass.donated_input_ids_sets) == 1 + assert len(donated_info_pass.donated_input_ids_sets[0]) == expected_donated + + # Verify expected number of clones after cleanup + actual_clones = backend.op_count(torch.ops.aten.clone.default, before=False) + assert actual_clones == expected_clones, ( + f"Expected {expected_clones} clones, got {actual_clones}:" + f"{backend.print_graphs()}" + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Only test on cuda and rocm platform", +) +def test_donated_buffer_context_propagation(default_vllm_config): + """Test that donated_input_ids propagates correctly through pass_context.""" + torch.set_default_device(current_platform.device_type) + + # Create a custom backend that inspects pass_context in cleanup pass + functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config) + lowering_pass = VllmIRLoweringPass(default_vllm_config) + + donation_info_pass = StoreDonationInfoPass() + cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config) + + backend = TestBackend(lowering_pass, donation_info_pass, cleanup_pass) + backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass + + model = MaybeInplaceModel() + x = torch.randn(8, 16, dtype=torch.bfloat16) + residual1 = torch.randn(8, 16, dtype=torch.bfloat16) + residual2 = torch.randn(8, 16, dtype=torch.bfloat16) + + compiled_model = torch.compile(model, backend=backend, fullgraph=True) + compiled_model(x.clone(), residual1.clone(), residual2.clone()) + + donated_ids_seen = donation_info_pass.donated_input_ids_sets + # Verify donated_input_ids was set and propagated + assert len(donated_ids_seen) == 1 + # Should have donated inputs (exact indices depend on AOTAutograd) + assert len(donated_ids_seen[0]) == 3 + # All donated ids should be valid non-negative integers + for idx in donated_ids_seen[0]: + assert isinstance(idx, int) and idx >= 0, f"Invalid donated index: {idx}" + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Only test on cuda and rocm platform", +) +def test_maybe_inplace_reuse_error(default_vllm_config): + """Test that reusing a donated activation input raises ValueError.""" + torch.set_default_device(current_platform.device_type) + + class ReuseModel(nn.Module): + """Model that incorrectly reuses a donated activation input.""" + + def __init__(self, hidden_size=16): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + + def forward(self, x: torch.Tensor, residual: torch.Tensor): + # x is donated to maybe_inplace + x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace( + x, residual, self.weight, 1e-5 + ) + # ERROR: x is used again after being donated + return x_normed + x # This should raise ValueError + + functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config) + lowering_pass = VllmIRLoweringPass(default_vllm_config) + cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config) + + backend = TestBackend(lowering_pass, cleanup_pass) + backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass + + model = ReuseModel() + x = torch.randn(8, 16, dtype=torch.bfloat16) + residual = torch.randn(8, 16, dtype=torch.bfloat16) + + # Compilation should raise BackendCompilerFailed wrapping ValueError + with pytest.raises( + torch._dynamo.exc.BackendCompilerFailed, + match="is used again after the node", + ): + compiled_model = torch.compile(model, backend=backend, fullgraph=True) + compiled_model(x.clone(), residual.clone()) + + +# Piecewise compilation tests with graph splitting + + +@torch.library.custom_op("vllm::test_split_marker", mutates_args=()) +def test_split_marker(x: torch.Tensor) -> torch.Tensor: + """Identity op that marks a split point for piecewise compilation.""" + return x.clone() + + +@test_split_marker.register_fake +def _fake_split_marker(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + +class TransformerBlockWithSplits(nn.Module): + """Transformer block with explicit split points for piecewise compilation.""" + + def __init__(self, hidden_size=32, intermediate_size=128): + super().__init__() + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + + # Attention-like projection + self.attn_proj = nn.Linear( + hidden_size, hidden_size, bias=False, dtype=torch.bfloat16 + ) + + # Post-attention norm + self.post_attn_norm = nn.Parameter( + torch.ones(hidden_size, dtype=torch.bfloat16) + ) + + # MLP + self.gate_proj = nn.Linear( + hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16 + ) + self.up_proj = nn.Linear( + hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16 + ) + self.down_proj = nn.Linear( + intermediate_size, hidden_size, bias=False, dtype=torch.bfloat16 + ) + + # Post-MLP norm + self.post_mlp_norm = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16)) + + def forward(self, x: torch.Tensor): + # Attention block with residual + residual1 = x + attn_out = self.attn_proj(x) + + # Fused add + norm (maybe_inplace: residual1 is donated) + normed1, residual1 = ops.fused_add_rms_norm.maybe_inplace( + attn_out, residual1, self.post_attn_norm, 1e-5 + ) + + # Force a graph split here + normed1 = torch.ops.vllm.test_split_marker(normed1) + + # MLP block + gate = self.gate_proj(normed1) + up = self.up_proj(normed1) + mlp_out = self.down_proj(gate * torch.nn.functional.silu(up)) + + # Fused add + norm (maybe_inplace: residual1 is donated) + normed2, residual2 = ops.fused_add_rms_norm.maybe_inplace( + mlp_out, residual1, self.post_mlp_norm, 1e-5 + ) + + return normed2, residual2 + + +def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int): + def inner(*args): + torch._dynamo.mark_dynamic(args[arg_index], dim_index) + return fn(*args) + + return inner + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Only test on cuda and rocm platform", +) +def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache): + """ + Test piecewise compilation with donated buffers across graph splits. + Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching. + """ + torch.set_default_device(current_platform.device_type) + + # Disable compilation cache to avoid serialization issues + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + + from vllm.compilation.backends import VllmBackend + from vllm.config import CompilationConfig, VllmConfig + + # Create config with custom splitting op + store_donation_info = StoreDonationInfoPass() + vllm_config = VllmConfig( + compilation_config=CompilationConfig( + custom_ops=["all"], + splitting_ops=["vllm::test_split_marker"], + inductor_compile_config={"post_grad_custom_post_pass": store_donation_info}, + ) + ) + + backend = VllmBackend(vllm_config) + + model = TransformerBlockWithSplits() + x = torch.randn(8, 32, dtype=torch.bfloat16) + + # Reference output + ref_output = with_dyn_arg(model, 0, 0)(x.clone()) + + # Compile with piecewise compilation (graph will split at split_marker) + compiled_model = torch.compile(model, backend=backend, fullgraph=False) + output = with_dyn_arg(compiled_model, 0, 0)(x.clone()) + + # Verify correctness (relaxed tolerance for bfloat16) + torch.testing.assert_close(output[0], ref_output[0], rtol=1e-2, atol=1e-2) + torch.testing.assert_close(output[1], ref_output[1], rtol=1e-2, atol=1e-2) + + # Verify the model was split into multiple submodules + assert hasattr(backend, "split_gm"), "Backend should have split graph module" + + # Should have at least 2 submodules (split by test_split_marker op) + submodules = list(backend.split_gm.named_children()) + num_submodules = len(submodules) + assert num_submodules >= 2, ( + f"Expected at least 2 submodules (split), got {num_submodules}" + ) + + # Check that donation info was propagated correctly + donated_inputs_sets = store_donation_info.donated_input_ids_sets + assert len(donated_inputs_sets) == 2 + assert len(donated_inputs_sets[0]) == 1 + assert len(donated_inputs_sets[1]) == 1 diff --git a/tests/compile/passes/test_functionalization.py b/tests/compile/passes/test_functionalization.py index 9a03a6988763..31bf225d4135 100644 --- a/tests/compile/passes/test_functionalization.py +++ b/tests/compile/passes/test_functionalization.py @@ -126,7 +126,7 @@ def ops_in_model(self, do_fusion): if TEST_FP8 and do_fusion: return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default] else: - return [torch.ops._C.fused_add_rms_norm.default] + return [] def ops_not_in_model(self): return [] diff --git a/tests/compile/passes/test_fuse_act_padding.py b/tests/compile/passes/test_fuse_act_padding.py index f3f3bda47277..bfbe762abdb6 100644 --- a/tests/compile/passes/test_fuse_act_padding.py +++ b/tests/compile/passes/test_fuse_act_padding.py @@ -59,7 +59,7 @@ def forward(self, x): def ops_in_model_before(self): return [ - rocm_aiter_ops.get_rmsnorm_fused_add_op(), + torch.ops.vllm_ir.fused_add_rms_norm, torch.ops.aten.constant_pad_nd, ] diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 32803aad8c1c..2feb0bc4f787 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -17,7 +17,6 @@ FusedRMSQuantKey, RMSNormQuantFusionPass, ) -from vllm.compilation.passes.fx_utils import find_op_nodes from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass from vllm.config import ( @@ -243,9 +242,10 @@ def ops_in_model_after(self): ] def ops_in_model_before_partial(self): - return [torch.ops.vllm_ir.rms_norm] + ( - [RMS_ADD_OP] if self.enable_rms_norm_custom_op else [torch.ops.aten.rsqrt] - ) + return [ + torch.ops.vllm_ir.rms_norm, + torch.ops.vllm_ir.fused_add_rms_norm.default, + ] def _run_fusion_test( @@ -383,17 +383,6 @@ def test_fusion_rmsnorm_quant( model.ops_in_model_before_partial(), fully_replaced=False ) - # If RMSNorm custom op is disabled (native/torch impl used), - # there's a risk that the fused add doesn't get included in the - # replacement and only the rms part gets fused with quant. - # Hence, we check only 2 add nodes are left (final fused rmsnorm add). - if not enable_rms_norm_custom_op: - n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g)) - # rms_norm is IR, not included - # 6 = 3x2 (3xRMS_ADD, 2 each) - assert n_add_nodes(backend.graph_pre_pass) == 6 - assert n_add_nodes(backend.graph_post_pass) == 2 - @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("hidden_size", [256]) diff --git a/tests/ir/test_inplace_op.py b/tests/ir/test_inplace_op.py new file mode 100644 index 000000000000..decc4f51c777 --- /dev/null +++ b/tests/ir/test_inplace_op.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +from torch import Tensor +from torch.fx.experimental.proxy_tensor import make_fx + +import vllm.ir.op +from vllm.ir.op import IrOp, IrOpInplaceOverload + + +@vllm.ir.register_op(allow_inplace=True) +def _custom_mm2(x: Tensor, w: Tensor) -> Tensor: + return x @ w + + +@_custom_mm2.register_impl("regular") +def _custom_mm2_regular(x: Tensor, w: Tensor) -> Tensor: + return x @ w + 1 + + +@_custom_mm2.register_impl("inplace", inplace=True) +def _custom_mm2_inplace(x: Tensor, w: Tensor) -> Tensor: + x.copy_(x @ w + 2) + return x + + +class TestInplaceOp: + def test_registration(self): + # Test that the inplace op is registered correctly. + assert "_custom_mm2" in IrOp.registry + assert IrOp.registry["_custom_mm2"] is _custom_mm2 + assert _custom_mm2.torch_op is torch.ops.vllm_ir._custom_mm2.default + assert isinstance(_custom_mm2.maybe_inplace, IrOpInplaceOverload) + assert ( + _custom_mm2.maybe_inplace.torch_op + is torch.ops.vllm_ir._custom_mm2.maybe_inplace + ) + + def test_inplace_dispatching(self): + # check that the correct implementation is dispatched based on priority, + # and inplace semantics hold + w = torch.randn(3, 3) + x = torch.randn(2, 3) + x1 = x.clone() + + with _custom_mm2.set_priority(["regular"]): + result_regular = _custom_mm2.maybe_inplace(x, w) + + # check that the regular op does not modify x + torch.testing.assert_close(x, x1, atol=0, rtol=0) + + with _custom_mm2.set_priority(["inplace"]): + result_inplace: Tensor = _custom_mm2.maybe_inplace(x, w) + + # check that the inplace op returns x directly + assert result_inplace.data_ptr() == x.data_ptr() + + torch.testing.assert_close(result_inplace, x1 @ w + 2) + torch.testing.assert_close(result_regular, x1 @ w + 1) + + def test_default_dispatching(self): + # check that the correct implementation is dispatched, + # and ops do not modify inputs when using the default overload + w = torch.randn(3, 3) + x = torch.randn(2, 3) + x1 = x.clone() + + with _custom_mm2.set_priority(["regular"]): + result_regular = _custom_mm2(x, w) + + with _custom_mm2.set_priority(["inplace"]): + result_inplace = _custom_mm2(x, w) + + # check that x was not modified by either impl + torch.testing.assert_close(x, x1, atol=0, rtol=0) + + torch.testing.assert_close(result_inplace, x1 @ w + 2) + torch.testing.assert_close(result_regular, x1 @ w + 1) + + def test_trace(self): + # Test that the inplace op can be used in a graph. + def func(x: Tensor, y: Tensor) -> Tensor: + return _custom_mm2.maybe_inplace(x, y) + + x = torch.randn(2, 3) + y = torch.randn(3, 4) + graph = make_fx(func)(x, y) + assert any( + node.target == torch.ops.vllm_ir._custom_mm2.maybe_inplace + for node in graph.graph.nodes + ) diff --git a/tests/ir/test_op.py b/tests/ir/test_op.py index 524497916b6c..3576e5aef8bd 100644 --- a/tests/ir/test_op.py +++ b/tests/ir/test_op.py @@ -21,7 +21,7 @@ class CustomError(Exception): pass -@vllm.ir.register_op +@vllm.ir.register_op(allow_inplace=True) def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y @@ -129,11 +129,15 @@ def test_schema_contains_tensor_signature(self): @pytest.mark.parametrize("enable_torch_wrap", [True, False]) @pytest.mark.parametrize("symbolic_trace", [True, False]) + @pytest.mark.parametrize("overload", ["default", "maybe_inplace"]) def test_trace_sees_single_custom_op( - self, symbolic_trace: bool, enable_torch_wrap: bool + self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str ): + op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace + torch_op = getattr(torch.ops.vllm_ir._custom_add, overload) + def fn(x, y): - return _custom_add(x, y) + return op_fn(x, y) def find_fn(target: Any, gm: fx.GraphModule): return gm.graph.find_nodes(op="call_function", target=target) @@ -155,7 +159,7 @@ def find_fn(target: Any, gm: fx.GraphModule): torch.testing.assert_close(out_fx, out_eager) # check that IR nodes only appear if enable_torch_wrap=True - ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm) + ir_nodes = find_fn(torch_op, gm) if enable_torch_wrap: assert len(ir_nodes) == 1, gm.code else: @@ -167,7 +171,7 @@ def find_fn(target: Any, gm: fx.GraphModule): else: gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2)) - ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm) + ir_nodes = find_fn(torch_op, gm) assert len(ir_nodes) == 1, gm.code @@ -176,9 +180,12 @@ def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y + 10 -@_custom_add.register_impl("impl_b") +@_custom_add.register_impl("impl_b", inplace=True) def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return x + y + 20 + """Computes x+y+20""" + x.add_(y) + x.add_(20) + return x @_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0) @@ -243,19 +250,23 @@ def test_set_priority_scoped(self): # Restored to empty assert _custom_add.get_priority() == [] - def test_dispatch_priority_order(self): + @pytest.mark.parametrize("overload", ["default", "maybe_inplace"]) + def test_dispatch_priority_order(self, overload: str): + op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace + torch_op = getattr(torch.ops.vllm_ir._custom_add, overload) + x = torch.tensor(1, dtype=torch.int32) y = torch.tensor(2, dtype=torch.int32) with _custom_add.set_priority(["impl_b", "impl_a"]): assert _custom_add.dispatch(x, y) is impl_b - out1 = _custom_add(x, y) - out2 = torch.ops.vllm_ir._custom_add(x, y) + out1 = op_fn(x.clone(), y) + out2 = torch_op(x.clone(), y) with _custom_add.set_priority(["impl_a"]): assert _custom_add.dispatch(x, y) is impl_a - out3 = _custom_add(x, y) - out4 = torch.ops.vllm_ir._custom_add(x, y) + out3 = op_fn(x.clone(), y) + out4 = torch_op(x.clone(), y) # impl_b assert out1.item() == 1 + 2 + 20 @@ -265,18 +276,18 @@ def test_dispatch_priority_order(self): assert out4.item() == 1 + 2 + 10 def test_unsupported_impl_filtered(self): - @_custom_add.register_impl("unsupported", supported=False) - def impl_bad(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + @_custom_add.register_impl("impl_unsupported", supported=False) + def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y + 999 x = torch.tensor(1, dtype=torch.int32) y = torch.tensor(2, dtype=torch.int32) - with _custom_add.set_priority(["unsupported", "impl_a"]): + with _custom_add.set_priority(["impl_unsupported", "impl_a"]): assert _custom_add.get_priority() == ["impl_a"] out = _custom_add(x, y) - # impl_bad skipped → impl_a + # impl_unsupported skipped → impl_a assert out.item() == 1 + 2 + 10 def test_supports_args_runtime_dispatch_and_warning( diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index 7510ae5010fa..e9661f5202f9 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -28,7 +28,9 @@ def test_rms_norm_registration(): "native": True, "vllm_c": current_platform.is_cuda_alike(), "aiter": current_platform.is_rocm(), - "oink": False, + "oink": current_platform.has_device_capability(100) + and hasattr(torch.ops, "oink") + and hasattr(torch.ops.oink, "rmsnorm"), "xpu_kernels": current_platform.is_xpu(), } @@ -67,6 +69,14 @@ def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon) torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3) + # Mean square should be approximately 1 (ignoring epsilon and weight scaling) + combined_norm = out.float() / weight.float() + variance = combined_norm.pow(2).mean(dim=-1) + # After RMS normalization, variance should be close to 1 + torch.testing.assert_close( + variance, torch.ones_like(variance), rtol=1e-2, atol=1e-2 + ) + # Check behavior with and without weight weight1 = torch.ones_like(weight) out3 = rms_norm_native(x, weight1, epsilon=epsilon) @@ -129,3 +139,197 @@ def test_aiter_rejects_unsupported_dtypes(): num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5 ) assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}" + + +fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike() and not current_platform.is_xpu(), + reason="Currently only kernels on CUDA, ROCm and XPU", +) +def test_fused_add_rms_norm_registration(): + expected = { + "native": True, + "vllm_c": current_platform.is_cuda_alike(), + "aiter": current_platform.is_rocm(), + "oink": current_platform.has_device_capability(100) + and hasattr(torch.ops, "oink") + and hasattr(torch.ops.oink, "fused_add_rms_norm"), + "xpu_kernels": current_platform.is_xpu(), + } + + actual = { + provider: impl.supported + for provider, impl in ir.ops.fused_add_rms_norm.impls.items() + } + + assert actual == expected + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("n_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES) +@pytest.mark.parametrize("epsilon", [1e-6, 1e-5]) +@pytest.mark.skipif( + not current_platform.is_cuda_alike() and not current_platform.is_xpu(), + reason="Currently only kernels on CUDA, ROCm and XPU", +) +class TestFusedAddRMSNorm: + @classmethod + def setup_class(cls, **kwargs): + torch.set_default_device(current_platform.device_type) + + def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon): + x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( + num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon + ) + out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps) + + # Check shape, dtype, device + assert out.shape == x.shape + assert out.dtype == x.dtype + assert out.device == x.device + assert residual_out.shape == x_residual.shape + assert residual_out.dtype == x_residual.dtype + assert residual_out.device == x_residual.device + + # Check that residual_out = x + x_residual + expected_residual = (x.float() + x_residual.float()).to(dtype) + torch.testing.assert_close( + residual_out, expected_residual, rtol=1e-3, atol=1e-3 + ) + + # Verify that the output is RMS normalized version of (x + x_residual) + expected_out = rms_norm_native(expected_residual, weight, epsilon) + assert_close( + ir.ops.fused_add_rms_norm, + (out, residual_out), + (expected_out, expected_residual), + ) + + # Check the scaling property of rms norm + out1, _ = fused_add_rms_norm_native( + x, torch.zeros_like(x), weight, epsilon=epsilon + ) + out2, _ = fused_add_rms_norm_native( + x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon + ) + torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3) + + # Check behavior with and without weight + weight1 = torch.ones_like(weight) + out3, _ = fused_add_rms_norm_native(x, x_residual, weight1, eps) + out4, _ = fused_add_rms_norm_native(x, x_residual, None, eps) + torch.testing.assert_close(out3, out4) + + @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm)) + def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider): + impl = ir.ops.fused_add_rms_norm.impls[provider] + x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( + num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + ) + args = (x, x_residual, weight, eps, None) + + if not impl.supports_args(*args): + pytest.skip(f"{provider} does not support args") + + ref_output, ref_residual = fused_add_rms_norm_native(*clone_args(args)) + output, residual = impl.impl_fn(*clone_args(args)) + assert_close(ir.ops.fused_add_rms_norm, output, ref_output) + assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual) + + # check that dispatched call matches direct call + with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]): + out_dispatched, residual_dispatched = ir.ops.fused_add_rms_norm(*args[:4]) + out_direct, residual_direct = impl.impl_fn(*clone_args(args)) + torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0) + torch.testing.assert_close( + residual_dispatched, residual_direct, rtol=0.0, atol=0.0 + ) + + # none of these support variance_size override + assert not impl.supports_args(x, x_residual, weight, epsilon, 4) + assert not impl.supports_args(x, x_residual, weight, epsilon, variance_size=4) + + # test weight=None behavior + out_no_weight, residual_no_weight = impl.impl_fn( + x.clone(), x_residual.clone(), None, epsilon + ) + out_unit_weight, residual_unit_weight = impl.impl_fn( + x.clone(), x_residual.clone(), torch.ones_like(weight), epsilon + ) + assert_close(ir.ops.fused_add_rms_norm, out_no_weight, out_unit_weight) + assert_close( + ir.ops.fused_add_rms_norm, residual_no_weight, residual_unit_weight + ) + + @pytest.mark.parametrize("provider", ["vllm_c"]) + def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider): + """Test that inplace implementations reuse inputs, + for maybe_inplace overload but not for default overload.""" + impl = ir.ops.fused_add_rms_norm.impls[provider] + if not impl.supported: + pytest.skip(f"{provider} impl not supported on this platform") + + x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs( + num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + ) + + # Test default overload - should NOT modify inputs even with inplace impl + x_default = x.clone() + x_residual_default = x_residual.clone() + x_default_ptr = x_default.data_ptr() + x_residual_default_ptr = x_residual_default.data_ptr() + + with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]): + out_default, residual_default = ir.ops.fused_add_rms_norm( + x_default, x_residual_default, weight, eps + ) + + # Default should NOT be inplace (even with inplace implementation) + assert out_default.data_ptr() != x_default_ptr + assert residual_default.data_ptr() != x_residual_default_ptr + torch.testing.assert_close(x, x_default, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_residual, x_residual_default, rtol=0.0, atol=0.0) + + # Test maybe_inplace overload - should modify inputs with inplace impl + x_inplace = x.clone() + x_residual_inplace = x_residual.clone() + x_inplace_ptr = x_inplace.data_ptr() + x_residual_inplace_ptr = x_residual_inplace.data_ptr() + + with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]): + out_inplace, residual_inplace = ir.ops.fused_add_rms_norm.maybe_inplace( + x_inplace, x_residual_inplace, weight, eps + ) + + # maybe_inplace should be inplace + assert out_inplace.data_ptr() == x_inplace_ptr + assert residual_inplace.data_ptr() == x_residual_inplace_ptr + + # Both should produce same results + torch.testing.assert_close(out_default, out_inplace, atol=0.0, rtol=0.0) + torch.testing.assert_close( + residual_default, residual_inplace, atol=0.0, rtol=0.0 + ) + + @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm)) + def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider): + args = ir.ops.fused_add_rms_norm.generate_inputs( + num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon + ) + args = args + (None,) # Add variance_size parameter + + # When checking the torch op, we have to set priority and use dispatch + with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]): + torch.library.opcheck(torch.ops.vllm_ir.fused_add_rms_norm.default, args) + + # Only test maybe_inplace with non-inplace implementations + # Inplace implementations return aliases of inputs which is not allowed. + # We break this invariant, but we also convert maybe_inplace to the default + # overload during compilation, so maybe_inplace never reaches Inductor. + if not ir.ops.fused_add_rms_norm.impls[provider].inplace: + torch.library.opcheck( + torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace, args + ) diff --git a/tests/model_executor/test_enabled_custom_ops.py b/tests/model_executor/test_enabled_custom_ops.py index fc4f6f6b63f9..490284f43954 100644 --- a/tests/model_executor/test_enabled_custom_ops.py +++ b/tests/model_executor/test_enabled_custom_ops.py @@ -23,11 +23,7 @@ vllm_topk_sigmoid, vllm_topk_softmax, ) -from vllm.model_executor.layers.layernorm import ( - RMSNorm, - dispatch_rocm_rmsnorm_func, - fused_add_rms_norm, -) +from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] @@ -153,26 +149,3 @@ def test_topk_sigmoid_dispatch(use_rocm_aiter: bool): assert topk_func == rocm_aiter_ops.topk_sigmoid else: assert topk_func == vllm_topk_sigmoid - - -@pytest.mark.parametrize("add_residual", [False]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("use_rocm_aiter", [True, False]) -@pytest.mark.skipif( - not current_platform.is_rocm(), reason="AITER is a feature exclusive for ROCm" -) -def test_rms_norm_dispatch( - add_residual: bool, dtype: torch.dtype, use_rocm_aiter: bool -): - rms_norm_func = dispatch_rocm_rmsnorm_func(dtype, use_rocm_aiter) - - should_use_rocm_aiter = ( - current_platform.is_rocm() - and use_rocm_aiter - and dtype in RMS_NORM_SUPPORTED_DTYPES - ) - - if should_use_rocm_aiter: - assert rms_norm_func == rocm_aiter_ops.rms_norm2d_with_add - else: - assert rms_norm_func == fused_add_rms_norm diff --git a/tests/model_executor/test_oink_integration.py b/tests/model_executor/test_oink_integration.py index d7f38fdd5158..2f37472b73ef 100644 --- a/tests/model_executor/test_oink_integration.py +++ b/tests/model_executor/test_oink_integration.py @@ -1,60 +1,97 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - +import multiprocessing import types import pytest -import torch - - -def _load_oink_ops_module(): - # Import the module normally (vllm is installed as an editable package in CI). - from vllm import _oink_ops - - return _oink_ops - -def test_oink_availability_checks(monkeypatch: pytest.MonkeyPatch): - _oink_ops = _load_oink_ops_module() - - # Ensure the ops namespace exists and is mutable for tests. - monkeypatch.setattr( - torch.ops, - "oink", - types.SimpleNamespace(rmsnorm=lambda x, w, eps: x), - raising=False, - ) - - # Case 1: CUDA not available. - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - assert _oink_ops.is_oink_available_for_device(0) is False - - # Case 2: CUDA available but < SM100. - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (9, 0)) - assert _oink_ops.is_oink_available_for_device(0) is False - - # Case 3: CUDA available and SM100, rmsnorm op registered. - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (10, 0)) - assert _oink_ops.is_oink_available_for_device(0) is True - - # fused op presence probe - assert _oink_ops.has_fused_add_rms_norm() is False - monkeypatch.setattr( - torch.ops, - "oink", - types.SimpleNamespace( - rmsnorm=lambda x, w, eps: x, - fused_add_rms_norm=lambda x, residual, w, eps: None, +from vllm.platforms import current_platform + + +def _test_oink_availability_impl( + device_capability: tuple[int, int], + has_rmsnorm: bool, + has_fused_add_rms_norm: bool, + expected_available: bool, + expected_fused: bool, +) -> None: + """Test OINK support detection with mocked state.""" + import torch + + from vllm import platforms + + # Mock device capability (class method, override on class) + dc = platforms.interface.DeviceCapability(*device_capability) + platforms.current_platform.__class__.get_device_capability = lambda device_id=0: dc + + # Mock oink ops + oink_ops = types.SimpleNamespace() + if has_rmsnorm: + oink_ops.rmsnorm = lambda x, w, eps: x + if has_fused_add_rms_norm: + oink_ops.fused_add_rms_norm = lambda x, residual, w, eps: None + + torch.ops.oink = oink_ops + + # Now import vllm modules with mocks in place (fresh import with mocked platform) + import vllm.kernels.oink_ops # noqa: F401 + from vllm.ir.ops import fused_add_rms_norm, rms_norm + + # Verify support checks + assert rms_norm.impls["oink"].supported is expected_available + assert fused_add_rms_norm.impls["oink"].supported is expected_fused + + +@pytest.mark.parametrize( + "device_capability,has_rmsnorm,has_fused_add_rms_norm,expected_available,expected_fused", + [ + # Case 1: < SM100, ops not supported + ((9, 0), True, False, False, False), + # Case 2: CUDA available and SM100, rmsnorm op registered + ((10, 0), True, False, True, False), + # Case 3: SM100 with both rmsnorm and fused_add_rms_norm + ((10, 0), True, True, True, True), + ], +) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA") +def test_oink_availability_checks( + device_capability: tuple[int, int], + has_rmsnorm: bool, + has_fused_add_rms_norm: bool, + expected_available: bool, + expected_fused: bool, +): + """Test OINK support detection with clean import state for each parameter set.""" + + # Use spawn to run function in fresh process with clean imports + # TODO migrate to spawn utility: + # https://github.com/vllm-project/vllm/issues/41415 + ctx = multiprocessing.get_context("spawn") + process = ctx.Process( + target=_test_oink_availability_impl, + args=( + device_capability, + has_rmsnorm, + has_fused_add_rms_norm, + expected_available, + expected_fused, ), - raising=False, ) - assert _oink_ops.has_fused_add_rms_norm() is True + process.start() + process.join() + + if process.exitcode != 0: + raise AssertionError( + f"Subprocess test failed with exit code {process.exitcode}" + ) def test_can_view_as_2d_stride_guard(): - # Import the helper from the layernorm module. - from vllm.model_executor.layers.layernorm import _can_view_as_2d + # No global import + import torch + + # Import the helper from the kernels module. + from vllm.kernels.oink_ops import _can_view_as_2d x = torch.zeros((2, 3, 4)) assert _can_view_as_2d(x) is True diff --git a/tests/test_config.py b/tests/test_config.py index 41d34a6cb06b..02e4d1d5d77b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1295,11 +1295,14 @@ def test_ir_op_priority_default(): # Assert default is applied to ops priority_config = IrOpPriorityConfig.with_default(["vllm_c", "native"]) assert priority_config.rms_norm == ["vllm_c", "native"] + assert priority_config.fused_add_rms_norm == ["vllm_c", "native"] # Assert single ops override the default - assert IrOpPriorityConfig.with_default( - ["vllm_c", "native"], rms_norm=["oink", "native"] - ) == IrOpPriorityConfig(rms_norm=["oink", "native"]) + priority_config = IrOpPriorityConfig.with_default( + ["native"], rms_norm=["oink", "native"] + ) + assert priority_config.rms_norm == ["oink", "native"] + assert priority_config.fused_add_rms_norm == ["native"] def test_ir_op_priority_str(): @@ -1318,3 +1321,34 @@ def test_ir_op_priority_str(): with pytest.raises(pydantic.ValidationError): # must be list of only strings priority_config = IrOpPriorityConfig(rms_norm=["vllm_c", 4, "native"]) + + +def test_ir_op_priority_ctx(): + """Test that the priority-setting context sets priority correctly.""" + from vllm import ir + from vllm.config.kernel import IrOpPriorityConfig + + priority = IrOpPriorityConfig.with_default(["native"], rms_norm=["vllm_c"]) + priority2 = IrOpPriorityConfig.with_default( + ["native"], fused_add_rms_norm=["vllm_c"] + ) + with priority.set_priority(): + assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"] + assert ir.ops.fused_add_rms_norm.get_priority() == ["native"] + with priority2.set_priority(): + assert ir.ops.rms_norm.get_priority() == ["native"] + assert ir.ops.fused_add_rms_norm.get_priority() == ["vllm_c", "native"] + + # context restored + assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"] + assert ir.ops.fused_add_rms_norm.get_priority() == ["native"] + + with pytest.raises(ValueError), priority2.set_priority(): + assert ir.ops.rms_norm.get_priority() == ["native"] + assert ir.ops.fused_add_rms_norm.get_priority() == ["vllm_c", "native"] + + raise ValueError + + # context restored even after exception + assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"] + assert ir.ops.fused_add_rms_norm.get_priority() == ["native"] diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 5c036c1b3802..2e9f77881273 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -12,7 +12,7 @@ from utils import skip_unsupported from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm -from vllm.model_executor.layers.layernorm import RMSNorm, fused_add_rms_norm +from vllm.model_executor.layers.layernorm import RMSNorm from vllm.platforms import current_platform DEVICE_TYPE = current_platform.device_type @@ -105,6 +105,12 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( dim=0, ) + def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]: + import vllm._custom_ops as ops + + ops.fused_add_rms_norm(x, residual, w, e) + return x, residual + out_single, residual_out_single = fused_add_rms_norm( x_single.clone(), residual_single.clone(), diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index b11fc21975ca..45defc6926ba 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -647,58 +647,6 @@ def _rocm_aiter_gemm_a8w8_blockscale_fake( return Y -def _rocm_aiter_rms_norm_impl( - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float -) -> torch.Tensor: - from aiter import rms_norm - - if x.dim() > 2: - x_original_shape = x.shape - x = x.reshape(-1, x_original_shape[-1]) - x = rms_norm(x, weight, variance_epsilon) - return x.reshape(x_original_shape) - - return rms_norm(x, weight, variance_epsilon) - - -def _rocm_aiter_rms_norm_fake( - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float -) -> torch.Tensor: - return torch.empty_like(x) - - -def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - variance_epsilon: float, -) -> tuple[torch.Tensor, torch.Tensor]: - from aiter import rmsnorm2d_fwd_with_add - - residual_out = torch.empty_like(residual) - out = torch.empty_like(x) - rmsnorm2d_fwd_with_add( - out, # output - x, # input - residual, # residual input - residual_out, # residual output - weight, - variance_epsilon, - ) - return out, residual_out - - -def _rocm_aiter_rmsnorm2d_fwd_with_add_fake( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - variance_epsilon: float, -) -> tuple[torch.Tensor, torch.Tensor]: - residual_out = torch.empty_like(residual) - out = torch.empty_like(x) - return out, residual_out - - def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl( x: torch.Tensor, residual: torch.Tensor, @@ -1229,10 +1177,9 @@ class rocm_aiter_ops: # Check if aiter is enabled before using operations if rocm_aiter_ops.is_enabled(): - result = rocm_aiter_ops.rms_norm(x, weight, epsilon) + result = rocm_aiter_ops.per_token_quant(x, FP8_DTYPE) Operations: - - RMS normalization: rms_norm, rms_norm2d_with_add - GEMM operations: gemm_a8w8, gemm_a8w8_blockscale - Fused MoE: fused_moe, asm_moe_tkw1 - Routing: topk_softmax, biased_grouped_topk, grouped_topk @@ -1244,7 +1191,6 @@ class rocm_aiter_ops: # Check if the env variable is set _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR - _RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA _MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA @@ -1275,7 +1221,6 @@ def refresh_env_variables(cls): """ cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR - cls._RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA @@ -1367,11 +1312,6 @@ def is_linear_enabled(cls) -> bool: def is_linear_fp8_enabled(cls) -> bool: return cls.is_linear_enabled() - @classmethod - @if_aiter_supported - def is_rmsnorm_enabled(cls) -> bool: - return cls._AITER_ENABLED and cls._RMSNORM_ENABLED - @classmethod @if_aiter_supported def is_fused_moe_enabled(cls) -> bool: @@ -1560,19 +1500,6 @@ def register_ops_once() -> None: fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake, ) - direct_register_custom_op( - op_name="rocm_aiter_rms_norm", - op_func=_rocm_aiter_rms_norm_impl, - fake_impl=_rocm_aiter_rms_norm_fake, - ) - - direct_register_custom_op( - op_name="rocm_aiter_rmsnorm2d_fwd_with_add", - op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl, - fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake, - dispatch_key=current_platform.dispatch_key, - ) - direct_register_custom_op( op_name="rocm_aiter_rmsnorm_fused_dynamic_quant", op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl, @@ -1672,14 +1599,6 @@ def register_ops_once() -> None: _OPS_REGISTERED = True - @staticmethod - def get_rmsnorm_fused_add_op() -> OpOverload: - return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add.default - - @staticmethod - def get_rmsnorm_op() -> OpOverload: - return torch.ops.vllm.rocm_aiter_rms_norm.default - @staticmethod def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload: return torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant.default @@ -1724,23 +1643,6 @@ def get_fused_allreduce_rmsnorm_op() -> OpOverload: def get_fused_mla_dual_rms_norm_op() -> OpOverload: return torch.ops.vllm.fused_mla_dual_rms_norm.default - @staticmethod - def rms_norm( - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float - ) -> torch.Tensor: - return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon) - - @staticmethod - def rms_norm2d_with_add( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - variance_epsilon: float, - ) -> tuple[torch.Tensor, torch.Tensor]: - return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add( - x, residual, weight, variance_epsilon - ) - @staticmethod def w8a8_gemm( A: torch.Tensor, diff --git a/vllm/_oink_ops.py b/vllm/_oink_ops.py deleted file mode 100644 index c7a055410b71..000000000000 --- a/vllm/_oink_ops.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Small helper wrappers for external Oink Blackwell custom ops. - -vLLM does not depend on the external Oink repository/package. When an external -plugin registers torch.library.custom_op entrypoints under the `oink::` -namespace (e.g. via vLLM's general_plugins mechanism) and -`VLLM_USE_OINK_OPS=1` is set, vLLM can route eligible calls to those ops. - -This module provides: -- A single place to probe Oink op availability at module init time - (outside torch.compile tracing), and -- Thin wrappers around the torch.ops entrypoints for use in CUDA fast paths, - without introducing graph breaks. - -Important: - Do not call the availability helpers in a compiled region. They may call - functions decorated with `torch._dynamo.disable` to safely check - conditions that should not be traced. -""" - -from __future__ import annotations - -from collections.abc import Callable - -import torch - -try: - from torch._dynamo import disable as _dynamo_disable # type: ignore[attr-defined] -except Exception: # pragma: no cover - - def _dynamo_disable(fn: Callable): # type: ignore[misc] - return fn - - -def _has_oink_op(op_name: str) -> bool: - """Check if a specific oink op is registered.""" - return hasattr(torch.ops, "oink") and hasattr(torch.ops.oink, op_name) - - -@_dynamo_disable -def is_oink_available_for_device(device_index: int) -> bool: - """Return True if Oink ops are registered and device is SM100+. - - This function is intended to be called during module initialization - (e.g., in RMSNorm.__init__), not in the forward path. - - External plugins are expected to gate registration on SM100+ and - VLLM_USE_OINK_OPS=1, so if the ops are present they should be usable. - """ - if not torch.cuda.is_available(): - return False - - try: - major, minor = torch.cuda.get_device_capability(device_index) - sm = 10 * major + minor - if sm < 100: - return False - except Exception: - return False - - return _has_oink_op("rmsnorm") - - -def has_fused_add_rms_norm() -> bool: - """Return True if the in-place fused op is registered.""" - return _has_oink_op("fused_add_rms_norm") - - -def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """Call `torch.ops.oink.rmsnorm`. - - This wrapper is safe to call in torch.compile regions. - """ - return torch.ops.oink.rmsnorm(x, weight, eps) - - -def fused_add_rms_norm_( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float, -) -> None: - """Call `torch.ops.oink.fused_add_rms_norm` (mutates x and residual).""" - torch.ops.oink.fused_add_rms_norm(x, residual, weight, eps) - - -def fused_add_rms_norm( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Convenience wrapper returning (x, residual) after in-place mutation.""" - fused_add_rms_norm_(x, residual, weight, eps) - return x, residual diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 353567fd96d3..4e8a485c01ab 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -50,6 +50,7 @@ should_split, ) from .passes.inductor_pass import InductorPass, pass_context +from .passes.ir.inplace_functionalization import VllmIRInplaceFunctionalizationPass from .passes.pass_manager import PostGradPassManager logger = init_logger(__name__) @@ -926,6 +927,24 @@ def collect_standalone_compile_artifacts( return standalone_compile_artifacts, sym_shape_indices_map, returns_tuple_map def configure_post_pass(self) -> None: + # TODO proper PassManager? + pre_grad_pass_key = "pre_grad_custom_pass" + assert self.pass_key != pre_grad_pass_key + assert pre_grad_pass_key not in self.inductor_config + self.inductor_config[pre_grad_pass_key] = VllmIRInplaceFunctionalizationPass( + self.vllm_config + ) + + # Make sure pre_grad_custom_pass is not pickled + # as part of AOTAutograd built-in cache key + # TODO(luka) is there a cleaner way to do this + import torch._inductor.config as inductor_config + + ignore = inductor_config._cache_config_ignore_prefix + [pre_grad_pass_key] + assert "_cache_config_ignore_prefix" not in self.inductor_config + self.inductor_config["_cache_config_ignore_prefix"] = ignore + + # Configure the (nominally post-grad) pass manager self.pass_manager.configure(self.vllm_config) # Post-grad custom passes are run using the post_grad_custom_post_pass diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index e683b1dfa69f..6cb0c8f49f3d 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -40,7 +40,7 @@ VllmPatternMatcherPass, VllmPatternReplacement, ) -from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8 +from .matcher_utils import MatcherQuantFP8 FP8_DTYPE = current_platform.fp8_dtype() @@ -356,10 +356,11 @@ def __init__( super().__init__(dtype, device) self.epsilon = epsilon self.allreduce_params = allreduce_params - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon) def get_inputs(self) -> list[torch.Tensor]: - input, residual, weight = self.rmsnorm_matcher.inputs() + input = self.empty(5, 16) + residual = self.empty(5, 16) + weight = self.empty(16) # input goes through allreduce first, always 16-bit return [residual, input.to(self.dtype), weight] @@ -369,7 +370,9 @@ def pattern( residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) - rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual) + rms, residual = vllm.ir.ops.fused_add_rms_norm( + allreduce_output, residual, weight, self.epsilon + ) return rms, residual def replacement( @@ -503,11 +506,12 @@ def __init__( self.allreduce_params = allreduce_params self.quant_dtype = torch.float8_e4m3fn - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon) self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) def get_inputs(self) -> list[torch.Tensor]: - input, residual, weight = self.rmsnorm_matcher.inputs() + input = self.empty(5, 16) + residual = self.empty(5, 16) + weight = self.empty(16) _, scale = self.quant_matcher.inputs() # input goes through allreduce first, always 16-bit @@ -521,7 +525,9 @@ def pattern( scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) - rms, res = self.rmsnorm_matcher(allreduce_output, weight, residual) + rms, res = vllm.ir.ops.fused_add_rms_norm( + allreduce_output, residual, weight, self.epsilon + ) quant, _ = self.quant_matcher(rms, scale) return quant, res @@ -668,7 +674,6 @@ def __init__( super().__init__(dtype, device) self.epsilon = epsilon self.allreduce_params = allreduce_params - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon) def get_inputs(self) -> list[torch.Tensor]: input = torch.empty([16, 16], device=self.device, dtype=self.dtype) @@ -700,7 +705,9 @@ def pattern( input_global_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) - rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual) + rms, residual = vllm.ir.ops.fused_add_rms_norm( + allreduce_output, residual, weight, self.epsilon + ) quant_out_tuple = auto_functionalized( STATIC_FP4_QUANT_OP, input=rms, @@ -955,15 +962,11 @@ def __init__( super().__init__(dtype, device) self.epsilon = epsilon self.dtype = dtype - self.rmsnorm_matcher = MatcherFusedAddRMSNorm( - epsilon, match_rocm_aiter=use_aiter_rmsnorm - ) self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op() def get_inputs(self) -> list[torch.Tensor]: - input, residual, weight = self.rmsnorm_matcher.inputs() - - return [residual, input.to(self.dtype), weight] + # input, residual, weight + return [self.empty(5, 16), self.empty(5, 16), self.empty(16)] @property def pattern(self): @@ -971,8 +974,9 @@ def _pattern( residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: allreduce_output = tensor_model_parallel_all_reduce(input) - rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual) - + rms, residual = vllm.ir.ops.fused_add_rms_norm( + allreduce_output, residual, weight, self.epsilon + ) return rms, residual return _pattern diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index c2490d8a21f5..e5130c19c392 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -10,7 +10,6 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import get_current_vllm_config from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, @@ -159,72 +158,6 @@ def forward_native( return result -class MatcherFusedAddRMSNorm(MatcherCustomOp): - def __init__( - self, - epsilon: float, - enabled: bool | None = None, - match_rocm_aiter: bool = False, - ) -> None: - if enabled is None: - enabled = RMSNorm.enabled() - - super().__init__(enabled) - self.epsilon = epsilon - self.match_rocm_aiter = match_rocm_aiter - - self._rmsnorm_op = RMS_ADD_OP - - if match_rocm_aiter: - self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_fused_add_op() - - def inputs(self) -> list[torch.Tensor]: - input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16) - weight = self.empty(16) - residual = self.empty(5, 16) - return [input, weight, residual] - - def forward_rocm_aiter( - self, - input: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - return self._rmsnorm_op( # type: ignore[no-any-return] - x=input, residual=residual, weight=weight, variance_epsilon=self.epsilon - ) - - def forward_custom( - self, - input: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.match_rocm_aiter: - return self.forward_rocm_aiter(input, weight, residual) - - _, result, residual = auto_functionalized( - self._rmsnorm_op, - input=input, - residual=residual, - weight=weight, - epsilon=self.epsilon, - ) - - return result, residual - - def forward_native( - self, - input: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - result: tuple[torch.Tensor, torch.Tensor] = RMSNorm.forward_static( - input, self.epsilon, input.size(-1), self.model_dtype, weight, residual - ) - return result - - class MatcherQuantFP8(MatcherCustomOp): def __init__( self, diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index 850e434a3e73..cc986595d436 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -29,7 +29,6 @@ from ..inductor_pass import enable_fake_mode from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass from .matcher_utils import ( - MatcherFusedAddRMSNorm, MatcherQuantFP8, ) @@ -146,9 +145,6 @@ def __init__( assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}" self.FUSED_OP = FUSED_OPS[key] - if key.fused_add: - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon) - self.quant_matcher = MatcherQuantFP8( key.quant, has_col_major_scales=has_col_major_scales, @@ -231,7 +227,9 @@ def pattern( residual: torch.Tensor, scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - result_rms, residual = self.rmsnorm_matcher(input, weight, residual) + result_rms, residual = vllm.ir.ops.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) result, _ = self.quant_matcher(result_rms, scale) return result, residual @@ -261,8 +259,9 @@ def replacement( return at[1], at[2] inputs = [ - # input, weight, residual - *self.rmsnorm_matcher.inputs(), + empty_bf16(5, 16), # input + empty_bf16(16), # weight + empty_bf16(5, 16), # residual self.quant_matcher.inputs()[1], # scale ] @@ -311,7 +310,9 @@ def pattern( residual: torch.Tensor, scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - result_rms, residual = self.rmsnorm_matcher(input, weight, residual) + result_rms, residual = vllm.ir.ops.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) result = torch.empty( result_rms.shape, device=result_rms.device, @@ -366,12 +367,17 @@ def replacement( # result, residual, scale return at[1], at[3], at[2] - scale = self.quant_matcher.empty_f32(1, 1) + inputs = [ + empty_bf16(5, 16), # input + empty_bf16(16), # weight + empty_bf16(5, 16), # residual + self.quant_matcher.empty_f32(1, 1), # scale + ] pm.register_replacement( pattern, replacement, - self.rmsnorm_matcher.inputs() + [scale], + inputs, pm.fwd_only, pm_pass, extra_check=_rms_input_weight_dtype_match, @@ -552,7 +558,9 @@ def register(self, pm_pass: PatternMatcherPass) -> None: def pattern( input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - result_rms, residual = self.rmsnorm_matcher(input, weight, residual) + result_rms, residual = vllm.ir.ops.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) result, scale = self.quant_matcher(result_rms) return result, residual, scale @@ -580,10 +588,16 @@ def replacement( # result, residual, scale return at[1], at[3], at[2] + inputs = [ + empty_bf16(5, 16), # input + empty_bf16(16), # weight + empty_bf16(5, 16), # residual + ] + pm.register_replacement( pattern, replacement, - self.rmsnorm_matcher.inputs(), + inputs, pm.fwd_only, pm_pass, extra_check=_rms_input_weight_dtype_match, diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index cdd0e23773d6..28159dbe0872 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -29,7 +29,6 @@ VllmPatternReplacement, ) from .matcher_utils import ( - MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherSiluAndMul, ) @@ -49,10 +48,6 @@ def __init__( self.quant_dtype = key.quant.dtype self.device = torch.device("cuda") - if key.fused_add: - self.rmsnorm_matcher = MatcherFusedAddRMSNorm( - epsilon, match_rocm_aiter=True - ) self.quant_matcher = MatcherQuantFP8( key.quant, match_rocm_aiter=match_aiter_quant, @@ -145,7 +140,9 @@ def pattern( weight: torch.Tensor, residual: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual) + result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) result, scale = self.quant_matcher(result_rms) return result, residual_out, scale @@ -163,10 +160,16 @@ def replacement( return result[0], result[1], result[2] + inputs = [ + self.empty(5, 16), # input + self.empty(16), # weight + self.empty(5, 16), # residual + ] + pm.register_replacement( pattern, replacement, - self.rmsnorm_matcher.inputs(), + inputs, pm.fwd_only, pm_pass, ) @@ -258,7 +261,9 @@ def pattern( weight: torch.Tensor, residual: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual) + result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) result, scale = self.quant_matcher(result_rms) return result, residual_out, scale @@ -279,9 +284,13 @@ def replacement( # result, scale, residual return at[0], at[1], at[2] - pm.register_replacement( - pattern, replacement, self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass - ) + inputs = [ + self.empty(5, 16), # input + self.empty(16), # weight + self.empty(5, 16), # residual + ] + + pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass) class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): @@ -420,12 +429,15 @@ def __init__( self.epsilon = epsilon self.hidden_size = hidden_size self.x_pad_to_multiple = x_pad_to_multiple - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon, match_rocm_aiter=True) def get_inputs(self) -> list[torch.Tensor]: - input, weight, residual = self.rmsnorm_matcher.inputs() - router_weight = torch.empty([8, 16], dtype=weight.dtype, device=weight.device) - router_bias = torch.empty([8], dtype=weight.dtype, device=weight.device) + device = torch.device("cuda") + dtype = torch.bfloat16 + input = torch.empty(5, 16, dtype=dtype, device=device) + weight = torch.empty(16, dtype=dtype, device=device) + residual = torch.empty(5, 16, dtype=dtype, device=device) + router_weight = torch.empty([8, 16], dtype=dtype, device=device) + router_bias = torch.empty([8], dtype=dtype, device=device) return [input, weight, residual, router_weight, router_bias] def register(self, pm_pass: PatternMatcherPass) -> None: @@ -439,7 +451,9 @@ def pattern( pad_size = self.x_pad_to_multiple - ( self.hidden_size % self.x_pad_to_multiple ) - result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual) + result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) router_logits = torch.ops.vllm.rocm_unquantized_gemm( result_rms, router_weight, router_bias ) diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index 35885eeb0b8e..2c7a1390bdb8 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/vllm/compilation/passes/fusion/sequence_parallelism.py @@ -23,7 +23,7 @@ from ..inductor_pass import enable_fake_mode from ..utility.noop_elimination import NoOpEliminationPass from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass -from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8 +from .matcher_utils import MatcherQuantFP8 logger = init_logger(__name__) @@ -174,7 +174,6 @@ def replacement( class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper): def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: super().__init__(epsilon, dtype, device) - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon) def get_inputs(self) -> list[torch.Tensor]: mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype) @@ -195,7 +194,9 @@ def pattern( rms_norm_weights: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: all_reduce = self._all_reduce(mm_1) - rmsnorm = self.rmsnorm_matcher(all_reduce, rms_norm_weights, residual) + rmsnorm = vllm.ir.ops.fused_add_rms_norm( + all_reduce, residual, rms_norm_weights, self.epsilon + ) return rmsnorm[0], rmsnorm[1] def replacement( @@ -208,7 +209,9 @@ def replacement( # once the seqpar pattern with the previous rmsnorm is replaced reduce_scatter = self._reduce_scatter(mm_1) residual = residual[0 : reduce_scatter.size(0), ...] - rmsnorm = self.rmsnorm_matcher(reduce_scatter, rms_norm_weights, residual) + rmsnorm = vllm.ir.ops.fused_add_rms_norm( + reduce_scatter, residual, rms_norm_weights, self.epsilon + ) all_gather = self._all_gather(rmsnorm[0]) # shape of residual changes but that's fine, # next node is already slicing it, now becomes a noop @@ -271,7 +274,6 @@ def replacement( class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper): def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None: super().__init__(epsilon, dtype, device) - self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon) self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) def get_inputs(self) -> list[torch.Tensor]: @@ -290,8 +292,8 @@ def pattern( scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: all_reduce = self._all_reduce(mm_1) - rms, residual_out = self.rmsnorm_matcher( - all_reduce, rms_norm_weights, residual + rms, residual_out = vllm.ir.ops.fused_add_rms_norm( + all_reduce, residual, rms_norm_weights, self.epsilon ) quant, _ = self.quant_matcher(rms, scale) return quant, residual_out @@ -308,8 +310,8 @@ def replacement( # once the seqpar pattern with the previous rmsnorm is replaced reduce_scatter = self._reduce_scatter(mm_1) residual = residual[0 : reduce_scatter.size(0), ...] - rms, residual_out = self.rmsnorm_matcher( - reduce_scatter, rms_norm_weights, residual + rms, residual_out = vllm.ir.ops.fused_add_rms_norm( + reduce_scatter, residual, rms_norm_weights, self.epsilon ) quant, _ = self.quant_matcher(rms, scale) all_gather = self._all_gather(quant) diff --git a/vllm/compilation/passes/inductor_pass.py b/vllm/compilation/passes/inductor_pass.py index b54c7bfa14d0..8a0d5326dd92 100644 --- a/vllm/compilation/passes/inductor_pass.py +++ b/vllm/compilation/passes/inductor_pass.py @@ -30,6 +30,9 @@ class PassContext: def __init__(self, compile_range: Range): self.compile_range: Range = compile_range + # set of arg indices + self.donated_input_ids: set[int] = set() + def get_pass_context() -> PassContext: """Get the current pass context.""" diff --git a/vllm/compilation/passes/ir/clone_elimination.py b/vllm/compilation/passes/ir/clone_elimination.py new file mode 100644 index 000000000000..61ba750a6c4e --- /dev/null +++ b/vllm/compilation/passes/ir/clone_elimination.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +from torch import fx +from torch._higher_order_ops.auto_functionalize import auto_functionalized +from torch._higher_order_ops.triton_kernel_wrap import TritonKernelWrapperFunctional +from torch._ops import HigherOrderOperator, OpOverload + +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from ..fx_utils import is_func +from ..inductor_pass import get_pass_context +from ..vllm_inductor_pass import VllmInductorPass + +logger = init_logger(__name__) + + +def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool: + if user.op == "output": + return False + + if is_func(user, auto_functionalized): + # While autofunc writes to the node, + # this is a follow-up use we're not interested in. + # It is also guaranteed to be the final use, + # as auto_functionalized returns the tensor back for follow-up use. + return False + elif user.op == "call_function" and isinstance(user.target, HigherOrderOperator): + # By default, be conservative, assume this could be a write + # (except functional HOPs) + return not isinstance(user.target, TritonKernelWrapperFunctional) + + assert isinstance(user.target, OpOverload), ( + f"{node=} {user=} {user.op=} {user.target=}" + ) + schema = user.target._schema + assert len(user.args) <= len(schema.arguments) + for i, arg in enumerate(user.args): + # Only interested in writes to node + if arg is not node: + continue + + # If not a write, next arg could be + if schema.arguments[i].is_write: + return True + + # No writes found + return False + + +class UnsafeCloneEliminationPass(VllmInductorPass): + """ + This pass removes clone nodes that are no longer needed after vLLM IR lowering. + It uses donated_input_ids to eliminate clones of donated graph inputs, preserving + contents of non-donated graph inputs. + + It is "unsafe" because it does not (yet) take aliasing into account. Solving + aliasing is an open problem, so this pass intends to support known vLLM cases + and not guarantee soundness on general graphs. In the future, this pass will likely + support basic forms of aliasing to handle simple views (e.g. qkv -> q,k,v). + """ + + def __init__(self, vllm_config: VllmConfig) -> None: + super().__init__(vllm_config) + + @VllmInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + count = 0 + node_to_idx = {node: i for i, node in enumerate(graph.nodes)} + pass_context = get_pass_context() + donated_input_ids = pass_context.donated_input_ids + logger.debug("Donated input ids: %s", donated_input_ids) + + for node in graph.nodes: + if not is_func(node, torch.ops.aten.clone.default): + continue + + original_node = node.args[0] + assert isinstance(original_node, fx.Node) + + # Clone needs to be preserved if node is getting written to and + # the old value is used again. + # This could only happen if an inplace implementation was lowered. + # Then node (the clone) will have one write. + # TODO(luka) hopefully this can be removed once we lower functional graphs. + write_idxs = [ + node_to_idx[u] for u in node.users if user_writes_to_node(u, node) + ] + assert len(write_idxs) in (0, 1) + if write_idxs: + # Check if a user of original_node occurs after a write + write_idx = write_idxs[0] + if any( + node_to_idx[orig_user] > write_idx + for orig_user in original_node.users + ): + logger.debug( + "Clone removal not possible, " + "original_node=%s used after mutation on node=%s", + original_node, + node, + ) + continue + + # Check if a node is a (non-donated) graph input + if ( + original_node.op == "placeholder" + and node_to_idx[original_node] not in donated_input_ids + ): + logger.debug( + "Graph input %s not donated, cannot eliminate its clone", + original_node, + ) + continue + + logger.debug( + "Node %s is a redundant clone node of %s, removing it", + node, + original_node, + ) + node.replace_all_uses_with(original_node) + graph.erase_node(node) + count += 1 + + logger.debug("CloneCleanupPass removed %d clone nodes", count) diff --git a/vllm/compilation/passes/ir/inplace_functionalization.py b/vllm/compilation/passes/ir/inplace_functionalization.py new file mode 100644 index 000000000000..e69351075bca --- /dev/null +++ b/vllm/compilation/passes/ir/inplace_functionalization.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections import defaultdict + +from torch import fx +from torch._inductor.pattern_matcher import ( + PatternMatcherPass, +) + +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from ..inductor_pass import get_pass_context +from ..vllm_inductor_pass import VllmInductorPass +from .lowering_pass import get_ir_op +from .utils import overload_or_default + +logger = init_logger(__name__) + + +class VllmIRInplaceFunctionalizationPass(VllmInductorPass): + """ + This pass functionalizes maybe_inplace vLLM IR ops to the default overload. + The maybe_inplace overloads have the same signature as the default overload + so the pass simply replaces the called overload. + That makes the graph properly functional. + The pass also validates that activations passed to maybe_inplace have no later + uses in the graph: they are donated to the maybe_inplace op call, + and their contents are not defined afterward. + + This pass operates pre-AOTAutograd, + so it must handle non-normalized and non-functional IR. + """ + + def __init__(self, vllm_config: VllmConfig) -> None: + super().__init__(vllm_config) + self.patterns = PatternMatcherPass(self.pass_name) + self.functionalized_ops: dict[str, int] = defaultdict(lambda: 0) + + @VllmInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + # clear at the beginning instead of end, so that tests can inspect + self.functionalized_ops.clear() + assert graph.owning_module is not None + node_to_idx = {node: i for i, node in enumerate(graph.nodes)} + + # Pass donated input via vLLM's pass context + pass_context = get_pass_context() + pass_context.donated_input_ids = set[int]() + + for node in graph.nodes: + if (ir_op := get_ir_op(node)) is None: + continue + + op_overload = overload_or_default(node.target) + overload_name = op_overload._overloadname + if overload_name != "maybe_inplace": + assert overload_name == "default", ( + f"Found overload {overload_name} for op {ir_op.name}, " + f"expected maybe_inplace or default" + ) + continue + + # must have maybe_inplace overload and allow_inplace + assert ir_op.allow_inplace and hasattr(ir_op, "maybe_inplace") + + # Check that activation inputs are not used after this op + for arg_idx in ir_op.activation_indices: + arg = node.args[arg_idx] + assert isinstance(arg, fx.Node), "Activation inputs must be fx.Node" + for user in arg.users: + if node_to_idx[user] > node_to_idx[node]: + raise ValueError( + f"Input {arg} to maybe_inplace node {node} " + f"is used again after the node. " + f"This is not allowed; activation inputs to maybe_inplace " + f"ops are donated to the op, meaning their memory may be " + f"recycled for outputs.\n\n" + f"To preserve the inputs, use the default overload or " + f"clone them manually beforehand." + ) + + if arg.op == "placeholder": + # Graph input that maybe_inplace might modify. + # Mark it so downstream passes know it's donated. + # TODO(luka) store in placeholder node meta once supported + pass_context.donated_input_ids.add(node_to_idx[arg]) + + # Same signature, just replace the overload that's called. + node.target = ir_op.torch_op + self.functionalized_ops[ir_op.name] += 1 + + count = sum(self.functionalized_ops.values()) + ops = ",".join(self.functionalized_ops.keys()) + logger.debug("Donated input IDs: %s", pass_context.donated_input_ids) + logger.debug( + "%s functionalized %d vLLM IR nodes for op(s) %s", + self.pass_name, + count, + ops, + ) diff --git a/vllm/compilation/passes/ir/lowering_pass.py b/vllm/compilation/passes/ir/lowering_pass.py index 02acdd1a298b..f34f1c64b76e 100644 --- a/vllm/compilation/passes/ir/lowering_pass.py +++ b/vllm/compilation/passes/ir/lowering_pass.py @@ -10,7 +10,6 @@ PatternMatcherPass, register_graph_pattern, ) -from torch._ops import OpOverload, OpOverloadPacket from vllm.config import VllmConfig from vllm.ir.op import IrOp @@ -18,41 +17,11 @@ from vllm.logging_utils import lazy from ..vllm_inductor_pass import VllmInductorPass +from .utils import get_ir_op logger = init_logger(__name__) -def get_default_overload(op: OpOverload | OpOverloadPacket) -> OpOverload: - if isinstance(op, OpOverloadPacket): - return op.default - assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket" - return op - - -def get_ir_op(node: fx.Node) -> IrOp | None: - if node.op != "call_function": - return None - - if not isinstance(node.target, (OpOverload, OpOverloadPacket)): - return None - - op_overload = get_default_overload(node.target) - if op_overload.namespace != "vllm_ir": - return None - - op_name = op_overload._opname - if op_name not in IrOp.registry: - logger.warning( - "Unknown vLLM IR op %s, there's likely an issue with torch registration, " - "or a torch custom op was registered in the vllm_ir namespace by mistake.", - op_name, - ) - return None - - ir_op = IrOp.registry[op_name] - return ir_op - - class VllmIRLoweringPass(VllmInductorPass): """ This pass lowers vLLM IR ops to their implementations the priority list. @@ -76,7 +45,7 @@ def lower_matched_op(self, match: Match, *args, **kwargs): assert len(match.nodes) == 1, "Expected single node match" node = match.nodes[0] - ir_op = get_ir_op(node) + ir_op = get_ir_op(node) # TODO is node.target always an overload? assert ir_op is not None, "Expected vLLM IR op" assert not node.kwargs # I think there should never be kwargs here @@ -86,13 +55,18 @@ def lower_matched_op(self, match: Match, *args, **kwargs): self.selected_impls[ir_op.name][node.name] = ir_op_impl.provider # replace_by_example wants node args, not the fake tensors + # use func_impl_fn to properly handle in-place implementations # TODO(luka): Use aot_export_module to get functionalized graph # TODO(luka): Cache the fx_replacement to avoid re-tracing the same impl # Defaults not present on node.args but required for replacement tracing bound_args = ir_op._py_signature.bind(*node.args) bound_args.apply_defaults() - match.replace_by_example(ir_op_impl.impl_fn, bound_args.args) + # It is not safe to run functional passes (like DCE) on the replacements + # as they might not be functional. + match.replace_by_example( + ir_op_impl.func_impl_fn, bound_args.args, run_functional_passes=False + ) @VllmInductorPass.time_and_log def __call__(self, graph: fx.Graph) -> None: @@ -136,7 +110,7 @@ def print_count(counts: dict[str, int]) -> str: if failed_nodes or failed_ops: logger.warning("Failed to lower vLLM IR ops: %s", ",".join(failed_ops)) - logger.warning("Full node list: %s", failed_nodes) + logger.warning("Full node list: %s", ",".join(str(n) for n in failed_nodes)) def uuid(self) -> str: """ diff --git a/vllm/compilation/passes/ir/utils.py b/vllm/compilation/passes/ir/utils.py new file mode 100644 index 000000000000..50b4773ce523 --- /dev/null +++ b/vllm/compilation/passes/ir/utils.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from torch import fx +from torch._ops import OpOverload, OpOverloadPacket + +from vllm.ir.op import IrOp +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def overload_or_default(op: OpOverload | OpOverloadPacket) -> OpOverload: + if isinstance(op, OpOverloadPacket): + return op.default + assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket" + return op + + +def get_ir_op(node: fx.Node) -> IrOp | None: + if node.op != "call_function": + return None + + if not isinstance(node.target, (OpOverload, OpOverloadPacket)): + return None + + op_overload = overload_or_default(node.target) + if op_overload.namespace != "vllm_ir": + return None + + op_name = op_overload._opname + if op_name not in IrOp.registry: + logger.warning( + "Unknown vLLM IR op %s, there's likely an issue with torch registration, " + "or a torch custom op was registered in the vllm_ir namespace by mistake.", + op_name, + ) + return None + + ir_op = IrOp.registry[op_name] + return ir_op diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 3dc0d7b096ba..5d4355a5b2b4 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -14,6 +14,7 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import set_env_var +from .ir.clone_elimination import UnsafeCloneEliminationPass from .ir.lowering_pass import VllmIRLoweringPass from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass @@ -115,6 +116,8 @@ def __call__(self, graph: fx.Graph) -> None: # DCE handles mutating ops correctly as well. self.ir_lowering(graph) VllmInductorPass.dump_prefix += 1 + self.clone_elimination(graph) + VllmInductorPass.dump_prefix += 1 # clean up after lowering again self.post_cleanup(graph) @@ -149,11 +152,12 @@ def configure(self, config: VllmConfig) -> None: self.passes += [MiniMaxQKNormPass(config)] if self.pass_config.fuse_norm_quant: - self.passes += [RMSNormQuantFusionPass(config)] if rocm_aiter_ops.is_enabled(): self.passes += [ RocmAiterRMSNormQuantFusionPass(config), ] + self.passes += [RMSNormQuantFusionPass(config)] + if self.pass_config.fuse_act_quant: self.passes += [ActivationQuantFusionPass(config)] if rocm_aiter_ops.is_enabled(): @@ -179,6 +183,7 @@ def configure(self, config: VllmConfig) -> None: self.passes += [QKNormRoPEFusionPass(config)] self.ir_lowering = VllmIRLoweringPass(config) + self.clone_elimination = UnsafeCloneEliminationPass(config) self.post_cleanup = PostCleanupPass(config) self.fix_functionalization = FixFunctionalizationPass(config) @@ -200,6 +205,7 @@ def uuid(self) -> str: passes.append(self.post_cleanup.uuid()) passes.append(self.ir_lowering.uuid()) + passes.append(self.clone_elimination.uuid()) passes.append(self.post_cleanup.uuid()) passes.append(self.fix_functionalization.uuid()) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 93fb4c54b7f1..e5dcc20b5048 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -31,6 +31,9 @@ class IrOpPriorityConfig: rms_norm: list[str] = Field(default_factory=list) """Priority list for vllm.ir.ops.rms_norm""" + fused_add_rms_norm: list[str] = Field(default_factory=list) + """Priority list for vllm.ir.ops.fused_add_rms_norm""" + def compute_hash(self) -> str: """ Produces a hash unique to the pass configuration. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8d2c2608e56b..0146ee4c144a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -125,9 +125,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool: from vllm._aiter_ops import rocm_aiter_ops return ( - rocm_aiter_ops.is_enabled() - and rocm_aiter_ops.is_rmsnorm_enabled() - and cfg.parallel_config.tensor_parallel_size > 1 + rocm_aiter_ops.is_enabled() and cfg.parallel_config.tensor_parallel_size > 1 ) return ( @@ -159,10 +157,9 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool: def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool: """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss.""" - from vllm._aiter_ops import rocm_aiter_ops return ( - rocm_aiter_ops.is_rmsnorm_enabled() + cfg.kernel_config.ir_op_priority.fused_add_rms_norm[0] == "aiter" and cfg.model_config is not None and cfg.model_config.get_hidden_size() == 2880 ) diff --git a/vllm/envs.py b/vllm/envs.py index 24ec92c3d755..8378c9762ae7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -143,7 +143,7 @@ VLLM_DP_RANK_LOCAL: int = -1 VLLM_DP_SIZE: int = 1 VLLM_USE_STANDALONE_COMPILE: bool = True - VLLM_ENABLE_PREGRAD_PASSES: bool = False + VLLM_ENABLE_PREGRAD_PASSES: bool = True VLLM_DP_MASTER_IP: str = "" VLLM_DP_MASTER_PORT: int = 0 VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False @@ -620,9 +620,10 @@ def _get_or_set_default() -> str: # The pre-grad passes get run even on cache-hit and negatively impact # vllm cold compile times by O(1s) # Can remove this after the following issue gets fixed + # TODO(luka): maybe_inplace requires this # https://github.com/pytorch/pytorch/issues/174502 "VLLM_ENABLE_PREGRAD_PASSES": lambda: ( - os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "0") == "1" + os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "1") == "1" ), # Debug pattern matching inside custom passes. # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3'). diff --git a/vllm/ir/op.py b/vllm/ir/op.py index 5d7c01be1bbc..841df0f9adf4 100644 --- a/vllm/ir/op.py +++ b/vllm/ir/op.py @@ -4,7 +4,7 @@ import inspect from collections.abc import Callable from pathlib import Path -from typing import Any, ClassVar, overload +from typing import Any, ClassVar, Literal, overload import torch from torch.library import Library, infer_schema @@ -46,35 +46,51 @@ def enable_torch_wrap(enable: bool = True): _ENABLE_TORCH_WRAP = old -# 0-param decorator overload +# 0-param decorator overload (no inplace) @overload def register_op(f: Callable[..., Any]) -> "IrOp": ... -# parametrized decorator overload +# parametrized decorator with allow_inplace=False (default) @overload def register_op( *, name: str | None = None, + activations: list[str] | None = None, + allow_inplace: Literal[False] = False, ) -> Callable[[Callable[..., Any]], "IrOp"]: ... +# parametrized decorator with allow_inplace=True +@overload +def register_op( + *, + name: str | None = None, + activations: list[str] | None = None, + allow_inplace: Literal[True], +) -> Callable[[Callable[..., Any]], "IrOpInplace"]: ... + + def register_op( f: Callable | None = None, *, name: str | None = None, + activations: list[str] | None = None, + allow_inplace: bool = False, ) -> "IrOp | Callable[[Callable], IrOp]": """ Register a new vLLM IR op. :param f: the native implementation of the op :param name: the name of the op, defaults to the function name + :param activations: list of activation params, defaults to params starting with 'x' + :param allow_inplace: add a maybe_inplace overload that allows inplace impls :return: the IrOp object if f is provided, otherwise a decorator Example usage: ```python @vllm.ir.register_op - def my_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y @@ -85,7 +101,10 @@ def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: def decorator(_f: Callable): op_name: str = _f.__name__ if name is None else name assert op_name not in IrOp.registry - op = IrOp(op_name, _f) + if allow_inplace: + op: IrOp = IrOpInplace(op_name, _f, activations) + else: + op = IrOp(op_name, _f, activations) IrOp.registry[op_name] = op return op @@ -100,8 +119,14 @@ class IrOp: name: str impls: dict[str, "IrOpImpl"] + allow_inplace: bool = False - def __init__(self, name: str, native_impl: Callable): + def __init__( + self, + name: str, + native_impl: Callable, + activations: list[str] | None = None, + ): self._py_signature = inspect.signature(native_impl) if any( p.kind == inspect.Parameter.KEYWORD_ONLY @@ -112,8 +137,22 @@ def __init__(self, name: str, native_impl: Callable): f"supported. That's because kwargs are not allowed during lowering." ) + # By convention, we consider parameters starting with 'x' as activations. + if activations is None: + activations = [ + p.name + for p in self._py_signature.parameters.values() + if p.name.startswith("x") + ] + self.name = name self.impls: dict[str, IrOpImpl] = {} + self.activations = activations + self.activation_indices = [ + i + for i, p in enumerate(self._py_signature.parameters.values()) + if p.name in activations + ] self._priority_impls: list[IrOpImpl] = [] self._schema_str = infer_schema(native_impl, mutates_args=[]) self._input_generator: InputGenerator | None = None @@ -121,7 +160,12 @@ def __init__(self, name: str, native_impl: Callable): # native implementation self.impls["native"] = IrOpImpl( - self, "native", native_impl, supported=True, supports_args=None + self, + "native", + native_impl, + # always supported + supported=True, + supports_args=None, ) # By default, fake routes directly to native, @@ -161,12 +205,14 @@ def register_impl( *, supported: bool = True, supports_args: Callable[..., bool] | None = None, + inplace: bool = False, ): """ Register an implementation for this custom op. :param provider: The name of the provider, must be unique. :param supported: Static support check, use this to check platform support. :param supports_args: Dynamic arg support check, used for types and shapes. + :param inplace: Does this op reuse activation input memory for outputs :return: A decorator that registers the implementation. The decorated function must have the same semantics and signature as @@ -193,7 +239,7 @@ def my_provider_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: ... ) def _register_impl(f: Callable): - impl = IrOpImpl(self, provider, f, supported, supports_args) + impl = IrOpImpl(self, provider, f, supported, supports_args, inplace) self.impls[provider] = impl if self.get_priority(): @@ -213,7 +259,10 @@ def _inner_call(self, *args, **kwargs) -> Any: __call__ routes straight here instead of going through torch op dispatching. """ impl = self.dispatch(*args, **kwargs) - return impl.impl_fn(*args, **kwargs) + + # Default overload must be functional, + # use func_impl_fn to correctly handle inplace impls. + return impl.func_impl_fn(*args, **kwargs) def apply_arg_defaults(self, args) -> tuple: """ @@ -314,6 +363,11 @@ def filter_priority_impls(p_list: list[str]) -> list[IrOpImpl]: old_priority_impls = self._priority_impls try: self._priority_impls = filter_priority_impls(priority) + logger.debug( + "Priority for vllm.ir.%s set to %s", + self.name, + lazy(lambda: [p.provider for p in self._priority_impls]), + ) yield finally: self._priority_impls = old_priority_impls @@ -354,6 +408,66 @@ def get_tolerance(self, dtype: torch.dtype) -> dict[str, float]: ) +class IrOpInplace(IrOp): + """IR op with inplace support via maybe_inplace.""" + + maybe_inplace: "IrOpInplaceOverload" + allow_inplace: bool = True + + def __init__( + self, + name: str, + native_impl: Callable, + activations: list[str] | None = None, + ): + super().__init__(name, native_impl, activations) + + # Create the inplace overload + self.maybe_inplace = IrOpInplaceOverload(self) + + +class IrOpInplaceOverload: + def __init__(self, op: IrOp): + params, returns = op._schema_str.split(" -> ") + n_outputs = returns.count("Tensor") + + assert returns.count("Tensor") == len(op.activations), ( + "Inplace overload requires the same number of outputs as activations." + ) + + assert returns.count(",") == n_outputs - 1, ( + "Inplace overload only supports Tensor outputs for now." + ) + + self.op = op + self.name = f"{op.name}.maybe_inplace" + self._schema_str = infer_schema( + op.impls["native"].impl_fn, mutates_args=op.activations + ) + + # torch registration + vllm_ir_lib.define(self.name + self._schema_str) + vllm_ir_lib.impl( + self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd" + ) + # fake goes to default overload for now + vllm_ir_lib._register_fake(self.name, self.op._fake_call) + + assert hasattr(getattr(torch.ops.vllm_ir, self.op.name), "maybe_inplace") + self.torch_op = getattr(torch.ops.vllm_ir, self.op.name).maybe_inplace + + def __call__(self, *args, **kwargs) -> Any: + if not _ENABLE_TORCH_WRAP: + return self._inner_call(*args, **kwargs) + + return self.torch_op(*args, **kwargs) + + def _inner_call(self, *args, **kwargs) -> Any: + # Calling the maybe_inplace overload means we can use inplace impls directly. + impl = self.op.dispatch(*args, **kwargs) + return impl.impl_fn(*args, **kwargs) + + class IrOpImpl: def __init__( self, @@ -362,6 +476,7 @@ def __init__( impl_fn: Callable, supported: bool, supports_args: Callable[..., bool] | None, + inplace: bool = False, ): assert provider not in op.impls, ( f"Implementation for provider {provider} already registered." @@ -420,11 +535,18 @@ def __init__( f"native default {op_p.default}'" ) + if inplace: + assert op.allow_inplace, ( + f"Inplace implementation cannot be registered for op {op.name}" + f" that does not allow inplace." + ) + self.op = op self.provider = provider self.impl_fn = impl_fn self.supported = supported self._supports_args = supports_args + self.inplace = inplace @property def supports_all_args(self) -> bool: @@ -449,3 +571,19 @@ def uuid(self): """ sources = [Path(inspect.getfile(self.impl_fn))] return hash_source(*sources) + + def func_impl_fn(self, *args, **kwargs) -> Any: + """ + Copy any inputs in activations if this is an inplace impl, + to ensure functional semantics. + """ + if not self.inplace: + return self.impl_fn(*args, **kwargs) + + # copy activations to ensure functional semantics + new_args = list(args) + for i in self.op.activation_indices: + assert isinstance(args[i], torch.Tensor) + new_args[i] = args[i].clone() + + return self.impl_fn(*new_args, **kwargs) diff --git a/vllm/ir/ops/__init__.py b/vllm/ir/ops/__init__.py index 25ad27c8a078..d4d71afef723 100644 --- a/vllm/ir/ops/__init__.py +++ b/vllm/ir/ops/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .layernorm import rms_norm +from .layernorm import fused_add_rms_norm, rms_norm -__all__ = ["rms_norm"] +__all__ = ["rms_norm", "fused_add_rms_norm"] diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py index 981d5e3bd836..33a71b8f853f 100644 --- a/vllm/ir/ops/layernorm.py +++ b/vllm/ir/ops/layernorm.py @@ -27,10 +27,46 @@ def _rms_norm_input_generator( ) -> tuple: x = torch.randn(num_tokens, hidden_size, dtype=dtype) weight = torch.randn(hidden_size, dtype=dtype) - return (x, weight, epsilon) + return x, weight, epsilon # Reductions in rms_norm accumulate rounding error at large shapes # (e.g. 32768x16384), causing a few elements out of millions to exceed # the default float16 tolerance. rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3) + + +@register_op(allow_inplace=True) +def fused_add_rms_norm( + x: Tensor, + x_residual: Tensor, + weight: Tensor | None, + epsilon: float, + variance_size: int | None = None, +) -> tuple[Tensor, Tensor]: + """Fused add and weighted root-mean-square layer normalization""" + orig_dtype = x.dtype + x = x.to(torch.float32) + x = x + x_residual.to(torch.float32) + x_residual = x.to(orig_dtype) + + x_var = x if variance_size is None else x[..., :variance_size] + variance = x_var.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + epsilon) + if weight is not None: + x = x.to(weight.dtype) * weight + return x.to(orig_dtype), x_residual + + +# fused_add_rms_norm has similar rounding error accumulation as rms_norm +fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3) + + +@fused_add_rms_norm.register_input_generator +def _fused_add_rms_norm_input_generator( + num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5 +) -> tuple: + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype) + weight = torch.randn(hidden_size, dtype=dtype) + return x, x_residual, weight, epsilon diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 14c2b87fbbdb..273bc58935b7 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -75,3 +75,72 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor direct_register_aiter_op( op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake ) + +rms_add_no_var_16bit_only = ( + lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None + and x.dtype in (torch.float16, torch.bfloat16) + and (weight is None or weight.dtype == x.dtype) +) +""" +AITER fused_add_rms_norm only supports 16-bit activations and no var_size override. +Requires weight dtype to match x dtype. +""" + + +@ir.ops.fused_add_rms_norm.register_impl( + "aiter", supports_args=rms_add_no_var_16bit_only, supported=AITER_SUPPORTED +) +def fused_add_rms_norm( + x: Tensor, + x_residual: Tensor, + weight: Tensor | None, + epsilon: float, + variance_size: int | None = None, +) -> tuple[Tensor, Tensor]: + assert variance_size is None + assert x.dtype in (torch.float16, torch.bfloat16) + if weight is None: + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) + return torch.ops.vllm_aiter.fused_add_rms_norm(x, x_residual, weight, epsilon) + + +def _rocm_aiter_rmsnorm2d_fwd_with_add_impl( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, +) -> tuple[torch.Tensor, torch.Tensor]: + from aiter import rmsnorm2d_fwd_with_add + + # TODO can out = x and residual_out = residual to save memory? + # Need to check if the kernel supports in-place residual output + # (if yes set mutates_args and inplace) + residual_out = torch.empty_like(residual) + out = torch.empty_like(x) + rmsnorm2d_fwd_with_add( + out, # output + x, # input + residual, # residual input + residual_out, # residual output + weight, + variance_epsilon, + ) + return out, residual_out + + +def _rocm_aiter_rmsnorm2d_fwd_with_add_fake( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, +) -> tuple[torch.Tensor, torch.Tensor]: + residual_out = torch.empty_like(residual) + out = torch.empty_like(x) + return out, residual_out + + +direct_register_aiter_op( + op_name="fused_add_rms_norm", + op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl, + fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake, +) diff --git a/vllm/kernels/oink_ops.py b/vllm/kernels/oink_ops.py index e8e3cb91f857..835cd062d037 100644 --- a/vllm/kernels/oink_ops.py +++ b/vllm/kernels/oink_ops.py @@ -1,6 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""This file registers Oink implementations for vLLM IR ops. + +vLLM does not depend on the external Oink repository/package. When an external +plugin registers torch.library.custom_op entrypoints under the `oink::` +namespace (e.g. via vLLM's general_plugins mechanism), these ops will be marked + as supported. To dispatch to those ops, set kernel_config.ir_op_priority. to oink. +Alternatively, `VLLM_USE_OINK_OPS=1` will add this to priority by default. +""" + import torch +from torch import Tensor from vllm import ir from vllm.platforms import current_platform @@ -15,7 +25,7 @@ def has_oink_op(name: str) -> bool: return OINK_AVAILABLE and hasattr(torch.ops.oink, name) -def _can_view_as_2d(x: torch.Tensor) -> bool: +def _can_view_as_2d(x: Tensor) -> bool: """Return True if x.view(-1, x.shape[-1]) is viewable (no copy).""" if x.dim() < 2: return False @@ -32,7 +42,7 @@ def _can_view_as_2d(x: torch.Tensor) -> bool: return True -def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool: +def _is_oink_stride_compatible_2d(x_2d: Tensor) -> bool: """Return True if x_2d meets Oink's pointer-path stride constraints.""" if x_2d.dim() != 2: return False @@ -67,11 +77,51 @@ def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool: "oink", supports_args=oink_rms_supported, supported=has_oink_op("rmsnorm") ) def rms_norm( - x: torch.Tensor, - weight: torch.Tensor | None, + x: Tensor, + weight: Tensor | None, epsilon: float, variance_size: int | None = None, -) -> torch.Tensor: +) -> Tensor: assert variance_size is None x_2d = x.view(-1, x.shape[-1]) return torch.ops.oink.rmsnorm(x_2d, weight, epsilon).view_as(x) + + +oink_add_rms_supported = ( + lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None + and weight is not None + and x.dim() >= 2 + and x.dtype == weight.dtype + and weight.is_contiguous() + and _can_view_as_2d(x) + and _is_oink_stride_compatible_2d(x.view(-1, x.shape[-1])) + # residual must have 2d-compatible strides and match x shape/dtype + and x.dtype == x_residual.dtype + and x.shape == x_residual.shape + and _can_view_as_2d(x_residual) + and _is_oink_stride_compatible_2d(x_residual.view(-1, x_residual.shape[-1])) +) +""" +Oink fused_add_rms_norm has the same constraints as rms_norm, +and residual must be 2d-like with compatible strides. +""" + + +@ir.ops.fused_add_rms_norm.register_impl( + "oink", + supports_args=oink_add_rms_supported, + supported=has_oink_op("fused_add_rms_norm"), + inplace=True, +) +def fused_add_rms_norm( + x: Tensor, + x_residual: Tensor, + weight: Tensor | None, + epsilon: float, + variance_size: int | None = None, +) -> tuple[Tensor, Tensor]: + assert variance_size is None + x_2d = x.view(-1, x.shape[-1]) + residual_2d = x_residual.view(-1, x_residual.shape[-1]) + torch.ops.oink.fused_add_rms_norm(x_2d, residual_2d, weight, epsilon) + return x, x_residual diff --git a/vllm/kernels/vllm_c.py b/vllm/kernels/vllm_c.py index 124b02e4e27a..5c602c39843b 100644 --- a/vllm/kernels/vllm_c.py +++ b/vllm/kernels/vllm_c.py @@ -31,3 +31,33 @@ def rms_norm( output = torch.empty(x.shape, device=x.device, dtype=x.dtype) torch.ops._C.rms_norm(output, x, weight, epsilon) return output + + +rms_add_no_var_size = ( + lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None + and (weight is None or weight.dtype == x.dtype) +) +"""vLLM Kernel does not support variance_size parameter and requires +matching input/weight dtype.""" + + +@ir.ops.fused_add_rms_norm.register_impl( + "vllm_c", + supports_args=rms_add_no_var_size, + supported=CUDA_ALIKE, + inplace=True, +) +def fused_add_rms_norm( + x: Tensor, + x_residual: Tensor, + weight: Tensor | None, + epsilon: float, + variance_size: int | None = None, +) -> tuple[Tensor, Tensor]: + if weight is None: + # Kernel requires weight tensor, pass ones + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) + + assert variance_size is None + torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) + return x, x_residual diff --git a/vllm/kernels/xpu_ops.py b/vllm/kernels/xpu_ops.py index c680c542c1df..5e7f90f70868 100644 --- a/vllm/kernels/xpu_ops.py +++ b/vllm/kernels/xpu_ops.py @@ -36,3 +36,31 @@ def rms_norm( output = torch.empty(x.shape, device=x.device, dtype=x.dtype) torch.ops._C.rms_norm(output, x, weight, epsilon) return output + + +rms_add_no_var_size = ( + lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None + and (weight is None or weight.dtype == x.dtype) +) + + +@ir.ops.fused_add_rms_norm.register_impl( + "xpu_kernels", + supports_args=rms_add_no_var_size, + supported=XPU_KERNELS_SUPPORTED, + inplace=True, +) +def fused_add_rms_norm( + x: Tensor, + x_residual: Tensor, + weight: Tensor | None, + epsilon: float, + variance_size: int | None = None, +) -> tuple[Tensor, Tensor]: + if weight is None: + # Kernel requires weight tensor, pass ones + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype) + + assert variance_size is None + torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) + return x, x_residual diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index d9184bb77070..a5d4e4db79fe 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -8,68 +8,15 @@ # Import kernels import vllm.kernels # noqa: F401 -from vllm import _oink_ops, envs, ir -from vllm._aiter_ops import rocm_aiter_ops +from vllm import envs, ir +from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp -from vllm.model_executor.layers.batch_invariant import ( - rms_norm_batch_invariant, -) -from vllm.platforms import current_platform +from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant logger = init_logger(__name__) -def _can_view_as_2d(x: torch.Tensor) -> bool: - """Return True if x.view(-1, x.shape[-1]) is viewable (no copy).""" - if x.dim() < 2: - return False - if x.dim() == 2: - return True - # For a view(-1, N) to be valid, all leading dims must be contiguous with - # respect to each other (size-1 dims are ignored). - for dim in range(x.dim() - 1): - # Strides for size-1 dims are irrelevant and can be arbitrary. - if x.size(dim + 1) != 1 and x.stride(dim) != x.stride(dim + 1) * x.size( - dim + 1 - ): - return False - return True - - -def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool: - """Return True if x_2d meets Oink's pointer-path stride constraints.""" - if x_2d.dim() != 2: - return False - if x_2d.stride(1) != 1: - return False - # Match Oink's vectorization constraint: stride(0) divisible by 256b. - if x_2d.dtype in (torch.float16, torch.bfloat16): - divby = 16 - elif x_2d.dtype == torch.float32: - divby = 8 - else: - return False - return (x_2d.stride(0) % divby) == 0 - - -def fused_add_rms_norm( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - variance_epsilon: float, -) -> tuple[torch.Tensor, torch.Tensor]: - from vllm import _custom_ops as ops - - ops.fused_add_rms_norm( - x, - residual, - weight, - variance_epsilon, - ) - return x, residual - - def poly_norm( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, variance_epsilon: float ) -> torch.Tensor: @@ -86,18 +33,6 @@ def poly_norm( return out -def dispatch_rocm_rmsnorm_func(dtype: torch.dtype, use_aiter: bool = False): - use_aiter = use_aiter and dtype in [ - torch.float16, - torch.bfloat16, - ] - - if use_aiter: - return rocm_aiter_ops.rms_norm2d_with_add - else: - return fused_add_rms_norm - - # --8<-- [start:rms_norm] @CustomOp.register("rms_norm") class RMSNorm(CustomOp): @@ -130,105 +65,19 @@ def __init__( if self.has_weight: self.weight = nn.Parameter(self.weight) - if current_platform.is_rocm(): - aiter_rmsnorm_enabled = rocm_aiter_ops.is_rmsnorm_enabled() - self.rocm_norm_func_with_add = dispatch_rocm_rmsnorm_func( - dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled - ) - - # Optional: enable Oink Blackwell RMSNorm custom-op fast path on - # compatible CUDA devices (e.g., SM100) when the external Oink - # package is available. This is detected once at construction time - # to avoid per-call device queries in the hot path. - self._use_oink_fused_add_rmsnorm = False - if ( - not current_platform.is_rocm() - and torch.cuda.is_available() - and bool(getattr(envs, "VLLM_USE_OINK_OPS", False)) - ): - # NOTE: vLLM disables custom ops by default when using Inductor. - # If this op is disabled, CustomOp will dispatch to forward_native, - # and the Oink path in forward_cuda will never run. - if getattr(self._forward_method, "__func__", None) is getattr( - self.forward_native, "__func__", None - ): - try: - from vllm.config import get_cached_compilation_config - - custom_ops = get_cached_compilation_config().custom_ops - except Exception: - custom_ops = [""] - logger.warning_once( - "VLLM_USE_OINK_OPS=1 but the `rms_norm` custom op is " - "disabled (CompilationConfig.custom_ops=%s). Enable it via " - "`compilation_config={'custom_ops': ['none', '+rms_norm']}` " - "(or `['all']`) to let vLLM call into torch.ops.oink.*.", - custom_ops, - ) - # Custom op disabled => forward_cuda won't run. Avoid doing any - # external Oink initialization work in this case. - else: - try: - device_index = torch.accelerator.current_device_index() - if _oink_ops.is_oink_available_for_device(device_index): - self._use_oink_fused_add_rmsnorm = ( - _oink_ops.has_fused_add_rms_norm() - ) - except Exception as e: - # If anything goes wrong (no Oink install, CPU-only env, etc.), - # silently fall back to the built-in RMSNorm path. - logger.warning_once( - "VLLM_USE_OINK_OPS=1 but failed to initialize Oink " - "RMSNorm; falling back to vLLM RMSNorm. Error: %s", - e, - ) - self._use_oink_fused_add_rmsnorm = False - - @staticmethod - def forward_static( - x: torch.Tensor, - variance_epsilon: float, - hidden_size: int, - orig_dtype: torch.dtype, - weight: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - variance_size_override: int | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """PyTorch-native implementation equivalent to forward().""" - x = x.to(torch.float32) - if residual is not None: - # residual promoted f16->f32 automatically, - # otherwise Inductor eliminates the casts to and from f16, - # increasing memory usage (and complicating pattern matching) - x = x + residual - residual = x.to(orig_dtype) - - if x.shape[-1] != hidden_size: - raise ValueError( - f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}" - ) - - if variance_size_override is None: - x_var = x - else: - if hidden_size < variance_size_override: - raise ValueError( - "Expected hidden_size to be at least " - f"{variance_size_override}, but found: {hidden_size}" - ) - - x_var = x[:, :, :variance_size_override] - - variance = x_var.pow(2).mean(dim=-1, keepdim=True) - - x = x * torch.rsqrt(variance + variance_epsilon) - x = x.to(orig_dtype) - if weight is not None: - x = x * weight - if residual is None: - return x - else: - return x, residual + # Do not pass identity weight to native implementation (causes issue on TPU). + # Other implementations require weight to be passed even if all ones. + # Cheat and predict if native will be dispatched to: + # 1) if native is first in priority list + # 2) if variance_size_override is given (only supported by native impl) + # TODO(luka): address weight passing inconsistency: + # https://github.com/vllm-project/vllm/issues/39370 + priority = get_current_vllm_config().kernel_config.ir_op_priority + var_override = self.variance_size_override is not None + native_rms_norm = priority.rms_norm[0] == "native" or var_override + native_add_rms_norm = priority.fused_add_rms_norm[0] == "native" or var_override + self.pass_weight = self.has_weight or not native_rms_norm + self.pass_weight_add = self.has_weight or not native_add_rms_norm def forward_native( self, @@ -237,106 +86,34 @@ def forward_native( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """PyTorch-native implementation equivalent to forward().""" if residual is None: - # TODO(luka): address the weight=None passing issue more generally return ir.ops.rms_norm( x, - self.weight.data if self.has_weight else None, + self.weight.data if self.pass_weight else None, + self.variance_epsilon, + self.variance_size_override, + ) + else: + return ir.ops.fused_add_rms_norm.maybe_inplace( + x, + residual, + self.weight.data if self.pass_weight_add else None, self.variance_epsilon, self.variance_size_override, ) - - return self.forward_static( - x, - self.variance_epsilon, - self.hidden_size, - x.dtype, - self.weight.data if self.has_weight else None, - residual, - self.variance_size_override, - ) def forward_cuda( self, x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if residual is None and not envs.VLLM_BATCH_INVARIANT: - return ir.ops.rms_norm( - x, self.weight.data, self.variance_epsilon, self.variance_size_override - ) - - if self.variance_size_override is not None: - return self.forward_native(x, residual) - - # Optional Oink SM100 fast path (fused residual-add + RMSNorm, in-place). - # This mirrors vLLM's fused_add_rms_norm semantics by mutating both - # `x` (normalized output) and `residual` (residual-out buffer). if ( - residual is not None - and getattr(self, "_use_oink_fused_add_rmsnorm", False) - and x.is_cuda - and residual.is_cuda - and x.shape == residual.shape - and x.dtype == residual.dtype - and x.dim() >= 2 - and self.has_weight - and not envs.VLLM_BATCH_INVARIANT - and self.weight.data.dtype == x.dtype - and self.weight.data.is_contiguous() + envs.VLLM_BATCH_INVARIANT + and residual is None + and self.variance_size_override is None ): - orig_shape = x.shape - hidden_size = orig_shape[-1] - if _can_view_as_2d(x) and _can_view_as_2d(residual): - x_2d = x.view(-1, hidden_size) - res_2d = residual.view(-1, hidden_size) - - # The Oink in-place pointer path supports the common vLLM - # layout where: - # - `x` may be strided/padded row-major (stride(1) == 1), and - # - `residual` is contiguous row-major ([M, N] with stride(0) == N). - # If these conditions are not met, fall back to vLLM's built-in - # fused kernel. - if ( - _is_oink_stride_compatible_2d(x_2d) - and _is_oink_stride_compatible_2d(res_2d) - and res_2d.is_contiguous() - ): - _oink_ops.fused_add_rms_norm_( - x_2d, - res_2d, - self.weight.data, - self.variance_epsilon, - ) - return x, residual - - if residual is not None: - return fused_add_rms_norm( - x, residual, self.weight.data, self.variance_epsilon - ) - else: - assert envs.VLLM_BATCH_INVARIANT return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon) - def forward_hip( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if residual is None and not envs.VLLM_BATCH_INVARIANT: - return ir.ops.rms_norm( - x, self.weight.data, self.variance_epsilon, self.variance_size_override - ) - - if self.variance_size_override is not None: - return self.forward_native(x, residual) - - if residual is not None: - return self.rocm_norm_func_with_add( - x, residual, self.weight.data, self.variance_epsilon - ) - else: - assert envs.VLLM_BATCH_INVARIANT - return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon) + return self.forward_native(x, residual) def forward_xpu( self, diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 4f9b9d7bf234..9f04bf11660a 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -578,7 +578,9 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf if envs.VLLM_USE_OINK_OPS: rms_norm = ["oink"] + default - return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm) + return IrOpPriorityConfig.with_default( + default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + ) # NVML utils diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 866b9ffd1a6d..7200a7698d6c 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -693,21 +693,11 @@ def get_device_total_memory(cls, device_id: int = 0) -> int: @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: from vllm._aiter_ops import rocm_aiter_ops - from vllm.config.compilation import CUDAGraphMode compilation_config = vllm_config.compilation_config - is_eager_execution = compilation_config.cudagraph_mode == CUDAGraphMode.NONE use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled() - use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled() use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled() use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - # Aiter rms norm perform best when CUDA Graph capture is enabled. - if ( - use_aiter_rms_norm - and not is_eager_execution - and "-rms_norm" not in compilation_config.custom_ops - ): - compilation_config.custom_ops.append("+rms_norm") if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops: compilation_config.custom_ops.append("+quant_fp8") @@ -939,7 +929,7 @@ def use_custom_op_collectives(cls) -> bool: def get_default_ir_op_priority( cls, vllm_config: "VllmConfig" ) -> "IrOpPriorityConfig": - from vllm.config.compilation import CompilationMode + from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.config.kernel import IrOpPriorityConfig # Native used by default when compiling, @@ -949,12 +939,10 @@ def get_default_ir_op_priority( using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE default = ["native"] if using_inductor else ["vllm_c", "native"] - # This (mostly) preserves previous CustomOp behavior - # Necessary on ROCm because it's common that users - # enable rms_norm to use the aiter kernel. + # Aiter rms norm perform best when CUDA Graph capture is enabled. # TODO(luka/TJ) remove env vars completely if ( - cc.is_custom_op_enabled("rms_norm") + cc.cudagraph_mode != CUDAGraphMode.NONE and envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_RMSNORM ): @@ -962,7 +950,9 @@ def get_default_ir_op_priority( else: rms_norm = default - return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm) + return IrOpPriorityConfig.with_default( + default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + ) @classmethod @with_amdsmi_context From c293ccc58ef6e1a0976a62f79f57bc045108073d Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Date: Fri, 1 May 2026 21:13:15 -0700 Subject: [PATCH 0012/1083] [ROCm][Bugfix] Fix init-time bias dtype cast when gate.out_dtype is None (#41405) Signed-off-by: Rita Brugarolas Brufau --- vllm/model_executor/models/deepseek_v2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 15913c418b05..26a1903fa7b9 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -351,8 +351,9 @@ def __init__( self.is_rocm_aiter_moe_enabled and self.gate.e_score_correction_bias is not None ): + gate_out_dtype = self.gate.out_dtype or self.gate.weight.dtype self.gate.e_score_correction_bias.data = ( - self.gate.e_score_correction_bias.data.to(self.gate.out_dtype) + self.gate.e_score_correction_bias.data.to(gate_out_dtype) ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: From ae3b4deb8a5987759d4732e67767146a46ee72ed Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sat, 2 May 2026 13:27:43 +0800 Subject: [PATCH 0013/1083] [Doc] Add Codex usage example (#41358) Signed-off-by: chaunceyjiang --- docs/serving/integrations/codex.md | 88 ++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 docs/serving/integrations/codex.md diff --git a/docs/serving/integrations/codex.md b/docs/serving/integrations/codex.md new file mode 100644 index 000000000000..48148acfd51f --- /dev/null +++ b/docs/serving/integrations/codex.md @@ -0,0 +1,88 @@ +# Codex + +[Codex](https://github.com/openai/codex) is OpenAI's official agentic coding tool that lives in your terminal. It can understand your codebase, edit files, run commands, and help you write code more efficiently. + +By pointing Codex at a vLLM server, you can use your own models as the backend instead of the OpenAI API. This is useful for: + +- Running fully local/private coding assistance +- Using open-weight models with tool calling capabilities +- Testing and developing with custom models + +## How It Works + +vLLM implements the OpenAI-Responses API, which is the same API that Codex uses to communicate with OpenAI's servers. By configuring Codex to point at your vLLM server, Codex sends its requests to vLLM instead of OpenAI. vLLM then translates these requests to work with your local model and returns responses in the format Codex expects. + +This means any model served by vLLM with proper tool calling support can act as a drop-in replacement for OpenAI models in Codex. + +## Requirements + +Codex requires a model with strong tool calling capabilities. The model must support the OpenAI-Responses tool calling API. See [Tool Calling](../../features/tool_calling.md) for details on enabling tool calling for your model. + +## Installation + +First, install Codex by following the [official installation guide](https://github.com/openai/codex). + +## Starting the vLLM Server + +Start vLLM with a tool-calling capable model - here's an example using `Qwen/Qwen3-27B`: + +```bash +vllm serve Qwen/Qwen3.6-27B --port 8000 --tensor-parallel-size 8 --max-model-len 262144 --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_coder + +``` + +For other models, you'll need to enable tool calling explicitly with `--enable-auto-tool-choice` and the right `--tool-call-parser`. Refer to the [Tool Calling documentation](../../features/tool_calling.md) for the correct flags for your model. + +## Configuring Codex + +Codex is configured via a TOML file located at `~/.codex/config.toml`. Create or edit this file to point Codex at your vLLM server: + +```toml +model = "my-model" +model_provider = "vllm" + +[model_providers.vllm] +name = "vLLM" +env_key = "VLLM_API_KEY" +base_url = "http://localhost:8000/v1" +wire_api = "responses" +``` + +The configuration fields: + +| Field | Description | +| ----- | ----------- | +| `model` | The model name to use. Must match the `--served-model-name` you passed to vLLM. | +| `model_provider` | Set to `"vllm"` to use your local vLLM server. | +| `[model_providers.vllm]` | Configuration section for the vLLM provider. | +| `name` | A display name for your vLLM provider. | +| `env_key` | The name of an environment variable that Codex will read for the API key. vLLM does not require authentication by default, so this can be any value. | +| `base_url` | The URL of your vLLM server's OpenAI-compatible API endpoint (default is `http://localhost:8000/v1`). | +| `wire_api` | The API style to use. Set to `"responses"` for the OpenAI Responses API | + +!!! tip + You can set the `env_key` to any dummy environment variable since vLLM doesn't require authentication by default: + ```bash + export VLLM_API_KEY=dummy + ``` + +!!! warning + When using the `responses` API, ensure your vLLM version supports the OpenAI Responses API. + +## Testing the Setup + +Once Codex is configured, launch it in your project directory: + +```bash +codex +``` + +Try a simple prompt to verify the connection, such as asking it to explain a file in your project. If the model responds correctly, your setup is working. You can now use Codex with your vLLM-served model for coding tasks. + +## Troubleshooting + +**Connection refused**: Ensure vLLM is running and accessible at the specified URL. Check that the port matches and that `base_url` includes the `/v1` path suffix. + +**Tool calls not working**: Verify that your model supports tool calling and that you've enabled it with the correct `--tool-call-parser` flag. See [Tool Calling](../../features/tool_calling.md). + +**Model not found**: Ensure the `model` field in `~/.codex/config.toml` matches the `--served-model-name` you passed to vLLM. From 8586369f617a964235d0d9d32d6ebb1076a4581d Mon Sep 17 00:00:00 2001 From: Matthew Santiago Date: Sat, 2 May 2026 01:22:14 -0500 Subject: [PATCH 0014/1083] Refactor Step3Text loading to use AutoWeightsLoader (#41492) Signed-off-by: Matthew Santiago --- vllm/model_executor/models/step3_text.py | 103 ++++++++++++----------- 1 file changed, 54 insertions(+), 49 deletions(-) diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index 8f08f6c60713..a0e7e16a9bbf 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -41,6 +41,7 @@ from .interfaces import SupportsPP from .utils import ( + AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -382,55 +383,6 @@ def forward( hidden_states, _ = self.norm(hidden_states, residual) return hidden_states - -class Step3TextForCausalLM(nn.Module, SupportsPP): - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__() - config = vllm_config.model_config.hf_config - - self.config = config - self.vllm_config = vllm_config - - self.model = Step3TextModel(vllm_config=vllm_config, prefix=prefix) - - if get_pp_group().is_last_rank: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - self.logits_processor = LogitsProcessor(config.vocab_size) - else: - self.lm_head = PPMissingLayer() - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ): - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: qkv_params_mapping = [ # (param_name, shard_name, relative_start_idx, relative_end_idx) @@ -553,3 +505,56 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + + +class Step3TextForCausalLM(nn.Module, SupportsPP): + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + config = vllm_config.model_config.hf_config + + self.config = config + self.vllm_config = vllm_config + + self.model = Step3TextModel(vllm_config=vllm_config, prefix=prefix) + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + else: + self.lm_head = PPMissingLayer() + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ): + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) From c3ad791e1a9ad9a2bff082ebbdf2b03749deef3e Mon Sep 17 00:00:00 2001 From: Hoang Nguyen <118159510+hnt2601@users.noreply.github.com> Date: Sat, 2 May 2026 13:34:59 +0700 Subject: [PATCH 0015/1083] [Bugfix][Gemma 4] Clamp soft-token estimate to max_soft_tokens (#40796) Signed-off-by: Hoang Nguyen <118159510+hnt2601@users.noreply.github.com> Signed-off-by: Isotr0py Co-authored-by: Claude Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py --- .../multimodal/processing/test_gemma4.py | 54 +++++++++++++++++++ vllm/model_executor/models/gemma4_mm.py | 9 +++- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index 808fab6a030f..bd1f2bb86779 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -12,6 +12,60 @@ GEMMA4_MODEL_ID = "google/gemma-4-E2B-it" +@pytest.mark.parametrize( + "image_width,image_height,max_soft_tokens", + [ + # Production repro: a 3x900 image (extreme aspect ratio) made the + # prompt-side estimator return 289 while the HF Gemma 4 image + # processor's vision tower output capped at 280, producing the + # "Attempted to assign 280 multimodal tokens to 289 placeholders" + # mismatch that crashed EngineCore. + (900, 3, 280), + (3, 900, 280), + # Same pathology should hold for the video-frame budget (70 tokens). + (900, 3, 70), + # And for any other supported budget. + (4000, 2, 1120), + ], +) +@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) +def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens( + model_id: str, + image_width: int, + image_height: int, + max_soft_tokens: int, +): + """Regression for the Gemma 3/4 multimodal crash. + + `_compute_num_soft_tokens` must never return a value larger than + `max_soft_tokens`. The HF Gemma 4 image processor clamps its vision + tower output to that value; if the prompt-side estimator returns more, + the prompt has more `image` placeholder tokens than the encoder will + fill, and `_merge_multimodal_embeddings` raises `ValueError` deep in + the model forward. + """ + ctx = build_model_context( + model_id, + mm_processor_kwargs={"do_pan_and_scan": True}, + limit_mm_per_prompt={"image": 1}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + num_soft_tokens = processor.info._compute_num_soft_tokens( + image_width=image_width, + image_height=image_height, + max_soft_tokens=max_soft_tokens, + ) + + assert num_soft_tokens <= max_soft_tokens, ( + f"_compute_num_soft_tokens returned {num_soft_tokens} for " + f"image_width={image_width}, image_height={image_height}, " + f"max_soft_tokens={max_soft_tokens} — exceeds the cap that the HF " + f"image processor enforces on its vision tower output. This is " + f"the placeholder/encoder count mismatch that crashes EngineCore." + ) + + @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) def test_limit_mm_per_prompt( image_assets: ImageTestAssets, diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index cdc54609a652..9b2c54e27354 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -265,7 +265,14 @@ def _compute_num_soft_tokens( target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit) target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit) num_patches = (target_h // patch_size) * (target_w // patch_size) - return num_patches // (pooling_kernel_size**2) + # Clamp to ``max_soft_tokens``: extreme aspect ratios (e.g. 3x900) + # cause the floor() above to round one dim up to ``unit`` while the + # other scales freely, which over-shoots ``max_patches``. The HF + # Gemma 4 image processor caps its vision-tower output at + # ``max_soft_tokens``, so without this clamp the prompt-side + # placeholder count exceeds the encoder output and + # ``_merge_multimodal_embeddings`` crashes. + return min(num_patches // (pooling_kernel_size**2), max_soft_tokens) def get_image_repl( self, From cfd2573f239f8e5be370b0148809e25ea5bb0a3e Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Sat, 2 May 2026 08:51:28 -0400 Subject: [PATCH 0016/1083] [Build] Switch CUDA 13.0 wheel builds to PyTorch manylinux_2_28 base (#41416) Signed-off-by: mgoin Co-authored-by: Claude --- .buildkite/release-pipeline.yaml | 4 +- docker/Dockerfile | 93 +++++++++++++++++++++++--------- docker/versions.json | 3 ++ 3 files changed, 72 insertions(+), 28 deletions(-) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 74227da45c71..cdb5b00d4143 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -37,7 +37,7 @@ steps: agents: queue: arm64_cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -76,7 +76,7 @@ steps: agents: queue: cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" diff --git a/docker/Dockerfile b/docker/Dockerfile index a6b291407713..fd0622e2416a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -41,6 +41,13 @@ ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 # Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels) ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION} +# OS family of BUILD_BASE_IMAGE. Controls package manager (apt vs dnf) and +# Python bootstrap. Set to "manylinux" alongside a manylinux build base such +# as pytorch/manylinux2_28-builder:cuda13.0 to produce wheels with a glibc +# 2.28 floor (matches PyTorch's own published wheels). Default stays on +# Ubuntu for backwards compatibility. +ARG BUILD_OS=ubuntu + # By parameterizing the Deadsnakes repository URL, we allow third-party to use # their own mirror. When doing so, we don't benefit from the transparent # installation of the GPG key of the PPA, as done by add-apt-repository, so we @@ -94,35 +101,64 @@ FROM ${BUILD_BASE_IMAGE} AS base ARG CUDA_VERSION ARG PYTHON_VERSION +ARG BUILD_OS ENV DEBIAN_FRONTEND=noninteractive -# Install system dependencies including build tools -RUN apt-get update -y \ - && apt-get install -y --no-install-recommends \ - ccache \ - software-properties-common \ - git \ - curl \ - sudo \ - python3-pip \ - libibverbs-dev \ - # Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519 - # as it was causing spam when compiling the CUTLASS kernels - gcc-10 \ - g++-10 \ - && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \ - # Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04 - # which ships cmake 3.28 and requires Development.SABIModule; silently skipped on - # Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA) - && (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \ - && rm -rf /var/lib/apt/lists/* \ - && curl -LsSf https://astral.sh/uv/install.sh | sh \ - && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \ +# Install system dependencies including build tools. +# The Ubuntu path uses apt + deadsnakes-via-uv for Python; the manylinux path +# (AlmaLinux 8, e.g. pytorch/manylinux2_28-builder) uses dnf and the Python +# interpreters pre-installed at /opt/python/cpXY-cpXY/. +RUN if [ "${BUILD_OS}" = "manylinux" ]; then \ + # rdma-core-devel provides libibverbs headers; ccache lives in EPEL, + # which the pytorch manylinux image already enables. git/curl/sudo + # are typically pre-installed but listed defensively. + dnf install -y --setopt=install_weak_deps=False \ + ccache \ + git \ + curl \ + sudo \ + rdma-core-devel \ + && dnf clean all \ + && rm -rf /var/cache/dnf; \ + else \ + apt-get update -y \ + && apt-get install -y --no-install-recommends \ + ccache \ + software-properties-common \ + git \ + curl \ + sudo \ + python3-pip \ + libibverbs-dev \ + # Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519 + # as it was causing spam when compiling the CUTLASS kernels + gcc-10 \ + g++-10 \ + && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \ + # Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04 + # which ships cmake 3.28 and requires Development.SABIModule; silently skipped on + # Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA) + && (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \ + && rm -rf /var/lib/apt/lists/*; \ + fi + +# Install uv and bootstrap /opt/venv. Both paths converge on /opt/venv so all +# downstream stages stay distro-agnostic. +RUN curl -LsSf https://astral.sh/uv/install.sh | sh \ + && if [ "${BUILD_OS}" = "manylinux" ]; then \ + # manylinux images ship Python at /opt/python/cpXY-cpXY/; point uv + # at the matching interpreter rather than letting it fetch one. + PYV_NODOT=$(echo ${PYTHON_VERSION} | tr -d '.') \ + && MANYLINUX_PY=/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin/python${PYTHON_VERSION} \ + && $HOME/.local/bin/uv venv /opt/venv --python "$MANYLINUX_PY"; \ + else \ + $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION}; \ + fi \ && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \ - && ln -s /opt/venv/bin/python3 /usr/bin/python3 \ - && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \ - && ln -s /opt/venv/bin/pip /usr/bin/pip \ + && ln -sf /opt/venv/bin/python3 /usr/bin/python3 \ + && ln -sf /opt/venv/bin/python3-config /usr/bin/python3-config \ + && ln -sf /opt/venv/bin/pip /usr/bin/pip \ && python3 --version && python3 -m pip --version # Activate virtual environment and add uv to PATH @@ -433,6 +469,7 @@ FROM base AS dev ARG PIP_INDEX_URL UV_INDEX_URL ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL ARG PYTORCH_CUDA_INDEX_BASE_URL +ARG BUILD_OS # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out # Reference: https://github.com/astral-sh/uv/pull/1694 @@ -442,7 +479,11 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE=copy # Install libnuma-dev, required by fastsafetensors (fixes #20384) -RUN apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/* +RUN if [ "${BUILD_OS}" = "manylinux" ]; then \ + dnf install -y numactl-devel && dnf clean all && rm -rf /var/cache/dnf; \ + else \ + apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*; \ + fi # We can specify the standard or nightly build of PyTorch diff --git a/docker/versions.json b/docker/versions.json index b6b555790d2a..75652823db0b 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -16,6 +16,9 @@ "FINAL_BASE_IMAGE": { "default": "nvidia/cuda:13.0.2-base-ubuntu22.04" }, + "BUILD_OS": { + "default": "ubuntu" + }, "GET_PIP_URL": { "default": "https://bootstrap.pypa.io/get-pip.py" }, From 0a9362d6ab88eed6fe7b52ed8424794cd492f888 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Sat, 2 May 2026 12:42:41 -0400 Subject: [PATCH 0017/1083] Revert "[Build] Make bundled DeepGEMM wheel portable across Python versions" (#41512) --- cmake/external_projects/deepgemm.cmake | 27 ++++---------------------- 1 file changed, 4 insertions(+), 23 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index b821b90ec8e9..0d7ea43fb7d0 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -59,26 +59,11 @@ if(DEEPGEMM_ARCHS) # Build the _C pybind11 extension from DeepGEMM's C++ source. # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime. # - # Free-threaded Python doesn't yet support the stable ABI, so skip USE_SABI - # there. (The other vLLM extensions get this guard for free via - # define_extension_target; this target uses raw Python_add_library.) - run_python(IS_FREETHREADED_PYTHON - "import sysconfig; print(1 if sysconfig.get_config_var(\"Py_GIL_DISABLED\") else 0)" - "Failed to determine whether interpreter is free-threaded") - if (NOT IS_FREETHREADED_PYTHON) - Python_add_library(_deep_gemm_C MODULE WITH_SOABI USE_SABI 3 - "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") - else() - Python_add_library(_deep_gemm_C MODULE WITH_SOABI - "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") - endif() + Python_add_library(_deep_gemm_C MODULE WITH_SOABI + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") # The pybind11 module name must be _C to match DeepGEMM's Python imports. - # Place the build artifact in a subdir so it doesn't collide with vLLM's own - # `_C.abi3.so` in the build tree (the install destination still differs). - set_target_properties(_deep_gemm_C PROPERTIES - OUTPUT_NAME "_C" - LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm") + set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C") target_compile_definitions(_deep_gemm_C PRIVATE "-DTORCH_EXTENSION_NAME=_C") @@ -90,15 +75,11 @@ if(DEEPGEMM_ARCHS) "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include" "${deepgemm_SOURCE_DIR}/third-party/fmt/include") - # Keep Stable ABI for the module, but *not* for CUDA/C++ files. - # This prevents Py_LIMITED_API from affecting nvcc and C++ compiles. target_compile_options(_deep_gemm_C PRIVATE $<$:-std=c++17> $<$:-O3> $<$:-Wno-psabi> - $<$:-Wno-deprecated-declarations> - $<$:-UPy_LIMITED_API> - $<$:-UPy_LIMITED_API>) + $<$:-Wno-deprecated-declarations>) # torch_python is required because DeepGEMM uses pybind11 type casters # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which From 4f7309fcc05d614b78477e2f0e24581fc09a2b3b Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Sat, 2 May 2026 18:23:59 -0400 Subject: [PATCH 0018/1083] [CI] Add ci-fetch-log.sh helper for Buildkite job logs (#41517) Signed-off-by: mgoin Co-authored-by: Claude Opus 4.7 (1M context) --- .buildkite/scripts/ci-fetch-log.sh | 55 ++++++++++++++++++++++++++++++ docs/contributing/ci/failures.md | 14 ++++++-- 2 files changed, 67 insertions(+), 2 deletions(-) create mode 100755 .buildkite/scripts/ci-fetch-log.sh diff --git a/.buildkite/scripts/ci-fetch-log.sh b/.buildkite/scripts/ci-fetch-log.sh new file mode 100755 index 000000000000..02798b56f4a9 --- /dev/null +++ b/.buildkite/scripts/ci-fetch-log.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# Usage: ./ci-fetch-log.sh [output_file] +# ./ci-fetch-log.sh [output_file] +# +# Downloads the raw log for a Buildkite job from the public, unauthenticated +# /organizations//pipelines//builds//jobs//download +# endpoint, then strips ANSI/timestamps via ci-clean-log.sh. +# +# Find and via: +# gh pr checks --repo vllm-project/vllm +# Each failing row's URL is .../builds/#. + +set -euo pipefail + +ORG="vllm" +PIPELINE="ci" + +usage() { + echo "Usage: $0 [output_file]" + echo " $0 [output_file]" + exit 1 +} + +if [ $# -lt 1 ]; then usage; fi + +if [[ "$1" == https://* ]]; then + BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p') + JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1) + OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}" +else + if [ $# -lt 2 ]; then usage; fi + BUILD="$1" + JOB="$2" + OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}" +fi + +if [ -z "$BUILD" ] || [ -z "$JOB" ]; then + echo "Could not parse build number or job UUID from: $1" >&2 + usage +fi + +COOKIES=$(mktemp) +trap 'rm -f "$COOKIES"' EXIT + +# Buildkite issues a session cookie on first hit; subsequent /download needs it. +curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \ + "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null + +curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \ + "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \ + -o "$OUT" + +bash "$(dirname "$0")/ci-clean-log.sh" "$OUT" + +echo "$OUT" diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md index dad04e75fbb6..a0038f461a04 100644 --- a/docs/contributing/ci/failures.md +++ b/docs/contributing/ci/failures.md @@ -60,9 +60,19 @@ the failure? ## Logs Wrangling -Download the full log file from Buildkite locally. +Download a job's log (no Buildkite login required): -Strip timestamps and colorization: +[.buildkite/scripts/ci-fetch-log.sh](../../../.buildkite/scripts/ci-fetch-log.sh) + +```bash +# Find the failing job. Each row's URL is .../builds/#: +gh pr checks --repo vllm-project/vllm + +# Download + strip timestamps/ANSI in one step: +.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/#" +``` + +To clean an already-downloaded log: [.buildkite/scripts/ci-clean-log.sh](../../../.buildkite/scripts/ci-clean-log.sh) From 1c607d7b2cd4fb572b919c6053f19d0577203495 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Sat, 2 May 2026 19:41:40 -0400 Subject: [PATCH 0019/1083] [DSV4] Guard megamoe flag with Pure TP (#41522) Signed-off-by: Yongye Zhu --- vllm/model_executor/models/deepseek_v4.py | 26 ++++++++++++++--------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v4.py b/vllm/model_executor/models/deepseek_v4.py index 0b762d50fe72..01aa922f3c29 100644 --- a/vllm/model_executor/models/deepseek_v4.py +++ b/vllm/model_executor/models/deepseek_v4.py @@ -715,12 +715,15 @@ def __init__( config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.prefix = prefix - if vllm_config.parallel_config.enable_expert_parallel: - self.use_mega_moe = ( - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." ) - else: - self.use_mega_moe = False self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.hidden_size = config.hidden_size @@ -1223,12 +1226,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config - if vllm_config.parallel_config.enable_expert_parallel: - self.use_mega_moe = ( - vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." ) - else: - self.use_mega_moe = False self.vocab_size = config.vocab_size self.hc_eps = config.hc_eps self.hc_mult = config.hc_mult From 856ec4804a4e236bdd796e0954094a532c7d19bc Mon Sep 17 00:00:00 2001 From: Roger Wang Date: Sat, 2 May 2026 18:32:09 -0700 Subject: [PATCH 0020/1083] [DSv4] Tune default value of `VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD` (#41526) Co-authored-by: Copilot --- vllm/envs.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index 8378c9762ae7..0955894754fb 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -245,7 +245,7 @@ VLLM_DEBUG_WORKSPACE: bool = False VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 - VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 4096 + VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" VLLM_USE_V2_MODEL_RUNNER: bool = False VLLM_LOG_MODEL_INSPECTION: bool = False @@ -1686,10 +1686,10 @@ def _get_or_set_default() -> str: # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device # and the cross-stream sync becomes pure overhead. Set to 0 to disable - # the multi-stream path entirely. Empirical crossover on B300 (148 SMs) - # is ~4096; B200 (132 SMs) is expected ~3072. + # the multi-stream path entirely. See #PR 41526 for the empirical result + # for the default value of 1024 tokens. "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int( - os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "4096") + os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "1024") ), # Format for saving torch.compile cache artifacts # - "binary": saves as binary file From 08834cc3ceb86d77231666a5e07847f741637355 Mon Sep 17 00:00:00 2001 From: Jinzhen Lin Date: Sun, 3 May 2026 09:36:03 +0800 Subject: [PATCH 0021/1083] [Quantization] add humming mxfp4 moe backend (#41083) Signed-off-by: Jinzhen Lin --- vllm/config/kernel.py | 2 + vllm/envs.py | 4 +- .../layers/fused_moe/fused_humming_moe.py | 119 ++++++---- vllm/model_executor/layers/fused_moe/layer.py | 3 - .../layers/fused_moe/oracle/mxfp4.py | 64 +++++- .../layers/quantization/humming.py | 96 +++----- .../layers/quantization/mxfp4.py | 4 + .../quantization/utils/humming_moe_utils.py | 35 --- .../quantization/utils/humming_utils.py | 214 ++++++++++++++++++ 9 files changed, 384 insertions(+), 157 deletions(-) delete mode 100644 vllm/model_executor/layers/quantization/utils/humming_moe_utils.py create mode 100644 vllm/model_executor/layers/quantization/utils/humming_utils.py diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index e5dcc20b5048..f7d5f19e2388 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -118,6 +118,7 @@ def with_default( "flashinfer_cutlass", "flashinfer_cutedsl", "marlin", + "humming", "aiter", "emulation", ] @@ -148,6 +149,7 @@ class KernelConfig: - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) - "marlin": Use Marlin kernels (weight-only quantization) + - "humming": Use Humming Mixed Precision kernels - "aiter": Use AMD AITer kernels (ROCm only) - "emulation": use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations. diff --git a/vllm/envs.py b/vllm/envs.py index 0955894754fb..b2db5a8112bc 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1231,8 +1231,8 @@ def _get_or_set_default() -> str: # if 1, force use indexed gemm # if 0, force use grouped gemm # if None, choose better gemm type automatically - "VLLM_HUMMING_MOE_GEMM_TYPE": lambda: maybe_convert_bool( - os.environ.get("VLLM_HUMMING_MOE_GEMM_TYPE", None) + "VLLM_HUMMING_MOE_GEMM_TYPE": lambda: os.environ.get( + "VLLM_HUMMING_MOE_GEMM_TYPE", None ), # Whether to use DeepEPLL kernels for NVFP4 quantization and dispatch method # only supported on Blackwell GPUs and with diff --git a/vllm/model_executor/layers/fused_moe/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/fused_humming_moe.py index 6a2417cd4d31..64d12f9558b4 100644 --- a/vllm/model_executor/layers/fused_moe/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_humming_moe.py @@ -4,7 +4,7 @@ import json import math -from typing import TYPE_CHECKING, Any +from typing import Any import torch from humming import dtypes @@ -16,7 +16,11 @@ from vllm.forward_context import get_forward_context 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 FusedMoEParallelConfig +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) @@ -34,21 +38,16 @@ from vllm.platforms import current_platform from vllm.v1.worker.workspace import current_workspace_manager -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.humming import HummingMoEMethod - - logger = init_logger(__name__) def get_humming_moe_gemm_type() -> str: env_gemm_type: str = envs.VLLM_HUMMING_MOE_GEMM_TYPE or "" env_gemm_type = env_gemm_type.lower() - if env_gemm_type in ["indexed", "grouped"]: + if env_gemm_type == "indexed": gemm_type = env_gemm_type - elif current_platform.has_device_capability(90): - # for device that supports TMA, use grouped gemm - gemm_type = "grouped" + elif env_gemm_type in ["grouped_contiguous", "grouped"]: + gemm_type = "grouped_contiguous" else: gemm_type = "indexed" @@ -60,49 +59,44 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): def __init__( self, layer: torch.nn.Module, - quant_method: "HummingMoEMethod", - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular | None = None, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, ): self.layer = layer self.num_experts = self.layer.num_experts self.global_num_experts = self.layer.global_num_experts self.init_humming_moe() - if prepare_finalize is not None: - max_num_tokens: int | None = None - num_dispatchers: int | None = None - if self.is_batched: - max_num_tokens = prepare_finalize.max_num_tokens_per_rank() - num_dispatchers = prepare_finalize.num_dispatchers() - - assert quant_method.moe_quant_config is not None - super().__init__( - moe_config=quant_method.moe, - quant_config=quant_method.moe_quant_config, - max_num_tokens=max_num_tokens, - num_dispatchers=num_dispatchers, - ) - else: - assert not self.is_batched + if self.is_batched(): + assert max_num_tokens is not None and num_dispatchers is not None + + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) def init_humming_moe(self): self.compute_config = { "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, - "gemm_type": self.humming_gemm_type.value, + "gemm_type": self.humming_gemm_type().value, } self.w13_tuning_config = HummingMethod.get_default_tuning_configs( layer=self.layer, use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM, use_batch_invariant=envs.VLLM_BATCH_INVARIANT, - gemm_type=self.humming_gemm_type, + gemm_type=self.humming_gemm_type(), sublayer_name="w13", ) self.w2_tuning_config = HummingMethod.get_default_tuning_configs( layer=self.layer, use_f16_accum=envs.VLLM_HUMMING_USE_F16_ACCUM, use_batch_invariant=envs.VLLM_BATCH_INVARIANT, - gemm_type=self.humming_gemm_type, + gemm_type=self.humming_gemm_type(), sublayer_name="w2", ) self.compute_config_str = json.dumps(self.compute_config) @@ -124,13 +118,13 @@ def estimate_local_valid_shape_m(self, topk_ids: torch.Tensor): global_num_experts = self.global_num_experts return math.ceil(global_valid_shape_m * num_experts / global_num_experts) - @property - def humming_gemm_type(self) -> HummingGemmType: + @staticmethod + def humming_gemm_type() -> HummingGemmType: raise NotImplementedError - @property - def is_batched(self) -> bool: - return self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts + @classmethod + def is_batched(cls) -> bool: + return cls.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts @staticmethod def _supports_quant_scheme( @@ -189,7 +183,7 @@ def moe_problem_size( assert w1.size(0) == num_experts assert w2.size(0) == num_experts - if not self.is_batched: + if not self.is_batched(): num_tokens = a1.size(0) assert topk_ids.size(0) == num_tokens else: @@ -201,7 +195,7 @@ def moe_problem_size( def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): num_experts = self.num_experts - N = self.layer.intermediate_size + N = self.layer.intermediate_size_per_partition K = self.layer.hidden_size assert isinstance(num_experts, int) assert isinstance(N, int) @@ -218,7 +212,7 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): # The output must be derived from workspace1. output_shape: tuple[int, ...] - if self.is_batched: + if self.is_batched(): max_num_tokens = self.max_num_tokens num_dispatchers = self.num_dispatchers assert max_num_tokens is not None and num_dispatchers is not None @@ -227,7 +221,7 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): output_shape = (num_experts, max_num_tokens * num_dispatchers, K) else: input_shape_m = M - if self.humming_gemm_type != HummingGemmType.INDEXED: + if self.humming_gemm_type() != HummingGemmType.INDEXED: input_shape_m = M * topk real_shape_m = M * topk output_shape = (M, K) @@ -262,7 +256,7 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): "dtype": torch_dtype_map[a_dtype], }, "down_output": { - "shape": output_shape if self.is_batched else (real_shape_m, K), + "shape": output_shape if self.is_batched() else (real_shape_m, K), "dtype": torch_dtype_map[c_dtype], }, "output": { @@ -288,7 +282,7 @@ def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): ] # batched moe use down_output as output - if not self.is_batched: + if not self.is_batched(): required_buffers.append("output") return buffer_metas, required_buffers @@ -308,7 +302,7 @@ def _workspace_shapes(self, M: int, topk: int, activation: MoEActivation): else: workspace2_nbytes = max(workspace2_nbytes, nbytes) - output_key = "down_output" if self.is_batched else "output" + output_key = "down_output" if self.is_batched() else "output" output_shape = buffer_metas[output_key]["shape"] return (workspace1_nbytes // 2,), (workspace2_nbytes // 2,), output_shape @@ -395,6 +389,33 @@ def main_apply( ): raise NotImplementedError + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: + supported = cls.activation_format() == activation_format + reason = "activation_format mismatched" + elif activation_format == mk.FusedMoEActivationFormat.Standard: + if cls.activation_format() != mk.FusedMoEActivationFormat.Standard: + supported = False + reason = "activation_format mismatched" + else: + assert hasattr(cls, "humming_gemm_type") + gemm_type = cls.humming_gemm_type().value.lower() + preferred_gemm_type = get_humming_moe_gemm_type().lower() + supported = preferred_gemm_type == gemm_type + reason = "preferred gemm type mismatched" + else: + supported = False + reason = "unsupported activation_format" + + return supported, None if supported else reason + class HummingIndexedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -404,8 +425,8 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - @property - def humming_gemm_type(self) -> HummingGemmType: + @staticmethod + def humming_gemm_type() -> HummingGemmType: return HummingGemmType.INDEXED def prepare_humming_moe_kwargs( @@ -526,8 +547,8 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - @property - def humming_gemm_type(self) -> HummingGemmType: + @staticmethod + def humming_gemm_type() -> HummingGemmType: return HummingGemmType.GROUPED_CONTIGUOUS def main_apply( @@ -619,8 +640,8 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.BatchedExperts - @property - def humming_gemm_type(self) -> HummingGemmType: + @staticmethod + def humming_gemm_type() -> HummingGemmType: return HummingGemmType.GROUPED_MASKED def main_apply( diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 577f9a986790..3de05cd93d36 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1103,9 +1103,6 @@ def weight_loader( return_success: bool = False, ) -> bool | None: quant_config_name = self.quant_config and self.quant_config.get_name() - if quant_config_name == "humming": - assert hasattr(self.quant_method, "weight_schema") - quant_config_name = self.quant_method.weight_schema.quant_method if quant_config_name == "gpt_oss_mxfp4": # (FIXME) for gpt-oss all experts are combined if "bias" in weight_name: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index c1423362d737..3f2aca277160 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -68,6 +68,8 @@ class Mxfp4MoeBackend(Enum): XPU = "XPU" # Emulation EMULATION = "EMULATION" + # Humming + HUMMING = "HUMMING" # Backends that share the same TRTLLM weight format @@ -130,6 +132,19 @@ def backend_to_kernel_cls( return [UnfusedOAITritonExperts] + elif backend == Mxfp4MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == Mxfp4MoeBackend.MARLIN: from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( MarlinExperts, @@ -177,6 +192,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> Mxfp4MoeBackend: "flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, "triton": Mxfp4MoeBackend.TRITON, "triton_unfused": Mxfp4MoeBackend.TRITON_UNFUSED, + "humming": Mxfp4MoeBackend.HUMMING, "marlin": Mxfp4MoeBackend.MARLIN, "aiter": Mxfp4MoeBackend.AITER, "xpu": Mxfp4MoeBackend.XPU, @@ -573,7 +589,21 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( sf_block_size = 32 # mxfp4 block size - if mxfp4_backend in ( + if mxfp4_backend == Mxfp4MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + prepare_humming_moe_layer, + ) + + prepare_humming_moe_layer(layer, {"quant_method": "gpt_oss_mxfp4"}) + return ( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + getattr(layer, "w13_bias", None), + getattr(layer, "w2_bias", None), + ) + elif mxfp4_backend in ( Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, ): @@ -970,6 +1000,21 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) + if mxfp4_backend == Mxfp4MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + prepare_humming_moe_layer, + ) + + prepare_humming_moe_layer(layer, {"quant_method": "mxfp4"}) + return ( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + getattr(layer, "w13_bias", None), + getattr(layer, "w2_bias", None), + ) + if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN): from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_moe_mxfp4_layer_for_marlin, @@ -1175,6 +1220,7 @@ def make_mxfp4_moe_quant_config( swiglu_limit: float | None = None, w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig | None: """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: @@ -1234,6 +1280,14 @@ def make_mxfp4_moe_quant_config( gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, ) + elif mxfp4_backend == Mxfp4MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.layer import FusedMoE + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, FusedMoE) + return get_humming_moe_quant_config(layer) else: return ocp_mx_moe_quant_config( quant_dtype="mxfp4", @@ -1254,6 +1308,7 @@ def make_mxfp4_moe_kernel( mxfp4_backend: Mxfp4MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, shared_experts: torch.nn.Module | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: """Create a FusedMoEKernel for the given MXFP4 backend.""" is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) @@ -1269,6 +1324,11 @@ def make_mxfp4_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if mxfp4_backend == Mxfp4MoeBackend.HUMMING: + assert layer is not None + extra_kwargs["layer"] = layer + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -1278,11 +1338,13 @@ def make_mxfp4_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index 59f9c2ee9b97..79a1057c6003 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -9,11 +9,9 @@ import torch from vllm import envs -from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, - FusedMoEQuantDesc, ) from vllm.model_executor.layers.fused_moe.layer import ( FusedMoE, @@ -32,7 +30,6 @@ QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.parameter import ( BasevLLMParameter, @@ -215,6 +212,15 @@ def from_config(cls, config: dict[str, Any]) -> "HummingConfig": def override_quantization_method( cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: + if user_quant == "humming" and hf_config is not None: + model_type = hf_config.model_type + quant_method = hf_quant_cfg.get("quant_method", None) + if model_type == "gpt_oss" and quant_method == "mxfp4": + msg = ( + "For gpt-oss model, use '--moe-backend humming' " + "instead of '--quantization humming'." + ) + raise ValueError(msg) return "humming" if user_quant == "humming" else None def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): @@ -299,8 +305,6 @@ def get_quant_config_for_layer( force_weight_schema = schema if weight_schema is not None: - if weight_schema.quant_method == "gpt_oss_mxfp4" and layer_type != "moe": - return None input_schema = None force_input_schema = None @@ -335,12 +339,6 @@ def get_quant_method( elif isinstance(layer, LinearBase): layer_type = "linear" - # TODO: remove this after humming moe backend is ready - quant_method = self.full_config.get("quant_method", None) - moe_activation = getattr(layer, "activation", None) - if quant_method == "mxfp4" and moe_activation == MoEActivation.SWIGLUOAI: - self.full_config["quan_method"] = "gpt_oss_mxfp4" - quant_config = self.get_quant_config_for_layer(prefix, layer_type) if quant_config is None: if isinstance(layer, FusedMoE): @@ -760,62 +758,18 @@ def create_weights( layer.register_buffer("locks", locks) def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: - self.process_weights_after_loading(layer) - - input_schema = self.input_schemas["w13"] - weight_schema = self.weight_schemas["w13"] - - a_dtype = input_schema.a_dtype - if a_dtype is None or a_dtype.num_bits == 16: - a_quant_desc = FusedMoEQuantDesc(dtype=None) - else: - shape = GroupShape(row=1, col=-1) - a_quant_desc = FusedMoEQuantDesc(dtype=str(a_dtype), shape=shape) - - weight_scale_group_size = weight_schema.weight_scale_group_size - weight_scale_group_size_n = weight_schema.weight_scale_group_size_n - weight_group_shape: tuple[int, ...] = () - if weight_scale_group_size_n > 1: - weight_group_shape = GroupShape( - row=weight_scale_group_size, - col=weight_scale_group_size_n, - ) - elif weight_scale_group_size == 0: - weight_group_shape = GroupShape(row=-1, col=1) - else: - weight_group_shape = GroupShape(row=weight_scale_group_size, col=1) - - w1_quant_desc = FusedMoEQuantDesc( - dtype=str(weight_schema.b_dtype), - shape=weight_group_shape, - scale=getattr(layer, "w13_weight_scale", None), - alpha_or_gscale=getattr(layer, "w13_global_scale", None), - zp=getattr(layer, "w13_zero_point", None), - bias=getattr(layer, "w13_bias", None), - ) - - w2_quant_desc = FusedMoEQuantDesc( - dtype=str(weight_schema.b_dtype), - shape=weight_group_shape, - scale=getattr(layer, "w2_weight_scale", None), - alpha_or_gscale=getattr(layer, "w2_global_scale", None), - zp=getattr(layer, "w2_zero_point", None), - bias=getattr(layer, "w2_bias", None), + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, ) - return FusedMoEQuantConfig( - _a1=a_quant_desc, - _a2=a_quant_desc, - _w1=w1_quant_desc, - _w2=w2_quant_desc, - ) + return get_humming_moe_quant_config(layer) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if getattr(self, "processed", False): return self.processed = True - self.weight_schemas = {} - self.input_schemas = {} + layer.weight_schemas = {} + layer.input_schemas = {} for sublayer_name, configs in layer.sublayer_configs.items(): input_schema = self.input_schema weight_schema = self.weight_schema @@ -858,8 +812,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: param = torch.nn.Parameter(tensor, requires_grad=False) setattr(layer, name, param) - self.weight_schemas[sublayer_name] = weight_schema - self.input_schemas[sublayer_name] = input_schema + layer.weight_schemas[sublayer_name] = weight_schema + layer.input_schemas[sublayer_name] = input_schema # force requant (origin quant setting -> fp16/bf16 -> new_quant setting) assert isinstance(weight_schema, HummingWeightSchema) @@ -913,10 +867,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # use moe modular experts: HummingIndexedExperts | HummingGroupedExperts + assert self.moe_quant_config is not None if get_humming_moe_gemm_type() == "indexed": - experts = HummingIndexedExperts(layer, self) + experts = HummingIndexedExperts(layer, self.moe, self.moe_quant_config) else: - experts = HummingGroupedExperts(layer, self) + experts = HummingGroupedExperts(layer, self.moe, self.moe_quant_config) self.experts = experts def select_gemm_impl( @@ -927,12 +882,19 @@ def select_gemm_impl( from vllm.model_executor.layers.fused_moe import modular_kernel as mk activation_format = prepare_finalize.activation_format + assert self.moe_quant_config is not None if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: - return BatchedHummingGroupedExperts(layer, self, prepare_finalize) + return BatchedHummingGroupedExperts( + layer=layer, + moe_config=self.moe, + quant_config=self.moe_quant_config, + max_num_tokens=prepare_finalize.max_num_tokens_per_rank(), + num_dispatchers=prepare_finalize.num_dispatchers(), + ) elif get_humming_moe_gemm_type() == "indexed": - return HummingIndexedExperts(layer, self, prepare_finalize) + return HummingIndexedExperts(layer, self.moe, self.moe_quant_config) else: - return HummingGroupedExperts(layer, self, prepare_finalize) + return HummingGroupedExperts(layer, self.moe, self.moe_quant_config) def apply( self, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 0a516831c4ec..2be77f2b8b82 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -366,6 +366,7 @@ def _setup_kernel( experts_cls=self.experts_cls, routing_tables=layer._maybe_init_expert_routing_tables(), shared_experts=layer.shared_experts, + layer=layer, ) def process_weights_after_loading(self, layer): @@ -404,6 +405,7 @@ def get_fused_moe_quant_config( gemm1_alpha=1.702, gemm1_beta=1.0, swiglu_limit=7.0, + layer=layer, ) def select_gemm_impl( @@ -692,6 +694,7 @@ def _setup_kernel( experts_cls=self.experts_cls, routing_tables=layer._maybe_init_expert_routing_tables(), shared_experts=layer.shared_experts, + layer=layer, ) def process_weights_after_loading(self, layer): @@ -729,6 +732,7 @@ def get_fused_moe_quant_config( w1_bias=w1_bias, w2_bias=w2_bias, swiglu_limit=swiglu_limit, + layer=layer, ) def select_gemm_impl( diff --git a/vllm/model_executor/layers/quantization/utils/humming_moe_utils.py b/vllm/model_executor/layers/quantization/utils/humming_moe_utils.py deleted file mode 100644 index 82788a0e76e8..000000000000 --- a/vllm/model_executor/layers/quantization/utils/humming_moe_utils.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch - -from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( - moe_align_block_size, -) - - -def humming_moe_align( - configs: list[int], - topk_ids: torch.Tensor, - num_experts: int, - expert_map: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - assert len(configs) > 0 and len(configs) % 3 == 0 - # NOTE: we choose moe_block_size based on - # num_tokens * top_k (= topk_ids.nelement()) - shape_m = topk_ids.nelement() - - for i in range(len(configs) // 3): - if shape_m > configs[i * 3] and shape_m <= configs[i * 3 + 1]: - block_size = configs[i * 3 + 2] - break - else: - raise ValueError(f"Could not find a matching block_size for shape_m={shape_m}") - - return moe_align_block_size( - topk_ids=topk_ids, - block_size=block_size, - num_experts=num_experts, - expert_map=expert_map, - pad_sorted_ids=False, - ignore_invalid_experts=True, - ) diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py new file mode 100644 index 000000000000..f8c10bdcae16 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import regex as re +import torch +from humming.layer import HummingInputSchema, HummingMethod +from humming.schema import BaseWeightSchema + +from vllm import envs +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + FusedMoEQuantDesc, +) +from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape + + +def humming_is_layer_skipped(config: dict[str, Any], prefix: str): + if not config: + return True + + keys = ["ignored_layers", "ignore", "modules_to_not_convert"] + ignored_layers: list[str] = [] + for key in keys: + ignored_layers = config.get(key, []) or [] + if not ignored_layers: + break + + if any(module_name in prefix for module_name in ignored_layers): + return True + if "lm_head" in prefix: + return True + + for regex in config.get("dynamic", {}): + if regex[:1] != "-": + continue + if re.match(regex[2:], prefix): + return True + + return False + + +def prepare_humming_layer(layer: LinearBase, quant_config: dict): + weight_schema = BaseWeightSchema.from_config(quant_config) + input_schema = HummingInputSchema() + + shape_k_stacks = [layer.input_size_per_partition] + shape_n_stacks = layer.output_partition_sizes + + # Step 1: convert weight to humming standard format + weight_schema, tensors = weight_schema.convert_humming( + tensors=layer.named_parameters(), + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=layer.params_dtype, + ) + + layer.weight_schema = weight_schema + + for name, _ in list(layer.named_parameters()): + delattr(layer, name) + + for name, tensor in tensors.items(): + param = torch.nn.Parameter(tensor, requires_grad=False) + setattr(layer, name, param) + + # Step 2: transform weight (humming standard format) for forwarding + HummingMethod.prepare_layer_meta( + layer=layer, + shape_n=layer.output_partition_sizes_sum, + shape_k=layer.input_size_per_partition, + weight_schema=weight_schema, + input_schema=input_schema, + pad_n_to_multiple=256, + pad_k_to_multiple=128, + has_bias=layer.has_bias, + torch_dtype=layer.param_dtype, + ) + + HummingMethod.transform_humming_layer(layer) + + +def prepare_humming_moe_layer(layer: FusedMoE, quant_config: dict): + weight_schema = BaseWeightSchema.from_config(quant_config) + input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {} + if humming_is_layer_skipped(input_quant_config, layer.layer_name): + input_schema = HummingInputSchema() + else: + # TODO: read input_quant_config from quant_config + input_schema = HummingInputSchema.from_config(input_quant_config) + + is_gated = layer.activation.is_gated + shape_config = { + "w13": ( + layer.moe_config.intermediate_size_per_partition * 2, + layer.moe_config.hidden_dim, + ), + "w2": ( + layer.moe_config.hidden_dim, + layer.moe_config.intermediate_size_per_partition * (1 if is_gated else 2), + ), + } + + layer.weight_schemas = {} + layer.input_schemas = {} + + for sublayer_name in shape_config: + # Step 1: convert weight to humming standard format + tensors: dict[str, torch.Tensor] = dict( + (key.removeprefix(sublayer_name + "_"), value) + for key, value in layer.state_dict().items() + if key.startswith(sublayer_name + "_") + ) + + shape_n, shape_k = shape_config[sublayer_name] + shape_n_stacks = [shape_n] + shape_k_stacks = [shape_k] + if sublayer_name == "w13": + shape_n_stacks = [shape_n // 2] * 2 + + weight_schema_new, tensors = weight_schema.convert_humming( + tensors=tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + num_experts=layer.local_num_experts, + param_dtype=layer.params_dtype, + ) + + layer.weight_schemas[sublayer_name] = weight_schema_new + layer.input_schemas[sublayer_name] = input_schema + + for name, _ in list(layer.named_parameters()): + if not name.startswith(sublayer_name + "_"): + continue + delattr(layer, name) + + for name, tensor in tensors.items(): + name = f"{sublayer_name}_{name}" + param = torch.nn.Parameter(tensor, requires_grad=False) + setattr(layer, name, param) + + # Step 2: transform weight (humming standard format) for forwarding + HummingMethod.prepare_layer_meta( + layer=layer, + shape_n=shape_n, + shape_k=shape_k, + pad_n_to_multiple=256, + pad_k_to_multiple=128, + input_schema=input_schema, + weight_schema=weight_schema_new, + has_bias=layer.moe_config.has_bias, + num_experts=layer.num_experts, + torch_dtype=layer.params_dtype, + sublayer_name=sublayer_name, + ) + + HummingMethod.transform_humming_layer(layer, sublayer_name=sublayer_name) + + if not hasattr(layer, "locks"): + device = layer.w13_weight.device + locks = torch.zeros(1024, dtype=torch.int32, device=device) + layer.register_buffer("locks", locks) + + +def get_humming_moe_quant_config(layer: FusedMoE): + input_schema = layer.input_schemas["w13"] + weight_schema = layer.weight_schemas["w13"] + + a_dtype = input_schema.a_dtype + if a_dtype is None or a_dtype.num_bits == 16: + a_quant_desc = FusedMoEQuantDesc(dtype=None) + else: + shape = GroupShape(row=1, col=-1) + a_quant_desc = FusedMoEQuantDesc(dtype=str(a_dtype), shape=shape) + + weight_scale_group_size = weight_schema.weight_scale_group_size + weight_scale_group_size_n = weight_schema.weight_scale_group_size_n + weight_group_shape: tuple[int, ...] = () + if weight_scale_group_size_n > 1: + weight_group_shape = GroupShape( + row=weight_scale_group_size, + col=weight_scale_group_size_n, + ) + elif weight_scale_group_size == 0: + weight_group_shape = GroupShape(row=-1, col=1) + else: + weight_group_shape = GroupShape(row=weight_scale_group_size, col=1) + + w1_quant_desc = FusedMoEQuantDesc( + dtype=str(weight_schema.b_dtype), + shape=weight_group_shape, + scale=getattr(layer, "w13_weight_scale", None), + alpha_or_gscale=getattr(layer, "w13_global_scale", None), + zp=getattr(layer, "w13_zero_point", None), + bias=getattr(layer, "w13_bias", None), + ) + + w2_quant_desc = FusedMoEQuantDesc( + dtype=str(weight_schema.b_dtype), + shape=weight_group_shape, + scale=getattr(layer, "w2_weight_scale", None), + alpha_or_gscale=getattr(layer, "w2_global_scale", None), + zp=getattr(layer, "w2_zero_point", None), + bias=getattr(layer, "w2_bias", None), + ) + + return FusedMoEQuantConfig( + _a1=a_quant_desc, + _a2=a_quant_desc, + _w1=w1_quant_desc, + _w2=w2_quant_desc, + ) From e6ff3e9c83a6520c3793f4e0511ac8591a07c243 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 2 May 2026 21:06:30 -0700 Subject: [PATCH 0022/1083] [MRV2] Add shutdown() method (#41297) Signed-off-by: Woosuk Kwon --- vllm/v1/worker/gpu/model_runner.py | 19 +++++++++++++++++++ vllm/v1/worker/gpu/shutdown.py | 20 ++++++++++++++++++++ 2 files changed, 39 insertions(+) create mode 100644 vllm/v1/worker/gpu/shutdown.py diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 98d889cdbb88..bf882a8af311 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -91,6 +91,7 @@ from vllm.v1.worker.gpu.sample.output import SamplerOutput from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.shutdown import free_before_shutdown from vllm.v1.worker.gpu.spec_decode import init_speculator from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, @@ -1339,6 +1340,24 @@ def postprocess_pool(self, input_batch: InputBatch) -> None: input_batch.num_scheduled_tokens ) + def shutdown(self) -> None: + """Release GPU tensors (model weights, KV caches, workspace) so that + memory is reclaimable when running in the same process.""" + torch.accelerator.synchronize() + if hasattr(self, "kv_caches"): + self.kv_caches.clear() + if hasattr(self, "attn_groups"): + self.attn_groups.clear() + if hasattr(self, "kv_cache_config"): + del self.kv_cache_config + free_before_shutdown(self.vllm_config) + if hasattr(self, "model"): + del self.model + + gc.collect() + torch.accelerator.empty_cache() + logger.debug("Cleaned up model weights, KV caches, and workspace") + ########### EPLB methods start ########### @property def eplb_state(self): diff --git a/vllm/v1/worker/gpu/shutdown.py b/vllm/v1/worker/gpu/shutdown.py new file mode 100644 index 000000000000..830083962347 --- /dev/null +++ b/vllm/v1/worker/gpu/shutdown.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.config import VllmConfig +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def free_before_shutdown(vllm_config: VllmConfig) -> None: + from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT + from vllm.v1.worker.workspace import reset_workspace_manager + + cache_config = vllm_config.cache_config + cache_config.num_gpu_blocks = None + + compilation_config = vllm_config.compilation_config + compilation_config.static_forward_context.clear() + + _ROPE_DICT.clear() + reset_workspace_manager() From 54dc64d5d399d960b2f0bf7a00f1a90d16ca6178 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 3 May 2026 07:47:55 -0500 Subject: [PATCH 0023/1083] [Doc] Add Qwen3-30B-A3B-Thinking-2507-FP8 to batch invariance verified models (#41513) Signed-off-by: Taneem Ibrahim --- docs/features/batch_invariance.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 804cd905e3b1..b23631484508 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -105,7 +105,7 @@ Batch invariance has been tested and verified on the following models: - **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1` - **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ` -- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct` +- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8` - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct` - **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct` - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` From c51df43005726a09c6eb7348e8c1b00501c70a8e Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sun, 3 May 2026 12:19:59 -0400 Subject: [PATCH 0024/1083] Disable flashinfer autotune temporarily due to correctness issues (#41524) Signed-off-by: wzhao18 --- vllm/config/vllm.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0146ee4c144a..88e6660e2161 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -209,7 +209,9 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "use_inductor_graph_partition": False, }, "kernel_config": { - "enable_flashinfer_autotune": True, + # Disabled for now due to correctness issues: + # https://github.com/flashinfer-ai/flashinfer/issues/3197 + "enable_flashinfer_autotune": False, }, } OPTIMIZATION_LEVEL_02 = { @@ -229,7 +231,9 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "use_inductor_graph_partition": False, }, "kernel_config": { - "enable_flashinfer_autotune": True, + # Disabled for now due to correctness issues: + # https://github.com/flashinfer-ai/flashinfer/issues/3197 + "enable_flashinfer_autotune": False, }, } OPTIMIZATION_LEVEL_03 = { From cb03fee32b5c191b3ac5a248e47cd58a66c75591 Mon Sep 17 00:00:00 2001 From: tomeras91 <57313761+tomeras91@users.noreply.github.com> Date: Sun, 3 May 2026 23:00:41 +0300 Subject: [PATCH 0025/1083] [Bugfix][Ray] Fix RayExecutorV2 actor name collision with DP > 1 (#40398) Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> --- vllm/v1/engine/utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 7b0f00d14c8a..1f0b9bbb19d5 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -403,6 +403,11 @@ def __init__( range(dp_size), local_dp_ranks, placement_groups ): dp_vllm_config = copy.deepcopy(vllm_config) + if dp_size > 1: + # Append the DP rank to instance_id so that per-engine + # identifiers (e.g. Ray actor names in RayExecutorV2) are + # unique across DP replicas. + dp_vllm_config.instance_id = f"{dp_vllm_config.instance_id}_dp{index}" dp_vllm_config.parallel_config.placement_group = pg local_client = index < local_engine_count From db9a84e0cd0e17ab693467ff4a71103abd4b77bf Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Sun, 3 May 2026 14:30:04 -0600 Subject: [PATCH 0026/1083] [Bugfix] Fix FP8 Bias Loading (#41424) Signed-off-by: Alex Brooks --- .../model_loader/test_reload.py | 28 +++++++++++++++++++ .../model_loader/reload/meta.py | 2 +- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 6e3e2d63e144..cf3553bd57de 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -59,6 +59,34 @@ def test_reload_lifecycle(): assert tensor.__dict__ == materialized_tensor.__dict__ +def test_materialize_layer_preserves_non_meta_tensors(): + """Ensure that materialize_layer does not overwrite non meta tensors.""" + layer = torch.nn.Linear(2, 3, bias=True) + + # Create a non meta bias tensor and meta weight, which can happen with FP8 + bias_values = torch.ones(3) + layer.bias.data.copy_(bias_values) + layer.weight = torch.nn.Parameter(layer.weight.data.to("meta")) + + assert layer.weight.is_meta + assert not layer.bias.is_meta + + # materialize the layer weights after the bias is initialized + info = LayerReloadingInfo( + restore_metadata=({}, {}), + restore_device=torch.device("cpu"), + ) + materialize_layer(layer, info) + + # Ensure the weight materialized off meta + assert not layer.weight.is_meta + assert layer.weight.device.type == "cpu" + + # Ensure that the bias is (still) not meta and values are unchanged + assert not layer.bias.is_meta + assert torch.equal(layer.bias.data, bias_values) + + def test_model_cleanup(dist_init, default_vllm_config): layer = QKVParallelLinear(2, 3, 4) assert layer.weight.weight_loader.__self__ is layer diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 91fce6f57b3e..baa2081d58b2 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -102,7 +102,7 @@ def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo): with info.restore_device: for name, tensor in get_layer_tensors(layer).items(): - if name not in SKIP_TENSORS: + if name not in SKIP_TENSORS and tensor.is_meta: setattr(layer, name, materialize_meta_tensor(tensor)) From 66dfee7121dfcbdfcce04ce92c117e5e10b25b14 Mon Sep 17 00:00:00 2001 From: David Oy <58150256+the-david-oy@users.noreply.github.com> Date: Sun, 3 May 2026 16:52:18 -0700 Subject: [PATCH 0027/1083] [Bugfix] Fix degenerate KV cache stride causing TMA cudaErrorIllegalInstruction (#40737) Signed-off-by: David Oy Signed-off-by: David Oy <58150256+the-david-oy@users.noreply.github.com> Signed-off-by: David Oy Co-authored-by: David Oy Co-authored-by: Claude Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> --- .../test_kv_head_stride_canonicalization.py | 162 ++++++++++++++++++ vllm/utils/cpu_resource_utils.py | 2 +- vllm/utils/torch_utils.py | 26 +++ vllm/v1/attention/backends/flash_attn.py | 24 ++- .../attention/backends/flash_attn_diffkv.py | 25 ++- vllm/v1/attention/backends/flashinfer.py | 50 ++++-- 6 files changed, 267 insertions(+), 22 deletions(-) create mode 100644 tests/v1/attention/test_kv_head_stride_canonicalization.py diff --git a/tests/v1/attention/test_kv_head_stride_canonicalization.py b/tests/v1/attention/test_kv_head_stride_canonicalization.py new file mode 100644 index 000000000000..635f46390cfc --- /dev/null +++ b/tests/v1/attention/test_kv_head_stride_canonicalization.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for canonicalize_singleton_dim_strides. + +Background +---------- +When num_kv_heads_per_rank == 1 (e.g. Qwen3.5-397B with TP=8 → 1 KV head +per rank), PyTorch's is_contiguous() returns True for *any* stride on the +size-1 dimension. The KV cache allocator can therefore produce a tensor +where that singleton dim has stride = 1 element (2 bytes for bf16) instead +of the canonical product-of-remaining-dims value. + +CUDA TMA (used by FlashInfer XQA SM90 and Flash-Attention 3/4 on H100+) +requires all non-outermost strides to be multiples of 16 bytes. A 2-byte +stride triggers cudaErrorIllegalInstruction. + +canonicalize_singleton_dim_strides() patches degenerate strides on all +size-1 dimensions via torch.as_strided — zero-copy. + +The degenerate stride manifests at different positions in different backends: +- FlashInfer: stride(-3) after kv_cache.permute() → shape [..., 1, B, D] +- FlashAttention: stride(-2) after kv_cache.unbind(0) → shape [N, B, 1, D] +""" + +import torch + +from vllm.utils.torch_utils import canonicalize_singleton_dim_strides + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _inject_degenerate_stride(t: torch.Tensor, dim: int) -> torch.Tensor: + """Return a view of t with a degenerate (stride=1) on a size-1 dim.""" + assert t.shape[dim] == 1, f"dim {dim} must have size 1" + strides = list(t.stride()) + strides[dim] = 1 # inject the bug + return t.as_strided(t.shape, strides) + + +# --------------------------------------------------------------------------- +# Tests: canonicalize_singleton_dim_strides +# --------------------------------------------------------------------------- + + +class TestCanonicalizeSingletonDimStrides: + def test_flashinfer_layout_dim_neg3(self): + """FlashInfer path: degenerate stride at dim -3 (num_kv_heads).""" + # Shape after permute: [num_blocks, 2, num_kv_heads, block_size, head_size] + num_blocks, block_size, head_size = 64, 16, 128 + t = torch.zeros(num_blocks, 2, 1, block_size, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + + assert t_deg.stride(-3) == 1 # confirm degenerate + assert t_deg.is_contiguous() # PyTorch doesn't notice + + fixed = canonicalize_singleton_dim_strides(t_deg) + + assert fixed.stride(-3) == block_size * head_size # canonical = 2048 + assert fixed.stride(-2) == head_size # inner dims unchanged + assert fixed.stride(-1) == 1 + + def test_flash_attn_layout_dim_neg2(self): + """FlashAttention path: degenerate stride at dim -2 (num_kv_heads).""" + # Shape after unbind(0): [num_blocks, block_size, num_kv_heads, head_size] + num_blocks, block_size, head_size = 64, 16, 128 + t = torch.zeros(num_blocks, block_size, 1, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-2) + + assert t_deg.stride(-2) == 1 + assert t_deg.is_contiguous() + + fixed = canonicalize_singleton_dim_strides(t_deg) + + assert fixed.stride(-2) == head_size # canonical = 128 + assert fixed.stride(-1) == 1 + + def test_canonical_strides_returned_as_is(self): + """No degenerate strides → same object returned (no copy, no new view).""" + t = torch.zeros(64, 2, 1, 16, 128, dtype=torch.bfloat16) + result = canonicalize_singleton_dim_strides(t) + assert result is t + + def test_multi_kv_heads_unchanged(self): + """num_kv_heads > 1 → strides are already canonical → unchanged.""" + t = torch.zeros(16, 2, 4, 16, 128, dtype=torch.bfloat16) + original_strides = t.stride() + result = canonicalize_singleton_dim_strides(t) + assert result.stride() == original_strides + + def test_data_pointer_preserved(self): + """Fix is zero-copy: same underlying storage.""" + t = torch.zeros(8, 2, 1, 16, 128, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + fixed = canonicalize_singleton_dim_strides(t_deg) + assert fixed.data_ptr() == t_deg.data_ptr() + assert fixed.storage_offset() == t_deg.storage_offset() + + def test_multiple_singleton_dims(self): + """All size-1 dims with degenerate strides are fixed.""" + # Shape: [1, 1, 8, 32] — two size-1 dims + t = torch.zeros(1, 1, 8, 32, dtype=torch.float16) + # Both size-1 dims get degenerate strides + t_deg = t.as_strided(t.shape, (1, 1, 32, 1)) # both leading dims = 1 + + fixed = canonicalize_singleton_dim_strides(t_deg) + + assert fixed.stride(0) == 1 * 8 * 32 # canonical: 256 + assert fixed.stride(1) == 1 * 8 * 32 # canonical: 256 (same since size-1) + assert fixed.stride(2) == 32 + assert fixed.stride(3) == 1 + + def test_various_shapes_flashinfer(self): + """Correctness across different block_size / head_size for FlashInfer layout.""" + for block_size, head_size in [(16, 64), (16, 128), (32, 128), (16, 256)]: + t = torch.zeros(8, 2, 1, block_size, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + fixed = canonicalize_singleton_dim_strides(t_deg) + assert fixed.stride(-3) == block_size * head_size, ( + f"Failed for block_size={block_size}, head_size={head_size}: " + f"got stride(-3)={fixed.stride(-3)}" + ) + + def test_various_shapes_flash_attn(self): + """Correctness across different shapes for FlashAttention layout.""" + for block_size, head_size in [(16, 64), (16, 128), (32, 128)]: + t = torch.zeros(8, block_size, 1, head_size, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-2) + fixed = canonicalize_singleton_dim_strides(t_deg) + assert fixed.stride(-2) == head_size, ( + f"Failed for block_size={block_size}, head_size={head_size}: " + f"got stride(-2)={fixed.stride(-2)}" + ) + + def test_tma_alignment_satisfied_after_fix_bf16(self): + """After fix, all strides meet 16-byte TMA alignment for bf16.""" + t = torch.zeros(64, 2, 1, 16, 128, dtype=torch.bfloat16) + t_deg = _inject_degenerate_stride(t, dim=-3) + fixed = canonicalize_singleton_dim_strides(t_deg) + + element_size = fixed.element_size() # 2 bytes for bf16 + for i, s in enumerate(fixed.stride()): + assert (s * element_size) % 16 == 0 or i == len(fixed.stride()) - 1, ( + f"dim {i} stride {s} * {element_size} bytes not 16-byte aligned" + ) + + def test_non_contiguous_outer_dims_preserved(self): + """Outer (non-size-1) non-contiguous strides are left unchanged.""" + # Simulate cross-layer unified allocation: num_blocks stride is non-canonical + # but the inner dims should be fixed. + base = torch.zeros(200, 2, 1, 16, 128, dtype=torch.bfloat16) + # Slice every 2nd block → non-canonical outer stride + t_sliced = base[::2] # shape [100, 2, 1, 16, 128], stride[0] = 2*canonical + t_deg = _inject_degenerate_stride(t_sliced, dim=-3) + + fixed = canonicalize_singleton_dim_strides(t_deg) + + # Outer stride should be unchanged (not a size-1 dim) + assert fixed.stride(0) == t_sliced.stride(0) + # Inner degenerate stride should be fixed + assert fixed.stride(-3) == 16 * 128 diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index bbf554d0ccdd..25c299a0c0c1 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -125,7 +125,7 @@ def get_allowed_cpu_list() -> list[LogicalCPUInfo]: if platform.system() == "Darwin": return cpu_list - global_allowed_cpu_id_list = os.sched_getaffinity(0) + global_allowed_cpu_id_list = os.sched_getaffinity(0) # type: ignore[attr-defined] logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list] return logical_cpu_list diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 1eb9306ed4b1..798c136fc239 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -110,6 +110,32 @@ def is_strictly_contiguous(t: torch.Tensor) -> bool: return True +def canonicalize_singleton_dim_strides(t: torch.Tensor) -> torch.Tensor: + """Fix degenerate strides on size=1 dimensions for CUDA TMA compatibility. + + PyTorch allows any stride on a size=1 dim (is_contiguous() is always True + there), so a size=1 dim may have stride=1 (2 bytes for bf16) instead of + the canonical product(shape[i+1:]). CUDA TMA on H100+ requires all + non-outermost strides to be ≥16-byte aligned; stride=1 triggers + cudaErrorIllegalInstruction. Zero-copy: patches stride metadata only via + as_strided; returns t unchanged if all size=1 strides are already canonical. + """ + if 1 not in t.shape: + return t + strides = list(t.stride()) + shape = t.shape + prev_stride = 1 + changed = False + for i in range(len(shape) - 1, -1, -1): + if shape[i] == 1 and strides[i] != prev_stride: + strides[i] = prev_stride + changed = True + prev_stride = strides[i] * shape[i] + if not changed: + return t + return t.as_strided(t.shape, strides) + + @contextlib.contextmanager def set_default_torch_dtype(dtype: torch.dtype): """Sets the default torch dtype to the given dtype.""" diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 1c9ff3f79e43..e73954ee7478 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -11,7 +11,10 @@ from vllm.model_executor.layers.attention import Attention from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionImpl, @@ -747,6 +750,23 @@ def forward( # For decoder and cross-attention, use KV cache as before key_cache, value_cache = kv_cache.unbind(0) + # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). + # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. + # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides. + fixed_k = canonicalize_singleton_dim_strides(key_cache) + fixed_v = canonicalize_singleton_dim_strides(value_cache) + if fixed_k is not key_cache or fixed_v is not value_cache: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashAttention): " + "shape=%s, key strides before=%s after=%s, " + "value strides before=%s after=%s", + key_cache.shape, + key_cache.stride(), + fixed_k.stride(), + value_cache.stride(), + fixed_v.stride(), + ) + key_cache, value_cache = fixed_k, fixed_v if is_quantized_kv_cache(self.kv_cache_dtype): # queries are quantized in the attention layer @@ -861,6 +881,8 @@ def do_kv_cache_update( # we use direct Q, K, V tensors without caching return + # Scatter write into the KV cache using slot_mapping indices. + # No TMA kernel is invoked here, so stride canonicalization is not needed. key_cache, value_cache = kv_cache.unbind(0) # Reshape the input keys and values and store them in the cache. diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index d18054769711..82a9f07a4e59 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -4,7 +4,11 @@ import torch -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.logger import init_logger +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import AttentionType from vllm.v1.attention.backends.fa_utils import ( get_flash_attn_version, @@ -25,6 +29,8 @@ cascade_attention, ) +logger = init_logger(__name__) + class FlashAttentionDiffKVBackend(FlashAttentionBackend): # Default to 128 for this backend @@ -204,6 +210,23 @@ def forward( # Different head_size for K and V key_cache = kv_cache[..., : self.head_size] value_cache = kv_cache[..., self.head_size :] + # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP). + # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment. + # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides. + fixed_k = canonicalize_singleton_dim_strides(key_cache) + fixed_v = canonicalize_singleton_dim_strides(value_cache) + if fixed_k is not key_cache or fixed_v is not value_cache: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashAttentionDiffKV): " + "shape=%s, key strides before=%s after=%s, " + "value strides before=%s after=%s", + key_cache.shape, + key_cache.stride(), + fixed_k.stride(), + value_cache.stride(), + fixed_v.stride(), + ) + key_cache, value_cache = fixed_k, fixed_v if is_quantized_kv_cache(self.kv_cache_dtype): # queries are quantized in the attention layer diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 8f5cb6206bd0..2de61a2b1f28 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -43,6 +43,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, is_quantized_kv_cache, is_strictly_contiguous, nvfp4_kv_cache_full_dim, @@ -1479,6 +1480,21 @@ def forward( stride_order = FlashInferBackend.get_kv_cache_stride_order() kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous + # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1 + # with TP=8). PyTorch permits non-canonical strides on size-1 dims; + # CUDA TMA requires ≥16-byte alignment on all non-outermost strides. + # canonicalize_singleton_dim_strides patches metadata via as_strided — + # zero-copy. See vllm.utils.torch_utils. + fixed = canonicalize_singleton_dim_strides(kv_cache_permute) + if fixed is not kv_cache_permute: + logger.debug( + "Canonicalized degenerate KV cache strides (FlashInfer): " + "shape=%s, strides before=%s, strides after=%s", + kv_cache_permute.shape, + kv_cache_permute.stride(), + fixed.stride(), + ) + kv_cache_permute = fixed # For NVFP4, the kv_cache last dim is full_dim (data + scale packed). # Split into correctly-strided data and scale views. @@ -1568,10 +1584,11 @@ def forward( else: assert isinstance(attn_metadata.prefill, TRTLLMPrefill) # prefill_query may be non-contiguous or have degenerate strides - # First ensure memory contiguity, then fix degenerate strides - # with reshape. contiguous() alone doesn't fix degenerate - # strides when a dimension has size 1. - prefill_query = prefill_query.contiguous().reshape(prefill_query.shape) + # on size=1 dims. contiguous() ensures memory layout; then + # canonicalize_singleton_dim_strides fixes any remaining + # degenerate strides on size=1 dims for TMA alignment. + prefill_query = prefill_query.contiguous() + prefill_query = canonicalize_singleton_dim_strides(prefill_query) workspace_buffer = _get_trtllm_gen_workspace_buffer() block_tables_prefill = attn_metadata.prefill.block_tables seq_lens_prefill = attn_metadata.prefill.seq_lens @@ -1621,11 +1638,9 @@ def forward( # with fp8 kv cache, we can construct a mock block # and mock kv cache with BF16 KV involved in the prefill # - # The inner (block_size, head_size) dims must be - # contiguous; outer dims may have non-canonical strides - # (e.g. cross-layer unified allocation). - # Degenerate strides on outer dims break TMA descriptors - # (see flashinfer-ai/flashinfer#2232). + kv_cache_permute = canonicalize_singleton_dim_strides( + kv_cache_permute + ) kv_strides = kv_cache_permute.stride() assert ( kv_strides[-1] == 1 @@ -1732,12 +1747,13 @@ def forward( if needs_fp8_out: output[:num_decode_tokens].copy_(out_decode.to(output.dtype)) else: - # decode_query may be non-contiguous or have degenerate strides assert isinstance(attn_metadata.decode, TRTLLMDecode) - # First ensure memory contiguity, then fix degenerate strides - # with reshape. contiguous() alone doesn't fix degenerate - # strides when a dimension has size 1. - decode_query = decode_query.contiguous().reshape(decode_query.shape) + # decode_query may be non-contiguous or have degenerate strides + # on size=1 dims. contiguous() ensures memory layout; then + # canonicalize_singleton_dim_strides fixes any remaining + # degenerate strides on size=1 dims for TMA alignment. + decode_query = decode_query.contiguous() + decode_query = canonicalize_singleton_dim_strides(decode_query) workspace_buffer = _get_trtllm_gen_workspace_buffer() block_tables_decode = attn_metadata.decode.block_tables seq_lens_decode = attn_metadata.decode.seq_lens @@ -1748,11 +1764,7 @@ def forward( assert is_strictly_contiguous(workspace_buffer) assert is_strictly_contiguous(block_tables_decode) assert is_strictly_contiguous(seq_lens_decode) - # kv_cache outer dims may be non-contiguous (e.g. - # cross-layer unified allocation), but inner dims - # (block_size, head_size) must be contiguous and - # strides must be canonical to avoid TMA descriptor - # failures (see flashinfer-ai/flashinfer#2232). + kv_cache_permute = canonicalize_singleton_dim_strides(kv_cache_permute) kv_strides = kv_cache_permute.stride() assert ( kv_strides[-1] == 1 and kv_strides[-2] == kv_cache_permute.shape[-1] From 894a02500b16fceb77d92110d02c9b5127de0e5f Mon Sep 17 00:00:00 2001 From: Akim Tsvigun Date: Mon, 4 May 2026 02:39:10 +0200 Subject: [PATCH 0028/1083] [Bench] Forward --seed to CustomDataset and CustomMMDataset shuffle (#40788) Signed-off-by: akimtsvigun --- tests/benchmarks/test_custom_dataset_seed.py | 77 ++++++++++++++++++++ vllm/benchmarks/datasets/datasets.py | 8 +- 2 files changed, 83 insertions(+), 2 deletions(-) create mode 100644 tests/benchmarks/test_custom_dataset_seed.py diff --git a/tests/benchmarks/test_custom_dataset_seed.py b/tests/benchmarks/test_custom_dataset_seed.py new file mode 100644 index 000000000000..dac87e6e6d98 --- /dev/null +++ b/tests/benchmarks/test_custom_dataset_seed.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import argparse +import json +from pathlib import Path + +import pytest +from transformers import AutoTokenizer, PreTrainedTokenizerBase + +from vllm.benchmarks.datasets import get_samples + + +@pytest.fixture(scope="session") +def hf_tokenizer() -> PreTrainedTokenizerBase: + return AutoTokenizer.from_pretrained("gpt2") + + +def _write_jsonl(path: Path, n_rows: int) -> None: + with path.open("w") as f: + for i in range(n_rows): + f.write(json.dumps({"prompt": f"row {i}: unique prompt content."}) + "\n") + + +def _args_for_custom(dataset_path: str, seed: int) -> argparse.Namespace: + return argparse.Namespace( + dataset_name="custom", + dataset_path=dataset_path, + disable_shuffle=False, + num_prompts=30, + custom_output_len=32, + skip_chat_template=True, + no_oversample=False, + seed=seed, + request_id_prefix="", + ) + + +@pytest.mark.benchmark +def test_custom_dataset_seed_propagates( + hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path +) -> None: + """--seed must control the CustomDataset shuffle used by get_samples. + + Without the fix, CustomDataset was instantiated without random_seed, + so its load-time shuffle always used DEFAULT_SEED=0 regardless of + args.seed, causing every run with --dataset-name custom to pick the + same subset of rows from a larger file. + """ + jsonl = tmp_path / "data.jsonl" + _write_jsonl(jsonl, n_rows=60) + + samples_a = get_samples(_args_for_custom(str(jsonl), seed=0), hf_tokenizer) + samples_b = get_samples(_args_for_custom(str(jsonl), seed=42), hf_tokenizer) + + prompts_a = {s.prompt for s in samples_a} + prompts_b = {s.prompt for s in samples_b} + + assert len(prompts_a) == 30 + assert len(prompts_b) == 30 + assert prompts_a != prompts_b + + +@pytest.mark.benchmark +def test_custom_dataset_same_seed_is_deterministic( + hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path +) -> None: + """Same --seed must yield the same CustomDataset subset.""" + jsonl = tmp_path / "data.jsonl" + _write_jsonl(jsonl, n_rows=60) + + samples_a = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer) + samples_b = get_samples(_args_for_custom(str(jsonl), seed=7), hf_tokenizer) + + prompts_a = [s.prompt for s in samples_a] + prompts_b = [s.prompt for s in samples_b] + + assert prompts_a == prompts_b diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 419275d2e6ae..b032c0a0d613 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -1803,7 +1803,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: if args.dataset_name == "custom": dataset = CustomDataset( - dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle + dataset_path=args.dataset_path, + disable_shuffle=args.disable_shuffle, + random_seed=args.seed, ) input_requests = dataset.sample( num_requests=args.num_prompts, @@ -1816,7 +1818,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: elif args.dataset_name == "custom_mm": dataset = CustomMMDataset( - dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle + dataset_path=args.dataset_path, + disable_shuffle=args.disable_shuffle, + random_seed=args.seed, ) input_requests = dataset.sample( num_requests=args.num_prompts, From 67058ca326ac566db1aac6fc20ff0c43510422f8 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 3 May 2026 22:11:22 -0500 Subject: [PATCH 0029/1083] [CI] Clean up remote servers on pytest parent exit (#41570) Signed-off-by: Andreas Karatzas --- tests/utils.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 80 insertions(+), 3 deletions(-) diff --git a/tests/utils.py b/tests/utils.py index e4b6a6ff6e70..d2d07937c34e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import atexit import contextlib import copy import functools @@ -134,6 +135,11 @@ class RemoteVLLMServer: """ DUMMY_API_KEY = "token-abc123" # vLLM's OpenAI server does not need API key + _active_servers: set["RemoteVLLMServer"] = set() + _active_servers_lock = threading.RLock() + _cleanup_hooks_registered = False + _signal_hooks_registered = False + _previous_signal_handlers: dict[int, Any] = {} proc: subprocess.Popen def _create_cli_subcommand(self): @@ -209,6 +215,7 @@ def __init__( ) self._pre_download_model(model, args) + self._shutdown_complete = False # Record GPU memory before server start so we know what # "released" looks like. @@ -221,6 +228,7 @@ def __init__( ) self._start_server(model, vllm_serve_args, env_dict) + self._register_active_server() max_wait_seconds = max_wait_seconds or 480 try: self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds) @@ -246,8 +254,70 @@ def _shutdown(self) -> None: (when the server fails to start). Must be safe to call even if the process is already dead. """ - self._terminate_process_tree() - self._wait_for_gpu_memory_release() + if self._shutdown_complete: + return + + self._shutdown_complete = True + try: + self._terminate_process_tree() + self._wait_for_gpu_memory_release() + finally: + self._unregister_active_server() + + @classmethod + def _ensure_cleanup_hooks_registered(cls) -> None: + """Register process-exit cleanup for detached server subprocesses.""" + root_cls = RemoteVLLMServer + with root_cls._active_servers_lock: + if not root_cls._cleanup_hooks_registered: + atexit.register(root_cls._shutdown_active_servers) + root_cls._cleanup_hooks_registered = True + + if ( + threading.current_thread() is threading.main_thread() + and not root_cls._signal_hooks_registered + ): + for signum in (signal.SIGTERM, signal.SIGINT): + root_cls._previous_signal_handlers[signum] = signal.getsignal( + signum + ) + signal.signal(signum, root_cls._handle_parent_signal) + root_cls._signal_hooks_registered = True + + def _register_active_server(self) -> None: + """Track this server so parent-process exits still clean it up.""" + RemoteVLLMServer._ensure_cleanup_hooks_registered() + with RemoteVLLMServer._active_servers_lock: + RemoteVLLMServer._active_servers.add(self) + + def _unregister_active_server(self) -> None: + with RemoteVLLMServer._active_servers_lock: + RemoteVLLMServer._active_servers.discard(self) + + @classmethod + def _shutdown_active_servers(cls) -> None: + """Best-effort shutdown for all live RemoteVLLMServer instances.""" + with cls._active_servers_lock: + servers = list(cls._active_servers) + + for server in servers: + with contextlib.suppress(Exception): + server._shutdown() + + @classmethod + def _handle_parent_signal(cls, signum, frame) -> None: + """Clean up detached servers before letting the signal terminate pytest.""" + cls._shutdown_active_servers() + + previous_handler = cls._previous_signal_handlers.get(signum, signal.SIG_DFL) + if callable(previous_handler): + previous_handler(signum, frame) + elif previous_handler == signal.SIG_IGN: + return + elif signum == signal.SIGINT: + raise KeyboardInterrupt + else: + raise SystemExit(128 + signum) def _terminate_process_tree(self) -> None: """Kill the server process tree without waiting for GPU memory release. @@ -315,6 +385,9 @@ def shutdown_many(cls, servers: Sequence["RemoteVLLMServer"]) -> None: if not servers: return + for server in servers: + server._shutdown_complete = True + threads = [ threading.Thread( target=s._terminate_process_tree, @@ -339,7 +412,11 @@ def shutdown_many(cls, servers: Sequence["RemoteVLLMServer"]) -> None: else s._pre_server_gpu_memory ), ) - earliest._wait_for_gpu_memory_release() + try: + earliest._wait_for_gpu_memory_release() + finally: + for server in servers: + server._unregister_active_server() def _kill_process_group_survivors( self, pgid: int | None, timeout: float = 15.0 From c103c02a1a97819f1f83c03c60dcc532293f520f Mon Sep 17 00:00:00 2001 From: Fang Han Date: Sun, 3 May 2026 21:19:52 -0700 Subject: [PATCH 0030/1083] [Transformers v5] Vendor HCXVisionConfig for compatibility (#38447) Signed-off-by: Fang Han Co-authored-by: Claude Opus 4.7 (1M context) --- tests/models/registry.py | 7 -- vllm/transformers_utils/config.py | 6 +- vllm/transformers_utils/configs/__init__.py | 2 + .../transformers_utils/configs/hyperclovax.py | 73 +++++++++++++++++++ 4 files changed, 79 insertions(+), 9 deletions(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 21d3a50ce996..ab5869cd3fda 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -946,13 +946,6 @@ def check_available_online( "HCXVisionForCausalLM": _HfExamplesInfo( "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B", trust_remote_code=True, - max_transformers_version="4.57", - transformers_version_reason={ - "vllm": ( - "Custom config cannot be loaded with Transformers " - "v5 because `text_config` is not always set" - ) - }, ), "HCXVisionV2ForCausalLM": _HfExamplesInfo( "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B", diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 2f00178ba6ef..c95df9c1077c 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -101,6 +101,7 @@ def __getitem__(self, key): fireredlid="FireRedLIDConfig", funaudiochat="FunAudioChatConfig", granite4_vision="Granite4VisionConfig", + hyperclovax_vlm="HCXVisionConfig", hunyuan_vl="HunYuanVLConfig", hy_v3="HYV3Config", isaac="IsaacConfig", @@ -217,8 +218,9 @@ def parse( ) else: if model_type in _CONFIG_REGISTRY: - # Register the config class to AutoConfig to ensure it's used in future - # calls to `from_pretrained` + # Register the config class to AutoConfig to ensure it's used + # in future calls to `from_pretrained` (e.g. from + # AutoTokenizer or AutoProcessor). config_class = _CONFIG_REGISTRY[model_type] config_class.model_type = model_type AutoConfig.register(model_type, config_class, exist_ok=True) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 44abe32c916f..99f099adc786 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -37,6 +37,7 @@ "HunYuanVLConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLTextConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLVisionConfig": "vllm.transformers_utils.configs.hunyuan_vl", + "HCXVisionConfig": "vllm.transformers_utils.configs.hyperclovax", "HYV3Config": "vllm.transformers_utils.configs.hy_v3", "HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax", "IsaacConfig": "vllm.transformers_utils.configs.isaac", @@ -104,6 +105,7 @@ "HunYuanVLConfig", "HunYuanVLTextConfig", "HunYuanVLVisionConfig", + "HCXVisionConfig", "HYV3Config", "HyperCLOVAXConfig", "IsaacConfig", diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py index 9fa823743d66..839b8ed50b08 100644 --- a/vllm/transformers_utils/configs/hyperclovax.py +++ b/vllm/transformers_utils/configs/hyperclovax.py @@ -17,6 +17,7 @@ # limitations under the License. """HyperCLOVA X model configuration.""" +from transformers import AutoConfig from transformers.configuration_utils import PretrainedConfig @@ -275,3 +276,75 @@ def __init__( auto_map=auto_map, **kwargs, ) + + +class HCXVisionConfig(PretrainedConfig): + """Vendored HyperCLOVAX Vision config with transformers v5 fix. + + The original remote code config does not handle empty initialization + (text_config=None), which breaks transformers v5's @strict validation. + + TODO: Remove this class once HyperCLOVAX is upstreamed to transformers. + Tracking PR: https://github.com/huggingface/transformers/pull/44956 + """ + + model_type = "hyperclovax_vlm" + keys_to_ignore_at_inference = ["past_key_values"] + + text_config_attribute_map = { + "n_embd": "hidden_size", + "n_positions": "max_position_embeddings", + "n_head": "num_attention_heads", + "n_layer": "num_hidden_layers", + } + + def __init__( + self, + text_config=None, + vision_config=None, + use_nth_layer=-2, + img_start_id=100009, + decoder_max_length=4096, + anyres=False, + unpad=False, + max_num_grids=-1, + num_queries_vis_abstractor=-1, + ignore_index=-100, + proj_pos_emb=True, + proj_prenorm=False, + use_1x1_grid=False, + **kwargs, + ): + for key, val in self.text_config_attribute_map.items(): + if text_config is not None and key in text_config: + text_config[val] = text_config.pop(key) + + self.text_config = None + if text_config is not None: + _text_config = AutoConfig.for_model(text_config["model_type"]) + self.text_config = _text_config.from_dict(text_config) + self.hidden_size = self.text_config.hidden_size + + self.vision_config = None + if vision_config is not None: + _vision_config = AutoConfig.for_model( + vision_config["model_type"]) + self.vision_config = _vision_config.from_dict(vision_config) + + self.use_nth_layer = use_nth_layer + self.decoder_max_length = decoder_max_length + self.anyres = anyres + self.unpad = unpad + self.max_num_grids = max_num_grids + self.num_queries_vis_abstractor = num_queries_vis_abstractor + self.img_start_id = img_start_id + self.ignore_index = ignore_index + self.proj_pos_emb = proj_pos_emb + self.proj_prenorm = proj_prenorm + self.use_1x1_grid = use_1x1_grid + super().__init__(**kwargs) + + def get_text_config(self, decoder=False): + if self.text_config is not None: + return self.text_config + return self From 01d4d1ad375dc5854779c593eee093bcebb0cada Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 3 May 2026 23:33:29 -0500 Subject: [PATCH 0031/1083] [ROCm][CI] Align spec decode logprob test prefill settings (#41335) Signed-off-by: Andreas Karatzas --- tests/v1/sample/test_logprobs.py | 35 +++++++++++++++----------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 28fb2931b229..460e0d685649 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -33,11 +33,10 @@ SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT # On ROCm, floating-point reductions in attention and GEMM kernels are -# non-associative and sensitive to batch geometry. The ref LLM (no spec -# decode, default scheduling) and the spec-decode LLM (chunked prefill, -# different effective batch sizes) follow different reduction orders, -# producing numerically divergent logprobs that get misattributed to -# spec-decode incorrectness. +# non-associative and sensitive to batch geometry. If the ref LLM and +# spec-decode LLM use different scheduling or batch geometry, they can +# follow different reduction orders and produce numerically divergent +# logprobs that get misattributed to spec-decode incorrectness. # # Force LLM instances into an identical, deterministic execution # mode so the test isolates spec-decode correctness only: @@ -1086,18 +1085,25 @@ def test_spec_decode_logprobs( ) max_model_len = 256 - - # Run base LLM. - ref_llm = LLM( - model=model_name, + llm_kwargs = dict( max_logprobs=5, max_model_len=max_model_len, seed=42, logprobs_mode=logprobs_mode, gpu_memory_utilization=0.4, + # Force the same prefill chunking for both the base model and + # spec decode model so the comparison isolates spec decode. + enable_chunked_prefill=True, + max_num_batched_tokens=32, enable_prefix_caching=False, **ROCM_DETERMINISM_KWARGS, ) + + # Run base LLM. + ref_llm = LLM( + model=model_name, + **llm_kwargs, + ) ref_results = ref_llm.generate( [prompt, prompt], [sampling_params, penalty_sampling_params] ) @@ -1117,16 +1123,7 @@ def test_spec_decode_logprobs( spec_llm = LLM( model_name, speculative_config=spec_config_with_len, - max_logprobs=5, - max_model_len=max_model_len, - seed=42, - logprobs_mode=logprobs_mode, - gpu_memory_utilization=0.4, - # Force prefill chunking - enable_chunked_prefill=True, - max_num_batched_tokens=32, - enable_prefix_caching=False, - **ROCM_DETERMINISM_KWARGS, + **llm_kwargs, ) spec_results = spec_llm.generate( [prompt, prompt], [sampling_params, penalty_sampling_params] From 6ec9bbec384b14401f901af189754f7d8a6754e2 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 4 May 2026 00:22:42 -0500 Subject: [PATCH 0032/1083] [CI] Stabilize cpu offload compressed tensors test (#41102) Signed-off-by: Andreas Karatzas --- tests/quantization/test_cpu_offload.py | 1 + tests/utils.py | 71 +++++++++++++++++--------- 2 files changed, 47 insertions(+), 25 deletions(-) diff --git a/tests/quantization/test_cpu_offload.py b/tests/quantization/test_cpu_offload.py index 3b58614e58d4..151b5d97ddf3 100644 --- a/tests/quantization/test_cpu_offload.py +++ b/tests/quantization/test_cpu_offload.py @@ -70,4 +70,5 @@ def test_cpu_offload_compressed_tensors(monkeypatch): ["--enforce_eager"], ["--enforce_eager", "--cpu-offload-gb", "1"], max_wait_seconds=480, + include_seeded_sampling=False, ) diff --git a/tests/utils.py b/tests/utils.py index d2d07937c34e..41202aa19481 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -782,6 +782,7 @@ def _test_completion( model: str, prompt: str, token_ids: list[int], + include_seeded_sampling: bool = True, ): results = [] @@ -816,33 +817,40 @@ def _test_completion( } ) - # test seeded random sampling - completion = client.completions.create( - model=model, prompt=prompt, max_tokens=5, seed=33, temperature=1.0 - ) + if include_seeded_sampling: + # test seeded random sampling + completion = client.completions.create( + model=model, prompt=prompt, max_tokens=5, seed=33, temperature=1.0 + ) - results.append( - { - "test": "seeded_sampling", - "text": completion.choices[0].text, - "finish_reason": completion.choices[0].finish_reason, - "usage": completion.usage, - } - ) + results.append( + { + "test": "seeded_sampling", + "text": completion.choices[0].text, + "finish_reason": completion.choices[0].finish_reason, + "usage": completion.usage, + } + ) - # test seeded random sampling with multiple prompts - completion = client.completions.create( - model=model, prompt=[prompt, prompt], max_tokens=5, seed=33, temperature=1.0 - ) + # test seeded random sampling with multiple prompts + completion = client.completions.create( + model=model, + prompt=[prompt, prompt], + max_tokens=5, + seed=33, + temperature=1.0, + ) - results.append( - { - "test": "seeded_sampling", - "text": [choice.text for choice in completion.choices], - "finish_reason": [choice.finish_reason for choice in completion.choices], - "usage": completion.usage, - } - ) + results.append( + { + "test": "seeded_sampling", + "text": [choice.text for choice in completion.choices], + "finish_reason": [ + choice.finish_reason for choice in completion.choices + ], + "usage": completion.usage, + } + ) # test simple list batch = client.completions.create( @@ -1037,6 +1045,7 @@ def compare_two_settings( *, method: str = "generate", max_wait_seconds: float | None = None, + include_seeded_sampling: bool = True, ) -> None: """ Launch API server with two different sets of arguments/environments @@ -1048,6 +1057,8 @@ def compare_two_settings( arg2: The second set of arguments to pass to the API server. env1: The first set of environment variables to pass to the API server. env2: The second set of environment variables to pass to the API server. + include_seeded_sampling: Whether to include temperature=1.0 seeded + sampling checks in the default generate comparison. """ compare_all_settings( @@ -1056,6 +1067,7 @@ def compare_two_settings( [env1, env2], method=method, max_wait_seconds=max_wait_seconds, + include_seeded_sampling=include_seeded_sampling, ) @@ -1066,6 +1078,7 @@ def compare_all_settings( *, method: str = "generate", max_wait_seconds: float | None = None, + include_seeded_sampling: bool = True, ) -> None: """ Launch API server with several different sets of arguments/environments @@ -1074,6 +1087,8 @@ def compare_all_settings( model: The model to test. all_args: A list of argument lists to pass to the API server. all_envs: A list of environment dictionaries to pass to the API server. + include_seeded_sampling: Whether to include temperature=1.0 seeded + sampling checks in the default generate comparison. """ trust_remote_code = False @@ -1134,7 +1149,13 @@ def compare_all_settings( ) if method == "generate": - results += _test_completion(client, model, prompt, token_ids) + results += _test_completion( + client, + model, + prompt, + token_ids, + include_seeded_sampling=include_seeded_sampling, + ) elif method == "generate_close": results += _test_completion_close(client, model, prompt) elif method == "generate_chat": From 6f53753fc924c86982b028d241b8b912123b97f1 Mon Sep 17 00:00:00 2001 From: Stefano Castagnetta Date: Mon, 4 May 2026 12:37:16 +0200 Subject: [PATCH 0033/1083] [Bugfix] Apply ruff-format to hyperclovax.py (#41620) Signed-off-by: Stefano Castagnetta --- vllm/transformers_utils/configs/hyperclovax.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/transformers_utils/configs/hyperclovax.py b/vllm/transformers_utils/configs/hyperclovax.py index 839b8ed50b08..d1a3218fe4dd 100644 --- a/vllm/transformers_utils/configs/hyperclovax.py +++ b/vllm/transformers_utils/configs/hyperclovax.py @@ -327,8 +327,7 @@ def __init__( self.vision_config = None if vision_config is not None: - _vision_config = AutoConfig.for_model( - vision_config["model_type"]) + _vision_config = AutoConfig.for_model(vision_config["model_type"]) self.vision_config = _vision_config.from_dict(vision_config) self.use_nth_layer = use_nth_layer From 62ba7516e87a18c9d3407c1bb526384a7474cf44 Mon Sep 17 00:00:00 2001 From: Stefano Castagnetta Date: Mon, 4 May 2026 13:47:42 +0200 Subject: [PATCH 0034/1083] Revert "[Doc] Fix RTD build: pytorch.org/docs/stable/objects.inv returns 404" (#41618) Signed-off-by: Stefano Castagnetta --- mkdocs.yaml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mkdocs.yaml b/mkdocs.yaml index 6afc44d71af5..4b06b31ebe35 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -105,8 +105,7 @@ plugins: - https://docs.aiohttp.org/en/stable/objects.inv - https://pillow.readthedocs.io/en/stable/objects.inv - https://numpy.org/doc/stable/objects.inv - # TODO revert to stable once https://github.com/pytorch/pytorch/issues/182007 is fixed - - https://pytorch.org/docs/2.11/objects.inv + - https://pytorch.org/docs/stable/objects.inv - redirects: redirect_maps: features/spec_decode/README.md: features/speculative_decoding/README.md From 8decbfa02c9bbc1699b2136ecc72b1ef30c438a0 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Mon, 4 May 2026 16:31:37 +0300 Subject: [PATCH 0035/1083] Test nemotron nano-v2 and nemotron nano-v3 separately, disable super-omni redundant tests (#41616) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- tests/models/registry.py | 51 ++++++++-------------------------------- 1 file changed, 10 insertions(+), 41 deletions(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index ab5869cd3fda..ec6c3473a785 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1141,30 +1141,17 @@ def check_available_online( "NemotronH_Nano_VL_V2": _HfExamplesInfo( "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", max_model_len=4096, - # NemotronH layers are constructed via `hybrid_override_pattern`: + # NemotronH layers are constructed via `hybrid_override_pattern` use_original_num_layers=True, hf_overrides={ - "vision_config": PretrainedConfig( - args={ - "min_num_patches": 1, # Trigger image dynamic res - "max_num_patches": 12, - "model": "vit_huge_patch16_224", - }, - # Trigger conv3d: - video_temporal_patch_size=2, - ), - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "M*", - }, + "text_config": {"num_hidden_layers": 2, "hybrid_override_pattern": "M*"}, }, trust_remote_code=True, ), - # NemotronH_Nano_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 - # Use the same registry test as NemotronH_Nano_VL_V2 above "NemotronH_Nano_Omni_Reasoning_V3": _HfExamplesInfo( - "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", max_model_len=4096, + # NemotronH layers are constructed via `hybrid_override_pattern` use_original_num_layers=True, hf_overrides={ "vision_config": PretrainedConfig( @@ -1174,35 +1161,17 @@ def check_available_online( "model": "vit_huge_patch16_224", }, video_temporal_patch_size=2, + # TODO(nhaber): This is `true` in the official `config.json`, + # but this causes a processor exception in the tests due to a known bug + # with mixed-resolution video when `true`. To be resolved. + video_maintain_aspect_ratio=False, ), - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "M*", - }, + "text_config": {"num_hidden_layers": 2, "hybrid_override_pattern": "M*"}, }, trust_remote_code=True, ), - # NemotronH_Super_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 as well - # Use the same registry test as NemotronH_Nano_VL_V2 above "NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo( - "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", - max_model_len=4096, - use_original_num_layers=True, - hf_overrides={ - "vision_config": PretrainedConfig( - args={ - "min_num_patches": 1, - "max_num_patches": 12, - "model": "vit_huge_patch16_224", - }, - video_temporal_patch_size=2, - ), - "text_config": { - "num_hidden_layers": 2, - "hybrid_override_pattern": "M*", - }, - }, - trust_remote_code=True, + "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", From 3e1ad4435f7c205dcbd5b14d1a529cc328b22ce8 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 4 May 2026 12:22:07 -0400 Subject: [PATCH 0036/1083] [Bug] Fix `tests/compile/test_config.py` AttributeError: 'NoneType' object has no attribute 'dtype' (#41288) Signed-off-by: yewentao256 --- vllm/config/vllm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 88e6660e2161..52c04509b2fa 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1617,7 +1617,7 @@ def _set_compile_ranges(self): max_size = rocm_aiter_ops.get_aiter_allreduce_max_size() else: max_size = compilation_config.pass_config.flashinfer_max_size(tp_size) - if max_size is not None: + if max_size is not None and self.model_config is not None: assert isinstance(self.model_config.dtype, torch.dtype) max_token_num = max_size // ( self.model_config.get_hidden_size() From 321fa2d6d1644629ac39d173f6393f37e14bf7b4 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Mon, 4 May 2026 13:30:02 -0400 Subject: [PATCH 0037/1083] Limit gpu utils and lower max BS on test_transcription_api_correctness.py (#41649) Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- .../openai/correctness/test_transcription_api_correctness.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py index a3df30fb02b2..fedbd74795b5 100644 --- a/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/openai/correctness/test_transcription_api_correctness.py @@ -27,7 +27,8 @@ from ....utils import RemoteOpenAIServer # Tuned to prevent OOM on 18GB GPUs in transcription correctness tests. -MAX_SEQS_FOR_TRANSCRIPTION_TEST = 32 +MAX_SEQS_FOR_TRANSCRIPTION_TEST = 8 +GPU_UTIL_FOR_TRANSCRIPTION_TEST = 0.5 def to_bytes(y, sr): @@ -188,6 +189,7 @@ def test_wer_correctness( "--enforce-eager", f"--tokenizer_mode={model_info.tokenizer_mode}", f"--max_num_seqs={MAX_SEQS_FOR_TRANSCRIPTION_TEST}", + f"--gpu_memory_utilization={GPU_UTIL_FOR_TRANSCRIPTION_TEST}", ] if model_info.trust_remote_code: server_args.append("--trust-remote-code") From 712ad0286c9afe461cfce28283c2506d83178689 Mon Sep 17 00:00:00 2001 From: Keyi Li <94494390+JasonKeyiL@users.noreply.github.com> Date: Mon, 4 May 2026 10:42:05 -0700 Subject: [PATCH 0038/1083] [Bugfix] KimiK2ReasoningParser: guard against buffered end-token in streaming (#41068) Signed-off-by: Keyi Li Co-authored-by: Keyi Li Co-authored-by: Claude Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_kimi_k2_reasoning_parser.py | 63 +++++++++++++++++++ vllm/reasoning/kimi_k2_reasoning_parser.py | 7 +++ 2 files changed, 70 insertions(+) diff --git a/tests/reasoning/test_kimi_k2_reasoning_parser.py b/tests/reasoning/test_kimi_k2_reasoning_parser.py index 0f80bb8854a8..dfce2075c6a9 100644 --- a/tests/reasoning/test_kimi_k2_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k2_reasoning_parser.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -12,6 +14,20 @@ REASONING_MODEL_NAME = "moonshotai/Kimi-K2.5" +@pytest.fixture +def mock_kimi_k2_tokenizer(): + tokenizer = MagicMock() + tokenizer.get_vocab.return_value = { + "": 100, + "": 101, + "<|tool_calls_section_begin|>": 200, + "<|tool_calls_section_end|>": 201, + "<|tool_call_begin|>": 202, + "<|tool_call_end|>": 203, + } + return tokenizer + + @pytest.fixture(scope="module") def kimi_k2_tokenizer(): return get_tokenizer(tokenizer_name=REASONING_MODEL_NAME, trust_remote_code=True) @@ -153,3 +169,50 @@ def test_streaming_tool_section_ends_reasoning(kimi_k2_tokenizer): ) assert isinstance(result, DeltaMessage) assert result.content == "<|tool_calls_section_begin|>" + + +def test_streaming_end_token_id_buffered(mock_kimi_k2_tokenizer): + """When stop sequences buffer text, ID arrives before its text. + + The token ID is present in delta_token_ids but the actual string is not + yet in delta_text (still buffered). The parser must return None to wait + for the next delta, instead of calling find() which returns -1 and + silently corrupting the text split. + """ + parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer) + think_id = parser._start_token_id + end_think_id = parser._end_token_id + + # Simulate: ID arrived but text not yet flushed. + # Two token IDs in delta to bypass the single-special-token guard. + result = parser.extract_reasoning_streaming( + previous_text="some reasoning", + current_text="some reasoning extra", + delta_text="extra", # text not yet flushed + previous_token_ids=[think_id], + current_token_ids=[think_id, end_think_id, 999], + delta_token_ids=[end_think_id, 999], + ) + assert result is None + + +def test_streaming_tool_section_id_buffered(mock_kimi_k2_tokenizer): + """When stop sequences buffer text, tool section start ID arrives before its text. + + Same buffering scenario as above but for <|tool_calls_section_begin|>. + Without the guard, find() returns -1 and delta_text[:tool_index] silently + drops the last character of reasoning. + """ + parser = KimiK2ReasoningParser(mock_kimi_k2_tokenizer) + think_id = parser._start_token_id + tool_begin_id = parser._tool_section_start_token_id + + result = parser.extract_reasoning_streaming( + previous_text="some reasoning", + current_text="some reasoning extra", + delta_text="extra", # tool section text not yet flushed + previous_token_ids=[think_id], + current_token_ids=[think_id, tool_begin_id, 999], + delta_token_ids=[tool_begin_id, 999], + ) + assert result is None diff --git a/vllm/reasoning/kimi_k2_reasoning_parser.py b/vllm/reasoning/kimi_k2_reasoning_parser.py index 7a92703426fc..0b64c5c62ea1 100644 --- a/vllm/reasoning/kimi_k2_reasoning_parser.py +++ b/vllm/reasoning/kimi_k2_reasoning_parser.py @@ -221,6 +221,10 @@ def extract_reasoning_streaming( return None if self._end_token_id in delta_token_ids: + if self._end_token not in delta_text: + # Token ID arrived before text was flushed (stop-sequence buffering). + # Wait for the next delta when the text becomes visible. + return None end_index = delta_text.find(self._end_token) reasoning = delta_text[:end_index] content = delta_text[end_index + len(self._end_token) :] @@ -229,6 +233,9 @@ def extract_reasoning_streaming( ) if self._tool_section_start_token_id in delta_token_ids: + if self._tool_section_start_token not in delta_text: + # Token ID arrived before text was flushed (stop-sequence buffering). + return None tool_index = delta_text.find(self._tool_section_start_token) reasoning = delta_text[:tool_index] content = delta_text[tool_index:] From e724b0ea8d3b34e349a8635b1b49cb3005bf25f5 Mon Sep 17 00:00:00 2001 From: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Date: Mon, 4 May 2026 13:07:19 -0500 Subject: [PATCH 0039/1083] [ROCm] ROCm7.2.2 + profiler fix + AITER 0.1.12.post2 (#41386) Signed-off-by: Rohan138 Signed-off-by: Gregory Shtrasberg Co-authored-by: Rohan138 --- .buildkite/release-pipeline.yaml | 2 +- docker/Dockerfile.rocm_base | 28 ++++++++++++++++++++++++---- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index cdb5b00d4143..f122c423ba5c 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -723,7 +723,7 @@ steps: - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" env: S3_BUCKET: "vllm-wheels" - VARIANT: "rocm721" + VARIANT: "rocm722" # ROCm Job 6: Build ROCm Release Docker Image - label: ":docker: Build release image - x86_64 - ROCm" diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 5940a4ee564d..a21916d0b531 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.1-complete +ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete ARG TRITON_BRANCH="ba5c1517" ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17 @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.10.post3" +ARG AITER_BRANCH="v0.1.12.post2" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -104,6 +104,28 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} +# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062 +# will be removed once we move to ROCm 7.2.3 +RUN apt-get update && apt-get install -y rocm-llvm-dev +RUN pip install CppHeaderParser +RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \ + && cd /tmp/rocm-systems \ + && git sparse-checkout init --cone \ + && git sparse-checkout set projects/hip projects/clr \ + && git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \ + && export CLR_DIR=/tmp/rocm-systems/projects/clr \ + && export HIP_DIR=/tmp/rocm-systems/projects/hip \ + && mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \ + && cmake \ + -DHIP_COMMON_DIR=$HIP_DIR \ + -DCMAKE_PREFIX_PATH="/opt/rocm/" \ + -DCLR_BUILD_HIP=ON \ + -DCLR_BUILD_OCL=OFF \ + -DHIP_PLATFORM=amd \ + .. \ + && make -j$(nproc) \ + && make install \ + && rm -rf /tmp/rocm-systems ### ### Triton Build @@ -153,8 +175,6 @@ RUN git clone ${PYTORCH_REPO} pytorch RUN cd pytorch && git checkout ${PYTORCH_BRANCH} RUN cd pytorch \ && pip install -r requirements.txt && git submodule update --init --recursive -RUN cd pytorch/third_party/kineto \ - && git remote add rocm https://github.com/ROCm/kineto && git fetch rocm && git checkout 2d73be3 RUN cd pytorch && python3 tools/amd_build/build_amd.py \ && if [ "$USE_SCCACHE" = "1" ]; then \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ From 8c780943b492a26dc2032f40a18a77f92d1931e6 Mon Sep 17 00:00:00 2001 From: Baekpica <35071468+Baekpica@users.noreply.github.com> Date: Tue, 5 May 2026 03:43:07 +0900 Subject: [PATCH 0040/1083] Fix Nano Nemotron text-only weight loading (#41205) Signed-off-by: sunghoon.baek Signed-off-by: Baekpica <35071468+Baekpica@users.noreply.github.com> Signed-off-by: sunghoon.baek Co-authored-by: sunghoon.baek Co-authored-by: OpenAI Codex Co-authored-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- .../multimodal/test_nano_nemotron_vl.py | 114 ++++++++++++++++++ .../model_executor/models/nano_nemotron_vl.py | 18 ++- 2 files changed, 129 insertions(+), 3 deletions(-) create mode 100644 tests/models/multimodal/test_nano_nemotron_vl.py diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py new file mode 100644 index 000000000000..6922af79c08e --- /dev/null +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.model_executor.models.nano_nemotron_vl import NemotronH_Nano_VL_V2 + + +class _TextOnlyMultiModalConfig: + def get_limit_per_prompt(self, modality: str) -> int: + return 0 + + +class _ImageOnlyMultiModalConfig: + def get_limit_per_prompt(self, modality: str) -> int: + return 1 if modality == "image" else 0 + + +class _ModelConfig: + multimodal_config = _TextOnlyMultiModalConfig() + + +class _ImageOnlyModelConfig: + multimodal_config = _ImageOnlyMultiModalConfig() + + +class _LanguageModel: + def __init__(self) -> None: + self.loaded_weights: list[tuple[str, object]] = [] + + def load_weights(self, weights): + self.loaded_weights = list(weights) + + +class _MissingMultiModalModule: + def named_parameters(self): + raise AssertionError("multimodal weights should not be inspected") + + def load_weights(self, weights): + raise AssertionError("multimodal weights should not be loaded") + + +class _AdapterModule: + def named_parameters(self): + return [] + + +class _VisionModel: + def __init__(self) -> None: + self.loaded_weights: list[tuple[str, object]] = [] + + def load_weights(self, weights): + self.loaded_weights = list(weights) + + +def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode(): + model = object.__new__(NemotronH_Nano_VL_V2) + language_model = _LanguageModel() + object.__setattr__(model, "model_config", _ModelConfig()) + object.__setattr__(model, "language_model", language_model) + object.__setattr__(model, "mlp1", _AdapterModule()) + object.__setattr__(model, "vision_model", _MissingMultiModalModule()) + object.__setattr__(model, "sound_encoder", None) + + language_weight = object() + model.load_weights( + [ + ("language_model.layers.0.weight", language_weight), + ("mlp1.0.weight", object()), + ("vision_model.radio_model.encoder.weight", object()), + ("sound_encoder.encoder.weight", object()), + ] + ) + + assert language_model.loaded_weights == [("layers.0.weight", language_weight)] + + +def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder(): + model = object.__new__(NemotronH_Nano_VL_V2) + language_model = _LanguageModel() + vision_model = _VisionModel() + object.__setattr__(model, "model_config", _ImageOnlyModelConfig()) + object.__setattr__(model, "language_model", language_model) + object.__setattr__(model, "mlp1", _AdapterModule()) + object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "sound_encoder", None) + + language_weight = object() + vision_weight = object() + model.load_weights( + [ + ("language_model.layers.0.weight", language_weight), + ("vision_model.radio_model.encoder.weight", vision_weight), + ] + ) + + assert language_model.loaded_weights == [("layers.0.weight", language_weight)] + assert vision_model.loaded_weights == [ + ("radio_model.encoder.weight", vision_weight) + ] + + +def test_nano_nemotron_vl_requires_sound_encoder_for_sound_weights(): + model = object.__new__(NemotronH_Nano_VL_V2) + language_model = _LanguageModel() + vision_model = _VisionModel() + object.__setattr__(model, "model_config", _ImageOnlyModelConfig()) + object.__setattr__(model, "language_model", language_model) + object.__setattr__(model, "mlp1", _AdapterModule()) + object.__setattr__(model, "vision_model", vision_model) + object.__setattr__(model, "sound_encoder", None) + + with pytest.raises(AssertionError): + model.load_weights([("sound_encoder.encoder.weight", object())]) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 684ced0a6abd..994b52606b18 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1499,6 +1499,11 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + mm_config = self.model_config.multimodal_config + load_multimodal_weights = not all( + mm_config.get_limit_per_prompt(modality) == 0 + for modality in ("image", "video", "audio") + ) adapter_dict = dict(self.mlp1.named_parameters()) def is_llm(name: str) -> bool: @@ -1523,23 +1528,30 @@ def is_sound_weights(name: str) -> bool: # Strip 'language_model.' prefix for LLM weights llm_weights.append((".".join(name.split(".")[1:]), w)) elif is_adapter_weights((name, w)): + if not load_multimodal_weights: + continue # Load vision-language adapter weights directly trimmed_name = ".".join(name.split(".")[1:]) param = adapter_dict[trimmed_name] with torch.no_grad(): default_weight_loader(param, w) elif is_vision_weights(name): + if not load_multimodal_weights: + continue # Convert: vision_model.radio_model.* → radio_model.* hf_key = name[len("vision_model.") :] # Remove "vision_model." prefix vision_weights.append((hf_key, w)) elif is_sound_weights(name): + if not load_multimodal_weights: + continue assert self.sound_encoder is not None sound_weights.append((name, w)) self.language_model.load_weights(llm_weights) - self.vision_model.load_weights(vision_weights) - if self.sound_encoder is not None and len(sound_weights) > 0: - self.sound_encoder.load_weights(sound_weights) + if load_multimodal_weights: + self.vision_model.load_weights(vision_weights) + if self.sound_encoder is not None and len(sound_weights) > 0: + self.sound_encoder.load_weights(sound_weights) def get_vit_model_from_radio_config(self, hf_config): hf_config_vision = hf_config.vision_config From 422dd0259853e5ce2752341109b38293b9bbb5b8 Mon Sep 17 00:00:00 2001 From: Joachim Studnia Date: Mon, 4 May 2026 11:46:00 -0700 Subject: [PATCH 0041/1083] [bugfix] Fix prompt logprobs on request eviction during chunked prefill (#41411) Signed-off-by: Joachim Studnia Co-authored-by: Claude Opus 4.6 (1M context) --- tests/conftest.py | 2 ++ tests/v1/e2e/general/test_async_scheduling.py | 11 ++++++++++- vllm/v1/worker/gpu/sample/prompt_logprob.py | 4 +--- vllm/v1/worker/gpu_input_batch.py | 6 ++---- vllm/v1/worker/gpu_model_runner.py | 9 ++++----- 5 files changed, 19 insertions(+), 13 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 40adeda2bd50..779bd475f34b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -996,6 +996,8 @@ def generate( req_sample_output_ids: list[list[int]] = [] req_sample_output_strs: list[str] = [] req_logprobs = [] + if req_output.prompt_logprobs: + req_logprobs.extend(req_output.prompt_logprobs) for sample in req_output.outputs: output_str = sample.text output_ids = list(sample.token_ids) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 28a1bedbe0b2..c3c4970de382 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -57,6 +57,8 @@ def test_without_spec_decoding( dict(bad_words=["the", " the"]), dict(logprobs=2), dict(logprobs=2, frequency_penalty=-1.0), + dict(prompt_logprobs=2), + dict(prompt_logprobs=2, logprobs=2), dict(structured_outputs=struct_outputs), dict( structured_outputs=struct_outputs, @@ -126,6 +128,8 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke dict(bad_words=["the", " the"]), dict(logprobs=2), dict(logprobs=2, frequency_penalty=-1.0), + dict(prompt_logprobs=2), + dict(prompt_logprobs=2, logprobs=2), dict(structured_outputs=struct_outputs), dict( structured_outputs=struct_outputs, @@ -413,7 +417,12 @@ def _all_logprobs_match(req_a, req_b) -> bool: ) -def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool: +def _logprobs_match( + lps_a: dict[int, Logprob] | None, + lps_b: dict[int, Logprob] | None, +) -> bool: + if lps_a is None or lps_b is None: + return lps_a is lps_b rel_tol, abs_tol = 1e-3, 1e-6 return ( len(lps_a) == len(lps_b) diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py index 11dbf6985279..baa48ebf900c 100644 --- a/vllm/v1/worker/gpu/sample/prompt_logprob.py +++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py @@ -55,10 +55,8 @@ def compute_prompt_logprobs( num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np] prompt_lens = prompt_lens[idx_mapping_np] - # NOTE(woosuk): -1 because the last prompt token's hidden state is not - # needed for prompt logprobs. computed_prefill = num_computed_prefill_tokens[idx_mapping_np] - includes_prompt = computed_prefill < prompt_lens - 1 + includes_prompt = computed_prefill < prompt_lens # NOTE(woosuk): If the request was resumed after preemption, its prompt # logprobs must have been computed before preemption. Skip. resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np] diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 75898f463272..44e0efaaa2f2 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -49,6 +49,8 @@ class CachedRequestState: lora_request: LoRARequest | None = None prompt_embeds: torch.Tensor | None = None + # To accumulate prompt logprobs tensor chunks across prefill steps. + in_progress_prompt_logprobs_cpu: LogprobsTensors | None = None # Per-position mask for mixed-mode inputs (e.g chat completion with # prompt_embeds content parts). See `Request.prompt_is_token_ids`. @@ -255,9 +257,6 @@ def __init__( # More efficient than num_logprobs=-1 when only a few tokens are needed self.logprob_token_ids: dict[str, list[int]] = {} - # To accumulate prompt logprobs tensor chunks across prefill steps. - self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {} - # Internal representation of per-step batch state changes, used for # reordering persistent batch and generating logitsprocs batch state # updates. Should reset each step. @@ -552,7 +551,6 @@ def remove_request(self, req_id: str) -> int | None: self.generators.pop(req_index, None) self.num_logprobs.pop(req_id, None) self.logprob_token_ids.pop(req_id, None) - self.in_progress_prompt_logprobs_cpu.pop(req_id, None) if self.prev_req_id_to_index is not None: self.prev_req_id_to_index.pop(req_id, None) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bcab2ca2d4c2..0ca530c15bac 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5094,7 +5094,6 @@ def _get_prompt_logprobs_dict( if not num_prompt_logprobs_dict: return {} - in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {} # Since prompt logprobs are a rare feature, prioritize simple, @@ -5118,14 +5117,14 @@ def _get_prompt_logprobs_dict( ) # Set up target LogprobsTensors object. - logprobs_tensors = in_progress_dict.get(req_id) - if not logprobs_tensors: + logprobs_tensors = request.in_progress_prompt_logprobs_cpu + if logprobs_tensors is None: # Create empty logprobs CPU tensors for the entire prompt. # If chunked, we'll copy in slice by slice. logprobs_tensors = LogprobsTensors.empty_cpu( num_prompt_tokens - 1, num_prompt_logprobs + 1 ) - in_progress_dict[req_id] = logprobs_tensors + request.in_progress_prompt_logprobs_cpu = logprobs_tensors # Determine number of logits to retrieve. start_idx = request.num_computed_tokens @@ -5182,7 +5181,7 @@ def _get_prompt_logprobs_dict( # num_prompt_logprobs_dict. for req_id in completed_prefill_reqs: del num_prompt_logprobs_dict[req_id] - del in_progress_dict[req_id] + self.requests[req_id].in_progress_prompt_logprobs_cpu = None # Must synchronize the non-blocking GPU->CPU transfers. if prompt_logprobs_dict: From 844df542694089045589ceaad52cba69ab58527a Mon Sep 17 00:00:00 2001 From: Linzhang Li Date: Tue, 5 May 2026 03:45:24 +0800 Subject: [PATCH 0042/1083] feat: update xgrammar==0.2.0 to use structural tags for strict tool calling + reasoning for more models (#40894) Signed-off-by: Yuchuan Signed-off-by: Michael Goin Signed-off-by: mgoin Signed-off-by: Ubospica Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Michael Goin Co-authored-by: Ubospica Co-authored-by: sfeng33 <4florafeng@gmail.com> --- requirements/common.txt | 2 +- requirements/test/rocm.txt | 5 +- .../test_deepseekv4_tool_parser.py | 82 +++++ .../test_qwen3coder_tool_parser.py | 109 ++++++ vllm/entrypoints/openai/api_server.py | 15 + vllm/envs.py | 7 + vllm/tool_parsers/abstract_tool_parser.py | 40 ++- vllm/tool_parsers/deepseekv4_tool_parser.py | 15 + vllm/tool_parsers/qwen3coder_tool_parser.py | 14 + vllm/tool_parsers/structural_tag_registry.py | 330 ++++++++++++++++++ 10 files changed, 613 insertions(+), 6 deletions(-) create mode 100644 vllm/tool_parsers/structural_tag_registry.py diff --git a/requirements/common.txt b/requirements/common.txt index 5d4519204ee9..652738eebe74 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -24,7 +24,7 @@ outlines_core == 0.2.14 # required for outlines backend disk cache diskcache == 5.6.3 lark == 1.2.2 -xgrammar >= 0.1.32, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" +xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le" typing_extensions >= 4.10 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317 partial-json-parser # used for parsing partial JSON outputs diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 801af7db9db1..8445634ded40 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -42,6 +42,8 @@ anyio==4.13.0 # sse-starlette # starlette # watchfiles +apache-tvm-ffi==0.1.10 + # via xgrammar arctic-inference==0.1.1 # via -r requirements/test/rocm.in argcomplete==3.6.3 @@ -1264,6 +1266,7 @@ typing-extensions==4.15.0 # alembic # anthropic # anyio + # apache-tvm-ffi # azure-core # azure-identity # azure-storage-blob @@ -1345,7 +1348,7 @@ word2number==1.1 # via lm-eval wrapt==2.1.2 # via smart-open -xgrammar==0.1.33 +xgrammar==0.2.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/tests/tool_parsers/test_deepseekv4_tool_parser.py b/tests/tool_parsers/test_deepseekv4_tool_parser.py index 631d0fb97b33..cc77a1f77756 100644 --- a/tests/tool_parsers/test_deepseekv4_tool_parser.py +++ b/tests/tool_parsers/test_deepseekv4_tool_parser.py @@ -6,6 +6,15 @@ import json from unittest.mock import MagicMock +import pytest +from xgrammar import StructuralTag + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedFunction, + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, + ChatCompletionToolsParam, +) from vllm.tool_parsers import ToolParserManager from vllm.tool_parsers.deepseekv4_tool_parser import DeepSeekV4ToolParser @@ -20,6 +29,43 @@ PARAM_END = "" +@pytest.fixture +def sample_tools() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "The city name"}, + "state": {"type": "string", "description": "The state code"}, + "unit": {"type": "string", "enum": ["fahrenheit", "celsius"]}, + }, + "required": ["city", "state"], + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "calculate_area", + "description": "Calculate area of a shape", + "parameters": { + "type": "object", + "properties": { + "shape": {"type": "string"}, + "dimensions": {"type": "object"}, + "precision": {"type": "integer"}, + }, + }, + }, + ), + ] + + def make_parser(tools=None) -> DeepSeekV4ToolParser: return DeepSeekV4ToolParser(MOCK_TOKENIZER, tools=tools) @@ -121,3 +167,39 @@ def test_streaming_extracts_complete_invokes(): ] assert names == ["search"] assert json.loads(reconstruct_args(deltas)) == {"query": "deepseek v4"} + + +def test_get_vllm_registry_structural_tag_returns_structural_tag( + sample_tools: list[ChatCompletionToolsParam], +) -> None: + parser = make_parser() + req = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="auto", + ) + tag = parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + req = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="required", + ) + tag = parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + if sample_tools: + tool = sample_tools[0] + req = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + ) + req.tool_choice = ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=tool.function.name) + ) + tag = parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index c62e95830243..26bbf1a044bc 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -6,8 +6,11 @@ import pytest from openai.types.responses.function_tool import FunctionTool +from xgrammar import StructuralTag from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedFunction, + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionToolsParam, ) @@ -108,6 +111,27 @@ def sample_tools(request): ] +def _as_chat_completion_tools( + tools: list[ChatCompletionToolsParam | FunctionTool], +) -> list[ChatCompletionToolsParam]: + normalized: list[ChatCompletionToolsParam] = [] + for tool in tools: + if isinstance(tool, ChatCompletionToolsParam): + normalized.append(tool) + else: + normalized.append( + ChatCompletionToolsParam( + type="function", + function={ + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + }, + ) + ) + return normalized + + def assert_tool_calls( actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] ): @@ -1146,3 +1170,88 @@ def test_no_double_serialization_string_args(qwen3_tool_parser): args = json.loads(raw_arguments) assert args["message"] == "hello world" assert '\\"hello world\\"' not in raw_arguments + + +def test_get_vllm_registry_structural_tag_returns_structural_tag( + qwen3_tool_parser: Qwen3CoderToolParser, + sample_tools: list[ChatCompletionToolsParam], +) -> None: + request_tools = _as_chat_completion_tools(sample_tools) + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="auto", + ) + tag = qwen3_tool_parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="required", + ) + tag = qwen3_tool_parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + if request_tools: + tool = request_tools[0] + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + ) + req.tool_choice = ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=tool.function.name) + ) + tag = qwen3_tool_parser.get_structural_tag(req) + assert isinstance(tag, StructuralTag) + + +@pytest.mark.parametrize("include_reasoning", [True, False]) +def test_adjust_request_auto_uses_vllm_registry_structural_tag( + monkeypatch: pytest.MonkeyPatch, + qwen3_tool_parser: Qwen3CoderToolParser, + sample_tools: list[ChatCompletionToolsParam], + include_reasoning: bool, +) -> None: + monkeypatch.setattr( + "vllm.tool_parsers.abstract_tool_parser.VLLM_ENFORCE_STRICT_TOOL_CALLING", + True, + ) + request_tools = _as_chat_completion_tools(sample_tools) + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="auto", + include_reasoning=include_reasoning, + ) + out = qwen3_tool_parser.adjust_request(req) + assert out.structured_outputs is not None + assert out.structured_outputs.structural_tag is not None + assert isinstance(out.structured_outputs.structural_tag, str) + loaded = json.loads(out.structured_outputs.structural_tag) + assert isinstance(loaded, dict) + + +def test_adjust_request_required_prefers_structural_tag( + monkeypatch: pytest.MonkeyPatch, + qwen3_tool_parser: Qwen3CoderToolParser, + sample_tools: list[ChatCompletionToolsParam], +) -> None: + monkeypatch.setattr( + "vllm.tool_parsers.abstract_tool_parser.VLLM_ENFORCE_STRICT_TOOL_CALLING", + True, + ) + request_tools = _as_chat_completion_tools(sample_tools) + req = ChatCompletionRequest( + messages=[], + model="m", + tools=request_tools, + tool_choice="required", + ) + out = qwen3_tool_parser.adjust_request(req) + assert out.structured_outputs is not None + assert out.structured_outputs.structural_tag is not None diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 9aac19e2fda5..da2ec10284c5 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -321,6 +321,21 @@ async def init_app_state( supported_tasks: tuple["SupportedTask", ...] | None = None, ) -> None: vllm_config = engine_client.vllm_config + + # Propagate enable_in_reasoning to the API-server process. The engine core + # runs in a separate process, so the contextvar that backs + # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers + # call `get_enable_structured_outputs_in_reasoning()` during request + # handling and need to see the real flag, otherwise they silently fall + # back to False and mismatch the engine-side bitmask gating. + from vllm.tool_parsers.structural_tag_registry import ( + set_enable_structured_outputs_in_reasoning, + ) + + set_enable_structured_outputs_in_reasoning( + vllm_config.structured_outputs_config.enable_in_reasoning + ) + if supported_tasks is None: warnings.warn( "The 'supported_tasks' parameter was not provided to " diff --git a/vllm/envs.py b/vllm/envs.py index b2db5a8112bc..e456bec5bfb1 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -226,6 +226,7 @@ VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False VLLM_SYSTEM_START_DATE: str | None = None VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: bool = False + VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = False VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False VLLM_NVTX_SCOPES_FOR_PROFILING: bool = False VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True @@ -1593,6 +1594,12 @@ def _get_or_set_default() -> str: "VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY": lambda: bool( int(os.getenv("VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY", "0")) ), + # When 1,the model structural tags will be used to enforce the model + # output conforming to the model's tool-calling format and schema. + # Default 0 (off). + "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: bool( + int(os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "0")) + ), # Add optional custom scopes for profiling, disable to avoid overheads "VLLM_CUSTOM_SCOPES_FOR_PROFILING": lambda: bool( int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0")) diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 75181d8dfac6..c3438082a72d 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib +import json import os from collections.abc import Callable, Sequence from functools import cached_property @@ -13,6 +14,7 @@ from openai.types.responses.function_tool import FunctionTool from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionToolsParam, ) @@ -23,6 +25,7 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ) +from vllm.envs import VLLM_ENFORCE_STRICT_TOOL_CALLING from vllm.logger import init_logger from vllm.sampling_params import ( StructuredOutputsParams, @@ -83,13 +86,39 @@ def vocab(self) -> dict[str, int]: return self.model_tokenizer.get_vocab() def adjust_request( - self, request: ChatCompletionRequest | ResponsesRequest + self, + request: ChatCompletionRequest | ResponsesRequest, ) -> ChatCompletionRequest | ResponsesRequest: - """ - Static method that used to adjust the request parameters. - """ + # If there are no tools, return the request as is. if not request.tools: return request + + # Step 1 (highest priority for ChatCompletionRequest): apply + # vLLM-owned structural tag support for model-specific tool formats. + if ( + isinstance(request, ChatCompletionRequest) + and VLLM_ENFORCE_STRICT_TOOL_CALLING + ): + need_tool_calling = ( + request.tool_choice == "auto" + or request.tool_choice == "required" + or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) + ) + if need_tool_calling: + structure_tag = self.get_structural_tag(request) + if structure_tag is not None: + if request.structured_outputs is None: + request.structured_outputs = StructuredOutputsParams( + structural_tag=json.dumps(structure_tag.model_dump()), + ) + else: + request.structured_outputs.structural_tag = json.dumps( + structure_tag.model_dump() + ) + return request + + # Step 2: set structured output params when tool constraints are + # derived from the tool schema. json_schema_from_tool = get_json_schema_from_tools( tool_choice=request.tool_choice, tools=request.tools ) @@ -121,6 +150,9 @@ def adjust_request( return request + def get_structural_tag(self, request: ChatCompletionRequest): + return None + def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest ) -> ExtractedToolCallInformation: diff --git a/vllm/tool_parsers/deepseekv4_tool_parser.py b/vllm/tool_parsers/deepseekv4_tool_parser.py index 45a9c1302578..e32451cd8bbd 100644 --- a/vllm/tool_parsers/deepseekv4_tool_parser.py +++ b/vllm/tool_parsers/deepseekv4_tool_parser.py @@ -1,7 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser +from vllm.tool_parsers.structural_tag_registry import ( + get_enable_structured_outputs_in_reasoning, + get_model_structural_tag, +) class DeepSeekV4ToolParser(DeepSeekV32ToolParser): @@ -14,3 +21,11 @@ class DeepSeekV4ToolParser(DeepSeekV32ToolParser): tool_call_start_token: str = "<｜DSML｜tool_calls>" tool_call_end_token: str = "" + + def get_structural_tag(self, request: ChatCompletionRequest): + return get_model_structural_tag( + model="deepseek_v4", + tools=request.tools, + tool_choice=request.tool_choice, + reasoning=get_enable_structured_outputs_in_reasoning(), + ) diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 7b089ceffbc0..73850b2ab0c5 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -25,12 +25,18 @@ Tool, ToolParser, ) +from vllm.tool_parsers.structural_tag_registry import ( + get_enable_structured_outputs_in_reasoning, + get_model_structural_tag, +) from vllm.tool_parsers.utils import find_tool_properties logger = init_logger(__name__) class Qwen3CoderToolParser(ToolParser): + supports_required_and_named: bool = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) @@ -681,3 +687,11 @@ def extract_tool_calls_streaming( return result return None + + def get_structural_tag(self, request: ChatCompletionRequest): + return get_model_structural_tag( + model="qwen_3_5", + tools=request.tools, + tool_choice=request.tool_choice, + reasoning=get_enable_structured_outputs_in_reasoning(), + ) diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py new file mode 100644 index 000000000000..754cc52361c5 --- /dev/null +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Model-specific structural tag builders adapted from XGrammar's +# builtin structural tag implementations: +# https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/builtin_structural_tag.py + +from collections.abc import Callable +from typing import Any, Literal + +from xgrammar import StructuralTag +from xgrammar.structural_tag import ( + AnyTextFormat, + ConstStringFormat, + JSONSchemaFormat, + SequenceFormat, + TagFormat, + TagsWithSeparatorFormat, + TriggeredTagsFormat, +) + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionToolsParam, +) + +SimplifiedToolChoice = Literal["auto", "required", "forced"] +ToolChoice = ( + Literal["none", "auto", "required"] | ChatCompletionNamedToolChoiceParam | None +) +StructuralTagBuilder = Callable[ + [list[ChatCompletionToolsParam], SimplifiedToolChoice, bool], + StructuralTag, +] + +_structural_tag_registry: dict[str, StructuralTagBuilder] = {} + + +def register_model_structural_tag(name: str): + """Register a vLLM-owned model-specific structural tag builder.""" + + def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder: + _structural_tag_registry[name] = func + return func + + return decorator + + +def get_model_structural_tag( + model: str, + tools: list[ChatCompletionToolsParam] | None, + tool_choice: ToolChoice, + reasoning: bool, +) -> StructuralTag | None: + """Build a structural tag from vLLM-owned model-specific builders.""" + + builder = _structural_tag_registry.get(model) + if builder is None: + supported = list(_structural_tag_registry.keys()) + raise ValueError(f"Unknown format type: {model}, supported types: {supported}") + + normalized_tools, simplified_tool_choice = _normalize_tool_choice( + tools=tools, + tool_choice=tool_choice, + ) + if not normalized_tools: + return None + + return builder(normalized_tools, simplified_tool_choice, reasoning) + + +def _normalize_tool_choice( + tools: list[ChatCompletionToolsParam] | None, + tool_choice: ToolChoice, +) -> tuple[list[ChatCompletionToolsParam], SimplifiedToolChoice]: + """Normalize vLLM ChatCompletion tool_choice for structural tag builders.""" + + if not tools: + return [], "auto" + + if tool_choice is None or tool_choice == "none": + return [], "auto" + + if tool_choice == "auto": + return tools, "auto" + + if tool_choice == "required": + return tools, "required" + + if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): + tool_name = tool_choice.function.name + filtered_tools = [tool for tool in tools if tool.function.name == tool_name] + if not filtered_tools: + raise ValueError( + f"The tool with name '{tool_name}' is not found in the tools list." + ) + return filtered_tools, "forced" + + raise ValueError(f"Unsupported tool_choice for structural tag: {tool_choice}") + + +def _get_function_parameters(function: Any) -> dict[str, Any] | bool: + """Return the JSON schema used for constrained tool arguments.""" + + if getattr(function, "strict", None) is False: + return True + if function.parameters is None: + return True + return function.parameters + + +_enable_structured_outputs_in_reasoning: bool = False + + +def set_enable_structured_outputs_in_reasoning(enabled: bool) -> None: + """Publish the engine's ``enable_in_reasoning`` flag to tool parsers. + + Called once during APIServer startup so request-time parsers can read + it without going through the EngineCore-only contextvar. + """ + + global _enable_structured_outputs_in_reasoning + _enable_structured_outputs_in_reasoning = bool(enabled) + + +def get_enable_structured_outputs_in_reasoning() -> bool: + """Whether structured outputs are active during the reasoning phase. + + When ``True``, the structural tag will cover the reasoning part: + ``...`` prefix (if available); when ``False`` (default), the tag only + constrains the post-reasoning suffix. + """ + + return _enable_structured_outputs_in_reasoning + + +@register_model_structural_tag("deepseek_v4") +def get_deepseek_v4_structural_tag( + tools: list[ChatCompletionToolsParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + """Build DeepSeek V4 structural tags.""" + + invoke_begin_prefix = '<｜DSML｜invoke name="' + invoke_begin_suffix = '">\n' + invoke_end = "\n" + tool_calls_prefix = "\n\n" + function_calls_begin = "<｜DSML｜tool_calls>\n" + function_calls_end = "" + function_calls_trigger = "<｜DSML｜tool_calls>" + think_tag_end = "" + think_exclude_tokens = ["", ""] + xml_style = "deepseek_xml" + + if tool_choice == "auto": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=invoke_begin_prefix + function.name + invoke_begin_suffix, + content=JSONSchemaFormat( + json_schema=parameters, + style=xml_style, + ), + end=invoke_end, + ) + ) + + if tags: + function_calling_tags = TagsWithSeparatorFormat( + tags=tags, + separator="\n", + at_least_one=True, + ) + suffix_tag = TriggeredTagsFormat( + triggers=[function_calls_trigger], + tags=[ + TagFormat( + begin=function_calls_begin, + content=function_calling_tags, + end=function_calls_end, + ) + ], + excludes=think_exclude_tokens, + ) + else: + suffix_tag = AnyTextFormat(excludes=think_exclude_tokens) + + elif tool_choice == "forced": + if not tools: + raise ValueError("Forced tool choice must resolve to exactly one tool.") + function = tools[0].function + suffix_tag = SequenceFormat( + elements=[ + ConstStringFormat(value=tool_calls_prefix + function_calls_begin), + TagFormat( + begin=invoke_begin_prefix + function.name + invoke_begin_suffix, + content=JSONSchemaFormat( + json_schema=_get_function_parameters(function), + style=xml_style, + ), + end=invoke_end, + ), + ConstStringFormat(value=function_calls_end), + ] + ) + + elif tool_choice == "required": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=invoke_begin_prefix + function.name + invoke_begin_suffix, + content=JSONSchemaFormat( + json_schema=parameters, + style=xml_style, + ), + end=invoke_end, + ) + ) + assert len(tags) > 0 + suffix_tag = SequenceFormat( + elements=[ + ConstStringFormat(value=tool_calls_prefix + function_calls_begin), + TagsWithSeparatorFormat( + tags=tags, + separator="\n", + at_least_one=True, + ), + ConstStringFormat(value=function_calls_end), + ] + ) + + if not reasoning: + return StructuralTag(format=suffix_tag) + + prefix_tag = TagFormat(begin="", content=AnyTextFormat(), end=think_tag_end) + return StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) + + +@register_model_structural_tag("qwen_3_5") +def get_qwen_3_5_structural_tag( + tools: list[ChatCompletionToolsParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + """Build Qwen XML structural tags. + + This format is used for Qwen3-Coder/Qwen3.5/Qwen3.6 and is compatible with + Qwen variants that use the same XML tool-call format. + """ + tool_call_begin_prefix = "\n", ""] + + if tool_choice == "auto": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + content=JSONSchemaFormat(json_schema=parameters, style="qwen_xml"), + end=tool_call_end, + ) + ) + + if tags: + suffix_tag = TriggeredTagsFormat( + triggers=[tool_call_trigger], + tags=tags, + excludes=think_exclude_tokens, + ) + else: + suffix_tag = AnyTextFormat(excludes=think_exclude_tokens) + + elif tool_choice == "forced": + if not tools: + raise ValueError("Forced tool choice must resolve to exactly one tool.") + function = tools[0].function + suffix_tag = TagFormat( + begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + content=JSONSchemaFormat( + json_schema=_get_function_parameters(function), + style="qwen_xml", + ), + end=tool_call_end, + ) + + elif tool_choice == "required": + tags = [] + for tool in tools: + function = tool.function + parameters = _get_function_parameters(function) + tags.append( + TagFormat( + begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + content=JSONSchemaFormat(json_schema=parameters, style="qwen_xml"), + end=tool_call_end, + ) + ) + assert len(tags) > 0 + suffix_tag = TagsWithSeparatorFormat( + tags=tags, + separator="", + at_least_one=True, + ) + + if not reasoning: + result = StructuralTag(format=suffix_tag) + else: + prefix_tag = SequenceFormat( + elements=[ + TagFormat(begin="", content=AnyTextFormat(), end=think_tag_end), + ConstStringFormat(value=think_suffix), + ] + ) + result = StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) + + return result From 9c07342fdc556efd4b70a6c4183cac699bfcd1cd Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Mon, 4 May 2026 22:13:37 +0200 Subject: [PATCH 0043/1083] [NVFP4][fix] Fix `layer.weight` -> `w13` typo in NVFP4 MOE emulation kernel preparation (#41630) Signed-off-by: Felix Marty --- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 01ac5cfa9da7..f4796243e013 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -381,7 +381,7 @@ def convert_to_nvfp4_moe_kernel_format( elif nvfp4_backend == NvFp4MoeBackend.EMULATION: # Move the E2M1 lookup table to the device now, because # `.to(device)` is not allowed during CUDA graph capture. - kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device) + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(w13.device) if a13_scale is None or a2_scale is None: raise ValueError( From be5983b874bd483edef50b75ea5ceb0b28c1ea8a Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 4 May 2026 16:35:15 -0400 Subject: [PATCH 0044/1083] [Docs] Add non-causal support to attention backend docs (#41643) Signed-off-by: Matthew Bonanni --- docs/design/attention_backends.md | 59 ++++++++++--------- .../generate_attention_backend_docs.py | 9 +++ 2 files changed, 39 insertions(+), 29 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index dc4b5402cab6..83c5fc1b435f 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -155,6 +155,7 @@ Priority is **1 = highest** (tried first). | **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) | | **Head Sizes** | Supported attention head sizes | | **Sink** | Attention sink support (for StreamingLLM) | +| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models | | **Sparse** | Sparse attention support (MLA only) | | **MM Prefix** | Multimodal prefix full attention support | | **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) | @@ -165,22 +166,22 @@ Priority is **1 = highest** (tried first). ## Standard Attention (MHA, MQA, GQA) Backends -| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ | -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A | -| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x | -| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 | -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x | -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 | -| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any | -| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A | -| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | -| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any | -| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any | +| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | +| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | +| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A | +| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | +| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | +| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | +| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | +| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | +| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | +| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any | +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any | +| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. > @@ -211,16 +212,16 @@ hardware and configuration. MLA decode backends are selected using the standard `-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`). -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ | -| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | -| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | -| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | -| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any | +| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | +| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | +| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | +| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | +| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | +| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | +| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any | +| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 73ef8b915821..c0503fd69712 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -810,6 +810,9 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None "compute_capability": compute_cap, "is_mla": is_mla_backend or check_method_overrides(class_node, "is_mla"), "supports_sink": check_method_overrides(class_node, "supports_sink"), + "supports_non_causal": check_method_overrides( + class_node, "supports_non_causal" + ), "is_sparse": check_method_overrides(class_node, "is_sparse"), "supports_mm_prefix": check_method_overrides(class_node, "supports_mm_prefix"), "supports_dcp": supports_dcp, @@ -1311,6 +1314,10 @@ def _extract_priorities(body: list, priorities: dict[str, list[str]], prefix: st _COL_BLOCK_SIZES: TableColumn = ("Block Sizes", lambda b: b["block_sizes"]) _COL_HEAD_SIZES: TableColumn = ("Head Sizes", lambda b: b["head_sizes"]) _COL_SINK: TableColumn = ("Sink", lambda b: bool_to_emoji(b["supports_sink"])) +_COL_NON_CAUSAL: TableColumn = ( + "Non-Causal", + lambda b: bool_to_emoji(b["supports_non_causal"]), +) _COL_SPARSE: TableColumn = ("Sparse", lambda b: bool_to_emoji(b["is_sparse"])) _COL_MM_PREFIX: TableColumn = ( "MM Prefix", @@ -1344,6 +1351,7 @@ def _build_columns(is_mla: bool, has_versions: bool) -> list[TableColumn]: cols.append(_COL_VERSION) cols.extend([_COL_DTYPES, _COL_KV_DTYPES, _COL_BLOCK_SIZES, _COL_HEAD_SIZES]) cols.append(_COL_SINK) + cols.append(_COL_NON_CAUSAL) if is_mla: cols.append(_COL_SPARSE) cols.extend([_COL_MM_PREFIX, _COL_DCP, _COL_ATTN_TYPES, _COL_COMPUTE_CAP]) @@ -1554,6 +1562,7 @@ def generate_legend() -> str: | **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) | | **Head Sizes** | Supported attention head sizes | | **Sink** | Attention sink support (for StreamingLLM) | +| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models | | **Sparse** | Sparse attention support (MLA only) | | **MM Prefix** | Multimodal prefix full attention support | | **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) | From 1cb08387214fba82c2705429100c5942fab65f05 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 4 May 2026 18:32:55 -0500 Subject: [PATCH 0045/1083] [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 f91ea85779d5..8ab47b618957 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -22,6 +22,7 @@ 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 068eb43faf40..873cfb18701b 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 82eecc8cd49b..20981f60cd24 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -238,6 +238,9 @@ 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 @@ def __init__( 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 0046/1083] [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 7b963b99d565..1f18b92f95b4 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 02e4d1d5d77b..57d1e1bc686b 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 cfef8449ebf8..06e8e574a05d 100644 --- a/tests/v1/distributed/test_external_lb_dp.py +++ b/tests/v1/distributed/test_external_lb_dp.py @@ -14,7 +14,7 @@ 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 6ba392802e31..95fd8787afe5 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 cd9551003339..4c37d5a149c3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -962,7 +962,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "-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 @@ def create_engine_config( 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 0047/1083] [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 | 71 +++++++++++++-- vllm/platforms/interface.py | 36 ++++++++ 4 files changed, 186 insertions(+), 27 deletions(-) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index f074ce119ae8..b9567195b3a8 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -182,22 +182,100 @@ def test_all_presets_all_head_dims(self, preset, head_dim): # ---- 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 4c37d5a149c3..1b3803139217 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1699,29 +1699,15 @@ def create_engine_config( 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 f9cfc89c0c1d..50beb8d1d9bf 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 @@ def slot_size_aligned(self) -> int: 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). - - Returns first N and last N layer indices as strings, suitable for - kv_cache_dtype_skip_layers. + def get_boundary_skip_layers( + model_config: ModelConfig, + n: int = 2, + ) -> list[str]: + """Layer indices to skip TQ compression (boundary protection). + + 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 @@ def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]: 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 @@ def from_cache_dtype(cache_dtype: str, head_dim: int) -> "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 2753326755fb..80952ced73d1 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -545,6 +545,42 @@ def _align_hybrid_block_size( 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 0048/1083] [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 | 291 +++++++++++------- .../probabilistic_rejection_sampler_utils.py | 6 +- 3 files changed, 198 insertions(+), 121 deletions(-) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 62912491492e..a02dd62026ad 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 c6b0aa364f53..efe510f16e22 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 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): 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 @@ def _sample_draft( 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 @@ def _sample_draft( 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 @@ def prefill( 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 @@ def generate_draft( 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 - - 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 - ) + # 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, + ) + + # 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 @@ def _build_draft_attn_metadata( : 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 @@ def capture( 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 @@ def capture( 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 @@ def propose( # 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 @@ def propose( # 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 @@ def propose( 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, - hidden_states, - hidden_states.stride(0), input_buffers.seq_lens, - max_model_len, draft_tokens, - output_hidden_states, - output_hidden_states.stride(0), + current_draft_step, + hidden_states, + 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 9d86372e624b..10b29433efb2 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 0049/1083] [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 e10be4af8680..d6bd2b140087 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 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_activation(activation: MoEActivation) -> bool: @@ -70,6 +70,7 @@ def _supports_activation(activation: MoEActivation) -> bool: 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 3de05cd93d36..456f40bbf7a3 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -538,9 +538,11 @@ def _get_quant_method() -> FusedMoEMethodBase: # 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 0050/1083] [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 b967ff4ede7b..82170b5fbdc4 100644 --- a/vllm/model_executor/layers/pooler/seqwise/methods.py +++ b/vllm/model_executor/layers/pooler/seqwise/methods.py @@ -68,21 +68,23 @@ def forward( "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 d06663b5b947..ae5926cd62ff 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 @@ def forward( 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 @@ def forward( ) 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 d3fefb745cfe..59b7234661b5 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/vllm/model_executor/layers/pooler/tokwise/methods.py @@ -47,17 +47,12 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 0051/1083] [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 aefc35324d86..8ed7757f6553 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -28,6 +28,25 @@ 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 dequant_mxfp4_batches(mat_fp4: torch.Tensor, scale_tensor: torch.Tensor): 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 000000000000..3906a7e057ca --- /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 ac317ac7762c..e10514debd08 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 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 3f2aca277160..437da8e6438e 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 @@ 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.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 _return_or_raise( 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 _return_or_raise( ) 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 _return_or_raise( 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 _interleave_mxfp4_cutlass_sm90(w): 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 _interleave_mxfp4_cutlass_sm90(w): 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 2be77f2b8b82..d6fef0b3d3d5 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -24,7 +24,7 @@ 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 1eeca142343b..a14bfbc9c19b 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 @@ 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 @@ __all__ = [ "QuarkMoEMethod", "QuarkOCP_MX_MoEMethod", - "QuarkOCP_MX_MoEMethod_OSS", ] @@ -94,22 +93,9 @@ def get_moe_method( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( ) # 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( "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 @@ def create_weights( 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 @@ def process_weights_after_loading(self, layer): 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 @@ def process_weights_after_loading(self, layer): 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 @@ def _setup_kernel_via_oracle(self, layer: FusedMoE): 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 @@ def apply( 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 @@ def apply_monolithic( 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 0052/1083] [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 18afe4b7925c..4750dd78838d 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 f5b473bd262a..b9987fb26c19 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 61cae12d67da..c3983e0578a5 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 000000000000..c7e1502bcb05 --- /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