From c231d1f2906ed3f9a995d84a6c7ba67e6914fb16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:08:51 +0200 Subject: [PATCH 01/68] fix(security): bound tokenizer work when explicit truncation_side is set (#47007) Signed-off-by: jperezde Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/renderers/test_completions.py | 81 +++++++++++++++++++++++++++++ vllm/renderers/params.py | 22 ++++---- 2 files changed, 94 insertions(+), 9 deletions(-) diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 76e88f4213e0..d184eb8621ce 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -64,12 +64,14 @@ class DummyTokenizer: def __post_init__(self) -> None: self._captured_encode_kwargs: dict = {} + self._captured_text_len: int = 0 def decode(self, tokens: list[int]): return str(tokens) def encode(self, text: str, **kwargs): self._captured_encode_kwargs = kwargs + self._captured_text_len = len(text) in_length = len(text) truncation = kwargs.get("truncation") @@ -366,6 +368,85 @@ def test_tokens_input_with_needs_detokenization(self): assert results[0]["prompt_token_ids"] == tokens assert results[0]["prompt"] == "[1, 2, 3, 4]" + def test_explicit_side_tokenizer_unbounded(self): + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 500) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=4, + truncation_side="left", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 4 + + kwargs = renderer.tokenizer._captured_encode_kwargs + assert kwargs["truncation"] is False + + def test_explicit_side_left_text(self): + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 50) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=5, + truncation_side="left", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 5 + assert results[0]["prompt_token_ids"] == list(range(45, 50)) + + def test_explicit_side_right_text(self): + renderer = _build_renderer(MockModelConfig()) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 50) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=5, + truncation_side="right", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 5 + assert results[0]["prompt_token_ids"] == list(range(5)) + + def test_explicit_side_text_pretokenization_guard(self): + renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1) + + prompts = renderer.render_prompts( + _preprocess_prompt(renderer.model_config, "x" * 500) + ) + results = renderer.tokenize_prompts( + prompts, + TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=4, + truncation_side="left", + ), + ) + + assert len(results) == 1 + assert len(results[0]["prompt_token_ids"]) == 4 + + assert renderer.tokenizer._captured_text_len <= 100 + class TestRenderEmbedPrompt: def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes: diff --git a/vllm/renderers/params.py b/vllm/renderers/params.py index 7e3670c738d2..a07a49230676 100644 --- a/vllm/renderers/params.py +++ b/vllm/renderers/params.py @@ -314,10 +314,11 @@ def get_encode_kwargs(self) -> dict[str, Any]: # while still failing `self._token_len_check` as expected by users max_length = self.max_input_tokens + 1 - # Explicit truncation-side overrides require the full token sequence so - # we can slice from the requested side in _token_truncation. Disable - # tokenizer-level truncation because generation tokenizers default to - # left truncation while callers may request right truncation. + # Explicit truncation-side overrides require the full token sequence + # so we can slice from the requested side in _token_truncation. + # Disable tokenizer-level truncation because its default side may + # differ from the requested side. The defense against unbounded + # tokenization lives in _text_len_check (character-level pre-trim). if self.truncation_side is not None and self.truncate_prompt_tokens is not None: return dict( truncation=False, @@ -333,15 +334,13 @@ def get_encode_kwargs(self) -> dict[str, Any]: def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: """Apply length checks to prompt text if necessary.""" max_input_tokens = self.max_input_tokens - if max_input_tokens is None: + if max_input_tokens is None or tokenizer is None: return text - if self.truncate_prompt_tokens is None and tokenizer is not None: - max_input_chars = max_input_tokens * tokenizer.max_chars_per_token + max_input_chars = max_input_tokens * tokenizer.max_chars_per_token + if self.truncate_prompt_tokens is None: if len(text) > max_input_chars: - # To save resources, fail the request outright without even - # attempting tokenization raise VLLMValidationError( f"This model's maximum context length is " f"{self.max_total_tokens} tokens. However, you requested " @@ -354,6 +353,11 @@ def _text_len_check(self, tokenizer: TokenizerLike | None, text: str) -> str: parameter="input_text", value=len(text), ) + elif self.truncation_side is not None and len(text) > max_input_chars: + if self.truncation_side == "left": + text = text[-max_input_chars:] + else: + text = text[:max_input_chars] return text From 7cf7cbcd9500028f230deff5d194da2f00a2728b Mon Sep 17 00:00:00 2001 From: tc-mb <157115220+tc-mb@users.noreply.github.com> Date: Tue, 30 Jun 2026 23:12:44 +0800 Subject: [PATCH 02/68] [Bugfix] MiniCPM-V 4.6: fix grid rows/cols swap in placeholder generation (#45918) Signed-off-by: tc-mb --- vllm/model_executor/models/minicpmv4_6.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 79ac79d709ac..6cdc3624151b 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -140,7 +140,7 @@ def get_video_prompt_texts( per_frame = image_start + video_token * source_tokens + image_end if grids[0] > 0 and grids[1] > 0 and patch_tokens > 0: slice_ph = slice_start + video_token * patch_tokens + slice_end - rows = [slice_ph * grids[0] for _ in range(grids[1])] + rows = [slice_ph * grids[1] for _ in range(grids[0])] per_frame += "\n".join(rows) body = per_frame * num_frames @@ -596,7 +596,7 @@ def get_slice_image_placeholder( if use_image_id: placeholder = f"{id_start}{image_idx}{id_end}" + placeholder - num_cols, num_rows = grids[0], grids[1] + num_rows, num_cols = grids[0], grids[1] if num_cols > 0 and num_rows > 0 and patch_tokens > 0: slice_ph = slice_start + image_token * patch_tokens + slice_end slices = [slice_ph * num_cols for _ in range(num_rows)] From dc148dc4d7633f749409f4b724a25669fd455d3d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 30 Jun 2026 17:14:13 +0200 Subject: [PATCH 03/68] [CI][Bugfix] Fix `Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)` (#47157) Signed-off-by: NickLucche --- .../nixl_integration/run_mamba_prefix_cache_test.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index c7e65972004a..a34f07edc974 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -51,6 +51,7 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ + --attention-backend FLASHINFER \ --kv-transfer-config "$KV_CONFIG" & # Start decode instance @@ -68,6 +69,7 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ + --attention-backend FLASHINFER \ --kv-transfer-config "$KV_CONFIG" & echo "Waiting for prefill instance on port $PREFILL_PORT..." From d8f483dc30a9a74b8fdf1c6dcb11cff503fe40a3 Mon Sep 17 00:00:00 2001 From: Igor Margulis Date: Tue, 30 Jun 2026 18:19:51 +0300 Subject: [PATCH 04/68] [Spec Decode] Fix hidden-state extraction block size for hybrid verifiers (#46301) Signed-off-by: Igor Margulis Signed-off-by: mgoin Co-authored-by: Cursor Co-authored-by: mgoin --- .../unit/test_hidden_states_connector.py | 126 ++++++++++++++++++ .../v1/example_hidden_states_connector.py | 73 +++++++--- .../models/extract_hidden_states.py | 3 + 3 files changed, 187 insertions(+), 15 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_hidden_states_connector.py diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py new file mode 100644 index 000000000000..ffd648289134 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only unit tests for ExampleHiddenStatesConnector KV-cache-group logic.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 + ExampleHiddenStatesConnector, +) +from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheGroupSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) + + +def _full(block_size: int) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ) + + +def _hidden(block_size: int) -> HiddenStateCacheSpec: + return HiddenStateCacheSpec( + block_size=block_size, num_kv_heads=6, head_size=2048, dtype=torch.bfloat16 + ) + + +def _config(*specs): + """Minimal stand-in exposing only ``kv_cache_groups`` (all the helpers read).""" + return SimpleNamespace( + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=[f"layer.{i}"], kv_cache_spec=spec) + for i, spec in enumerate(specs) + ] + ) + + +# ---- _find_cache_kv_group_id ------------------------------------------------ + + +def test_find_group_id_none_config_returns_zero(): + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(None) == 0 + + +def test_find_group_id_single_non_hidden_group_returns_zero(): + # Uniform (dense) model: one group, no HiddenStateCacheSpec -> group 0. + cfg = _config(_full(16)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 0 + + +def test_find_group_id_locates_hidden_group_when_not_first(): + # Hybrid layout: the hidden-states group is not group 0. + cfg = _config(_full(528), _hidden(22), _full(528)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1 + + +def test_find_group_id_locates_hidden_group_last(): + cfg = _config(_full(528), _full(528), _hidden(22)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 2 + + +def test_find_group_id_raises_when_no_hidden_group_and_multiple_groups(): + cfg = _config(_full(16), _full(16)) + with pytest.raises(ValueError, match="Could not uniquely identify"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) + + +def test_find_group_id_raises_when_multiple_hidden_groups(): + cfg = _config(_hidden(22), _hidden(22)) + with pytest.raises(ValueError, match="Could not uniquely identify"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) + + +# ---- _get_cache_block_size -------------------------------------------------- + + +def test_get_block_size_reads_hidden_group_spec_not_global(): + # Hidden group keeps block size 22; the global is bumped to 528 for hybrids. + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528)) + cfg = _config(_full(528), _hidden(22)) + block_size = ExampleHiddenStatesConnector._get_cache_block_size( + vllm_config, cfg, cache_kv_group_id=1 + ) + assert block_size == 22 + + +def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config(): + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16)) + block_size = ExampleHiddenStatesConnector._get_cache_block_size( + vllm_config, None, cache_kv_group_id=0 + ) + assert block_size == 16 + + +# ---- MLA-verifier absorption ------------------------------------------------ + + +def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier(): + # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding- + # window MLA verifier absorbs it into the MLA group instead of isolating it. + dt = torch.bfloat16 + spec = { + "layers.0.mla": MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=dt + ), + "layers.1.swa": SlidingWindowMLASpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=dt, sliding_window=512 + ), + "cache_only_layers.61": _hidden(64), + } + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + ) + groups = get_kv_cache_groups(vllm_config, spec) + assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups) + cfg = SimpleNamespace(kv_cache_groups=groups) + with pytest.raises(ValueError, match="MLA verifiers are unsupported"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index 7e6c95bf8fb1..299ff037ad2c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -109,6 +109,49 @@ def prefer_cross_layer_blocks(self) -> bool: # Must be False so that drafter kv cache isn't merged with verifier's return False + @classmethod + def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int: + """Index of the KV cache group holding the extracted hidden states. + + Located by spec type so it resolves on both scheduler and worker side. + """ + if kv_cache_config is None: + return 0 + + from vllm.v1.kv_cache_interface import HiddenStateCacheSpec + + groups = kv_cache_config.kv_cache_groups + group_ids = [ + gid + for gid, group in enumerate(groups) + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ] + if len(group_ids) == 1: + return group_ids[0] + if not group_ids and len(groups) == 1: + return 0 + raise ValueError( + "Could not uniquely identify the extract-hidden-states KV cache " + f"group among {len(groups)} groups; the hidden-states layer must be " + "isolated in its own group (MLA verifiers are unsupported)." + ) + + @staticmethod + def _get_cache_block_size( + vllm_config: "VllmConfig", + kv_cache_config: "KVCacheConfig | None", + cache_kv_group_id: int, + ) -> int: + """Block size of the hidden-states group, read from its own spec. + + cache_config.block_size is bumped to a common multiple for hybrid + verifiers; the page-aligned hidden-states group keeps a smaller one. + """ + if kv_cache_config is None: + return vllm_config.cache_config.block_size + cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id] + return cache_group.kv_cache_spec.block_size + def __init__( self, vllm_config: "VllmConfig", @@ -120,7 +163,12 @@ def __init__( role=role, kv_cache_config=kv_cache_config, ) - self._block_size = vllm_config.cache_config.block_size + # Read the hidden-states group and its block size from the group spec; + # cache_config.block_size is bumped (wrong) for hybrid verifiers. + self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config) + self._block_size = self._get_cache_block_size( + vllm_config, kv_cache_config, self._cache_kv_group_id + ) self._storage_path = self._kv_transfer_config.get_from_extra_config( "shared_storage_path", "/tmp" ) @@ -151,13 +199,6 @@ def __init__( # Worker-side state (set by register_kv_caches). self._kv_cache: torch.Tensor | None = None - # Identify which KV cache group holds the hidden-states layer. - self._hs_group_idx: int = 0 - if self._kv_cache_config is not None: - for i, group in enumerate(self._kv_cache_config.kv_cache_groups): - if any("cache_only_layers" in n for n in group.layer_names): - self._hs_group_idx = i - break # Only TP rank 0 writes hidden states to disk; other TP ranks no-op. # Set in register_kv_caches (after distributed init). self._is_tp_rank_zero: bool = True @@ -267,12 +308,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) self._kv_cache = kv_caches[self.cache_layers[0]] - # Find the KV cache group index for hidden states - if self._kv_cache_config is not None: - for i, group in enumerate(self._kv_cache_config.kv_cache_groups): - if self.cache_layers[0] in group.layer_names: - self._hs_group_idx = i - break + # Block size must match the indexed buffer, else reads hit the wrong + # slots. Raise (not assert) so the check survives `python -O`. + if self._block_size != self._kv_cache.shape[1]: + raise ValueError( + f"Hidden-states block-size mismatch: derived {self._block_size} " + f"but buffer block size is {self._kv_cache.shape[1]}; read slots " + "would be wrong (likely a hybrid block-size resolution bug)." + ) @staticmethod def _write_tensors( @@ -543,7 +586,7 @@ def request_finished_all_groups( request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - return self.request_finished(request, block_ids[self._hs_group_idx]) + return self.request_finished(request, block_ids[self._cache_kv_group_id]) @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 8df4823b6973..87dba7b5c984 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -83,7 +83,10 @@ def basic_cache( kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size] slot_mapping: torch.Tensor, # shape: [seq_len] ): + # Padding slots are -1; redirect them to the null block (block 0, never + # allocated to a request) so the scatter stays branch-free and sync-free. block_size = kv_cache.shape[1] + slot_mapping = slot_mapping.clamp_min(0) kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache From 9e84ec86486df4aa5872c43a813e70772ab15f3c Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 30 Jun 2026 11:29:21 -0400 Subject: [PATCH 05/68] [Refactor] Remove dead minimax allreduce rms kernel (#46842) Signed-off-by: yewentao256 --- .../minimax_reduce_rms_kernel.cu | 29 ------------------- csrc/libtorch_stable/ops.h | 4 --- csrc/libtorch_stable/torch_bindings.cpp | 5 ---- vllm/_custom_ops.py | 14 --------- 4 files changed, 52 deletions(-) diff --git a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu index d9af0f5efe0f..da7fee5670f5 100644 --- a/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu +++ b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu @@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) { } // namespace tensorrt_llm } // namespace vllm -torch::stable::Tensor minimax_allreduce_rms( - torch::stable::Tensor const& input, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, - int64_t const rank, int64_t const nranks, double const eps) { - const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); - - allreduce_params.nranks = static_cast(nranks); - allreduce_params.rank = static_cast(rank); - allreduce_params.dtype = input.scalar_type(); - allreduce_params.size_q = static_cast(input.numel()); - allreduce_params.hidden_dim = static_cast(input.size(-1)); - allreduce_params.stride_q = allreduce_params.hidden_dim; - allreduce_params.workspace = - reinterpret_cast(workspace.mutable_data_ptr()); - allreduce_params.allreduce_in = const_cast(input.const_data_ptr()); - allreduce_params.rms_gamma = const_cast(norm_weight.const_data_ptr()); - allreduce_params.rms_eps = static_cast(eps); - allreduce_params.stream = get_current_cuda_stream(input.get_device_index()); - - torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input); - allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr(); - - vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params); - - return rms_norm_out; -} - std::tuple minimax_allreduce_rms_qk(torch::stable::Tensor qkv, torch::stable::Tensor const& norm_weight_q, diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index d60b68a5868d..7cf34d8b03ab 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -288,10 +288,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( int64_t cache_block_size); #ifndef USE_ROCM -torch::stable::Tensor minimax_allreduce_rms( - torch::stable::Tensor const& input, - torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, - int64_t const rank, int64_t const nranks, double const eps); std::tuple minimax_allreduce_rms_qk(torch::stable::Tensor qkv, torch::stable::Tensor const& norm_weight_q, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 1be7217ce789..158999a6633b 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -449,10 +449,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int cache_block_size) -> ()"); #ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms(" - "Tensor input, Tensor norm_weight, Tensor workspace, " - "int rank, int nranks, float eps) -> Tensor"); ops.def( "minimax_allreduce_rms_qk(" "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " @@ -705,7 +701,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); #ifndef USE_ROCM - ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms)); ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 02404a2f5173..7dcb890aecef 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3785,20 +3785,6 @@ def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor: return torch.empty_like(x) if not inplace else x -if hasattr(torch.ops._C, "minimax_allreduce_rms"): - - @register_fake("_C::minimax_allreduce_rms") - def _minimax_allreduce_rms_fake( - input: torch.Tensor, - norm_weight: torch.Tensor, - workspace: torch.Tensor, - rank: int, - nranks: int, - eps: float, - ) -> torch.Tensor: - return torch.empty_like(input) - - if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): @register_fake("_C::minimax_allreduce_rms_qk") From fcaa84efa7a980cee3681976ecfe36133c48251a Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 30 Jun 2026 16:31:27 +0100 Subject: [PATCH 06/68] [BugFix] Gate MRV2 mixed sparse-MLA warmup on `max_num_seqs` > 1 (#47050) Signed-off-by: Nick Hill Co-authored-by: ziminghuang --- tests/v1/worker/test_mixed_warmup_gate.py | 30 +++++++++++++++++++ .../warmup/flashinfer_sparse_mla_warmup.py | 6 ++-- vllm/v1/worker/gpu/warmup.py | 2 +- 3 files changed, 34 insertions(+), 4 deletions(-) create mode 100644 tests/v1/worker/test_mixed_warmup_gate.py diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py new file mode 100644 index 000000000000..6941df6773c1 --- /dev/null +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the max_num_reqs gate on the V2 mixed prefill+decode warmup.""" + +from types import SimpleNamespace + +import pytest + +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup + + +def _fail(*args, **kwargs): + raise AssertionError("worker callback must not run when warmup is skipped") + + +@pytest.mark.parametrize("max_num_reqs", [1, 0]) +def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): + """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 + the warmup must be skipped without touching the worker callbacks.""" + runner = SimpleNamespace(is_pooling_model=False, max_num_reqs=max_num_reqs) + + assert ( + run_mixed_prefill_decode_warmup( + runner, + worker_execute_model=_fail, + worker_sample_tokens=_fail, + num_tokens=128, + ) + is False + ) diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 44be769e246e..80dfd5199614 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -130,7 +130,7 @@ def _run_flashinfer_sparse_mla_decode_autotune( with torch.inference_mode(): warmup_executed = True if is_leader: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) warmup_executed = run_mixed_prefill_decode_warmup( v2_runner, @@ -144,7 +144,7 @@ def _run_flashinfer_sparse_mla_decode_autotune( with flashinfer_autotune(True, cache=str(cache_path)): runner._dummy_run(**dummy_run_kwargs) else: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) warmup_executed = run_mixed_prefill_decode_warmup( v2_runner, @@ -236,7 +236,7 @@ def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None: ) mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens) if not mixed_warmup_done: - if _uses_v2_model_runner(runner): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: v2_runner = cast("V2GPUModelRunner", runner) run_mixed_prefill_decode_warmup( v2_runner, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index ff9a75f0ef1a..de47cb8880a3 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -34,7 +34,7 @@ def run_mixed_prefill_decode_warmup( req_id_prefix: str = "_v2_mixed_warmup", ) -> bool: """Run a V2 mixed prefill+decode step through normal scheduler inputs.""" - if model_runner.is_pooling_model or num_tokens < 3: + if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3: return False decode_req_id = f"{req_id_prefix}_decode_" From b5633f9ec2beb210bdaf08d9180b5b0fc068d305 Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Tue, 26 May 2026 21:03:47 -0500 Subject: [PATCH 07/68] Apply gfx1250_wip_dllehr gfx1250 enablement onto upstream main Squashed application of rocm/vllm gfx1250_wip_dllehr (19 commits, 3 internal merges) onto current origin/main (~2666 commits newer). Conflict resolutions: - docker/Dockerfile.rocm, docker/Dockerfile.rocm_base, .buildkite/ release-pipeline.yaml, .buildkite/scripts/annotate-rocm-release.sh: kept WIP versions (TheRock-based ROCm build automation). - CMakeLists.txt: upstream HIP_SUPPORTED_ARCHS list + gfx1250. - vllm/platforms/rocm.py: kept upstream _ON_GFX12X/_ON_GFX90A AND added gfx1250 to _ON_MI3XX/_ON_GFX9 (WIP intent). - vllm/model_executor/layers/quantization/utils/mxfp4_utils.py: re-applied the live gfx1250 scale_layout change (on_gfx950() or on_gfx1250()); dropped the stale _can_support_mxfp4/get_padding_alignment (upstream removed them + callers). - vllm/_aiter_ops.py (new), vllm/envs.py, vllm/v1/worker/gpu_worker.py, csrc/quickreduce/base.h: applied cleanly. - vllm/model_executor/layers/attention/mla_attention.py: resolved to upstream; the WIP's ROCm flash_attn fallback was relocated upstream to vllm/v1/attention/backends/fa_utils.py. DEFERRED: porting the aiter.ops.triton.mha.flash_attn_varlen_func substitution there. - vllm/model_executor/layers/quantization/quark/quark_moe.py: resolved to upstream (it was rewritten +747/-304 into a backend/kernel-factory model). DEFERRED: the VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4 shuffle-skip guardrail now belongs in fused_moe/oracle/mxfp4.py (AITER_MXFP4_MXFP4 W4A4 backend selection), an area being reworked separately. AI-assisted (Claude Code); human review required before any PR, especially the two DEFERRED ports above. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- .buildkite/release-pipeline.yaml | 693 ++++++++---------- .buildkite/scripts/annotate-rocm-release.sh | 15 +- CMakeLists.txt | 2 +- csrc/quickreduce/base.h | 4 +- docker/Dockerfile.rocm | 343 +++++---- docker/Dockerfile.rocm_base | 159 ++-- vllm/_aiter_ops.py | 162 ++++ vllm/envs.py | 6 + .../layers/quantization/utils/mxfp4_utils.py | 4 +- vllm/platforms/rocm.py | 11 +- vllm/v1/worker/gpu_worker.py | 5 +- 11 files changed, 744 insertions(+), 660 deletions(-) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 897c98145341..7029c39f994b 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -38,13 +38,14 @@ steps: depends_on: ~ id: build-wheel-arm64-cuda-12-9 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge 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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: + # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 + - "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=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.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" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -52,13 +53,14 @@ steps: depends_on: ~ id: build-wheel-arm64-cuda-13-0 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge 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_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: + # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 + - "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.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --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" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" @@ -66,13 +68,12 @@ steps: depends_on: ~ id: build-wheel-arm64-cpu agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." - "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" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" @@ -80,13 +81,12 @@ steps: depends_on: ~ id: build-wheel-x86-cuda-12-9 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge 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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --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=12.9.1 --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" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31" env: DOCKER_BUILDKIT: "1" @@ -94,13 +94,12 @@ steps: depends_on: ~ id: build-wheel-x86-cuda-13-0 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge 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_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 ." + - "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.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --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" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" @@ -108,265 +107,80 @@ steps: depends_on: ~ id: build-wheel-x86-cpu agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." - "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" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' + - "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35" env: DOCKER_BUILDKIT: "1" - - label: "Generate and upload wheel indices" - depends_on: "build-wheels" - allow_dependency_failure: true - agents: - queue: cpu_queue_release - commands: - - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" - - - block: "Unblock to build release Docker images" - depends_on: ~ - key: block-build-release-images - if: build.env("NIGHTLY") != "1" - - group: "Build release Docker images" key: "build-release-images" - depends_on: block-build-release-images - allow_dependency_failure: true steps: - - label: "Build release image - x86_64 - CUDA 13.0" + - label: "Build release image - x86_64 - CUDA 12.9" depends_on: ~ id: build-release-image-x86 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh) \ - --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 INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ - --target vllm-openai \ - --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=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" # re-tag to default image tag and push, just in case arm64 build fails - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"' - - label: "Build release image - aarch64 - CUDA 13.0" + - label: "Build release image - aarch64 - CUDA 12.9" depends_on: ~ id: build-release-image-arm64 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh) \ - --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 INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ - --target vllm-openai \ - --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=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"' - - label: "Build release image - x86_64 - CUDA 12.9" + - label: "Build release image - x86_64 - CUDA 13.0" depends_on: ~ - id: build-release-image-x86-cuda-12-9 + id: build-release-image-x86-cuda-13-0 agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" + - "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.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" # re-tag to default image tag and push, just in case arm64 build fails - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"' + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" - - label: "Build release image - aarch64 - CUDA 12.9" - depends_on: ~ - id: build-release-image-arm64-cuda-12-9 - agents: - queue: arm64_cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"' - - - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04" - depends_on: ~ - id: build-release-image-x86-ubuntu2404 - agents: - queue: cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \ - --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 UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"' - - - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04" - depends_on: ~ - id: build-release-image-arm64-ubuntu2404 - agents: - queue: arm64_cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \ - --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 UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"' - - - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04" - depends_on: ~ - id: build-release-image-x86-cuda-12-9-ubuntu2404 - agents: - queue: cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' - - - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04" + - label: "Build release image - aarch64 - CUDA 13.0" depends_on: ~ - id: build-release-image-arm64-cuda-12-9-ubuntu2404 + id: build-release-image-arm64-cuda-13-0 agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - | - DOCKER_BUILDKIT=1 docker build \ - $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \ - --build-arg max_jobs=16 \ - --build-arg USE_SCCACHE=1 \ - --build-arg GIT_REPO_CHECK=1 \ - --build-arg CUDA_VERSION=12.9.1 \ - --build-arg UBUNTU_VERSION=24.04 \ - --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ - --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ - --build-arg INSTALL_KV_CONNECTORS=true \ - --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ - --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ - --target vllm-openai \ - --progress plain \ - -f docker/Dockerfile . - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' + # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark + - "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.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" - block: "Build release image for x86_64 CPU" key: block-cpu-release-image-build depends_on: ~ - label: "Build release image - x86_64 - CPU" - key: build-cpu-release-image-x86 depends_on: - block-cpu-release-image-build - input-release-version agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest" - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"' env: DOCKER_BUILDKIT: "1" @@ -375,82 +189,61 @@ steps: depends_on: ~ - label: "Build release image - arm64 - CPU" - key: build-cpu-release-image-arm64 - depends_on: + depends_on: - block-arm64-cpu-release-image-build - input-release-version agents: - queue: arm64_cpu_queue_release + queue: arm64_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest" - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"' env: DOCKER_BUILDKIT: "1" - group: "Publish release images" key: "publish-release-images" steps: - - label: "Create multi-arch manifest - CUDA 13.0" + - label: "Create multi-arch manifest - CUDA 12.9" depends_on: - build-release-image-x86 - build-release-image-arm64 id: create-multi-arch-manifest agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend" - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"' - - - label: "Create multi-arch manifest - CUDA 12.9" - depends_on: - - build-release-image-x86-cuda-12-9 - - build-release-image-arm64-cuda-12-9 - id: create-multi-arch-manifest-cuda-12-9 - agents: - queue: small_cpu_queue_release - commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"' - - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04" + - label: "Annotate release workflow - CUDA 12.9" depends_on: - - build-release-image-x86-ubuntu2404 - - build-release-image-arm64-ubuntu2404 - id: create-multi-arch-manifest-ubuntu2404 + - create-multi-arch-manifest + id: annotate-release-workflow agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"' + - "bash .buildkite/scripts/annotate-release.sh" - - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04" + - label: "Create multi-arch manifest - CUDA 13.0" depends_on: - - build-release-image-x86-cuda-12-9-ubuntu2404 - - build-release-image-arm64-cuda-12-9-ubuntu2404 - id: create-multi-arch-manifest-cuda-12-9-ubuntu2404 + - build-release-image-x86-cuda-13-0 + - build-release-image-arm64-cuda-13-0 + id: create-multi-arch-manifest-cuda-13-0 agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"' + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" - label: "Publish nightly multi-arch image to DockerHub" depends_on: - create-multi-arch-manifest if: build.env("NIGHTLY") == "1" agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - "bash .buildkite/scripts/push-nightly-builds.sh" # Clean up old nightly builds (keep only last 14) @@ -463,16 +256,16 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" - - label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9" + - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0" depends_on: - - create-multi-arch-manifest-cuda-12-9 + - create-multi-arch-manifest-cuda-13-0 if: build.env("NIGHTLY") == "1" agents: - queue: small_cpu_queue_release + queue: small_cpu_queue_postmerge commands: - - "bash .buildkite/scripts/push-nightly-builds.sh cu129" + - "bash .buildkite/scripts/push-nightly-builds.sh cu130" # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-" plugins: - docker-login#v3.0.0: username: vllmbot @@ -481,6 +274,24 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" + - group: "Publish wheels" + key: "publish-wheels" + steps: + - block: "Confirm update release wheels to PyPI (experimental, use with caution)?" + key: block-upload-release-wheels + depends_on: + - input-release-version + - build-wheels + + - label: "Upload release wheels to PyPI" + depends_on: + - block-upload-release-wheels + id: upload-release-wheels + agents: + queue: small_cpu_queue_postmerge + commands: + - "bash .buildkite/scripts/upload-release-wheels-pypi.sh" + # ============================================================================= # ROCm Release Pipeline (x86_64 only) # ============================================================================= @@ -489,112 +300,184 @@ steps: # To build a specific version, trigger the build from that branch/tag. # # Environment variables for ROCm builds (set via Buildkite UI or schedule): + # ROCM_PYTHON_VERSION: Python version (default: 3.12) + # PYTORCH_ROCM_ARCH: GPU architectures (default: gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151) + # ROCM_UPLOAD_WHEELS: Upload to S3 (default: false for nightly, true for releases) + # ROCM_FORCE_REBUILD: Force rebuild base wheels, ignore S3 cache (default: false) # # Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base + # (currently rocm/dev-ubuntu-22.04:7.1-complete) # # ============================================================================= + # ROCm Input Step - Collect build configuration (manual trigger only) + - input: "ROCm Wheel Release Build Configuration" + key: input-rocm-config + depends_on: ~ + if: build.source == "ui" + fields: + - text: "Python Version" + key: "rocm-python-version" + default: "3.12" + hint: "Python version (e.g., 3.12)" + - text: "GPU Architectures" + key: "rocm-pytorch-rocm-arch" + default: "gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" + hint: "Semicolon-separated GPU architectures" + - select: "Upload Wheels to S3" + key: "rocm-upload-wheels" + default: "true" + options: + - label: "No - Build only (nightly/dev)" + value: "false" + - label: "Yes - Upload to S3 (release)" + value: "true" + - select: "Force Rebuild Base Wheels" + key: "rocm-force-rebuild" + default: "false" + hint: "Ignore S3 cache and rebuild base wheels from scratch" + options: + - label: "No - Use cached wheels if available" + value: "false" + - label: "Yes - Rebuild even if cache exists" + value: "true" + # ROCm Job 1: Build ROCm Base Wheels (with S3 caching) - - label: ":rocm: Build ROCm Base Image & Wheels" + - label: ":rocm: Build ROCm Base Wheels" id: build-rocm-base-wheels - depends_on: ~ + depends_on: + - step: input-rocm-config + allow_failure: true # Allow failure so non-UI builds can proceed (input step is skipped) agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: + # Set configuration and check cache - | set -euo pipefail - # Generate cache key - CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base" + # Get values from meta-data (set by input step) or use defaults + PYTHON_VERSION="$$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo '')" + export PYTHON_VERSION="$${PYTHON_VERSION:-3.12}" + + PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" + export PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" + + # Check for force rebuild flag + ROCM_FORCE_REBUILD="$${ROCM_FORCE_REBUILD:-}" + if [ -z "$${ROCM_FORCE_REBUILD}" ]; then + ROCM_FORCE_REBUILD="$$(buildkite-agent meta-data get rocm-force-rebuild 2>/dev/null || echo '')" + fi echo "========================================" - echo "ROCm Base Build Configuration" + echo "ROCm Base Wheels Build Configuration" echo "========================================" - echo " CACHE_KEY: $${CACHE_KEY}" - echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}" + echo " PYTHON_VERSION: $${PYTHON_VERSION}" + echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}" + echo " ROCM_FORCE_REBUILD: $${ROCM_FORCE_REBUILD:-false}" echo "========================================" - - # Login to ECR - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - IMAGE_EXISTS=false - WHEELS_EXIST=false - - # Check ECR for Docker image - if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then - IMAGE_EXISTS=true - echo "ECR image cache HIT" - fi - - # Check S3 for wheels - WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check) - if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then - WHEELS_EXIST=true - echo "S3 wheels cache HIT" + # Save resolved config for later jobs + buildkite-agent meta-data set "rocm-python-version" "$${PYTHON_VERSION}" + buildkite-agent meta-data set "rocm-pytorch-rocm-arch" "$${PYTORCH_ROCM_ARCH}" + + # Check S3 cache for pre-built wheels + CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + CACHE_PATH=$$(.buildkite/scripts/cache-rocm-base-wheels.sh path) + echo "" + echo "Cache key: $${CACHE_KEY}" + echo "Cache path: $${CACHE_PATH}" + + # Save cache key for downstream jobs + buildkite-agent meta-data set "rocm-cache-key" "$${CACHE_KEY}" + + CACHE_STATUS="miss" + if [ "$${ROCM_FORCE_REBUILD}" != "true" ]; then + CACHE_STATUS=$$(.buildkite/scripts/cache-rocm-base-wheels.sh check) + else + echo "Force rebuild requested, skipping cache check" fi - - # Scenario 1: Both cached (best case) - if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then + if [ "$${CACHE_STATUS}" = "hit" ]; then echo "" - echo "FULL CACHE HIT - Reusing both image and wheels" + echo "CACHE HIT! Downloading pre-built wheels..." echo "" - - # Download wheels .buildkite/scripts/cache-rocm-base-wheels.sh download - - # Save ECR tag for downstream jobs - buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" - - # Scenario 2: Full rebuild needed + + # Set the S3 path for the cached Docker image (for Job 2 to download) + S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}" + buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + + # Mark that we used cache (for Docker image handling) + buildkite-agent meta-data set "rocm-used-cache" "true" + + echo "" + echo "Cache download complete. Skipping Docker build." + echo "Docker image will be downloaded from: $${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" else echo "" - echo " CACHE MISS - Building from scratch..." + echo "CACHE MISS. Building from scratch..." echo "" - - # Build full base image and push to ECR + + # Build full base image (for later vLLM build) DOCKER_BUILDKIT=1 docker buildx build \ --file docker/Dockerfile.rocm_base \ - --tag "$${ECR_CACHE_TAG}" \ + --tag rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} \ + --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ + --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --push \ + --load \ . - - # Build wheel extraction stage + + # Build debs_wheel_release stage for wheel extraction DOCKER_BUILDKIT=1 docker buildx build \ --file docker/Dockerfile.rocm_base \ --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \ --target debs_wheel_release \ + --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ + --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ --load \ . - - # Extract and upload wheels + + # Extract wheels from Docker image mkdir -p artifacts/rocm-base-wheels - cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) - docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/ - docker rm $${cid} - + container_id=$$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) + docker cp $${container_id}:/app/debs/. artifacts/rocm-base-wheels/ + docker rm $${container_id} + echo "Extracted base wheels:" + ls -lh artifacts/rocm-base-wheels/ + + # Upload wheels to S3 cache for future builds + echo "" + echo "Uploading wheels to S3 cache..." .buildkite/scripts/cache-rocm-base-wheels.sh upload - # Cache base docker image to ECR - docker push "$${ECR_CACHE_TAG}" - - buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" - + # Export base Docker image for reuse in vLLM build + mkdir -p artifacts/rocm-docker-image + docker save rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} | gzip > artifacts/rocm-docker-image/rocm-base-image.tar.gz + echo "Docker image size:" + ls -lh artifacts/rocm-docker-image/ + + # Upload large Docker image to S3 (also cached by cache key) + S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}" + echo "Uploading Docker image to $${S3_ARTIFACT_PATH}/" + aws s3 cp artifacts/rocm-docker-image/rocm-base-image.tar.gz "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + + # Save the S3 path for downstream jobs + buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + + # Mark that we did NOT use cache + buildkite-agent meta-data set "rocm-used-cache" "false" + echo "" - echo " Build complete - Image and wheels cached" + echo "Build complete. Wheels cached for future builds." fi - artifact_paths: - "artifacts/rocm-base-wheels/*.whl" env: @@ -608,7 +491,7 @@ steps: - step: build-rocm-base-wheels allow_failure: false agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge timeout_in_minutes: 180 commands: # Download artifacts and prepare Docker image @@ -638,25 +521,31 @@ steps: echo "Downloading wheel artifacts from current build" buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . - # Get ECR image tag from metadata (set by build-rocm-base-wheels) - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rocm-base-image-tag metadata not found" + # Download Docker image from S3 (too large for Buildkite artifacts) + DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')" + if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then + echo "ERROR: rocm-docker-image-s3-path metadata not found" echo "This should have been set by the build-rocm-base-wheels job" exit 1 fi - - echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" - - # Login to ECR - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Pull base Docker image from ECR - docker pull "$${ECR_IMAGE_TAG}" - - echo "Loaded base image: $${ECR_IMAGE_TAG}" - + echo "Downloading Docker image from $${DOCKER_IMAGE_S3_PATH}" + mkdir -p artifacts/rocm-docker-image + aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz + + # Load base Docker image and capture the tag + echo "Loading base Docker image..." + LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load) + echo "$${LOAD_OUTPUT}" + # Extract the actual loaded image tag from "Loaded image: " output + # This avoids picking up stale images (like rocm/vllm-dev:nightly) already on the agent + BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //') + if [ -z "$${BASE_IMAGE_TAG}" ]; then + echo "ERROR: Failed to extract image tag from docker load output" + echo "Load output was: $${LOAD_OUTPUT}" + exit 1 + fi + echo "Loaded base image: $${BASE_IMAGE_TAG}" + # Prepare base wheels for Docker build context mkdir -p docker/context/base-wheels touch docker/context/base-wheels/.keep @@ -664,11 +553,16 @@ steps: echo "Base wheels for vLLM build:" ls -lh docker/context/base-wheels/ + # Get GPU architectures from meta-data + PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" + PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" + echo "========================================" echo "Building vLLM wheel with:" echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" + echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}" + echo " BASE_IMAGE: $${BASE_IMAGE_TAG}" echo "========================================" # Build vLLM wheel using local checkout (REMOTE_VLLM=0) @@ -676,7 +570,8 @@ steps: --file docker/Dockerfile.rocm \ --target export_vllm_wheel_release \ --output type=local,dest=rocm-dist \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ + --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \ + --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ --build-arg REMOTE_VLLM=0 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg USE_SCCACHE=1 \ @@ -684,8 +579,10 @@ steps: --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ . + echo "Built vLLM wheel:" ls -lh rocm-dist/*.whl + # Copy wheel to artifacts directory mkdir -p artifacts/rocm-vllm-wheel cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/ @@ -704,13 +601,35 @@ steps: - step: build-rocm-vllm-wheel allow_failure: false agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge timeout_in_minutes: 60 commands: # Download all wheel artifacts and run upload - | set -euo pipefail + # Check if upload is enabled (from env var, meta-data, or release branch) + ROCM_UPLOAD_WHEELS="$${ROCM_UPLOAD_WHEELS:-}" + if [ -z "$${ROCM_UPLOAD_WHEELS}" ]; then + # Try to get from meta-data (input form) + ROCM_UPLOAD_WHEELS="$$(buildkite-agent meta-data get rocm-upload-wheels 2>/dev/null || echo '')" + fi + + echo "========================================" + echo "Upload check:" + echo " ROCM_UPLOAD_WHEELS: $${ROCM_UPLOAD_WHEELS}" + echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" + echo "========================================" + + # Skip upload if not enabled + if [ "$${ROCM_UPLOAD_WHEELS}" != "true" ]; then + echo "Skipping S3 upload (ROCM_UPLOAD_WHEELS != true, NIGHTLY != 1, not a release branch)" + echo "To enable upload, set 'Upload Wheels to S3' to 'Yes' in the build configuration" + exit 0 + fi + + echo "Upload enabled, proceeding..." + # Download artifacts from current build echo "Downloading artifacts from current build" buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . @@ -726,9 +645,12 @@ steps: - label: ":memo: Annotate ROCm wheel release" id: annotate-rocm-release depends_on: - - upload-rocm-wheels + - step: upload-rocm-wheels + allow_failure: true + - step: input-release-version + allow_failure: true agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "bash .buildkite/scripts/annotate-rocm-release.sh" env: @@ -745,60 +667,61 @@ steps: depends_on: block-generate-root-index-rocm-wheels id: generate-root-index-rocm-wheels agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge commands: - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" env: S3_BUCKET: "vllm-wheels" VARIANT: "rocm723" - # ROCm Job 6: Build ROCm Release Docker Image + # ROCm Job 5: Build ROCm Release Docker Image - label: ":docker: Build release image - x86_64 - ROCm" id: build-rocm-release-image depends_on: - - step: block-build-release-images - allow_failure: true - step: build-rocm-base-wheels allow_failure: false agents: - queue: cpu_queue_release + queue: cpu_queue_postmerge timeout_in_minutes: 60 commands: - | set -euo pipefail - + # Login to ECR aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Get ECR image tag from metadata (set by build-rocm-base-wheels) - ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" - if [ -z "$${ECR_IMAGE_TAG}" ]; then - echo "ERROR: rocm-base-image-tag metadata not found" - echo "This should have been set by the build-rocm-base-wheels job" + + # Download Docker image from S3 (set by build-rocm-base-wheels) + DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')" + if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then + echo "ERROR: rocm-docker-image-s3-path metadata not found" exit 1 fi - - echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" - - # Pull base Docker image from ECR - docker pull "$${ECR_IMAGE_TAG}" - - echo "Loaded base image: $${ECR_IMAGE_TAG}" - - # Pass the base image ECR tag to downstream steps (nightly publish) - buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}" - - echo "========================================" - echo "Building vLLM ROCm release image with:" - echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" - echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" - echo "========================================" - + + echo "Downloading base image from $${DOCKER_IMAGE_S3_PATH}" + mkdir -p artifacts/rocm-docker-image + aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz + + # Load base Docker image + echo "Loading base Docker image..." + LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load) + BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //') + echo "Loaded base image: $${BASE_IMAGE_TAG}" + + # Tag and push the base image to ECR + docker tag "$${BASE_IMAGE_TAG}" public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base + docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base + echo "Pushed base image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base" + + # Get GPU architectures from meta-data + PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" + PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" + # Build vLLM ROCm release image using cached base DOCKER_BUILDKIT=1 docker build \ --build-arg max_jobs=16 \ - --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ + --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \ + --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ @@ -807,14 +730,10 @@ steps: --target vllm-openai \ --progress plain \ -f docker/Dockerfile.rocm . - + # Push to ECR docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm - - echo "" - echo " Successfully built and pushed ROCm release image" - echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" - echo "" + echo "Pushed: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" env: DOCKER_BUILDKIT: "1" S3_BUCKET: "vllm-wheels" diff --git a/.buildkite/scripts/annotate-rocm-release.sh b/.buildkite/scripts/annotate-rocm-release.sh index d66129722748..f42f77fc9a85 100755 --- a/.buildkite/scripts/annotate-rocm-release.sh +++ b/.buildkite/scripts/annotate-rocm-release.sh @@ -5,21 +5,20 @@ # Generate Buildkite annotation for ROCm wheel release set -ex -# Extract build configuration from Dockerfile.rocm_base (single source of truth) +# Get build configuration from meta-data # Extract ROCm version dynamically from Dockerfile.rocm_base # BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0" ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown") -PYTHON_VERSION=$(grep '^ARG PYTHON_VERSION=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTHON_VERSION=//') -PYTORCH_ROCM_ARCH=$(grep '^ARG PYTORCH_ROCM_ARCH=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTORCH_ROCM_ARCH=//') +PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12") +PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151") +# TODO: Enable the nightly build for ROCm # Get release version, default to 1.0.0.dev for nightly/per-commit builds RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "") if [ -z "${RELEASE_VERSION}" ]; then RELEASE_VERSION="1.0.0.dev" fi -ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - # S3 URLs S3_BUCKET="${S3_BUCKET:-vllm-wheels}" S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}" @@ -69,7 +68,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl . +aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl . \`\`\` @@ -81,7 +80,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash- - **torchvision**: TorchVision for ROCm PyTorch - **torchaudio**: Torchaudio for ROCm PyTorch - **amdsmi**: AMD SMI Python bindings -- **amd_aiter**: Aiter for ROCm +- **aiter**: Aiter for ROCm - **flash-attn**: Flash Attention for ROCm ### :warning: Notes @@ -97,7 +96,7 @@ To download and upload the image: docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base docker push vllm/vllm-openai-rocm:latest-base diff --git a/CMakeLists.txt b/CMakeLists.txt index 1ef9d596aec5..5d7672fd9222 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -49,7 +49,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") # ROCm installation prefix. Default to /opt/rocm but allow override via # -DROCM_PATH=/your/rocm/path when invoking cmake. diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index 6c3456d06f20..86aaf14c64a4 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -81,11 +81,11 @@ union BufferResource { __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( int32x4_t srsrc, int32_t voffset, int32_t soffset, - int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32"); + int32_t aux) {} __quickreduce_device_inline__ static void buffer_store_dwordx4( int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, - int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32"); + int32_t aux) {} __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) { #if defined(__gfx942__) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index c17444217f02..046ec3c7090f 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -33,6 +33,8 @@ ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ + software-properties-common build-essential git curl sudo vim less \ + libopenmpi-dev libpci-dev python3-venv python3-dev \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ libnuma-dev ccache mold @@ -44,9 +46,9 @@ RUN --mount=type=cache,target=/root/.cache/pip \ # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE RUN if [ "$USE_SCCACHE" != "1" ]; then \ - apt-get purge -y sccache || true; \ - python3 -m pip uninstall -y sccache || true; \ - rm -f "$(which sccache)" || true; \ + apt-get purge -y sccache || true; \ + python3 -m pip uninstall -y sccache || true; \ + rm -f "$(which sccache)" || true; \ fi # Install UV — download first, then run, so a curl failure is not masked by the pipe @@ -76,21 +78,21 @@ ARG SCCACHE_BUCKET_NAME ARG SCCACHE_REGION_NAME ARG SCCACHE_S3_NO_CREDENTIALS RUN if [ "$USE_SCCACHE" = "1" ]; then \ - if command -v sccache >/dev/null 2>&1; then \ - echo "sccache already installed, skipping installation"; \ - sccache --version; \ - else \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi; \ + if command -v sccache >/dev/null 2>&1; then \ + echo "sccache already installed, skipping installation"; \ + sccache --version; \ + else \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi; \ fi # Set sccache environment variables only when USE_SCCACHE=1 @@ -101,9 +103,44 @@ 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} +# Install uv for faster pip installs (required by build_rixl, test, and final stages) +# Base images may not include uv; install here so all stages that use uv have it +RUN python3 -m pip install uv + ARG COMMON_WORKDIR WORKDIR ${COMMON_WORKDIR} +# Ensure core build tools are installed (idempotent if base image already has them) +RUN uv pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 + +# ROCm SDK install + init (idempotent -- skips if already satisfied) +ARG PIP_EXTRA_INDEX_URL=https://rocm.genesis.amd.com/whl/gfx1250/ +RUN uv pip install --index-url ${PIP_EXTRA_INDEX_URL} "rocm[devel]" \ + && rocm-sdk init + +# ROCm SDK environment setup +ENV SITE_PACKAGES=${SITE_PACKAGES:-/opt/venv/lib/python3.12/site-packages} +ENV ROCM_PATH=$SITE_PACKAGES/_rocm_sdk_devel +ENV HIP_DEVICE_LIB_PATH=$SITE_PACKAGES/_rocm_sdk_core/lib/llvm/amdgcn/bitcode + +# Install amd_smi from SDK source and patch rtld_global (guarded for idempotency) +RUN if [ -d "$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi" ]; then \ + cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install .; \ + else \ + echo "amd_smi source not found at $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi, skipping"; \ + fi \ + && if [ -f "$SITE_PACKAGES/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' $SITE_PACKAGES/rocm_sdk/__init__.py; \ + else \ + echo "rocm_sdk/__init__.py not found at $SITE_PACKAGES/rocm_sdk/__init__.py, skipping patch"; \ + fi +ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi:$PYTHONPATH + +ENV CMAKE_PREFIX_PATH=$SITE_PACKAGES/torch/share/cmake +RUN if [ -d "$SITE_PACKAGES/_rocm_sdk_core/lib" ] && [ ! -e "$SITE_PACKAGES/_rocm_sdk_core/lib/libamdhip64.so" ]; then \ + ln -s libamdhip64.so.7 $SITE_PACKAGES/_rocm_sdk_core/lib/libamdhip64.so; \ + fi +ENV LD_LIBRARY_PATH=$SITE_PACKAGES/_rocm_sdk_core/lib:$LD_LIBRARY_PATH # ----------------------- # vLLM fetch stages @@ -115,12 +152,12 @@ ARG VLLM_BRANCH="main" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ - && cd vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git checkout FETCH_HEAD \ - && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ - git remote add upstream "https://github.com/vllm-project/vllm.git" \ - && git fetch upstream ; fi + && cd vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git checkout FETCH_HEAD \ + && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ + git remote add upstream "https://github.com/vllm-project/vllm.git" \ + && git fetch upstream ; fi FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm # ----------------------- @@ -132,7 +169,7 @@ ARG COMMON_WORKDIR # protoc is used by tonic-build/prost-build. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - ca-certificates curl unzip \ + ca-certificates curl unzip \ && rm -rf /var/lib/apt/lists/* COPY tools/install_protoc.sh /tmp/install_protoc.sh @@ -238,10 +275,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages FROM base AS build_rixl -ARG RIXL_BRANCH="39be1de8" +ARG RIXL_BRANCH="f33a5599" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" -ARG UCX_BRANCH="bfb51733" -ARG UCX_REPO="https://github.com/openucx/ucx.git" +ARG UCX_BRANCH="da3fac2a" +ARG UCX_REPO="https://github.com/ROCm/ucx.git" ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl @@ -276,16 +313,16 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ mkdir build && cd build && \ CC="ccache gcc" CXX="ccache g++" \ ../configure \ - --prefix=/usr/local/ucx \ - --enable-shared \ - --disable-static \ - --disable-doxygen-doc \ - --enable-optimizations \ - --enable-devel-headers \ - --with-rocm=${ROCM_PATH} \ - --with-verbs \ - --with-dm \ - --enable-mt && \ + --prefix=/usr/local/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-devel-headers \ + --with-rocm=${ROCM_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt && \ make -j$(nproc) && \ make install @@ -298,25 +335,19 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ git checkout ${RIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ - -Ducx_path=${UCX_HOME} \ - -Drocm_path=${ROCM_PATH} && \ + -Ducx_path=${UCX_HOME} \ + -Drocm_path=${ROCM_PATH} && \ cd build && \ ninja -j$(nproc) && \ ninja install # Generate RIXL wheel -# Exclude libcore and libpull from auditwheel: transitive dependencies -# that are not shipped in the wheel and vary across base images. -RUN cd /opt/rixl && \ - sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ - contrib/build-wheel.sh && \ - mkdir -p /app/install && \ - _ucx_install_dir=${UCX_HOME} \ +RUN cd /opt/rixl && mkdir -p /app/install && \ ./contrib/build-wheel.sh \ - --output-dir /app/install \ - --rocm-dir ${ROCM_PATH} \ - --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ - --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + --output-dir /app/install \ + --rocm-dir ${ROCM_PATH} \ + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ + --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins # ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not # invalidate the slow ROCShmem build. @@ -331,16 +362,16 @@ ENV ROCSHMEM_DIR=/opt/rocshmem RUN --mount=type=cache,target=/root/.cache/ccache \ git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ - && cd rocm-systems \ - && git sparse-checkout set --cone projects/rocshmem \ - && git checkout ${ROCSHMEM_BRANCH} \ - && mkdir -p projects/rocshmem/build \ - && cd projects/rocshmem/build \ - && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ + && git checkout ${ROCSHMEM_BRANCH} \ + && mkdir -p projects/rocshmem/build \ + && cd projects/rocshmem/build \ + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ bash ../scripts/build_configs/all_backends \ - -DROCM_PATH=${ROCM_PATH} \ - -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ - -DUSE_EXTERNAL_MPI=OFF + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF # DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. FROM build_rocshmem AS build_deepep @@ -352,10 +383,10 @@ ARG DEEPEP_NIC="cx7" # DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. RUN --mount=type=cache,target=/root/.cache/ccache \ export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ - && git clone ${DEEPEP_REPO} \ - && cd DeepEP \ - && git checkout ${DEEPEP_BRANCH} \ - && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + && git clone ${DEEPEP_REPO} \ + && cd DeepEP \ + && git checkout ${DEEPEP_BRANCH} \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do # not force users to rebuild the long-lived Dockerfile.rocm_base image. @@ -364,52 +395,52 @@ ARG NIC_BACKEND ARG AINIC_VERSION ARG UBUNTU_CODENAME RUN /bin/bash -lc 'set -euo pipefail; \ - \ - install_ainic() { \ - apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ - rm -rf /var/lib/apt/lists/*; \ - mkdir -p /etc/apt/keyrings; \ - curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ - echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ - > /etc/apt/sources.list.d/amdainic.list; \ - apt-get update && apt-get install -y --no-install-recommends \ - libionic-dev \ - ionic-common \ - ; \ - rm -rf /var/lib/apt/lists/*; \ - }; \ - \ - # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ - # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ - install_bnxt() { \ - install -m 0755 -d /etc/apt/keyrings; \ - curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ - -o /etc/apt/keyrings/broadcom-nic.asc; \ - chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ - echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ - > /etc/apt/sources.list.d/broadcom-nic.list; \ - apt-get update && apt-get install -y --no-install-recommends \ - bnxt-rocelib=235.2.86.0 \ - ; \ - cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ - ldconfig; \ - rm -rf /var/lib/apt/lists/*; \ - }; \ - \ - echo "[MORI] Install MoRI proxy deps"; \ - pip install --quiet --ignore-installed blinker && \ - pip install --quiet quart msgpack aiohttp pyzmq; \ - echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ - \ - # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ - # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ - case "${NIC_BACKEND}" in \ - none) ;; \ - all) install_ainic; install_bnxt ;; \ - ainic) install_ainic ;; \ - bnxt) install_bnxt ;; \ - *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ - esac' + \ + install_ainic() { \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ + rm -rf /var/lib/apt/lists/*; \ + mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ + # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ + install_bnxt() { \ + install -m 0755 -d /etc/apt/keyrings; \ + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ + -o /etc/apt/keyrings/broadcom-nic.asc; \ + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ + > /etc/apt/sources.list.d/broadcom-nic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + bnxt-rocelib=235.2.86.0 \ + ; \ + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ + ldconfig; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + \ + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ + # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ + case "${NIC_BACKEND}" in \ + none) ;; \ + all) install_ainic; install_bnxt ;; \ + ainic) install_ainic ;; \ + bnxt) install_bnxt ;; \ + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ + esac' # ----------------------- # vLLM wheel release build stage (for building distributable wheels) @@ -434,22 +465,22 @@ COPY docker/context/base-wheels/ /tmp/base-wheels/ # If there are not wheels found there, we are not building for a wheel release. # So we exit with an error. To skip this stage. RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ - echo "Found custom wheels - copying to /install"; \ - cp /tmp/base-wheels/*.whl /install/ && \ - echo "Copied custom wheels:"; \ - ls -lh /install/; \ + echo "Found custom wheels - copying to /install"; \ + cp /tmp/base-wheels/*.whl /install/ && \ + echo "Copied custom wheels:"; \ + ls -lh /install/; \ else \ - echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ - echo "Wheel releases require pre-built ROCm wheels."; \ - exit 1; \ + echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ + echo "Wheel releases require pre-built ROCm wheels."; \ + exit 1; \ fi # GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) # This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm ARG GIT_REPO_CHECK=0 RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ - echo "Running repository checks..."; \ - cd vllm && bash tools/check_repo.sh; \ + echo "Running repository checks..."; \ + cd vllm && bash tools/check_repo.sh; \ fi # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) @@ -468,22 +499,22 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \ RUN echo "Checking for git-based packages in requirements files..." \ && echo "Checking common.txt for git-based packages:" \ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ - echo "ERROR: Git-based packages found in common.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in common.txt"; \ - fi \ + echo "ERROR: Git-based packages found in common.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in common.txt"; \ + fi \ && echo "Checking rocm.txt for git-based packages:" \ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ - echo "ERROR: Git-based packages found in rocm.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in rocm.txt"; \ - fi \ + echo "ERROR: Git-based packages found in rocm.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in rocm.txt"; \ + fi \ && echo "All requirements files are clean - no git-based packages found" # Pin vLLM dependencies to exact versions of custom ROCm wheels @@ -536,7 +567,7 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ mkdir -p build && cd build && \ cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ -fi + fi # Install RIXL + DeepEP wheels. RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ @@ -589,7 +620,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ && test -n "${FASTSAFETENSORS_REQ}" \ && python3 -m pip install --force-reinstall --no-deps \ - --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ && rm /tmp/rocm-test-reqs.txt # Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. @@ -636,7 +667,7 @@ RUN mkdir src && mv vllm src/vllm # ----------------------- # Final vLLM image -FROM mori_base AS final +FROM base AS final RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* @@ -656,40 +687,44 @@ ENV SCCACHE_IDLE_TIMEOUT= # Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. # Manually remove it so that later steps of numpy upgrade can continue RUN case "$(which python3)" in \ - *"/opt/conda/envs/py_3.9"*) \ - rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ - *) ;; esac + *"/opt/conda/envs/py_3.9"*) \ + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ + *) ;; esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --upgrade huggingface-hub[cli] + uv pip install --upgrade huggingface-hub[cli,hf_transfer] scipy # Install vLLM using uv (inherited from base stage) # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ - && uv pip install --system -r requirements/rocm.txt \ + && uv pip install -r requirements/rocm.txt \ && pip uninstall -y vllm \ - && uv pip install --system *.whl - -# Install RIXL wheel -RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ - uv pip install --system /rixl_install/*.whl + && uv pip install *.whl + +RUN pip install amdsmi + +# Install amd-aiter and Triton from repo/branch (for gfx1250; base image may not include them) +ARG AITER_REPO="https://github.com/ROCm/aiter" +ARG AITER_BRANCH="shared/triton-gfx12" +# TODO: Triton repo requires github.amd.com authentication; re-enable once credentials are configured +# ARG TRITON_REPO="https://github.amd.com/GFX-IP-Arch/triton" +# ARG TRITON_BRANCH="shared/gfx1250" +RUN apt-get update -q -y && apt-get install -q -y git && rm -rf /var/lib/apt/lists/* \ + && git clone --depth 1 --branch "$AITER_BRANCH" "$AITER_REPO" /tmp/aiter \ + && pip install --no-build-isolation /tmp/aiter && rm -rf /tmp/aiter +# && git clone --depth 1 --branch "$TRITON_BRANCH" "$TRITON_REPO" /tmp/triton \ +# && pip install --no-build-isolation /tmp/triton && rm -rf /tmp/triton ARG COMMON_WORKDIR ARG BASE_IMAGE -ARG NIC_BACKEND -ARG AINIC_VERSION # Copy over the benchmark scripts as well COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc -# See: https://github.com/ROCm/rocm-libraries/issues/6266 -ENV HSA_ENABLE_IPC_MODE_LEGACY=1 - ENV TOKENIZERS_PARALLELISM=false # ENV that can improve safe tensor loading, and end-to-end time @@ -704,9 +739,7 @@ ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt CMD ["/bin/bash"] diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index fbd5e1e60e3b..ca3573c87947 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -11,8 +11,6 @@ ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" ARG AITER_BRANCH="v0.1.16.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" # Sccache configuration (only used in release pipeline) ARG USE_SCCACHE @@ -27,38 +25,37 @@ FROM ${BASE_IMAGE} AS base ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin ENV ROCM_PATH=/opt/rocm ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx942;gfx950 -ENV MORI_GPU_ARCHS=gfx942;gfx950 +ENV AITER_ROCM_ARCH=gfx942;gfx950;gfx1250 +ENV MORI_GPU_ARCHS=gfx942;gfx950;gfx1250 # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 -ARG PYTHON_VERSION=3.12 -ENV PYTHON_VERSION=${PYTHON_VERSION} - RUN mkdir -p /app WORKDIR /app ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ - && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ - done \ - && apt-get update -y \ - && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 \ - && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ - && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ - && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version && python3 -m pip --version - -RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython + && apt-get install -y software-properties-common build-essential git curl sudo vim less libopenmpi-dev libpci-dev python3-venv python3-dev + +RUN python3 -m venv /opt/python --system-site-packages +ENV PATH=/opt/python/bin:$PATH + +# Install UV +RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="/usr/local/bin" sh +ENV UV_PYTHON=/opt/python/bin/python +# 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 +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +# Use copy mode to avoid hardlink failures with Docker cache mounts +ENV UV_LINK_MODE=copy + + +RUN uv pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* # Install sccache if USE_SCCACHE is enabled (for release builds) @@ -69,16 +66,16 @@ ARG SCCACHE_BUCKET_NAME ARG SCCACHE_REGION_NAME ARG SCCACHE_S3_NO_CREDENTIALS RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ fi # Setup sccache for HIP compilation via HIP_CLANG_PATH @@ -87,13 +84,13 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ fi # Set sccache environment variables only when USE_SCCACHE=1 @@ -153,21 +150,21 @@ RUN cd pytorch \ && pip install -r requirements.txt && git submodule update --init --recursive RUN cd pytorch && python3 tools/amd_build/build_amd.py \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ - && sccache --show-stats; \ - fi \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ + && sccache --show-stats; \ + fi \ && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && pip install dist/*.whl RUN git clone ${PYTORCH_VISION_REPO} vision RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ && python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && pip install dist/*.whl @@ -176,10 +173,10 @@ RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ && git submodule update --init --recursive \ && pip install -r requirements.txt \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ && python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && pip install dist/*.whl @@ -211,16 +208,14 @@ FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl RUN git clone ${FA_REPO} RUN cd flash-attention \ && git checkout ${FA_BRANCH} \ && git submodule update --init \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install @@ -238,12 +233,12 @@ RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} RUN cd aiter \ && git submodule update --init --recursive \ - && pip install -r requirements.txt -RUN pip install pyyaml && cd aiter \ + && uv pip install -r requirements.txt +RUN uv pip install pyyaml && cd aiter \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl @@ -258,64 +253,28 @@ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install # only includes dependencies used by wheel release pipeline FROM base AS debs_wheel_release RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs # Full debs stage - includes Mori (used by Docker releases) FROM base AS debs RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs FROM base AS final RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ - pip install /install/*.whl + uv pip install /install/*.whl ARG BASE_IMAGE -ARG TRITON_BRANCH -ARG TRITON_REPO -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH ARG AITER_REPO -ARG MORI_BRANCH -ARG MORI_REPO RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ - && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ - && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ - && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4a8b4209d874..f826c1ee3834 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -6,6 +6,7 @@ from typing import Protocol import torch +import torch.nn.functional as F from torch._ops import OpOverload from torch.distributed import ProcessGroup @@ -149,6 +150,143 @@ def wrapper(*args, **kwargs): return wrapper +def _mxfp4_moe_w1_triton_gemm_view(w1: torch.Tensor, w2: torch.Tensor) -> torch.Tensor: + """Map expert w1 to moe_gemm_a4w4 layout (stride(-2)==1 on [E, K, N]). + + vLLM stores w13 as [E, 2*inter, hidden/2] (K along dim2). aiter downcast uses the + same axis order with K-packed stride-1 on dim1; do not compose transpose+as_strided + (that overruns storage). Standard layout only needs transpose(1, 2). + """ + if w1.stride(-2) == 1: + return w1 + h = w2.shape[1] + e, a, b = w1.shape + if b * 2 == h: + return w1.transpose(1, 2) + if a * 2 == h: + kp, n_out = a, b + return w1.as_strided((e, kp, n_out), (kp * n_out, 1, kp)) + return w1.transpose(1, 2) + + +def _mxfp4_moe_w2_triton_gemm_view(w2: torch.Tensor) -> torch.Tensor: + """Map vLLM w2 [E, model_dim, inter/2] to moe_gemm_a4w4 layout.""" + if w2.stride(-2) == 1: + return w2 + return w2.transpose(1, 2) + + +def _silu_and_mul_glu(x: torch.Tensor) -> torch.Tensor: + """SiLU(gate) * up for g1u1 tensors with last dim 2 * inter (matches CK ``silu_and_mul``).""" + d = x.shape[-1] // 2 + gate, up = x[..., :d], x[..., d:] + return (F.silu(gate.to(torch.float32)) * up.to(torch.float32)).to(dtype=x.dtype) + + +def _mxfp4_moe_weight_as_uint8(w: torch.Tensor) -> torch.Tensor: + """Triton moe_gemm_a4w4 JIT cannot canonicalize torch float4 dtypes (KeyError).""" + if w.dtype == torch.uint8: + return w + fp4 = getattr(torch, "float4_e2m1fn_x2", None) + if fp4 is not None and w.dtype == fp4: + return w.view(torch.uint8) + name = getattr(w.dtype, "name", str(w.dtype)) + if "float4" in name or "e2m1" in name: + return w.view(torch.uint8) + return w + + +def _rocm_aiter_fused_moe_triton_gemm_a4w4( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weight: torch.Tensor, + topk_ids: torch.Tensor, + activation_method: int, + w1_scale: torch.Tensor | None, + w2_scale: torch.Tensor | None, + output_dtype: torch.dtype | None, +) -> torch.Tensor: + from aiter import ActivationType + from aiter.fused_moe import get_inter_dim + from aiter.ops.triton.moe.moe_op_gemm_a4w4 import moe_gemm_a4w4, mxfp4_quant + from aiter.ops.triton.moe.moe_routing.routing import routing + + if activation_method != int(ActivationType.Silu): + raise RuntimeError( + "Triton moe_gemm_a4w4 path supports only Silu/SwiGLU (ActivationType.Silu)." + ) + + device = hidden_states.device + out_dtype = output_dtype or hidden_states.dtype + # FP32 matmul outputs for numerical headroom before the second mxfp4_quant. + gemm_out_dt = torch.float32 + m, _ = topk_ids.shape + num_experts = w1.shape[0] + + logits = torch.full( + (m, num_experts), -1e9, device=device, dtype=torch.float32 + ) + tid = topk_ids.long().clamp(min=0, max=num_experts - 1) + logits.scatter_( + 1, + tid, + torch.log(topk_weight.to(torch.float32).clamp(min=1e-20)), + ) + + rdata, gather_idx, scatter_idx = routing(logits, topk_ids.shape[1]) + gate_scal = rdata.gate_scal + + e, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape) + is_g1u1 = inter_dim != w1.shape[1] + if not is_g1u1: + raise RuntimeError( + "Triton moe_gemm_a4w4 path expects SwiGLU (w1 dim1 == 2 * inter)." + ) + + w1_v = _mxfp4_moe_weight_as_uint8(_mxfp4_moe_w1_triton_gemm_view(w1, w2)) + w2_v = _mxfp4_moe_weight_as_uint8(_mxfp4_moe_w2_triton_gemm_view(w2)) + + x_q, x_s = mxfp4_quant(hidden_states.to(out_dtype)) + # Raw W1@x per (token, expert), then CK-style silu_and_mul (not Triton ``apply_swiglu``). + # ``scatter_indx=None`` keeps expanded rows until after activation; higher VRAM than fused swiglu. + mid_raw = moe_gemm_a4w4( + x_q, + w1_v, + x_s, + w1_scale, + x_static_scale=None, + quant_static_scale=None, + bias=None, + routing_data=rdata, + gather_indx=gather_idx, + scatter_indx=None, + gammas=None, + swizzle_mx_scale=None, + out_dtype=gemm_out_dt, + apply_swiglu=False, + ) + mid = _silu_and_mul_glu(mid_raw) + mid_q, mid_s = mxfp4_quant(mid.to(out_dtype)) + out = moe_gemm_a4w4( + mid_q, + w2_v, + mid_s, + w2_scale, + x_static_scale=None, + quant_static_scale=None, + bias=None, + routing_data=rdata, + gather_indx=None, + scatter_indx=scatter_idx, + gammas=gate_scal, + swizzle_mx_scale=None, + out_dtype=gemm_out_dt, + apply_swiglu=False, + ) + return out.to(out_dtype) + + def _rocm_aiter_fused_moe_impl( hidden_states: torch.Tensor, w1: torch.Tensor, @@ -179,6 +317,30 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) + m_tokens = hidden_states.shape[0] + use_triton_a4w4 = ( + envs.VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4 + and quant_type == QuantType.per_1x32 + and expert_mask is None + and not doweight_stage1 + and num_local_tokens is None + and a1_scale is None + and a2_scale is None + and m_tokens >= envs.VLLM_ROCM_AITER_TRITON_MOE_MIN_TOKENS + ) + if use_triton_a4w4: + return _rocm_aiter_fused_moe_triton_gemm_a4w4( + hidden_states, + w1, + w2, + topk_weight, + topk_ids, + activation_method, + w1_scale, + w2_scale, + output_dtype, + ) + extra_kwargs: dict = {} if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): extra_kwargs["gate_mode"] = gate_mode diff --git a/vllm/envs.py b/vllm/envs.py index 1f94be8ac5fa..f880858b711d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -134,6 +134,7 @@ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True + VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4: bool = False VLLM_ROCM_USE_SKINNY_GEMM: bool = True VLLM_ROCM_FP8_PADDING: bool = True VLLM_ROCM_MOE_PADDING: bool = True @@ -1227,6 +1228,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1") ), + # ROCm AITER fused MoE: use Triton moe_gemm_a4w4 (MXFP4) instead of aiter.fused_moe. + "VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4": lambda: ( + os.getenv("VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4", "False").lower() + in ("true", "1") + ), # use rocm skinny gemms "VLLM_ROCM_USE_SKINNY_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1") diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index db88ba273cd2..c70dcaa3379b 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -58,8 +58,10 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): value_layout = StridedLayout scale_layout = StridedLayout elif current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + value_layout = StridedLayout - if should_use_cdna4_mx_scale_swizzle(): + if should_use_cdna4_mx_scale_swizzle() or on_gfx1250(): try: # triton < 3.6 from triton_kernels.tensor_details.layout import GFX950MXScaleLayout diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 6c3a0fe96ecc..e65ee9cf1fe4 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -213,11 +213,12 @@ def _get_gcn_arch() -> str: _ON_GFX1100 = "gfx1100" in _GCN_ARCH _ON_GFX1151 = "gfx1151" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) -_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) -_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) +_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950", "gfx1250"]) +_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950", "gfx1250"]) _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH +_ON_GFX1250 = "gfx1250" in _GCN_ARCH def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: @@ -331,6 +332,10 @@ def on_gfx950() -> bool: return _ON_GFX950 +def on_gfx1250() -> bool: + return _ON_GFX1250 + + # Enable HIP online tuning early, before hipBLASLt initializes. # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM. if ( @@ -864,7 +869,7 @@ def get_device_communicator_cls(cls) -> str: @classmethod def supports_mx(cls) -> bool: - return any(gfx in _GCN_ARCH for gfx in ["gfx95"]) + return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"]) @classmethod def supports_fp8(cls) -> bool: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 0c5512d5e15f..3db3fbf961c8 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -66,7 +66,6 @@ from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling -from vllm.utils.torch_utils import set_random_seed from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.outputs import ( @@ -348,7 +347,7 @@ def init_device(self): logger.info_once("Using V2 Model Runner") # Set random seed. - set_random_seed(self.model_config.seed) + # set_random_seed(self.model_config.seed) # Now take memory snapshot after NCCL is initialized gc.collect() @@ -848,7 +847,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. - set_random_seed(self.model_config.seed) + # set_random_seed(self.model_config.seed) # Eagerly trigger inductor's once-per-process lazy inits during # warmup (rather than on a later compile cache-miss at runtime). From b54e405b6009ec9819ba5a8af7db2b2b7660a51b Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Tue, 26 May 2026 21:11:46 -0500 Subject: [PATCH 08/68] [ROCm] Route flash_attn_varlen_func to aiter.ops.triton.mha Completes deferred port #1 from the gfx1250 squash. FlashAttention (the 'flash_attn' pip package) is not installed on ROCm, so fa_utils.py's ROCm branch now imports flash_attn_varlen_func from aiter.ops.triton.mha (same source as aiter_triton_mla.py) instead of 'flash_attn'. _ROCM_FLASH_ATTN_AVAILABLE is set True when AITER is present, so is_flash_attn_varlen_func_available() reports a working impl. All ROCm consumers that import flash_attn_varlen_func from fa_utils (mla/prefill/flash_attn.py, flash_attn_diffkv.py, turboquant_attn.py, the vit_attn_wrappers fallback) inherit this. The only other 'import flash_attn' (qwen2_5_omni_thinker.py) is already guarded (-> None). Verified in the gfx1250 FFM image: resolves to aiter.ops.triton.attention.mha.flash_attn_varlen_func, availability True. AI-assisted (Claude Code); human review + lint pending before any PR. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- vllm/v1/attention/backends/fa_utils.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 474523780ff7..b2b173858e07 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -30,17 +30,22 @@ flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[assignment] get_scheduler_metadata = xpu_ops.get_scheduler_metadata # type: ignore[assignment] elif current_platform.is_rocm(): + # On ROCm we use AITER's Triton flash-attention; the upstream flash-attn + # package is not installed/available. (Same source as aiter_triton_mla.py.) try: - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] + from aiter.ops.triton.mha import ( # type: ignore[no-redef] + flash_attn_varlen_func, + ) - # Mark that upstream flash-attn is available on ROCm + # A working flash_attn_varlen_func is available on ROCm (via AITER). _ROCM_FLASH_ATTN_AVAILABLE = True except ImportError: def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc] raise ImportError( - "ROCm platform requires upstream flash-attn " - "to be installed. Please install flash-attn first." + "ROCm platform requires AITER's Triton MHA " + "(aiter.ops.triton.mha.flash_attn_varlen_func). " + "Please install aiter." ) # ROCm doesn't use scheduler metadata (FA3 feature), provide stub @@ -237,7 +242,7 @@ def is_flash_attn_varlen_func_available() -> bool: Platform-specific sources: - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func - XPU: xpu_ops.flash_attn_varlen_func - - ROCm: upstream flash_attn.flash_attn_varlen_func (if available) + - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py) which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is From 6a7d9bdfe70a1298fec17189dd03ab14d04df192 Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Tue, 26 May 2026 22:05:17 -0500 Subject: [PATCH 09/68] [ROCm][gfx1250] Make vLLM C++/HIP extensions build for gfx1250 Enables a from-source vLLM build with PYTORCH_ROCM_ARCH=gfx1250 (verified in the FFM gfx1250 simulator image; _C and _rocm_C carry gfx1250 code objects, import OK): - csrc/rocm/attention.cu: guard the gfx12 WMMA builtins (__builtin_amdgcn_wmma_f32_16x16x16_{f16,bf16}_w32_gfx12) for gfx1250, which lacks them (needs wmma-128b-insts); trap if launched. - CMakeLists.txt + csrc/rocm/torch_bindings.cpp: skinny_gemms.cu is pervasively gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy s_waitcnt asm) unsupported on gfx1250. When VLLM_GPU_ARCHES matches gfx1250, exclude skinny_gemms.cu from _rocm_C and define VLLM_SKIP_SKINNY_GEMMS to skip its op registrations (LLMM1/wvSplitK/ wvSplitKrc/wvSplitKQ). vLLM falls back to default/Triton GEMM for these on gfx1250. Non-gfx1250 ROCm builds are unaffected. Caveat: these custom ROCm GEMM ops are unavailable on gfx1250 and the custom attention WMMA path traps; intended for functional bring-up on the simulator. AI-assisted (Claude Code); human review + real gfx1250 run needed before any PR. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- CMakeLists.txt | 13 ++++++++++++- csrc/rocm/attention.cu | 10 ++++++++++ csrc/rocm/torch_bindings.cpp | 5 +++++ 3 files changed, 27 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 5d7672fd9222..c0e4698e78fb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1370,6 +1370,17 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/torch_bindings.cpp" "csrc/rocm/skinny_gemms.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS}) + + # skinny_gemms.cu is built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy + # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude it from the + # gfx1250 build and disable its op registrations (VLLM_SKIP_SKINNY_GEMMS); + # vLLM falls back to default/Triton GEMM for those ops on gfx1250. + if(VLLM_GPU_ARCHES MATCHES "gfx1250") + message(STATUS "gfx1250 detected: excluding csrc/rocm/skinny_gemms.cu from _rocm_C") + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS") + endif() set(VLLM_ROCM_HAS_GFX1100 OFF) if(VLLM_GPU_ARCHES MATCHES "gfx1100") @@ -1385,7 +1396,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP") DESTINATION vllm LANGUAGE ${VLLM_GPU_LANG} SOURCES ${VLLM_ROCM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} + COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 4ac255d0a75f..86c5044ce981 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -2405,6 +2405,15 @@ template __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, const bit16x8& inpB, const floatx8& inpC) { +#if defined(__gfx1250__) + // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by + // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is + // unsupported on gfx1250; trap if ever launched (fail loud, not silent-wrong). + (void)inpA; + (void)inpB; + __builtin_trap(); + return inpC; +#else if constexpr (std::is_same::value) { return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); } else if constexpr (std::is_same::value) { @@ -2412,6 +2421,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, } else { static_assert(false, "unsupported 16b dtype"); } +#endif } template diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 03de6dcd1576..53943848f3f7 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -14,6 +14,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { // vLLM custom ops for rocm +// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250 +// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid +// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250. +#ifndef VLLM_SKIP_SKINNY_GEMMS // Custom gemm op for matrix-vector multiplication rocm_ops.def( "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> " @@ -38,6 +42,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "Tensor scale_a, " " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#endif // VLLM_SKIP_SKINNY_GEMMS #ifdef VLLM_ROCM_GFX1100 // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). From b816c7c088d629a9b82b41c7580b40c15a89345b Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Tue, 26 May 2026 22:35:57 -0500 Subject: [PATCH 10/68] [ROCm][gfx1250] Runtime fixes to run vLLM on the FFM simulator Lets vLLM initialize and run on the gfx1250 FFM pre-silicon simulator (verified: facebook/opt-125m generates end-to-end with attention_backend=TRITON_ATTN): - platforms/__init__.py: ROCm platform detection falls back to torch.version.hip when amdsmi is unavailable. amdsmi queries the driver/sysfs and does not see GPUs exposed only via HSA (the FFM model); torch's HIP runtime does. - platforms/rocm.py: _get_gcn_arch() uses logger.debug (not warning_once) on the amdsmi-failure path; warning_once imports vllm.distributed at module load, causing a circular import before current_platform is bound. The torch.cuda gcnArchName fallback then correctly reports gfx1250 (amdsmi would report the host's real gfx950 cards, not the sim). - model_executor/layers/utils.py: skip the skinny-GEMM (wvSplitK/LLMM1) path on gfx1250 (those kernels are excluded from _rocm_C); fall back to torch GEMM. Note: on gfx1250 the ROCm custom-attention backends route paged decode through _rocm_C.paged_attention; use attention_backend=TRITON_ATTN (pure Triton) for now. AI-assisted (Claude Code); human review needed before any PR. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- vllm/model_executor/layers/utils.py | 5 ++++- vllm/platforms/__init__.py | 14 ++++++++++++++ vllm/platforms/rocm.py | 12 ++++++++---- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 6ca42c0e7f0e..a01acec2ac2e 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950 + from vllm.platforms.rocm import on_gfx1250, on_gfx1x, on_gfx9, on_gfx950 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -172,6 +172,9 @@ def rocm_unquantized_gemm_impl( use_skinny = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) + # gfx1250 excludes the skinny-GEMM kernels (wvSplitK/LLMM1) from the + # build (gfx9/gfx11 ISA); fall back to torch GEMM there. + and not on_gfx1250() and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index ac536aff00c5..5f5bd78793e5 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -125,6 +125,20 @@ def rocm_platform_plugin() -> str | None: except Exception as e: logger.debug("ROCm platform is not available because: %s", str(e)) + # Fallback: amdsmi queries the kernel driver/sysfs, which does not see GPUs + # exposed only through the HSA layer (e.g. the FFM pre-silicon simulator). + # Trust torch's HIP runtime detection in that case. + if not is_rocm: + try: + import torch + + if torch.version.hip is not None and torch.cuda.device_count() > 0: + is_rocm = True + logger.debug("ROCm platform detected via torch.version.hip " + "(amdsmi unavailable).") + except Exception as e: + logger.debug("torch HIP ROCm fallback failed: %s", str(e)) + return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index e65ee9cf1fe4..3df9298c15b7 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -194,10 +194,14 @@ def _get_gcn_arch() -> str: return _query_gcn_arch_from_amdsmi() except Exception as e: logger.debug("Failed to get GCN arch via amdsmi: %s", e) - logger.warning_once( - "Failed to get GCN arch via amdsmi, falling back to torch.cuda. " - "This will initialize CUDA and may cause " - "issues if CUDA_VISIBLE_DEVICES is not set yet." + # NOTE: use logger.debug, not warning_once, here. This runs at module + # load while resolving the platform; warning_once imports + # vllm.distributed, which causes a circular import before + # current_platform is bound. This path is taken on the FFM simulator, + # where amdsmi is unavailable (and would report the host's real gfx950 + # cards rather than the simulated gfx1250) — torch.cuda below is correct. + logger.debug( + "Failed to get GCN arch via amdsmi, falling back to torch.cuda." ) # Ultimate fallback: use torch.cuda (will initialize CUDA) return torch.cuda.get_device_properties("cuda").gcnArchName From fe2642b3e9aa1945eac4fcc1475f77e8796b0d3e Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Tue, 26 May 2026 23:13:16 -0500 Subject: [PATCH 11/68] [ROCm][gfx1250] FFM-sim enablement: amdsmi torch fallbacks + DSV4 layer-prune loader More fixes to bring vLLM up on the gfx1250 FFM simulator (no amdsmi; amdsmi would report the host's real gfx950 cards, not the sim): - platforms/rocm.py: add _AMDSMI_AVAILABLE flag; with_amdsmi_context skips amdsmi init/shutdown when unavailable; get_device_name falls back to torch.cuda.get_device_name. (Complements the earlier _get_gcn_arch fix.) - models/deepseek_v4/amd/model.py: load_weights skips checkpoint weights for transformer layers pruned via a num_hidden_layers hf_override (was raising KeyError), enabling small-layer-count bring-up runs. Probe result (DeepSeek-V4-Flash, 4 layers, TP=1, AITER off, triton_unfused MoE): loads and runs the FP8 attention via the Triton GEMM, but the MXFP4 MoE fails to compile -- triton_kernels matmul_ogs uses unswizzle_mx_scale_cdna4 (gfx950 scale layout) for gfx1250; reshape mismatch. gfx1250 needs its own MXFP4 scale-unswizzle (cf. aiter swizzle_scales_gfx1250). MoE is the gfx1250 enablement gap. AI-assisted (Claude Code); human review needed. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- vllm/models/deepseek_v4/amd/model.py | 17 +++++++++++++++++ vllm/platforms/rocm.py | 13 +++++++++++++ 2 files changed, 30 insertions(+) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index edb923511501..2c6501452f87 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -626,6 +626,16 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: expert_mapping = self.get_expert_mapping() for name, loaded_weight in weights: + # Skip checkpoint weights for transformer layers pruned by a + # num_hidden_layers hf_override (keeps load + run tractable on the + # FFM simulator). Names here are already vLLM-mapped ("layers.N."). + if "layers." in name: + try: + _li = int(name.split("layers.", 1)[1].split(".", 1)[0]) + if _li >= self.config.num_hidden_layers: + continue + except (IndexError, ValueError): + pass for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -682,6 +692,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: elif "attn_sink" in name: if is_pp_missing_parameter(name, self): continue + if name not in params_dict: + # Layer pruned via a num_hidden_layers hf_override; skip + # its checkpoint weights instead of KeyError-ing. + continue narrow_weight = loaded_weight[head_rank_start:head_rank_end] n = narrow_weight.shape[0] params_dict[name][:n].copy_(narrow_weight) @@ -690,6 +704,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: else: if is_pp_missing_parameter(name, self): continue + if name not in params_dict: + # Layer pruned via a num_hidden_layers hf_override; skip. + continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 3df9298c15b7..613a28be27c3 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -37,8 +37,14 @@ amdsmi_topo_get_link_type, amdsmi_topo_get_numa_node_number, ) + + _AMDSMI_AVAILABLE = True except ImportError as e: logger.warning("Failed to import from amdsmi with %r", e) + # amdsmi is unavailable on the FFM pre-silicon simulator (and would report + # the host's real GPUs, not the sim). amdsmi-decorated methods fall back to + # torch when this is False. + _AMDSMI_AVAILABLE = False try: import vllm._C # noqa: F401 @@ -154,6 +160,10 @@ def _sync_hip_cuda_env_vars(): def with_amdsmi_context(fn): @wraps(fn) def wrapper(*args, **kwargs): + if not _AMDSMI_AVAILABLE: + # amdsmi unavailable (e.g. FFM simulator): skip init/shutdown and let + # the wrapped function use its own non-amdsmi (torch) fallback. + return fn(*args, **kwargs) amdsmi_init() try: return fn(*args, **kwargs) @@ -728,6 +738,9 @@ def is_fully_connected(cls, physical_device_ids: list[int]) -> bool: @with_amdsmi_context @lru_cache(maxsize=8) def get_device_name(cls, device_id: int = 0) -> str: + if not _AMDSMI_AVAILABLE: + # FFM simulator: amdsmi can't see the simulated GPU; use torch. + return torch.cuda.get_device_name(device_id) physical_device_id = cls.device_id_to_physical_device_id(device_id) handle = amdsmi_get_processor_handles()[physical_device_id] asic_info = amdsmi_get_gpu_asic_info(handle) From e7f5b7ce42875dbb2e890d41740390a6279b4768 Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Wed, 27 May 2026 09:54:44 -0500 Subject: [PATCH 12/68] [ROCm][gfx1250] Enable DeepSeek-V4 MXFP4 MoE + sparse-attn compressor on gfx1250 Routes the DeepSeek-V4-Flash MXFP4 MoE through aiter's gfx1250 W4A8 `moe_gemm_a8w4` instead of the vendored triton_kernels `matmul_ogs`, whose `unswizzle_mx_scale_cdna4` has no gfx1250 variant. - gpt_oss_triton_kernels_moe.py: add `UnfusedOAITritonExperts._try_apply_aiter_w4a8` plus an early-return gate in `apply()` (fires on use_mxfp4_w4a8/w4a16, no router-weight-on-input, no LoRA, rocm_aiter_ops enabled). Does a manual expert-sorted gather (in-kernel gather is numerically broken on gfx1250), builds aiter-native routing from topk_ids/topk_weights, resolves dynamic FP8 LHS scales, runs gemm1(swiglu, fp8 out) -> gemm2, and restores DeepSeek-V4's routed_scaling_factor via a per-token output rescale. - mxfp4_utils.py: store a plain StridedLayout MXFP4 scale on gfx1250 (only gfx950 keeps the CDNA4/GFX950 swizzle); the W4A8 path reads it un-swizzled. - deepseek_v4/compressor.py: on ROCm, select the pure-Triton sparse-attention compressor (head_dim==512) instead of the NVIDIA-only CuTeDSL path. - test_modular_oai_triton_moe.py: coverage for the W4A8 MoE path. Verified end-to-end on the gfx1250 FFM simulator: DeepSeek-V4-Flash loads and runs MLA + sparse attention + MXFP4 MoE to generation. MoE kernel math validated standalone to maxrel ~5e-3. Diagnostic logger.info calls in `_try_apply_aiter_w4a8` are retained intentionally for sim bring-up; drop them before an upstream PR. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- .../moe/test_modular_oai_triton_moe.py | 144 +++++ .../experts/gpt_oss_triton_kernels_moe.py | 508 +++++++++++++++++- .../layers/quantization/utils/mxfp4_utils.py | 10 +- 3 files changed, 638 insertions(+), 24 deletions(-) diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 0315d8d89e56..585aa34d5966 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -4,6 +4,12 @@ Test modular OAI Triton MoE """ +from __future__ import annotations + +import json +import os +from pathlib import Path + import pytest import torch @@ -48,6 +54,52 @@ ] +def deepseek_v4_flash_moe_topology(): + """MoE sizes for the MODEL path used by ``launch_dsv4.sh``. + + Default weights path: ``/data/deepseek-ai/DeepSeek-V4-Flash/config.json``. + When that file is readable, values come from ``hidden_size``, + ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``. + Otherwise fall back to the same numeric constants. + """ + defaults = { + "hidden_size": 4096, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + } + cfg_path = Path( + os.environ.get( + "DEEPSEEK_V4_FLASH_CONFIG", + "/data/deepseek-ai/DeepSeek-V4-Flash/config.json", + ) + ) + if cfg_path.is_file(): + with cfg_path.open() as f: + cfg = json.load(f) + return { + "hidden_size": int(cfg["hidden_size"]), + "moe_intermediate_size": int(cfg["moe_intermediate_size"]), + "n_routed_experts": int(cfg["n_routed_experts"]), + "num_experts_per_tok": int(cfg["num_experts_per_tok"]), + } + return defaults + + +def scaled_deepseek_v4_flash_problem( + *, + dim_scale: int = 8, + expert_scale: int = 8, +): + """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1).""" + t = deepseek_v4_flash_moe_topology() + k = max(128, t["hidden_size"] // dim_scale) + n = max(64, t["moe_intermediate_size"] // dim_scale) + num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale) + topk = t["num_experts_per_tok"] + return k, n, num_experts, topk + + def unshuffle_weight(w: torch.Tensor): first = w[..., ::2] second = w[..., 1::2] @@ -261,3 +313,95 @@ def test_oai_triton_moe( ) assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): + """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. + + Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``, + with linear dimensions and expert count scaled down for test GPU memory. + """ + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) + set_random_seed(0) + + k, n, num_experts, topk = scaled_deepseek_v4_flash_problem() + m = 7 + dtype = torch.bfloat16 + + ( + w1, + w2, + w1_bias, + w2_bias, + w1_tri, + w2_tri, + w1_bias_tri, + w2_bias_tri, + w1_precision_config, + w2_precision_config, + ) = make_weights(dtype, k, n, num_experts) + + x = torch.randn((m, k), dtype=dtype, device="cuda") + router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) + + quant_config = mxfp4_w4a16_moe_quant_config( + w1_bias=w1_bias_tri, + w2_bias=w2_bias_tri, + w1_scale=w1_precision_config, + w2_scale=w2_precision_config, + ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size_per_partition=n, + ) + experts = UnfusedOAITritonExperts(moe_config, quant_config) + + if not UnfusedOAITritonExperts._supports_current_device(): + pytest.skip("UnfusedOAITritonExperts does not support this device") + + _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids) + assert top_k == topk + ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( + m, + N, + K, + topk, + num_experts, + num_experts, + None, + MoEActivation.SWIGLUOAI, + ) + workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda") + workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda") + output = torch.empty(out_shape, dtype=dtype, device="cuda") + + with set_current_vllm_config(VllmConfig()): + out_ref = torch_moe_impl( + x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids + ) + experts.apply( + output=output, + hidden_states=x, + w1=w1_tri, + w2=w2_tri, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SWIGLUOAI, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) 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 4b0a0b8ecad3..bd3d9aa31bd3 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 @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import replace + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -814,6 +818,117 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx +def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_states: torch.Tensor, + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int = 1, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). + + Uses ``aiter.ops.triton.quant_moe.downcast_to_static_fp8`` with a scalar scale: + for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when + single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or + ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static + FP8 path). + + Returns: + (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for + FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + if not quant_config.use_mxfp4_w4a16: + return hidden_states, None + try: + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return hidden_states, None + if not rocm_aiter_ops.is_enabled(): + return hidden_states, None + + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + + lhs_scale = None + if qp is not None and qp.flex_ctx.lhs_data.scale is not None: + s = qp.flex_ctx.lhs_data.scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = hidden_states.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + + return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale + + +def _mxfp4_w4a8_unpadded_dims( + moe_cfg: FusedMoEConfig, +) -> tuple[int | None, int | None, int | None, int | None]: + """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" + if ( + moe_cfg.intermediate_size_per_partition_unpadded is None + or moe_cfg.hidden_dim_unpadded is None + ): + return None, None, None, None + unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 + unpadded_k_w1 = moe_cfg.hidden_dim_unpadded + unpadded_n_w2 = moe_cfg.hidden_dim_unpadded + unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded + return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 + + +def _mxfp4_w4a8_resolve_lhs_scale( + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int, + activations: torch.Tensor, +) -> torch.Tensor: + """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. + + Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise + ``max(|activations|) / 448`` (e4m3), matching + ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + lhs_scale = None + if qp is not None and qp.flex_ctx is not None: + ld = qp.flex_ctx.lhs_data + if ld is not None and ld.scale is not None and ld.scale.numel() == 1: + lhs_scale = ld.scale + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = activations.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + return lhs_scale + + +def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + precision_config, + lhs_scale: torch.Tensor | None, +): + """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" + if precision_config is None or lhs_scale is None: + return precision_config + from triton_kernels.numerics import InFlexData + + new_flex = replace( + precision_config.flex_ctx, + lhs_data=InFlexData(scale=lhs_scale), + ) + return replace(precision_config, flex_ctx=new_flex) + + @triton.jit def _masked_topk_sum_kernel( inp_ptr, # (M, topk, K) contiguous @@ -1143,6 +1258,193 @@ def activation( else: super().activation(activation, output, input) + def _try_apply_aiter_w4a8( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + apply_router_weight_on_input: bool, + ) -> torch.Tensor | None: + """ROCm/gfx1250 MXFP4 MoE via aiter's W4A8 ``moe_gemm_a8w4``. + + The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls + ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and + fails to compile on gfx1250. aiter's gluon ``moe_gemm_a8w4`` (FP8 act x MXFP4 + weight) does run on gfx1250, so we route the MoE through it. + + Routing is built aiter-native directly from ``topk_ids``/``topk_weights`` so we + never touch the diverged vendored ``triton_kernels`` routing structs (whose + ``RoutingData``/``GatherIndx``/``ExptData`` field names and layout differ from + aiter's). This mirrors ``triton_kernel_fused_mxfp4_w4a8_experts`` + (aiter_mxfp4_w4a8_moe.py) but resolves the FP8 LHS scales dynamically because + DeepSeek-V4 is ``activation_scheme: dynamic`` (no calibrated static scale). + + aiter's ``routing`` always softmaxes the gate weights, so we feed + ``log(topk_weights)``: softmax-over-selected then recovers + ``topk_weights / sum(topk_weights)``. DeepSeek-V4 weights are + norm_topk_prob-normalized then scaled by ``routed_scaling_factor`` (so they sum + to that factor, not 1); we restore the exact weighting with a per-token output + rescale by the weight sum (== 1 for plain renormalize routing, a no-op there). + + Returns ``output`` on success, or ``None`` if aiter is unavailable (caller then + falls back to ``matmul_ogs``). + """ + try: + from aiter.ops.triton.moe.moe_routing.routing import ( + routing as aiter_routing, + ) + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + M = hidden_states.shape[0] + num_experts = w1.shape[0] + topk = topk_ids.size(1) + + # Reconstruct dense gating logits from the sparse topk selection. log(weight) + # padded with a large-negative sentinel makes aiter's softmax reproduce both + # this exact expert selection and weight = topk_weights / per-token-sum. + tw = topk_weights.to(torch.float32) + logits = torch.full( + (M, num_experts), -1e30, device=hidden_states.device, dtype=torch.float32 + ) + logits.scatter_( + 1, + topk_ids.long().clamp(min=0, max=num_experts - 1), + torch.log(tw.clamp(min=1e-20)), + ) + logger.info( + "[aiter-w4a8] ENTER M=%d K=%d num_experts=%d topk=%d " + "w1=%s w2=%s w1_wscale=%s w2_wscale=%s", + M, + hidden_states.shape[1], + num_experts, + topk, + tuple(w1.storage.data.shape), + tuple(w2.storage.data.shape), + tuple(quant_config.w1_precision.weight_scale.storage.data.shape), + tuple(quant_config.w2_precision.weight_scale.storage.data.shape), + ) + routing_data, gather_idx, scatter_idx = aiter_routing( + logits, topk, sm_first=False + ) + gammas = routing_data.gate_scal + logger.info( + "[aiter-w4a8] routing OK block_m=%s gate_scal=%s gather=%s scatter=%s", + getattr(routing_data, "block_m", None), + tuple(gammas.shape), + tuple(gather_idx.shape), + tuple(scatter_idx.shape), + ) + + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + # FP8 LHS scales (static if calibrated, else dynamic amax/448). GEMM1 emits FP8 + # directly (apply_swiglu=True, out_dtype fp8) quantized with GEMM2's LHS scale, + # so resolve GEMM2's scale up front (hidden states as the amax proxy) and reuse + # it as GEMM2's input scale so quantize/dequantize are consistent. + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=1, activations=hidden_states + ) + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=2, activations=hidden_states + ) + # aiter's in-kernel gather is numerically broken on gfx1250 (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, sorted + # row i reads source token gather_idx[i] // n_expts_act, so this reproduces the + # in-kernel gather exactly (validated: manual gather -> maxrel ~5e-3). + gather_src = gather_idx.to(torch.long) // topk + hidden_sorted = hidden_states[gather_src] + hidden_fp8 = downcast_to_static_fp8(hidden_sorted, gemm1_lhs_scale) + logger.info( + "[aiter-w4a8] scales g1_lhs=%s g2_lhs=%s hidden_fp8=%s; calling gemm1 " + "(unpadded N/K w1=%s/%s w2=%s/%s, alpha=%.4f limit=%.2f)", + float(gemm1_lhs_scale), + float(gemm2_lhs_scale), + tuple(hidden_fp8.shape), + unpadded_n_w1, + unpadded_k_w1, + unpadded_n_w2, + unpadded_k_w2, + swiglu_alpha, + swiglu_limit, + ) + + intermediate_cache1 = moe_gemm_a8w4( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale, + quant_config.w1_bias, + routing_data, + gather_indx=None, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + out_dtype=torch.float8_e4m3fn, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + logger.info( + "[aiter-w4a8] gemm1 OK ic1=%s %s; calling gemm2", + tuple(intermediate_cache1.shape), + intermediate_cache1.dtype, + ) + intermediate_cache3 = moe_gemm_a8w4( + intermediate_cache1, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + + # Restore the per-token weight sum (e.g. routed_scaling_factor) that the + # log+softmax normalization divided out. Output weighting (gammas on GEMM2) is + # linear in the gate, so a single post-scale is exact. + logger.info( + "[aiter-w4a8] gemm2 OK ic3=%s %s; rescaling + writing output=%s", + tuple(intermediate_cache3.shape), + intermediate_cache3.dtype, + tuple(output.shape), + ) + out = intermediate_cache3.to(output.dtype) * tw.sum(dim=1, keepdim=True) + output.copy_(out.reshape(output.shape)) + logger.info("[aiter-w4a8] DONE") + return output + def apply( self, output: torch.Tensor, @@ -1203,6 +1505,40 @@ def apply( if global_num_experts == -1: global_num_experts = E + # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from + # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs + # unswizzle. Output-side router weighting only (the post-scale below assumes + # it); LoRA still uses the matmul_ogs path further down. + if ( + (quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16) + and not apply_router_weight_on_input + and self._lora_context is None + and rocm_aiter_ops.is_enabled() + ): + if ( + self._try_apply_aiter_w4a8( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + apply_router_weight_on_input, + ) + is not None + ): + return + + # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids + # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for + # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). + # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not + # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA + # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below + # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA + # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. + # Note that the output tensor might be in workspace13 intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) @@ -1211,18 +1547,108 @@ def apply( gammas = routing_data.gate_scal if routing_data else None - matmul_ogs( - hidden_states, - w1, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=quant_config.w1_precision, - gammas=gammas if apply_router_weight_on_input else None, - fused_activation=None, - y=intermediate_cache1, + hidden_bf16_for_lora = hidden_states + + _moe_gemm_a8w4_fn = None + _downcast_static_fp8_fn = None + # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's + # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also + # serve the w4a16 case by dynamically quantizing the bf16 activations to + # FP8 here and passing the scale. This routes the MoE through aiter's + # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. + _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 + if _mxfp4_weights and rocm_aiter_ops.is_enabled(): + try: + from aiter.ops.triton.moe_op_gemm_a8w4 import ( + moe_gemm_a8w4 as _moe_gemm_a8w4_fn, + ) + from aiter.ops.triton.quant_moe import ( + downcast_to_static_fp8 as _downcast_static_fp8_fn, + ) + except ImportError: + pass + use_aiter_unfused_a8w4 = ( + _mxfp4_weights + and _moe_gemm_a8w4_fn is not None + and _downcast_static_fp8_fn is not None ) + if use_aiter_unfused_a8w4: + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=1, + activations=hidden_bf16_for_lora, + ) + # GEMM1 passes W2's static scale through to the kernel; resolve it before + # GEMM2 activations exist (fallback uses token hidden states for amax). + gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=hidden_bf16_for_lora, + ) + hidden_fp8 = _downcast_static_fp8_fn( + hidden_bf16_for_lora, + gemm1_lhs_scale, + ) + gemm1_out = _moe_gemm_a8w4_fn( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale_proxy, + 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.bfloat16, + apply_swiglu=False, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) + else: + hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_bf16_for_lora, + quant_config, + ) + ) + w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w1_precision, + gemm1_lhs_scale_fp8, + ) + + matmul_ogs( + hidden_for_gemm1, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w1_precision_for_gemm, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=None, + y=intermediate_cache1, + ) + # w13 LoRA: gather the activation input from expert-sorted # intermediate_cache1, then add the LoRA delta in-place on that copy # before passing it to activation — exactly mirroring the old @@ -1243,7 +1669,7 @@ def apply( ) = self.apply_w13_lora( lora_context, y=act_input, - x=hidden_states, + x=hidden_bf16_for_lora, topk_ids=global_topk_ids, topk_weights=topk_weights, expert_map=expert_map, @@ -1264,16 +1690,56 @@ def apply( # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. routing_data.n_expts_act = 1 - matmul_ogs( - intermediate_cache2[gather_indx.src_indx], - w2, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - precision_config=quant_config.w2_precision, - gammas=None if apply_router_weight_on_input else gammas, - y=intermediate_cache3, - ) + gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] + if use_aiter_unfused_a8w4: + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=gemm2_input_bf16, + ) + gemm2_fp8 = _downcast_static_fp8_fn( + gemm2_input_bf16, + gemm2_lhs_scale, + ) + gemm2_out = _moe_gemm_a8w4_fn( + gemm2_fp8, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_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, + ) + intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) + else: + gemm2_input, gemm2_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + gemm2_input_bf16, + quant_config, + gemm_num=2, + ) + ) + w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w2_precision, + gemm2_lhs_scale_fp8, + ) + + matmul_ogs( + gemm2_input, + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=w2_precision_for_gemm, + gammas=None if apply_router_weight_on_input else gammas, + y=intermediate_cache3, + ) # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index c70dcaa3379b..5110bcb7be70 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -58,10 +58,8 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): value_layout = StridedLayout scale_layout = StridedLayout elif current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx1250 - value_layout = StridedLayout - if should_use_cdna4_mx_scale_swizzle() or on_gfx1250(): + if should_use_cdna4_mx_scale_swizzle(): try: # triton < 3.6 from triton_kernels.tensor_details.layout import GFX950MXScaleLayout @@ -73,6 +71,12 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): scale_layout = CDNA4MXScaleLayout else: + # gfx1250 (and other ROCm archs): the vendored triton_kernels has no + # gfx1250 scale-swizzle, and the gfx1250 aiter `moe_gemm_a8w4` kernel + # reads a CDNA4-swizzled scale as garbage (validated on the FFM sim: + # CDNA4_SCALE -> maxrel ~7e4, plain/None -> ~6e-3). Keep the scale + # unswizzled (StridedLayout) and pass swizzle_mx_scale=None in the + # W4A8 path (gpt_oss_triton_kernels_moe._try_apply_aiter_w4a8). scale_layout = StridedLayout else: value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout( From e3c259c35e331700b049e03e2096cca54ee028ee Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Wed, 27 May 2026 11:02:23 -0500 Subject: [PATCH 13/68] [ROCm][gfx1250] gpt-oss: skip pruned-layer weights in mxfp4 + quark loaders Enables `hf_overrides={"num_hidden_layers": N}` layer-prune smoke tests for gpt-oss-120b on the FFM sim (mirrors the DSV4 loader fix). Both `_load_weights_mxfp4` and `_load_weights_quark` did raw `params_dict[name]` lookups and KeyError'd on checkpoint weights for layers the pruned model doesn't build. Skip any weight whose `layers..` index is >= num_hidden_layers. WIP bring-up: with this, gpt-oss-120b (mxfp4) loads but the forward still crashes on gfx1250 with an async "invalid device function" upstream of the MoE experts (kernel not pinpointed yet). The AMD quark w4a8 variant (/data/amd/gpt-oss-120b-w-mxfp4-a-fp8) is staged to try next. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- vllm/model_executor/models/gpt_oss.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 01f2752ac54b..863d8ea3ba7f 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -413,6 +413,14 @@ def _load_weights_mxfp4( if is_pp_missing_parameter(name, self): continue + # Skip checkpoint weights for layers dropped by an hf_overrides + # `num_hidden_layers` prune: the model only builds params for the + # layers it keeps, so extra-layer weights have no destination param. + if "layers." in name: + idx_str = name.split("layers.", 1)[1].split(".", 1)[0] + if idx_str.isdigit() and int(idx_str) >= self.config.num_hidden_layers: + continue + if ".w13_weight_scale" in name: # Handle MLP gate and up projection weights scale if use_ep: @@ -635,6 +643,13 @@ def _get_moe_weight_dtype(layer_id: int = 0) -> str | None: if is_pp_missing_parameter(name, self): continue + # Skip checkpoint weights for layers dropped by an hf_overrides + # `num_hidden_layers` prune (mirrors _load_weights_mxfp4). + if "layers." in name: + idx_str = name.split("layers.", 1)[1].split(".", 1)[0] + if idx_str.isdigit() and int(idx_str) >= self.config.num_hidden_layers: + continue + layer_id, expert_id, fused_name = None, None, None moe_quant_method = None if "experts" in name: From 980069ab70845df7fbc7ca3ad1fa617a5ab24997 Mon Sep 17 00:00:00 2001 From: Douglas Lehr Date: Wed, 27 May 2026 21:59:45 -0500 Subject: [PATCH 14/68] [ROCm][gfx1250] Enable AiterW4A8ExpertsMonolithic on gfx1250 (gpt-oss path) Generalizes the monolithic AITER MXFP4 W4A8 MoE experts class so it can fire on gfx1250 (RDNA4 sim), not just gfx950. This is the code path gpt-oss-120b selects when use_mxfp4_w4a8 is set -- the earlier DSV4 commit only enabled the gpt_oss_triton_kernels path, not AiterW4A8ExpertsMonolithic. aiter_triton_kernel_w4a8_moe_forward / triton_kernel_fused_mxfp4_w4a8_experts: - Routing: on gfx1250 import aiter's pure-torch `routing_torch` instead of the triton `routing`. The triton routing kernel compiles a TMA (TDM) descriptor whose last dim is `topk * 2` bytes, which is < 16 (the descriptor minimum) for power-of-2 topk like gpt-oss' topk=4, and fails to compile for the warmup dummy run plus every decode step. routing_torch is numerically identical to routing on the sim where the latter compiles. DSV4's topk=6 dodged this because next_power_of_2(6)==8!=6 disables the descriptor branch in aiter. - Try aiter's nested `moe.moe_routing.routing` module first, fall back to the legacy `moe_routing.routing` path (handles aiter version skew). - Gather: gfx1250 in-kernel gather is numerically broken (validated on the sim: do_gather=True -> maxrel ~2.4). Replicate aiter's gather in torch (sorted row i reads token gather_idx[i] // topk) and pass gather_indx=None; reproduces the in-kernel gather to ~5e-3 maxrel. gfx950 path unchanged. - Scale swizzle: gfx1250 stores the MXFP4 weight scale unswizzled (mxfp4_utils._swizzle_mxfp4 keeps StridedLayout on gfx1250), and the gfx1250 moe_gemm_a8w4 reads CDNA4_SCALE as garbage (CDNA4_SCALE -> maxrel ~7e4 on the sim, None -> ~6e-3). Pass swizzle_mx_scale=None on gfx1250; gfx950 still uses "CDNA4_SCALE". _supports_current_device(): accept on_gfx950() OR on_gfx1250(). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Douglas Lehr --- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 62 +++++++++++++++++-- 1 file changed, 58 insertions(+), 4 deletions(-) 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 index 7c3fe5831f3b..49791f88e24a 100644 --- 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 @@ -46,11 +46,52 @@ def aiter_triton_kernel_w4a8_moe_forward( and quant_config.use_mxfp4_w4a8 and rocm_aiter_ops.is_enabled() ) - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + from vllm.platforms.rocm import on_gfx1250 + + # aiter exposes its MoE routing under two module paths across versions; + # prefer the nested `moe.moe_routing` location, fall back to the legacy one. + # + # On gfx1250 use aiter's pure-torch ``routing_torch`` instead of the triton + # ``routing``: the triton routing kernel compiles a TMA (TDM) descriptor + # whose last dim is ``topk * 2`` bytes, which is < 16 bytes (the descriptor + # minimum) for a power-of-2 topk such as gpt-oss' topk=4 and fails to + # compile for small batches (warmup dummy run + every decode step). + # ``routing_torch`` avoids the TDM kernel and is numerically identical + # (validated on the FFM sim: gather/scatter/gate_scal match the triton path + # exactly where the latter compiles). DeepSeek-V4's topk=6 dodged this since + # ``next_power_of_2(6) == 8 != 6`` disables the descriptor branch. + if on_gfx1250(): + try: + from aiter.ops.triton.moe.moe_routing.routing import ( + routing_torch as aiter_routing, + ) + except ImportError: + from aiter.ops.triton.moe_routing.routing import ( + routing_torch as aiter_routing, + ) + else: + try: + from aiter.ops.triton.moe.moe_routing.routing import ( + routing as aiter_routing, + ) + except ImportError: + 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 ) + + # gfx1250: aiter's in-kernel gather is numerically broken (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, + # sorted row i reads source token gather_idx[i] // n_expts_act, so this + # reproduces the in-kernel gather exactly (manual gather -> maxrel ~5e-3). + # gfx950 keeps the (working) in-kernel gather. + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + return triton_kernel_fused_mxfp4_w4a8_experts( None, hidden_states, @@ -130,6 +171,15 @@ def triton_kernel_fused_mxfp4_w4a8_experts( hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale ) + # gfx1250 stores the MXFP4 weight scale unswizzled (StridedLayout, see + # mxfp4_utils._swizzle_mxfp4) because the gfx1250 moe_gemm_a8w4 reads a + # CDNA4-swizzled scale as garbage (validated on the FFM sim: CDNA4_SCALE -> + # maxrel ~7e4, plain/None -> ~6e-3); pass swizzle_mx_scale=None there. + # gfx950 uses the CDNA4 swizzle layout. + from vllm.platforms.rocm import on_gfx1250 + + mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" + intermediate_cache1 = moe_gemm_a8w4( hidden_states, w1.storage.data, @@ -199,12 +249,16 @@ def activation_format() -> mk.FusedMoEActivationFormat: @staticmethod def _supports_current_device() -> bool: - # Requires AITER and GFX950 + # Requires AITER and a supported AMD arch. gfx950 (CDNA4) uses the + # in-kernel gather + CDNA4 scale swizzle; gfx1250 routes through the + # same moe_gemm_a8w4 kernel with a manual gather and unswizzled scales + # (see aiter_triton_kernel_w4a8_moe_forward / the swizzle handling in + # triton_kernel_fused_mxfp4_w4a8_experts). if not rocm_aiter_ops.is_enabled(): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - return on_gfx950() + return on_gfx950() or on_gfx1250() @staticmethod def _supports_no_act_and_mul() -> bool: From 1a9af4012fb4be5cba2cde095941f7e5faad95be Mon Sep 17 00:00:00 2001 From: Daniel Date: Tue, 2 Jun 2026 12:44:16 -0400 Subject: [PATCH 15/68] docker/Dockerfile.rocm: port gfx1250_wip build chain - add libpciaccess0/libpciaccess-dev/libdrm-dev/ pkg-config/cmake to base apt; drop the standalone cmake install - install rocm[libraries,devel] (vs rocm[devel]) and pin torch/torchvision/torchaudio via genesis gfx1250 index - /opt/rocm symlinks for bin/include/lib and rocprofiler-sdk; PATH and LD_LIBRARY_PATH wired through. - AITER: switch branch shared/triton-gfx12 -> main. - ENABLE_CK toggle: AITER_ENABLE_CK=0 (default) disables Composable Kernel for the current build Signed-off-by: Daniel --- docker/Dockerfile.rocm | 42 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 046ec3c7090f..cabc66f24474 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -34,7 +34,8 @@ ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ software-properties-common build-essential git curl sudo vim less \ - libopenmpi-dev libpci-dev python3-venv python3-dev \ + libopenmpi-dev libpci-dev libpciaccess0 libpciaccess-dev libdrm-dev \ + pkg-config cmake python3-venv python3-dev \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ libnuma-dev ccache mold @@ -113,11 +114,35 @@ WORKDIR ${COMMON_WORKDIR} # Ensure core build tools are installed (idempotent if base image already has them) RUN uv pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 -# ROCm SDK install + init (idempotent -- skips if already satisfied) +# ROCm PyTorch stack + SDK from genesis gfx1250 index, then init CLI. +# Uninstall first (no-op if missing) then install without --upgrade so torch/torchvision/torchaudio +# stay on the same genesis build (avoids torchvision::nms / partial upgrade skew). ARG PIP_EXTRA_INDEX_URL=https://rocm.genesis.amd.com/whl/gfx1250/ -RUN uv pip install --index-url ${PIP_EXTRA_INDEX_URL} "rocm[devel]" \ +ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} +RUN uv pip uninstall -y torch torchvision torchaudio || true \ + && uv pip install --index-url ${PIP_EXTRA_INDEX_URL} \ + torch torchvision torchaudio "rocm[libraries,devel]" \ && rocm-sdk init +# /opt/rocm symlinks for tools and rocprofiler (matches rocm-sdk layout) +RUN set -eux; \ + echo "=== setting up ROCm symlinks ===" 1>&2; \ + mkdir -p /opt/rocm; \ + ROCM_SDK="$(dirname "$(rocm-sdk path --bin)")"; \ + ROCM_SDK_CORE="$(dirname "$ROCM_SDK")/_rocm_sdk_core"; \ + ln -sf "$ROCM_SDK/bin" /opt/rocm/bin; \ + ln -sf "$ROCM_SDK/include" /opt/rocm/include; \ + ln -sf "$ROCM_SDK/lib" /opt/rocm/lib; \ + ROCPROF_DIR="$ROCM_SDK_CORE/lib/rocprofiler-sdk"; \ + mkdir -p "$ROCPROF_DIR"; \ + if [ -f "$ROCPROF_DIR/librocprofv3-list-avail.so.1" ]; then \ + ln -sf librocprofv3-list-avail.so.1 "$ROCPROF_DIR/librocprofv3-list-avail.so"; \ + else \ + echo "Note: librocprofv3-list-avail.so.1 not found under $ROCPROF_DIR, skipping symlink"; \ + fi + +ENV PATH=/opt/rocm/bin:${PATH} + # ROCm SDK environment setup ENV SITE_PACKAGES=${SITE_PACKAGES:-/opt/venv/lib/python3.12/site-packages} ENV ROCM_PATH=$SITE_PACKAGES/_rocm_sdk_devel @@ -140,7 +165,7 @@ ENV CMAKE_PREFIX_PATH=$SITE_PACKAGES/torch/share/cmake RUN if [ -d "$SITE_PACKAGES/_rocm_sdk_core/lib" ] && [ ! -e "$SITE_PACKAGES/_rocm_sdk_core/lib/libamdhip64.so" ]; then \ ln -s libamdhip64.so.7 $SITE_PACKAGES/_rocm_sdk_core/lib/libamdhip64.so; \ fi -ENV LD_LIBRARY_PATH=$SITE_PACKAGES/_rocm_sdk_core/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=$SITE_PACKAGES/_rocm_sdk_core/lib:${LD_LIBRARY_PATH}:/opt/rocm/lib/rocprofiler-sdk:/opt/rocm/lib/rocm_sysdeps/lib # ----------------------- # vLLM fetch stages @@ -700,6 +725,8 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ && uv pip install -r requirements/rocm.txt \ + && uv pip uninstall -y torch torchvision torchaudio || true \ + && uv pip install --index-url ${PIP_EXTRA_INDEX_URL} torch torchvision torchaudio \ && pip uninstall -y vllm \ && uv pip install *.whl @@ -707,7 +734,12 @@ RUN pip install amdsmi # Install amd-aiter and Triton from repo/branch (for gfx1250; base image may not include them) ARG AITER_REPO="https://github.com/ROCm/aiter" -ARG AITER_BRANCH="shared/triton-gfx12" +ARG AITER_BRANCH="main" +# ENABLE_CK=0 disables Composable Kernel both at wheel build time +# Override at build time with `--build-arg AITER_ENABLE_CK=1` to re-enable CK +# (requires a recursive aiter clone with the composable_kernel submodule). +ARG AITER_ENABLE_CK=0 +ENV ENABLE_CK=${AITER_ENABLE_CK} # TODO: Triton repo requires github.amd.com authentication; re-enable once credentials are configured # ARG TRITON_REPO="https://github.amd.com/GFX-IP-Arch/triton" # ARG TRITON_BRANCH="shared/gfx1250" From 18ca4bad59830b7e0efa8b16d8e815bb51f32153 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Fri, 5 Jun 2026 17:14:05 -0400 Subject: [PATCH 16/68] Add gfx1250 Enabled docker rocm base (#987) * Add 1250 only base+vllm Dockerfile * Add upstream dockerfiles for rocm and rocm base * Enable rocm_base for gfx1250 --------- Co-authored-by: jpvillam Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm | 271 +++++++++++++++++++++--------------- docker/Dockerfile.rocm_1250 | 114 +++++++++++++++ docker/Dockerfile.rocm_base | 258 +++++++++++++++++----------------- 3 files changed, 404 insertions(+), 239 deletions(-) create mode 100644 docker/Dockerfile.rocm_1250 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index cabc66f24474..f8f2d4cd1251 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -17,6 +17,23 @@ ARG NIC_BACKEND=all ARG AINIC_VERSION=1.117.3-hydra ARG UBUNTU_CODENAME=jammy +# Sccache configuration. Release builds use this today; CI can opt in when a +# shared S3-compatible cache backend is available. +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base +# NIC backend for MoRI RDMA support. +# By default (all), drivers and userspace libraries for all supported NIC types +# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. +# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: +# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only +# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only +# --build-arg NIC_BACKEND=none # Install nothing. +ARG NIC_BACKEND=all +# AMD AINIC apt repo settings +# Users can specify a custom version compatible with their host drivers. +# The default version has been tested with ioinic-dkms=25.11.1.001 +ARG AINIC_VERSION=1.117.3-hydra +ARG UBUNTU_CODENAME=jammy + # Sccache configuration. Release builds use this today; CI can opt in when a # shared S3-compatible cache backend is available. ARG USE_SCCACHE @@ -31,11 +48,9 @@ FROM ${BASE_IMAGE} AS base ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} +# Install build dependencies and utilities # Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ - software-properties-common build-essential git curl sudo vim less \ - libopenmpi-dev libpci-dev libpciaccess0 libpciaccess-dev libdrm-dev \ - pkg-config cmake python3-venv python3-dev \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ libnuma-dev ccache mold @@ -44,6 +59,13 @@ RUN --mount=type=cache,target=/root/.cache/pip \ # Note: mold is installed but not set as the system default linker because # some packages use JIT compilation at runtime with flags mold does not support. # Build stages opt in via LDFLAGS="-fuse-ld=mold". +apt-transport-https ca-certificates wget curl \ + libnuma-dev ccache mold +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --upgrade pip +# Note: mold is installed but not set as the system default linker because +# some packages use JIT compilation at runtime with flags mold does not support. +# Build stages opt in via LDFLAGS="-fuse-ld=mold". # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE RUN if [ "$USE_SCCACHE" != "1" ]; then \ @@ -71,6 +93,25 @@ ENV CCACHE_COMPILERCHECK=content ARG max_jobs ENV MAX_JOBS=${max_jobs} +# Install UV — download first, then run, so a curl failure is not masked by the pipe +RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ + && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ + && rm -f /tmp/uv-install.sh \ + && uv --version + +# 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 +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +# Use copy mode to avoid hardlink failures with Docker cache mounts +ENV UV_LINK_MODE=copy +# ccache directory - persisted across layer rebuilds via cache mounts. +ENV CCACHE_DIR=/root/.cache/ccache +ENV CCACHE_COMPILERCHECK=content +# Empty by default so build steps fall back to $(nproc); CI can override. +ARG max_jobs +ENV MAX_JOBS=${max_jobs} + # Install sccache if USE_SCCACHE is enabled (for release builds) ARG USE_SCCACHE ARG SCCACHE_DOWNLOAD_URL @@ -104,68 +145,9 @@ 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} -# Install uv for faster pip installs (required by build_rixl, test, and final stages) -# Base images may not include uv; install here so all stages that use uv have it -RUN python3 -m pip install uv - ARG COMMON_WORKDIR WORKDIR ${COMMON_WORKDIR} -# Ensure core build tools are installed (idempotent if base image already has them) -RUN uv pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 - -# ROCm PyTorch stack + SDK from genesis gfx1250 index, then init CLI. -# Uninstall first (no-op if missing) then install without --upgrade so torch/torchvision/torchaudio -# stay on the same genesis build (avoids torchvision::nms / partial upgrade skew). -ARG PIP_EXTRA_INDEX_URL=https://rocm.genesis.amd.com/whl/gfx1250/ -ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -RUN uv pip uninstall -y torch torchvision torchaudio || true \ - && uv pip install --index-url ${PIP_EXTRA_INDEX_URL} \ - torch torchvision torchaudio "rocm[libraries,devel]" \ - && rocm-sdk init - -# /opt/rocm symlinks for tools and rocprofiler (matches rocm-sdk layout) -RUN set -eux; \ - echo "=== setting up ROCm symlinks ===" 1>&2; \ - mkdir -p /opt/rocm; \ - ROCM_SDK="$(dirname "$(rocm-sdk path --bin)")"; \ - ROCM_SDK_CORE="$(dirname "$ROCM_SDK")/_rocm_sdk_core"; \ - ln -sf "$ROCM_SDK/bin" /opt/rocm/bin; \ - ln -sf "$ROCM_SDK/include" /opt/rocm/include; \ - ln -sf "$ROCM_SDK/lib" /opt/rocm/lib; \ - ROCPROF_DIR="$ROCM_SDK_CORE/lib/rocprofiler-sdk"; \ - mkdir -p "$ROCPROF_DIR"; \ - if [ -f "$ROCPROF_DIR/librocprofv3-list-avail.so.1" ]; then \ - ln -sf librocprofv3-list-avail.so.1 "$ROCPROF_DIR/librocprofv3-list-avail.so"; \ - else \ - echo "Note: librocprofv3-list-avail.so.1 not found under $ROCPROF_DIR, skipping symlink"; \ - fi - -ENV PATH=/opt/rocm/bin:${PATH} - -# ROCm SDK environment setup -ENV SITE_PACKAGES=${SITE_PACKAGES:-/opt/venv/lib/python3.12/site-packages} -ENV ROCM_PATH=$SITE_PACKAGES/_rocm_sdk_devel -ENV HIP_DEVICE_LIB_PATH=$SITE_PACKAGES/_rocm_sdk_core/lib/llvm/amdgcn/bitcode - -# Install amd_smi from SDK source and patch rtld_global (guarded for idempotency) -RUN if [ -d "$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi" ]; then \ - cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install .; \ - else \ - echo "amd_smi source not found at $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi, skipping"; \ - fi \ - && if [ -f "$SITE_PACKAGES/rocm_sdk/__init__.py" ]; then \ - sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' $SITE_PACKAGES/rocm_sdk/__init__.py; \ - else \ - echo "rocm_sdk/__init__.py not found at $SITE_PACKAGES/rocm_sdk/__init__.py, skipping patch"; \ - fi -ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi:$PYTHONPATH - -ENV CMAKE_PREFIX_PATH=$SITE_PACKAGES/torch/share/cmake -RUN if [ -d "$SITE_PACKAGES/_rocm_sdk_core/lib" ] && [ ! -e "$SITE_PACKAGES/_rocm_sdk_core/lib/libamdhip64.so" ]; then \ - ln -s libamdhip64.so.7 $SITE_PACKAGES/_rocm_sdk_core/lib/libamdhip64.so; \ - fi -ENV LD_LIBRARY_PATH=$SITE_PACKAGES/_rocm_sdk_core/lib:${LD_LIBRARY_PATH}:/opt/rocm/lib/rocprofiler-sdk:/opt/rocm/lib/rocm_sysdeps/lib # ----------------------- # vLLM fetch stages @@ -200,25 +182,27 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ COPY tools/install_protoc.sh /tmp/install_protoc.sh RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh +# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ + sh -s -- -y --profile minimal --default-toolchain none +ENV PATH="/root/.cargo/bin:${PATH}" + # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 ENV CARGO_NET_RETRY=10 ENV RUSTUP_MAX_RETRIES=10 -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - cd ${COMMON_WORKDIR}/vllm \ - && uv pip install --system -r requirements/build/rust.txt - # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while -# keeping the cache benefit. Do not cache target/, because stale target metadata -# can outlive source updates across BuildKit cache reuse. +# keeping the cache benefit. Copy the binary out so it persists into the image +# layer for later COPY --from=rust-build. RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ - && bash build_rust.sh \ - && test -x vllm/vllm-rs + && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \ + && test -x /tmp/vllm-rs # ----------------------- # vLLM native build stages @@ -238,7 +222,6 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ # pyproject.toml is bind-mounted in the RUN step so metadata-only changes do # not invalidate the expensive native build layer. COPY setup.py CMakeLists.txt ./ -COPY tools/build_rust.py tools/build_rust.py COPY cmake cmake/ COPY csrc csrc/ COPY vllm/envs.py vllm/envs.py @@ -270,10 +253,9 @@ ENV VLLM_TARGET_DEVICE=rocm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels -# Drop the pre-built Rust artifacts into the source tree. setup.py detects -# them and ships them as-is, skipping the local Rust build. -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ +# Drop the pre-built rust frontend binary into the source tree. setup.py +# detects it and ships it as-is, skipping the local cargo build. +COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd vllm \ @@ -296,14 +278,15 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh / COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages FROM base AS build_rixl -ARG RIXL_BRANCH="f33a5599" +ARG RIXL_BRANCH="39be1de8" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" -ARG UCX_BRANCH="da3fac2a" -ARG UCX_REPO="https://github.com/ROCm/ucx.git" +ARG UCX_BRANCH="bfb51733" +ARG UCX_REPO="https://github.com/openucx/ucx.git" ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl @@ -326,10 +309,14 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ ibverbs-providers \ && rm -rf /var/lib/apt/lists/* +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system meson auditwheel patchelf tomlkit RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system meson auditwheel patchelf tomlkit RUN --mount=type=cache,target=/root/.cache/ccache \ + cd /usr/local/src && \ + RUN --mount=type=cache,target=/root/.cache/ccache \ cd /usr/local/src && \ git clone ${UCX_REPO} && \ cd ucx && \ @@ -337,6 +324,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ ./autogen.sh && \ mkdir build && cd build && \ CC="ccache gcc" CXX="ccache g++" \ + CC="ccache gcc" CXX="ccache g++" \ ../configure \ --prefix=/usr/local/ucx \ --enable-shared \ @@ -355,19 +343,29 @@ ENV PATH=/usr/local/ucx/bin:$PATH ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone ${RIXL_REPO} /opt/rixl && \ + RUN --mount=type=cache,target=/root/.cache/ccache \ git clone ${RIXL_REPO} /opt/rixl && \ cd /opt/rixl && \ git checkout ${RIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ + CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ ninja -j$(nproc) && \ + ninja -j$(nproc) && \ ninja install # Generate RIXL wheel -RUN cd /opt/rixl && mkdir -p /app/install && \ +# Exclude libcore and libpull from auditwheel: transitive dependencies +# that are not shipped in the wheel and vary across base images. +RUN cd /opt/rixl && \ + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ + contrib/build-wheel.sh && \ + mkdir -p /app/install && \ + _ucx_install_dir=${UCX_HOME} \ ./contrib/build-wheel.sh \ --output-dir /app/install \ --rocm-dir ${ROCM_PATH} \ @@ -474,10 +472,9 @@ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR -# Drop the pre-built Rust artifacts into the source tree. setup.py detects -# them and ships them as-is, skipping the local Rust build. -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ +# Drop the pre-built rust frontend binary into the source tree. setup.py +# detects it and ships it as-is, skipping the local cargo build. +COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs # Create /install directory for custom wheels RUN mkdir -p /install @@ -511,9 +508,11 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) # This ensures setuptools_scm sees clean repo state for version detection RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/uv \ cd vllm \ && uv pip install --system setuptools_scm regex \ + && uv pip install --system setuptools_scm regex \ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ && echo "Detected vLLM version: ${VLLM_VERSION}" \ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt @@ -550,19 +549,26 @@ RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ # Install dependencies using custom wheels from /install RUN --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + RUN --mount=type=cache,target=/root/.cache/uv \ cd vllm \ && echo "Building vLLM with custom wheels from /install" \ && uv pip install --system --find-links /install -r requirements/rocm.txt +&& uv pip install --system --find-links /install -r requirements/rocm.txt # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt # (setup.py auto-detects ccache/sccache in PATH) +# (setup.py auto-detects ccache/sccache in PATH) RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ && export CCACHE_BASEDIR="$PWD" \ + && export CCACHE_BASEDIR="$PWD" \ && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist +&& MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist FROM scratch AS export_vllm_wheel_release ARG COMMON_WORKDIR @@ -575,6 +581,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchc COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- @@ -598,10 +605,17 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ uv pip install --system /rixl_install/*.whl /deep_install/*.whl +--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ + uv pip install --system /rixl_install/*.whl /deep_install/*.whl # Copy ROCShmem runtime libraries. COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + # Copy ROCShmem runtime libraries. + COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem + # RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ librdmacm1 \ @@ -610,16 +624,24 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ ibverbs-utils \ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* +# Install torchcodec from source for ROCm/torch ABI compatibility. # Install torchcodec from source for ROCm/torch ABI compatibility. COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + bash /tmp/install_torchcodec.sh \ + RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/pip \ --mount=type=cache,target=/root/.cache/torchcodec-wheels \ bash /tmp/install_torchcodec.sh \ && rm /tmp/install_torchcodec.sh \ && apt-get clean && rm -rf /var/lib/apt/lists/* +&& apt-get clean && rm -rf /var/lib/apt/lists/* # Pre-install shared ROCm runtime dependencies. COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ @@ -631,9 +653,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ENV HF_XET_HIGH_PERFORMANCE=1 ENV HF_HUB_DOWNLOAD_TIMEOUT=60 -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - # Pre-install vLLM test dependencies. COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt RUN --mount=type=cache,target=/root/.cache/uv \ @@ -648,6 +667,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ && rm /tmp/rocm-test-reqs.txt +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # See: https://github.com/pytorch/pytorch/issues/169857 ENV MIOPEN_DEBUG_CONV_DIRECT=0 @@ -687,12 +707,47 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace # Copy in the v1 package (for python-only install test group). COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 +# Hide source under src/ so it won't shadow the installed package in tests. +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" + +# Install vllm_test_utils in ci_base for ci_base + wheel parity. +COPY tests/vllm_test_utils /tmp/vllm_test_utils +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /tmp/vllm_test_utils \ + && rm -rf /tmp/vllm_test_utils + +# ----------------------- +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. +FROM ${CI_BASE_IMAGE} AS test +ARG COMMON_WORKDIR + +# Install the vLLM wheel (--no-deps: all deps already in ci_base). +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system --no-deps *.whl + +# Store the vLLM wheel in the image for python-only install tests. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ + +WORKDIR /vllm-workspace +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# Copy in the v1 package (for python-only install test group). +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 + # Hide source under src/ so it won't shadow the installed package in tests. RUN mkdir src && mv vllm src/vllm # ----------------------- # Final vLLM image -FROM base AS final +FROM mori_base AS final RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* @@ -706,6 +761,7 @@ RUN rm -f /usr/bin/sccache || true \ ENV SCCACHE_BUCKET= ENV SCCACHE_REGION= ENV SCCACHE_ENDPOINT= +ENV SCCACHE_ENDPOINT= ENV SCCACHE_S3_NO_CREDENTIALS= ENV SCCACHE_IDLE_TIMEOUT= @@ -717,46 +773,35 @@ RUN case "$(which python3)" in \ *) ;; esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --upgrade huggingface-hub[cli,hf_transfer] scipy + uv pip install --system --upgrade huggingface-hub[cli] # Install vLLM using uv (inherited from base stage) # Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ - && uv pip install -r requirements/rocm.txt \ - && uv pip uninstall -y torch torchvision torchaudio || true \ - && uv pip install --index-url ${PIP_EXTRA_INDEX_URL} torch torchvision torchaudio \ + && uv pip install --system -r requirements/rocm.txt \ && pip uninstall -y vllm \ - && uv pip install *.whl - -RUN pip install amdsmi - -# Install amd-aiter and Triton from repo/branch (for gfx1250; base image may not include them) -ARG AITER_REPO="https://github.com/ROCm/aiter" -ARG AITER_BRANCH="main" -# ENABLE_CK=0 disables Composable Kernel both at wheel build time -# Override at build time with `--build-arg AITER_ENABLE_CK=1` to re-enable CK -# (requires a recursive aiter clone with the composable_kernel submodule). -ARG AITER_ENABLE_CK=0 -ENV ENABLE_CK=${AITER_ENABLE_CK} -# TODO: Triton repo requires github.amd.com authentication; re-enable once credentials are configured -# ARG TRITON_REPO="https://github.amd.com/GFX-IP-Arch/triton" -# ARG TRITON_BRANCH="shared/gfx1250" -RUN apt-get update -q -y && apt-get install -q -y git && rm -rf /var/lib/apt/lists/* \ - && git clone --depth 1 --branch "$AITER_BRANCH" "$AITER_REPO" /tmp/aiter \ - && pip install --no-build-isolation /tmp/aiter && rm -rf /tmp/aiter -# && git clone --depth 1 --branch "$TRITON_BRANCH" "$TRITON_REPO" /tmp/triton \ -# && pip install --no-build-isolation /tmp/triton && rm -rf /tmp/triton + && uv pip install --system *.whl + +# Install RIXL wheel +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + uv pip install --system /rixl_install/*.whl ARG COMMON_WORKDIR ARG BASE_IMAGE +ARG NIC_BACKEND +ARG AINIC_VERSION # Copy over the benchmark scripts as well COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + ENV TOKENIZERS_PARALLELISM=false # ENV that can improve safe tensor loading, and end-to-end time @@ -771,7 +816,9 @@ ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt CMD ["/bin/bash"] diff --git a/docker/Dockerfile.rocm_1250 b/docker/Dockerfile.rocm_1250 new file mode 100644 index 000000000000..697d92dda512 --- /dev/null +++ b/docker/Dockerfile.rocm_1250 @@ -0,0 +1,114 @@ +FROM ubuntu:24.04 +ENV DEBIAN_FRONTEND=noninteractive + +# --- ROCm + PyTorch versioning (override any of these at build time) ------- +ARG GFX_ARCH=gfx1250 +ARG ROCM_VERSION=7.14.0a20260604 +ARG ROCM_BUILD=a0 +ARG TORCH_VERSION=2.10.0 +ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250/ +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" +ARG AITER_BRANCH="jpvillam/gfx1250_0604" + +ENV VLLM_TARGET_DEVICE=rocm +ENV PYTORCH_ROCM_ARCH=${GFX_ARCH} +ENV CMAKE_BUILD_TYPE=Release +ENV GPU_TARGET=${GFX_ARCH} +ENV VIRTUAL_ENV=/opt/venv + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + build-essential clang-19 lld ccache ninja-build cmake \ + git curl libcurl4-openssl-dev rsync ssh wget ca-certificates \ + python3 python3-pip python3-dev python3-venv \ + less vim libzstd-dev numactl libelf1 m4 && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + +SHELL ["/bin/bash", "-e", "-u", "-o", "pipefail", "-c"] + +RUN python3 -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH + +# COPY ffm/ /root/x/ffm/ +# ENV PATH=${ROCM_PATH}/llvm/bin:${ROCM_PATH}/bin:$PATH +# ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${ROCM_PATH}/lib/rocm_sysdeps/lib:${ROCM_PATH}/llvm/lib +# ENV HSA_MODEL_LIB=/root/x/ffm/libhsakmtmodel.so +# ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" +# ENV HSA_MODEL_ARGS=ffm_enable_time_slicing +# ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 +# ENV HSA_MODEL_NUM_THREADS=64 + +ENV HSA_ENABLE_SDMA=0 +ENV HSA_ENABLE_INTERRUPT=0 + +# Install the TheRock PyTorch wheel and the matching rocm-sdk wheels into the venv +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + "torch==${TORCH_VERSION}+rocm${ROCM_VERSION}${ROCM_BUILD:+.${ROCM_BUILD}}" \ + "rocm[libraries,devel]==${ROCM_VERSION}${ROCM_BUILD:++${ROCM_BUILD}}" && \ + rocm-sdk init + +RUN pip install \ + filelock \ + "typing-extensions>=4.10.0" \ + "sympy>=1.13.3" \ + "networkx>=2.5.1" \ + jinja2 \ + "fsspec>=0.8.5" + +# Set ROCm, Python, Paths, and Environment Variables +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python3.12/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib:${SITE_PACKAGES}/_rocm_sdk_libraries_${GFX_ARCH}/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# NPI base-image fix-ups so the wheel ROCm behaves like a normal install: +# - build amd_smi from the bundled source (provides the `amdsmi` module), +# - flip rocm_sdk's library preload off RTLD_GLOBAL (avoids symbol clashes), +RUN if [ -d "${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi" ]; then \ + cd "${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi" && pip install .; \ + fi && \ + if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi + +# Install AITER (https://github.com/ROCm/aiter): plain clone WITHOUT submodules +# (so 3rdparty/composable_kernel is absent), built with Composable Kernel disabled. +RUN git clone -b ${AITER_BRANCH} --depth 1 https://github.com/ROCm/aiter.git /app/aiter && \ + cd /app/aiter && \ + pip install packaging "cmake<4" ninja wheel setuptools_scm vcs_versioning pybind11 numpy && \ + ENABLE_CK=0 AITER_USE_SYSTEM_TRITON=1 pip install --no-build-isolation -e . + + +RUN apt-get update && \ + apt-get install -y --no-install-recommends pkg-config libnuma-dev libdrm-dev && \ + apt-get clean && rm -rf /var/lib/apt/lists/* + +RUN pip install "setuptools>=77.0.3,<81.0.0" setuptools_scm setuptools_rust wheel packaging jinja2 + +# Build + install vLLM for gfx1250 (cloned from VLLM_REPO @ VLLM_BRANCH) +RUN git clone "${VLLM_REPO}" /app/vllm \ + && cd /app/vllm \ + && git fetch -v --prune -- origin "${VLLM_BRANCH}" \ + && git checkout FETCH_HEAD + +# Compile the C++/HIP extensions (_C, _rocm_C) for gfx1250 and editable-install. +# MAX_JOBS is a build-arg so it can be tuned per machine (docker build --build-arg). +RUN cd /app/vllm && \ + pip install -e . --no-build-isolation --no-deps -v + +# vLLM Python runtime deps. common.txt pins no torch, so the NPI ROCm torch +# installed above is left untouched. +RUN pip install -r /app/vllm/requirements/common.txt + +WORKDIR /app/vllm +ENTRYPOINT ["/bin/bash"] \ No newline at end of file diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index ca3573c87947..541d31ff3743 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,16 +1,15 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="0f380657" -ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="d0c8b1f3" # release/2.11 as of 6/09 -ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.24.1" -ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG PYTORCH_AUDIO_BRANCH="v2.9.0" -ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" +ARG BASE_IMAGE=ubuntu:24.04 +ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250 +ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260605 +ARG TORCHVISION_VERSION=0.28.0a0+rocm7.14.0a20260605 +ARG TORCHAUDIO_VERSION=2.11.0a0+rocm7.14.0a20260605 +ARG ROCM_SDK_VERSION=7.14.0a20260605 ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.16.post2" +ARG AITER_BRANCH="jpvillam/gfx1250_0604" ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG MORI_BRANCH="v1.1.0" +ARG MORI_REPO="https://github.com/ROCm/mori.git" # Sccache configuration (only used in release pipeline) ARG USE_SCCACHE @@ -22,40 +21,46 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -ENV ROCM_PATH=/opt/rocm -ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 +ARG PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx942;gfx950;gfx1250 -ENV MORI_GPU_ARCHS=gfx942;gfx950;gfx1250 +ENV AITER_ROCM_ARCH=gfx1250 +ENV MORI_GPU_ARCHS=gfx942;gfx950 + +# TODO: Unset these when support is available for gfx1250 +ENV ENABLE_CK=0 +ARG PREBUILD_KERNELS=0 # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 +ARG PYTHON_VERSION=3.12 +ENV PYTHON_VERSION=${PYTHON_VERSION} + RUN mkdir -p /app WORKDIR /app ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common build-essential git curl sudo vim less libopenmpi-dev libpci-dev python3-venv python3-dev - -RUN python3 -m venv /opt/python --system-site-packages -ENV PATH=/opt/python/bin:$PATH - -# Install UV -RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="/usr/local/bin" sh -ENV UV_PYTHON=/opt/python/bin/python -# 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 -ENV UV_HTTP_TIMEOUT=500 -ENV UV_INDEX_STRATEGY="unsafe-best-match" -# Use copy mode to avoid hardlink failures with Docker cache mounts -ENV UV_LINK_MODE=copy - - -RUN uv pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ + && for i in 1 2 3; do \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + done \ + && apt-get update -y \ + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ + && python3 --version + +ENV VIRTUAL_ENV=/opt/venv +RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH + +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* # Install sccache if USE_SCCACHE is enabled (for release builds) @@ -78,6 +83,51 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ && sccache --version; \ fi +### +### Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +### +ARG ROCM_WHEEL_INDEX +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG ROCM_SDK_VERSION +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + "torch==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream +# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm; + +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi + # Setup sccache for HIP compilation via HIP_CLANG_PATH # This creates wrapper scripts in a separate directory and points HIP to use them # This avoids modifying the original ROCm binaries which can break detection @@ -86,9 +136,9 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ RUN if [ "$USE_SCCACHE" = "1" ]; then \ echo "Setting up sccache wrappers for HIP compilation..." \ && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ && chmod +x /opt/sccache-wrappers/clang \ && echo "sccache wrappers created in /opt/sccache-wrappers"; \ fi @@ -102,87 +152,13 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} -### -### Triton Build -### -FROM base AS build_triton -ARG TRITON_BRANCH -ARG TRITON_REPO -RUN git clone ${TRITON_REPO} -# Cherry picking the following -# https://github.com/triton-lang/triton/pull/8991 -RUN cd triton \ - && git checkout ${TRITON_BRANCH} \ - && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ - && git cherry-pick 555d04f \ - && if [ ! -f setup.py ]; then cd python; fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install -RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \ - && python3 -m build --wheel && cp dist/*.whl /app/install; fi - - ### ### AMD SMI Build ### FROM base AS build_amdsmi -RUN cd /opt/rocm/share/amd_smi \ +RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install - - -### -### Pytorch build -### -FROM base AS build_pytorch -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_REPO -ARG USE_SCCACHE - -RUN apt-get update && apt-get install -y pkg-config liblzma-dev -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 && python3 tools/amd_build/build_amd.py \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ - && sccache --show-stats; \ - fi \ - && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN git clone ${PYTORCH_VISION_REPO} vision -RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN git clone ${PYTORCH_AUDIO_REPO} audio -RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ - && git submodule update --init --recursive \ - && pip install -r requirements.txt \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && pip install dist/*.whl -RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ - && cp /app/vision/dist/*.whl /app/install \ - && cp /app/audio/dist/*.whl /app/install +RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install ### @@ -191,14 +167,17 @@ RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${MORI_REPO} -RUN cd mori \ +RUN mkdir -p /app/install; \ + if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ + echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ + else \ + git clone ${MORI_REPO} \ + && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ - && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl -RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ + && cp /app/mori/dist/*.whl /app/install; \ + fi ### @@ -208,8 +187,12 @@ FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE -RUN git clone ${FA_REPO} -RUN cd flash-attention \ +RUN mkdir -p /app/install; \ + if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ + echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping FlashAttention build"; \ + else \ + git clone ${FA_REPO} \ + && cd flash-attention \ && git checkout ${FA_BRANCH} \ && git submodule update --init \ && if [ "$USE_SCCACHE" = "1" ]; then \ @@ -217,8 +200,9 @@ RUN cd flash-attention \ && sccache --show-stats; \ fi \ && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && cp dist/*.whl /app/install; \ + fi ### @@ -228,18 +212,16 @@ FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} RUN cd aiter \ && git submodule update --init --recursive \ - && uv pip install -r requirements.txt -RUN uv pip install pyyaml && cd aiter \ + && pip install -r requirements.txt +RUN pip install pyyaml && cd aiter \ && if [ "$USE_SCCACHE" = "1" ]; then \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ && sccache --show-stats; \ fi \ - && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install @@ -254,27 +236,49 @@ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install FROM base AS debs_wheel_release RUN mkdir /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs # Full debs stage - includes Mori (used by Docker releases) FROM base AS debs RUN mkdir /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi FROM base AS final RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ - uv pip install /install/*.whl + pip install /install/*.whl ARG BASE_IMAGE +ARG ROCM_WHEEL_INDEX +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH ARG AITER_REPO +ARG MORI_BRANCH +ARG MORI_REPO RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ + && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ + && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt From b1f4294738dc1179b33a6f28be423edbe68f0405 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Fri, 5 Jun 2026 18:26:10 -0400 Subject: [PATCH 17/68] gfx1250 enabled rocm dockerfile (#988) Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm | 47 ++++++++++++++---------------------------- 1 file changed, 15 insertions(+), 32 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index f8f2d4cd1251..99fe42907c1d 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -45,7 +45,7 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG ARG_PYTORCH_ROCM_ARCH +ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install build dependencies and utilities @@ -53,14 +53,7 @@ ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ - libnuma-dev ccache mold -RUN --mount=type=cache,target=/root/.cache/pip \ - python3 -m pip install --upgrade pip -# Note: mold is installed but not set as the system default linker because -# some packages use JIT compilation at runtime with flags mold does not support. -# Build stages opt in via LDFLAGS="-fuse-ld=mold". -apt-transport-https ca-certificates wget curl \ - libnuma-dev ccache mold + build-essential libnuma-dev ccache mold RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --upgrade pip # Note: mold is installed but not set as the system default linker because @@ -86,25 +79,10 @@ ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy -# ccache directory - persisted across layer rebuilds via cache mounts. -ENV CCACHE_DIR=/root/.cache/ccache -ENV CCACHE_COMPILERCHECK=content -# Empty by default so build steps fall back to $(nproc); CI can override. -ARG max_jobs -ENV MAX_JOBS=${max_jobs} - -# Install UV — download first, then run, so a curl failure is not masked by the pipe -RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ - && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ - && rm -f /tmp/uv-install.sh \ - && uv --version - -# 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 -ENV UV_HTTP_TIMEOUT=500 -ENV UV_INDEX_STRATEGY="unsafe-best-match" -# Use copy mode to avoid hardlink failures with Docker cache mounts -ENV UV_LINK_MODE=copy +# python binary fall back for non venv builds +ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 +# Expose paths from wheel installation +ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} # ccache directory - persisted across layer rebuilds via cache mounts. ENV CCACHE_DIR=/root/.cache/ccache ENV CCACHE_COMPILERCHECK=content @@ -154,8 +132,8 @@ WORKDIR ${COMMON_WORKDIR} FROM base AS fetch_vllm_0 ONBUILD COPY ./ vllm/ FROM base AS fetch_vllm_1 -ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" -ARG VLLM_BRANCH="main" +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ @@ -287,7 +265,7 @@ ARG RIXL_BRANCH="39be1de8" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" ARG UCX_BRANCH="bfb51733" ARG UCX_REPO="https://github.com/openucx/ucx.git" -ENV ROCM_PATH=/opt/rocm +# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl ENV RIXL_BENCH_HOME=/usr/local/rixl_bench @@ -351,6 +329,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ CC="ccache gcc" CXX="ccache g++" \ CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ + --force-fallback-for=abseil-cpp \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ @@ -364,6 +343,10 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ RUN cd /opt/rixl && \ sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ contrib/build-wheel.sh && \ + # The wheel build re-runs meson via meson-python; force the bundled abseil + sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ + pyproject.toml && \ + grep -q 'force-fallback-for' pyproject.toml && \ mkdir -p /app/install && \ _ucx_install_dir=${UCX_HOME} \ ./contrib/build-wheel.sh \ @@ -380,7 +363,7 @@ ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" # DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so # it can be linked against DeepEP without arch mismatches. ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" -ENV ROCM_PATH=/opt/rocm +# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image ENV ROCSHMEM_DIR=/opt/rocshmem RUN --mount=type=cache,target=/root/.cache/ccache \ From aa5b54edb67e102f1d505fe65949e4c46dcfa249 Mon Sep 17 00:00:00 2001 From: Daniel Date: Mon, 8 Jun 2026 12:09:29 -0400 Subject: [PATCH 18/68] adding ffm dockerfile layer Signed-off-by: Daniel --- docker/Dockerfile.rocm_1250_ffm | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 docker/Dockerfile.rocm_1250_ffm diff --git a/docker/Dockerfile.rocm_1250_ffm b/docker/Dockerfile.rocm_1250_ffm new file mode 100644 index 000000000000..8ebf0db80015 --- /dev/null +++ b/docker/Dockerfile.rocm_1250_ffm @@ -0,0 +1,11 @@ +# FFM overlay for gfx1250 builds +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +# extract tarball into /root/x/ffm/. +ADD ubuntu_24_04_rel.tar.gz /root/x/ffm/ +ENV HSA_MODEL_LIB=/root/x/ffm/libhsakmtmodel.so +ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" +ENV HSA_MODEL_ARGS=ffm_enable_time_slicing +ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 +ENV HSA_MODEL_NUM_THREADS=64 From 95f289d6e5c69bff51d70f78aee8ef3701664037 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Tue, 9 Jun 2026 17:45:34 -0500 Subject: [PATCH 19/68] Fix small merge error Signed-off-by: jpvillam --- .../layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 49791f88e24a..6f31dad9b75a 100644 --- 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 @@ -67,7 +67,7 @@ def aiter_triton_kernel_w4a8_moe_forward( ) except ImportError: from aiter.ops.triton.moe_routing.routing import ( - routing_torch as aiter_routing, + routing as aiter_routing, ) else: try: From 271bcd86c89e676c9289d040eee07cb939ff1c65 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Wed, 10 Jun 2026 11:22:07 -0400 Subject: [PATCH 20/68] Remove AMDSMI Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm_1250_ffm | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docker/Dockerfile.rocm_1250_ffm b/docker/Dockerfile.rocm_1250_ffm index 8ebf0db80015..35a66fc41e66 100644 --- a/docker/Dockerfile.rocm_1250_ffm +++ b/docker/Dockerfile.rocm_1250_ffm @@ -9,3 +9,6 @@ ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" ENV HSA_MODEL_ARGS=ffm_enable_time_slicing ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 ENV HSA_MODEL_NUM_THREADS=64 + +### Remove AMD SMI +RUN pip uninstall -y amd_smi && rm -rf ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/amd_smi From 737038da46c77b605d7d4e166c854d151f64798d Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Wed, 10 Jun 2026 12:36:49 -0400 Subject: [PATCH 21/68] Fixed uninstall amdsmi in ffm Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm_1250_ffm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.rocm_1250_ffm b/docker/Dockerfile.rocm_1250_ffm index 35a66fc41e66..6d4b6fed6bd7 100644 --- a/docker/Dockerfile.rocm_1250_ffm +++ b/docker/Dockerfile.rocm_1250_ffm @@ -11,4 +11,4 @@ ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 ENV HSA_MODEL_NUM_THREADS=64 ### Remove AMD SMI -RUN pip uninstall -y amd_smi && rm -rf ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/amd_smi +RUN pip uninstall -y amdsmi && rm -rf ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/amdsmi From 6ae337d5d0d4f144d5d3f2a24cb1fc6aaa0ea661 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Wed, 10 Jun 2026 19:02:17 +0000 Subject: [PATCH 22/68] Build process with no whls Signed-off-by: jpvillam --- docker/Dockerfile.rocm | 2 +- docker/Dockerfile.rocm_1250_ffm | 2 +- docker/Dockerfile.rocm_base | 32 ++++++++++++++++---------------- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 99fe42907c1d..789c4446a29d 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,7 +1,7 @@ # default base image ARG REMOTE_VLLM="0" ARG COMMON_WORKDIR=/app -ARG BASE_IMAGE=rocm/vllm-dev:base +ARG BASE_IMAGE=rocm/vllm-private:juan_455_npi_base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types diff --git a/docker/Dockerfile.rocm_1250_ffm b/docker/Dockerfile.rocm_1250_ffm index 6d4b6fed6bd7..686d3f7b0864 100644 --- a/docker/Dockerfile.rocm_1250_ffm +++ b/docker/Dockerfile.rocm_1250_ffm @@ -1,5 +1,5 @@ # FFM overlay for gfx1250 builds -ARG BASE_IMAGE +ARG BASE_IMAGE=rocm/vllm-private:juan_455_npi_test FROM ${BASE_IMAGE} # extract tarball into /root/x/ffm/. diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 541d31ff3743..8a4f276257a4 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=ubuntu:24.04 +ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.11.0-rocm7.14.0a20260605-7.14.0a20260605-nightly-ubuntu24.04-gfx1250 ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250 ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260605 ARG TORCHVISION_VERSION=0.28.0a0+rocm7.14.0a20260605 @@ -86,21 +86,21 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ ### ### Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index ### -ARG ROCM_WHEEL_INDEX -ARG TORCH_VERSION -ARG TORCHVISION_VERSION -ARG TORCHAUDIO_VERSION -ARG ROCM_SDK_VERSION -# torch/torchvision/torchaudio must be pinned to mutually-consistent builds -# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm -# sdk version is derived from torch's own dependency pin unless overridden, -# which keeps the set consistent and avoids pip backtracking. -RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ - "torch==${TORCH_VERSION}" \ - "torchvision==${TORCHVISION_VERSION}" \ - "torchaudio==${TORCHAUDIO_VERSION}" \ - "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ - rocm-sdk init +#ARG ROCM_WHEEL_INDEX +#ARG TORCH_VERSION +#ARG TORCHVISION_VERSION +#ARG TORCHAUDIO_VERSION +#ARG ROCM_SDK_VERSION +## torch/torchvision/torchaudio must be pinned to mutually-consistent builds +## (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +## sdk version is derived from torch's own dependency pin unless overridden, +## which keeps the set consistent and avoids pip backtracking. +#RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ +# "torch==${TORCH_VERSION}" \ +# "torchvision==${TORCHVISION_VERSION}" \ +# "torchaudio==${TORCHAUDIO_VERSION}" \ +# "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ +# rocm-sdk init # Torch runtime deps that may not be published on the ROCm wheel index; # install them from PyPI afterwards. From 671dac0ffb838fc2ece8a5734964c4b554e976fd Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 12 Jun 2026 14:33:54 -0400 Subject: [PATCH 23/68] switching build branches Signed-off-by: Daniel --- docker/Dockerfile.rocm_base | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 8a4f276257a4..aa5d7ec00a34 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.11.0-rocm7.14.0a20260605-7.14.0a20260605-nightly-ubuntu24.04-gfx1250 +ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.10.0-rocm7.14.0a20260611.a0-7.14.0a20260611-a0-nightly-ubuntu24.04-gfx1250 ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250 ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260605 ARG TORCHVISION_VERSION=0.28.0a0+rocm7.14.0a20260605 @@ -6,7 +6,7 @@ ARG TORCHAUDIO_VERSION=2.11.0a0+rocm7.14.0a20260605 ARG ROCM_SDK_VERSION=7.14.0a20260605 ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="jpvillam/gfx1250_0604" +ARG AITER_BRANCH="main" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" From 681df96d9e2941b0c29538158c239866fe5af35f Mon Sep 17 00:00:00 2001 From: jpvillam Date: Fri, 12 Jun 2026 08:52:44 -0500 Subject: [PATCH 24/68] Aiter routing patch for gptoss fp4 on 1250 Signed-off-by: jpvillam --- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 36 +++++-------------- 1 file changed, 9 insertions(+), 27 deletions(-) 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 index 6f31dad9b75a..adae262888d3 100644 --- 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 @@ -48,34 +48,16 @@ def aiter_triton_kernel_w4a8_moe_forward( ) from vllm.platforms.rocm import on_gfx1250 - # aiter exposes its MoE routing under two module paths across versions; - # prefer the nested `moe.moe_routing` location, fall back to the legacy one. - # - # On gfx1250 use aiter's pure-torch ``routing_torch`` instead of the triton - # ``routing``: the triton routing kernel compiles a TMA (TDM) descriptor - # whose last dim is ``topk * 2`` bytes, which is < 16 bytes (the descriptor - # minimum) for a power-of-2 topk such as gpt-oss' topk=4 and fails to - # compile for small batches (warmup dummy run + every decode step). - # ``routing_torch`` avoids the TDM kernel and is numerically identical - # (validated on the FFM sim: gather/scatter/gate_scal match the triton path - # exactly where the latter compiles). DeepSeek-V4's topk=6 dodged this since - # ``next_power_of_2(6) == 8 != 6`` disables the descriptor branch. + try: + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + # TODO: (JPVILLAM) This causes a tl compile error on 1250. + # Need to figure out why this is a problem and sync with triton team if on_gfx1250(): - try: - from aiter.ops.triton.moe.moe_routing.routing import ( - routing_torch as aiter_routing, - ) - except ImportError: - from aiter.ops.triton.moe_routing.routing import ( - routing as aiter_routing, - ) - else: - try: - from aiter.ops.triton.moe.moe_routing.routing import ( - routing as aiter_routing, - ) - except ImportError: - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize From bd4e24a5e2984af8545ca404b576c23bf2e36a68 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Fri, 12 Jun 2026 08:56:34 -0500 Subject: [PATCH 25/68] Initial non-accurate impl for dsr1 Signed-off-by: jpvillam --- vllm/platforms/rocm.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 613a28be27c3..78a9aa8b615f 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -210,9 +210,7 @@ def _get_gcn_arch() -> str: # current_platform is bound. This path is taken on the FFM simulator, # where amdsmi is unavailable (and would report the host's real gfx950 # cards rather than the simulated gfx1250) — torch.cuda below is correct. - logger.debug( - "Failed to get GCN arch via amdsmi, falling back to torch.cuda." - ) + logger.debug("Failed to get GCN arch via amdsmi, falling back to torch.cuda.") # Ultimate fallback: use torch.cuda (will initialize CUDA) return torch.cuda.get_device_properties("cuda").gcnArchName @@ -228,10 +226,14 @@ def _get_gcn_arch() -> str: _ON_GFX1151 = "gfx1151" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950", "gfx1250"]) -_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950", "gfx1250"]) +_ON_GFX9 = any( + arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950", "gfx1250"] +) # TODO(JPVILLAM): Bubblegum patch to unlock gptoss _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH -_ON_GFX950 = "gfx950" in _GCN_ARCH +_ON_GFX950 = any( + arch in _GCN_ARCH for arch in ["gfx950", "gfx1250"] +) # TODO(JPVILLAM): Bubblegum patch to unlock DSR1 _ON_GFX1250 = "gfx1250" in _GCN_ARCH From afe99c2a9962c9fb7978c48546b32f4e0625dbb7 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Fri, 12 Jun 2026 09:33:20 -0500 Subject: [PATCH 26/68] Add gpt bench script Signed-off-by: jpvillam --- benchmarks/curl_gpt.sh | 7 +++++++ benchmarks/gpt_fp4_serve.sh | 7 +++++++ 2 files changed, 14 insertions(+) create mode 100644 benchmarks/curl_gpt.sh create mode 100644 benchmarks/gpt_fp4_serve.sh diff --git a/benchmarks/curl_gpt.sh b/benchmarks/curl_gpt.sh new file mode 100644 index 000000000000..51a467bfbc2f --- /dev/null +++ b/benchmarks/curl_gpt.sh @@ -0,0 +1,7 @@ +MODEL="${MODEL:-/data/models/gpt-oss-120b-w-mxfp4-a-fp8}" \ +curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{ + "model": ${MODEL}, + "prompt": "What is the capital of France, and what is it known for?", + "temperature": 0.0, + "max_tokens": 100 +}' diff --git a/benchmarks/gpt_fp4_serve.sh b/benchmarks/gpt_fp4_serve.sh new file mode 100644 index 000000000000..ef23808e95ab --- /dev/null +++ b/benchmarks/gpt_fp4_serve.sh @@ -0,0 +1,7 @@ +#ROCR_VISIBLE_DEVICE=0 \ +MODEL="${MODEL:-/data/models/gpt-oss-120b-w-mxfp4-a-fp8}" \ +HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 \ +VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 \ +VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 \ +VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 \ +vllm serve --model ${MODEL} --host localhost --port 8000 --tensor-parallel-size 1 --gpu_memory_utilization 0.7 #--compilation-config '{"mode":"None","cudagraph_mode": "FULL", "cudagraph_capture_sizes": [1]}' From 8602a447a49d27b864e29e0c02a150fa05a5fa08 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Fri, 12 Jun 2026 12:21:35 -0500 Subject: [PATCH 27/68] Remove patch for 950 gate Signed-off-by: jpvillam --- vllm/platforms/rocm.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 78a9aa8b615f..32d27ab0ebc4 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -231,9 +231,10 @@ def _get_gcn_arch() -> str: ) # TODO(JPVILLAM): Bubblegum patch to unlock gptoss _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH -_ON_GFX950 = any( - arch in _GCN_ARCH for arch in ["gfx950", "gfx1250"] -) # TODO(JPVILLAM): Bubblegum patch to unlock DSR1 +_ON_GFX950 = "gfx950" in _GCN_ARCH +# any( +# arch in _GCN_ARCH for arch in ["gfx950", "gfx1250"] +# ) # TODO(JPVILLAM): Bubblegum patch to unlock DSR1 _ON_GFX1250 = "gfx1250" in _GCN_ARCH From 6f2135574b8cf354d48ff8a351f441d4a7b227b5 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Thu, 18 Jun 2026 14:15:35 -0500 Subject: [PATCH 28/68] Missing cmake gate for skinny gemms Signed-off-by: jpvillam --- CMakeLists.txt | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index c0e4698e78fb..1037faf0de3a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1400,9 +1400,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP") ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) - - if(VLLM_ROCM_HAS_GFX1100) - target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) + + # Needed to pass skip skinny flag to .cpp targets + if(VLLM_GPU_ARCHES MATCHES "gfx1250") + target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS) endif() endif() From 7bded231d8e70362009ef17efcb099b56b689a0d Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Tue, 23 Jun 2026 16:06:48 -0400 Subject: [PATCH 29/68] Update default args for rocm_base for gfx1250 Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm_base | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index aa5d7ec00a34..f93c71611bb9 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,12 +1,12 @@ -ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.10.0-rocm7.14.0a20260611.a0-7.14.0a20260611-a0-nightly-ubuntu24.04-gfx1250 +ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.10.0-rocm7.14.0a20260605.a0-7.14.0a20260605-a0-nightly-ubuntu24.04-gfx1250 ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250 -ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260605 +ARG TORCH_VERSION=2.10.0+rocm7.14.0a20260605.a0 ARG TORCHVISION_VERSION=0.28.0a0+rocm7.14.0a20260605 ARG TORCHAUDIO_VERSION=2.11.0a0+rocm7.14.0a20260605 ARG ROCM_SDK_VERSION=7.14.0a20260605 ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="main" +ARG AITER_BRANCH="jpvillam/gfx1250_0604" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" From c11b2f59c2c35756894a12084009be5094404eeb Mon Sep 17 00:00:00 2001 From: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:12:24 -0500 Subject: [PATCH 30/68] Initial fixes to get dsr4 working (#1017) Co-authored-by: root Signed-off-by: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> --- CMakeLists.txt | 4 + .../kernels/linear/scaled_mm/aiter.py | 2 +- .../experts/gpt_oss_triton_kernels_moe.py | 30 +- .../experts/gpt_oss_triton_kernels_moe.py.bak | 1548 +++++++++++++++++ vllm/models/deepseek_v4/compressor.py | 3 + vllm/v1/sample/ops/topk_topp_sampler.py | 2 +- 6 files changed, 1577 insertions(+), 12 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak diff --git a/CMakeLists.txt b/CMakeLists.txt index 1037faf0de3a..3fada1a6d2c7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1405,6 +1405,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP") if(VLLM_GPU_ARCHES MATCHES "gfx1250") target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS) endif() + + if(VLLM_ROCM_HAS_GFX1100) + target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) + endif() endif() # Must run after the last HIP `define_extension_target` so every extension diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 1b39491ab346..cb175acf45c9 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -373,7 +373,7 @@ def __init__(self, config: FP8ScaledMMLinearLayerConfig): self.use_triton = ( not current_platform.is_fp8_fnuz() and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) - ) + ) or True @classmethod def is_supported(cls, compute_capability=None): 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 bd3d9aa31bd3..adb24d03dc5d 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 @@ -1225,7 +1225,7 @@ def activation( alpha = ( quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None - else 1.702 + else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default ) limit = ( quant_config.gemm1_clamp_limit @@ -1351,7 +1351,7 @@ def _try_apply_aiter_w4a8( swiglu_alpha = ( quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None - else 1.702 + else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default ) swiglu_limit = ( quant_config.gemm1_clamp_limit @@ -1391,27 +1391,37 @@ def _try_apply_aiter_w4a8( swiglu_limit, ) - intermediate_cache1 = moe_gemm_a8w4( + # GEMM1 WITHOUT fused swiglu. aiter's fused _swiglu splits gate/up by + # INTERLEAVING, but DeepSeek-V4 (SILU) stores w13 half/half, so emit the raw + # gate_up (bf16) and apply the half/half SILU+clamp ourselves -- identical to + # the unfused activation() path (swiglu_limit_func) -- then requantize to fp8. + raw_gate_up = moe_gemm_a8w4( hidden_fp8, w1.storage.data, None, quant_config.w1_precision.weight_scale.storage.data, gemm1_lhs_scale, - gemm2_lhs_scale, + None, quant_config.w1_bias, routing_data, gather_indx=None, gammas=gammas if apply_router_weight_on_input else None, swizzle_mx_scale=None, - out_dtype=torch.float8_e4m3fn, - apply_swiglu=True, - alpha=swiglu_alpha, - limit=swiglu_limit, + out_dtype=torch.bfloat16, + apply_swiglu=False, unpadded_N=unpadded_n_w1, unpadded_K=unpadded_k_w1, ) + gate_up = raw_gate_up[:, :unpadded_n_w1] + _act = torch.empty( + (gate_up.shape[0], unpadded_n_w1 // 2), + dtype=torch.bfloat16, + device=gate_up.device, + ) + swiglu_limit_func(_act, gate_up, swiglu_limit) + intermediate_cache1 = downcast_to_static_fp8(_act, gemm2_lhs_scale) logger.info( - "[aiter-w4a8] gemm1 OK ic1=%s %s; calling gemm2", + "[aiter-w4a8] gemm1(raw)+halfsplit-silu OK ic1=%s %s; calling gemm2", tuple(intermediate_cache1.shape), intermediate_cache1.dtype, ) @@ -1582,7 +1592,7 @@ def apply( swiglu_alpha = ( quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None - else 1.702 + else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default ) swiglu_limit = ( quant_config.gemm1_clamp_limit diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak new file mode 100644 index 000000000000..189a5ca8fd9e --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak @@ -0,0 +1,1548 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import replace + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FUSED_MOE_UNQUANTIZED_CONFIG, + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( + LoRAExpertsMixin, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.fused_moe.utils import _resize_cache +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.import_utils import has_triton_kernels + +from ..utils import swiglu_limit_func + +logger = init_logger(__name__) + + +def _triton_kernel_moe_supports_current_device() -> bool: + # Shared device gate for the OAI Triton MoE expert classes. + # Platform-aware to avoid ROCm capability aliasing — cap (9, 0) + # matches both gfx90a (verified) and gfx906 (unverified), so we + # dispatch on gfx-string helpers instead of the cap tuple on ROCm. + p = current_platform + if p.is_cuda(): + cap = p.get_device_capability() + # Keep the original `(9, 0) <= cap < (11, 0)` window on + # CUDA (covers Hopper SM90 and Blackwell SM100, excludes + # SM120) — this PR is ROCm-scoped and the broader CUDA + # range was not validated. + return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) + if p.is_rocm(): + from vllm.platforms.rocm import on_gfx1x, on_gfx9 + + # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx); + # on_gfx9() already excludes gfx906/gfx908. + # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); + # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). + return on_gfx9() or on_gfx1x() + return False + + +def _patch_make_bitmatrix_metadata() -> None: + """Monkey-patch make_bitmatrix_metadata to support non-power-of-2 top_k. + + triton's tl.arange requires a power-of-2 range. The original kernel + computes BLOCK_SIZE = BLOCK_PER_TOK * TOKS_PER_ROW (= 32 * top_k). For + DeepSeek-V4 with top_k=6 this gives 192, which is not a power of 2 and + causes a compile error at the first forward pass. + + Fix: define a drop-in replacement kernel that accepts an extra constexpr + BLOCK_SIZE_PADDED (next power of 2 >= BLOCK_SIZE) and uses it for the + tl.arange call while keeping the actual BLOCK_SIZE as the stride between + thread-blocks so that all flat indices into NonzeroIndx stay correct. + Elements beyond BLOCK_SIZE are masked out (col_indx = 0xffff) and ignored. + + This function is called once at module load time and patches the function + inside the triton_kernels tensor module so that SparseMatrix.__post_init__ + picks up the fixed version transparently. + """ + import torch + import triton + import triton.language as tl + + try: + if current_platform.is_rocm(): + from triton_kernels.tensor_details import bitmatrix as _bm + from triton_kernels.tensor_details.bitmatrix import ( + BitmatrixMetadata, + _keyed_add, + cdiv, + ) + from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 + sum_bitmatrix_rows, + ) + else: + from vllm.third_party.triton_kernels.tensor_details import ( + bitmatrix as _bm, + ) + from vllm.third_party.triton_kernels.tensor_details.bitmatrix import ( + BitmatrixMetadata, + _keyed_add, + cdiv, + ) + from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 + sum_bitmatrix_rows, + ) + except ImportError: + return + + @triton.jit + def _stage2_pow2( + ColSortedIndx, + RowSortedIndx, + NonzeroIndx, + n_tokens, + ColPartialSum, + stride_pm, + stride_pn, + ColOffs, + TOKS_PER_ROW: tl.constexpr, + BLOCK_PER_TOK: tl.constexpr, + BLOCK_SIZE_PADDED: tl.constexpr, + ): + # Actual number of elements per block (may not be a power of 2). + BLOCK_SIZE: tl.constexpr = BLOCK_PER_TOK * TOKS_PER_ROW + tl.static_assert(BLOCK_SIZE_PADDED <= 32768) + if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr(): + n_tokens = tl.load(n_tokens) + nonzero_indx_size = n_tokens * TOKS_PER_ROW + pid_m = tl.program_id(0) + # Use BLOCK_SIZE_PADDED (a power of 2) for tl.arange, but stride by + # the actual BLOCK_SIZE so flat positions in NonzeroIndx are correct. + # Elements with offs_local >= BLOCK_SIZE have offs_global beyond the + # valid range, get col_indx = 0xffff, and are filtered by the mask + # below without producing any output. + offs_local = tl.arange(0, BLOCK_SIZE_PADDED) + offs_global = pid_m * BLOCK_SIZE + offs_local + mask = offs_global < nonzero_indx_size + col_indx = tl.load(NonzeroIndx + offs_global, mask=mask, other=-1).to(tl.uint32) + kv_pairs = ((col_indx << 16) | offs_local).to(tl.uint32) + kv_pairs = tl.sort(kv_pairs, 0) + col_indx = kv_pairs >> 16 + offs_global = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF) + mask = col_indx != 0xFFFF + x = kv_pairs & 0xFFFF0000 | 0x00000001 + cols_and_inclusive_run_lengths = tl.associative_scan(x, 0, _keyed_add) + exclusive_run_lengths = (cols_and_inclusive_run_lengths - 1) & 0xFFFF + row_sorted_indx = tl.load( + ColPartialSum + pid_m * stride_pm + col_indx * stride_pn, mask=mask + ) + row_sorted_indx += tl.load(ColOffs + col_indx, mask=mask) + row_sorted_indx += exclusive_run_lengths + tl.store(RowSortedIndx + offs_global, row_sorted_indx, mask=mask) + tl.store(ColSortedIndx + row_sorted_indx, offs_global, mask=mask) + + def _make_bitmatrix_metadata_pow2_safe(nonzero_indx, bitmatrix): + assert nonzero_indx.ndim == 2 + PARTIAL_BLOCK_M = 32 + col_sum, col_partial_sum = sum_bitmatrix_rows( + bitmatrix, partials_block_size=PARTIAL_BLOCK_M + ) + device = bitmatrix.device + n_indx = nonzero_indx.numel() + n_cols = bitmatrix.shape[1] + col_offs = torch.empty(n_cols, dtype=torch.int32, device=device) + combined_indx = torch.empty(n_indx * 2, dtype=torch.int32, device=device) + col_sorted_indx = combined_indx[:n_indx] + row_sorted_indx = combined_indx[n_indx:] + MEMSET_BLOCK = 1024 + memset_grid = (cdiv(n_indx * 2, MEMSET_BLOCK) + n_cols + 1,) + _bm._bitmatrix_metadata_compute_stage1[memset_grid]( + combined_indx, + n_indx * 2, + -1, + MEMSET_BLOCK, + col_sum, + col_offs, + col_sum.shape[0], + col_partial_sum, + col_partial_sum.shape[0], + col_partial_sum.stride(0), + col_partial_sum.stride(1), + BLOCK_M=512, + BLOCK_N=512, + ) + toks_per_row = nonzero_indx.shape[-1] + block_size = PARTIAL_BLOCK_M * toks_per_row + # Next power of 2 >= block_size (required by tl.arange). + block_size_padded = 1 << (max(block_size, 1) - 1).bit_length() + compute_grid = (cdiv(bitmatrix.shape_max[0], PARTIAL_BLOCK_M),) + _stage2_pow2[compute_grid]( + col_sorted_indx, + row_sorted_indx, + nonzero_indx, + bitmatrix.shape[0], + col_partial_sum, + col_partial_sum.stride(0), + col_partial_sum.stride(1), + col_offs, + TOKS_PER_ROW=toks_per_row, + BLOCK_PER_TOK=PARTIAL_BLOCK_M, + BLOCK_SIZE_PADDED=block_size_padded, + ) + return BitmatrixMetadata( + col_sum=col_sum, + col_sorted_indx=col_sorted_indx, + row_sorted_indx=row_sorted_indx, + ) + + # The most reliable patch point: SparseMatrix.__post_init__ looks up + # make_bitmatrix_metadata via its own __globals__ dict (the tensor.py + # module dict). Patching through __globals__ works regardless of how + # sys.modules maps "triton_kernels.tensor" vs + # "vllm.third_party.triton_kernels.tensor". + from triton_kernels.tensor import SparseMatrix as _SparseMatrix + + _SparseMatrix.__post_init__.__globals__["make_bitmatrix_metadata"] = ( + _make_bitmatrix_metadata_pow2_safe + ) + # Also patch the bitmatrix module itself in case it is imported directly. + _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe + + +# Two API generations of triton_kernels are supported: +# - v3.5.1 (the version bundled with vLLM): exposes `routing()` and +# `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix` +# constructor takes a `scratchpad` argument. +# - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix` +# based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only +# when the user has triton_kernels installed system-wide at v3.6.0+. +# +# `use_legacy_triton_kernels` selects between them at import time based on +# whether `SparseMatrix` is importable. +use_legacy_triton_kernels = False + +if has_triton_kernels(): + try: + import triton_kernels.swiglu + from triton_kernels.matmul_ogs import ( + FnSpecs, + FusedActivation, + GatherIndx, + RoutingData, + ScatterIndx, + matmul_ogs, + ) + from triton_kernels.tensor import ( + BIT, + Bitmatrix, + ) + + try: + from triton_kernels.tensor import ( + SparseMatrix, + make_ragged_tensor_metadata, + ) + except ImportError: + # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once + # the gpt-oss perf regression in v3.6.0+ is resolved upstream. + # Tracking: https://github.com/triton-lang/triton/issues/9969 + use_legacy_triton_kernels = True + if not use_legacy_triton_kernels: + _patch_make_bitmatrix_metadata() + except (AttributeError, ImportError) as e: + logger.error( + "Failed to import Triton kernels. Please make sure your triton " + "version is compatible. Error: %s", + e, + ) + + +@triton.jit +def pack_bitmatrix( + bitmatrix, + topk_ids, + n_rows, # n_rows in bitmatrix / topk_ids + bm_cols: tl.constexpr, # n int32_t bitpacks in bitmatrix + n_expts_act, # num_topk + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + """ + Packs topk_ids into a bitmatrix. + code reference: + https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264 + """ + pid_m = tl.program_id(0) + offsets_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offsets_k = tl.arange(0, BLOCK_SIZE_K) + offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :] + mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :] + indices = tl.load(topk_ids + offsets, mask=mask, other=-1) + valid = indices >= 0 + div = indices // 32 + rem = indices % 32 + one = tl.cast(1, tl.uint32) + + # Iterate through all the relevant bitmatrix columns. + for i in range(bm_cols): + # When BLOCK_SIZE_K=32, offs is just the column index. + offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) + # All topks that need to go into this column has the correct bit set. + # Other bits are 0. x is a 2D tensor. + # Guard with `valid` to prevent negative indices from producing + # spurious bits (on HIP, -1 // 32 == 0 and 1 << (-1 % 32) sets + # bit 31). + x = tl.where( + valid[:, :, None] & (div[:, :, None] == offs[None, None, :]), + (one << rem)[:, :, None], + 0, + ) + # Reduce x to get a single int32_t bitpack. + y = tl.reduce_or(x, axis=1) + bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] + tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) + + +def triton_kernel_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, +) -> torch.Tensor: + sm_first = not renormalize + + # When no expert map is provided (no EP), call the fused `routing()` + # kernel directly. It combines softmax, topk, bitmatrix packing, and + # routing-metadata construction in a single launch, instead of the + # three separate kernels used by the generic path below. + # Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines + # equivalent logic via SparseMatrix in `make_routing_data`. + if use_legacy_triton_kernels and expert_map is None: + from triton_kernels.routing import routing as fused_routing + + routing_data, gather_idx, scatter_idx = fused_routing( + gating_output, topk, sm_first=sm_first + ) + effective_expert_map = None + effective_global_num_experts = global_num_experts + else: + from triton_kernels.topk import topk as topk_fn + + logits = gating_output + if sm_first: + logits = torch.softmax(logits, dim=-1) + topk_result = topk_fn(logits, topk, apply_softmax=not sm_first) + # topk may return a tuple (vals, indx, bitmatrix) or a + # SparseMatrix depending on the triton_kernels version. + if isinstance(topk_result, tuple): + topk_weights, topk_ids_raw, _ = topk_result + else: + topk_weights = topk_result.vals + topk_ids_raw = topk_result.indx + + if expert_map is not None: + # topk_ids_raw contains global expert IDs - remap to local. + topk_ids = expert_map[topk_ids_raw.to(torch.long)] + local_num_experts = w1.shape[0] + routing_data, gather_idx, scatter_idx = make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + # expert_map already applied; pass None downstream. + effective_expert_map = None + effective_global_num_experts = local_num_experts + else: + topk_ids = topk_ids_raw.to(torch.long) + routing_data, gather_idx, scatter_idx = make_routing_data( + topk_ids, topk_weights, gating_output.shape[-1] + ) + effective_expert_map = expert_map + effective_global_num_experts = global_num_experts + + output = torch.empty_like(hidden_states) + effective_quant_config = ( + quant_config if quant_config is not None else FUSED_MOE_UNQUANTIZED_CONFIG + ) + + return triton_kernel_fused_experts( + output, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk=topk, + activation=activation, + quant_config=effective_quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=effective_global_num_experts, + expert_map=effective_expert_map, + ) + + +# This is a triton implementation of the fused_experts function +def triton_kernel_fused_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 + topk: int, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + 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, + intermediate_cache: torch.Tensor | None = None, + a1q_scale: torch.Tensor | None = None, +) -> torch.Tensor: + """Triton implementation of fused expert computation using OAI kernels.""" + assert activation == MoEActivation.SWIGLUOAI, ( + "Only SWIGLUOAI activation is supported" + ) + 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, only check non-mxfp4 + assert hidden_states.ndim == 2 + assert hidden_states.shape[-1] == w1.shape[-2] + assert w2.shape[-1] == w1.shape[1] + + batch_dim = 1 + M, K = hidden_states.shape[-2:] + E, _, N = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + if intermediate_cache is None: + intermediate_cache = torch.empty( + (batch_dim, M * topk, N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + # Add batch_dim to output buffer because matmul_ogs expects 3D output + intermediate_cache = _resize_cache( + intermediate_cache, (batch_dim, M * topk, N // 2) + ) + output_tensor = _resize_cache(output_tensor, (batch_dim, M, K)) + + act = ( + FusedActivation( + FnSpecs( + "swiglu", + triton_kernels.swiglu.swiglu_fn, + ("alpha", "limit"), + reduction_n=2, + ), + (swiglu_alpha, swiglu_limit), + ) + if not use_legacy_triton_kernels + else FusedActivation( + FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), + (swiglu_alpha, swiglu_limit), + 2, + ) + ) + gammas = routing_data.gate_scal if routing_data else None + + matmul_ogs( + hidden_states, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=act, + y=intermediate_cache, + ) + + matmul_ogs( + intermediate_cache.view(M * topk, N // 2), + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + y=output_tensor, + ) + output_tensor = output_tensor.view(M, K) + return output_tensor + + +def make_routing_data( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_local_experts: int, +) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: + topk_ids = topk_ids.to(torch.int16) + topk_weights = topk_weights.to(torch.bfloat16) + + n_rows, num_topk = topk_ids.size() + + BLOCK_SIZE_M = 512 + BLOCK_SIZE_K = 32 + + bm_cols = triton.cdiv(num_local_experts, BLOCK_SIZE_K) # n_bitpacks + bitmatrix = torch.zeros( + (n_rows, bm_cols), dtype=torch.uint32, device=topk_ids.device + ) + + grid = (triton.cdiv(n_rows, BLOCK_SIZE_M),) + pack_bitmatrix[grid]( + bitmatrix, + topk_ids, + n_rows, + bm_cols, + num_topk, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_K=BLOCK_SIZE_K, + ) + + bitmatrix_shape = [n_rows, bm_cols * 32] + bitmatrix_shape_max = [n_rows, None] + bitmatrix = ( + Bitmatrix( + bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max + ) + if not use_legacy_triton_kernels + else Bitmatrix( + bitmatrix, + shape=bitmatrix_shape, + shape_max=bitmatrix_shape_max, + scratchpad=None, + ) + ) + + # matmul_ogs expects invalid topk_weights to be -1s + topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights) + + if use_legacy_triton_kernels: + from triton_kernels.routing import routing_from_bitmatrix + + return routing_from_bitmatrix( + bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk + ) + + sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix) + dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx + combine_indx = sparse_logits.mask_metadata.col_sorted_indx + ragged_batch_metadata = make_ragged_tensor_metadata( + sparse_logits.mask_metadata.col_sum, + dispatch_indx.shape[0], + ) + gate_scal = sparse_logits.vals.flatten()[combine_indx] + routing_data = RoutingData( + gate_scal, + ragged_batch_metadata.block_sizes, + num_local_experts, + num_topk, + ragged_batch_metadata, + ) + gather_indx = GatherIndx(combine_indx, dispatch_indx) + scatter_indx = ScatterIndx(dispatch_indx, combine_indx) + return routing_data, gather_indx, scatter_indx + + +def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_states: torch.Tensor, + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int = 1, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). + + Uses ``aiter.ops.triton.quant_moe.downcast_to_static_fp8`` with a scalar scale: + for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when + single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or + ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static + FP8 path). + + Returns: + (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for + FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + if not quant_config.use_mxfp4_w4a16: + return hidden_states, None + try: + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return hidden_states, None + if not rocm_aiter_ops.is_enabled(): + return hidden_states, None + + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + + lhs_scale = None + if qp is not None and qp.flex_ctx.lhs_data.scale is not None: + s = qp.flex_ctx.lhs_data.scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = hidden_states.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + + return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale + + +def _mxfp4_w4a8_unpadded_dims( + moe_cfg: FusedMoEConfig, +) -> tuple[int | None, int | None, int | None, int | None]: + """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" + if ( + moe_cfg.intermediate_size_per_partition_unpadded is None + or moe_cfg.hidden_dim_unpadded is None + ): + return None, None, None, None + unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 + unpadded_k_w1 = moe_cfg.hidden_dim_unpadded + unpadded_n_w2 = moe_cfg.hidden_dim_unpadded + unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded + return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 + + +def _mxfp4_w4a8_resolve_lhs_scale( + quant_config: FusedMoEQuantConfig, + *, + gemm_num: int, + activations: torch.Tensor, +) -> torch.Tensor: + """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. + + Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise + ``max(|activations|) / 448`` (e4m3), matching + ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. + """ + if gemm_num not in (1, 2): + raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") + qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision + a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale + lhs_scale = None + if qp is not None and qp.flex_ctx is not None: + ld = qp.flex_ctx.lhs_data + if ld is not None and ld.scale is not None and ld.scale.numel() == 1: + lhs_scale = ld.scale + if lhs_scale is None and a_scale is not None: + s = a_scale + if s.numel() == 1: + lhs_scale = s + if lhs_scale is None: + amax = activations.abs().max().clamp(min=1e-12) + lhs_scale = (amax / 448.0).to(dtype=torch.float32) + return lhs_scale + + +def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + precision_config, + lhs_scale: torch.Tensor | None, +): + """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" + if precision_config is None or lhs_scale is None: + return precision_config + from triton_kernels.numerics import InFlexData + + new_flex = replace( + precision_config.flex_ctx, + lhs_data=InFlexData(scale=lhs_scale), + ) + return replace(precision_config, flex_ctx=new_flex) + + +class BaseOAITritonExperts(mk.FusedMoEExpertsModular): + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def _supports_current_device() -> bool: + return _triton_kernel_moe_supports_current_device() and has_triton_kernels() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + raise NotImplementedError + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + """ + Extract the MoE problem size from the given tensor arguments: + - a: The hidden states, input to the MoE layer. + - w1: The first set of expert weights. + - w2: The second set of expert weights. + - topk_ids: The topk ids. + Note: extracting the problem shape from the weight and activation + tensors is not obvious. It needs to be done this way specifically + due to subtle issues with particular kernels, e.g. the int4 kernels + divide the trailing dimension by two, so it's not "correct" to + extract N or K from the trailing dimension of w1 or w2. Similarly, + some kernels transpose the weights, so this needs to be kept in mind. + Note: This implementation covers most cases. However, if experts + require a specialized implementation, like MarlinExperts, they are free + to override this function. + """ + assert len(w1.shape) == 3 and len(w2.shape) == 3 + E, _, N = w1.shape + K = a1.size(-1) + + assert a1.dim() == 2 + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" + M = a1.size(0) + + assert topk_ids.dim() == 2 + topk = topk_ids.size(1) + + return E, M, N, K, topk + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # Weight application and reduction happens in the fused_experts kernel. + return TopKWeightAndReduceNoOP() + + def _make_routing_data( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_local_experts: int, + ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: + return make_routing_data(topk_ids, topk_weights, num_local_experts) + + +class OAITritonExperts(BaseOAITritonExperts): + """OAI Triton-based fused MoE expert implementation.""" + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SWIGLUOAI + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # workspace are allocated inside the kernel + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (0, 0) + workspace2 = (M * topk, activation_out_dim) + output = (M, K) + return (workspace1, workspace2, output) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + if self.quant_config is None: + self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG + + if expert_map is not None: + topk_ids = expert_map[topk_ids] + + local_num_experts = w1.shape[0] + if global_num_experts == -1: + global_num_experts = local_num_experts + + routing_data, gather_indx, scatter_indx = self._make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + + topk = topk_ids.size(1) + triton_kernel_fused_experts( + output, + hidden_states, + w1, + w2, + routing_data, + gather_indx, + scatter_indx, + topk=topk, + activation=activation, + quant_config=self.quant_config, + apply_router_weight_on_input=False, + global_num_experts=local_num_experts, + expert_map=None, # applied already + intermediate_cache=workspace2, + a1q_scale=a1q_scale, + ) + + +class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): + """ + A Triton based MoE expert class that operates on expert standard + format and explicitly keeps the activation and reduction (moe_sum) steps + unfused from the matmul_ogs kernel. This exposes injection points + for activation and moe_sum. + + One use case for it is to inject LoRA modules on the activation and moe_sum. + """ + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ] + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # workspace are allocated inside the kernel + activation_out_dim = self.adjust_N_for_activation(N, activation) + workspace1 = (M * topk, activation_out_dim) + workspace2 = (M * topk, max(N, K)) + output = (M, K) + return (workspace1, workspace2, output) + + def moe_sum(self, input: torch.Tensor, output: torch.Tensor): + ops.moe_sum(input, output) + + def activation( + self, + activation: MoEActivation, + output: torch.Tensor, + input: torch.Tensor, + **kwargs, + ) -> None: + quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG + if activation == MoEActivation.SWIGLUOAI: + alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + torch.ops._C.swigluoai_and_mul(output, input, alpha, limit) + elif ( + activation == MoEActivation.SILU + and quant_config.gemm1_clamp_limit is not None + ): + swiglu_limit_func( + output, + input, + quant_config.gemm1_clamp_limit, + ) + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + assert quant_config.gemm1_clamp_limit is not None + alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.0 + ) + beta = ( + quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0 + ) + torch.ops._C.silu_and_mul_with_clamp( + output, input, quant_config.gemm1_clamp_limit, alpha, beta + ) + else: + super().activation(activation, output, input) + + def _try_apply_aiter_w4a8( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + apply_router_weight_on_input: bool, + ) -> torch.Tensor | None: + """ROCm/gfx1250 MXFP4 MoE via aiter's W4A8 ``moe_gemm_a8w4``. + + The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls + ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and + fails to compile on gfx1250. aiter's gluon ``moe_gemm_a8w4`` (FP8 act x MXFP4 + weight) does run on gfx1250, so we route the MoE through it. + + Routing is built aiter-native directly from ``topk_ids``/``topk_weights`` so we + never touch the diverged vendored ``triton_kernels`` routing structs (whose + ``RoutingData``/``GatherIndx``/``ExptData`` field names and layout differ from + aiter's). This mirrors ``triton_kernel_fused_mxfp4_w4a8_experts`` + (aiter_mxfp4_w4a8_moe.py) but resolves the FP8 LHS scales dynamically because + DeepSeek-V4 is ``activation_scheme: dynamic`` (no calibrated static scale). + + aiter's ``routing`` always softmaxes the gate weights, so we feed + ``log(topk_weights)``: softmax-over-selected then recovers + ``topk_weights / sum(topk_weights)``. DeepSeek-V4 weights are + norm_topk_prob-normalized then scaled by ``routed_scaling_factor`` (so they sum + to that factor, not 1); we restore the exact weighting with a per-token output + rescale by the weight sum (== 1 for plain renormalize routing, a no-op there). + + Returns ``output`` on success, or ``None`` if aiter is unavailable (caller then + falls back to ``matmul_ogs``). + """ + try: + from aiter.ops.triton.moe.moe_routing.routing import ( + routing as aiter_routing, + ) + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + except ImportError: + return None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + M = hidden_states.shape[0] + num_experts = w1.shape[0] + topk = topk_ids.size(1) + + # Reconstruct dense gating logits from the sparse topk selection. log(weight) + # padded with a large-negative sentinel makes aiter's softmax reproduce both + # this exact expert selection and weight = topk_weights / per-token-sum. + tw = topk_weights.to(torch.float32) + logits = torch.full( + (M, num_experts), -1e30, device=hidden_states.device, dtype=torch.float32 + ) + logits.scatter_( + 1, + topk_ids.long().clamp(min=0, max=num_experts - 1), + torch.log(tw.clamp(min=1e-20)), + ) + logger.info( + "[aiter-w4a8] ENTER M=%d K=%d num_experts=%d topk=%d " + "w1=%s w2=%s w1_wscale=%s w2_wscale=%s", + M, + hidden_states.shape[1], + num_experts, + topk, + tuple(w1.storage.data.shape), + tuple(w2.storage.data.shape), + tuple(quant_config.w1_precision.weight_scale.storage.data.shape), + tuple(quant_config.w2_precision.weight_scale.storage.data.shape), + ) + routing_data, gather_idx, scatter_idx = aiter_routing( + logits, topk, sm_first=False + ) + gammas = routing_data.gate_scal + logger.info( + "[aiter-w4a8] routing OK block_m=%s gate_scal=%s gather=%s scatter=%s", + getattr(routing_data, "block_m", None), + tuple(gammas.shape), + tuple(gather_idx.shape), + tuple(scatter_idx.shape), + ) + + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + # FP8 LHS scales (static if calibrated, else dynamic amax/448). GEMM1 emits FP8 + # directly (apply_swiglu=True, out_dtype fp8) quantized with GEMM2's LHS scale, + # so resolve GEMM2's scale up front (hidden states as the amax proxy) and reuse + # it as GEMM2's input scale so quantize/dequantize are consistent. + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=1, activations=hidden_states + ) + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=2, activations=hidden_states + ) + # aiter's in-kernel gather is numerically broken on gfx1250 (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, sorted + # row i reads source token gather_idx[i] // n_expts_act, so this reproduces the + # in-kernel gather exactly (validated: manual gather -> maxrel ~5e-3). + gather_src = gather_idx.to(torch.long) // topk + hidden_sorted = hidden_states[gather_src] + hidden_fp8 = downcast_to_static_fp8(hidden_sorted, gemm1_lhs_scale) + logger.info( + "[aiter-w4a8] scales g1_lhs=%s g2_lhs=%s hidden_fp8=%s; calling gemm1 " + "(unpadded N/K w1=%s/%s w2=%s/%s, alpha=%.4f limit=%.2f)", + float(gemm1_lhs_scale), + float(gemm2_lhs_scale), + tuple(hidden_fp8.shape), + unpadded_n_w1, + unpadded_k_w1, + unpadded_n_w2, + unpadded_k_w2, + swiglu_alpha, + swiglu_limit, + ) + + intermediate_cache1 = moe_gemm_a8w4( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale, + quant_config.w1_bias, + routing_data, + gather_indx=None, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + out_dtype=torch.float8_e4m3fn, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + logger.info( + "[aiter-w4a8] gemm1 OK ic1=%s %s; calling gemm2", + tuple(intermediate_cache1.shape), + intermediate_cache1.dtype, + ) + intermediate_cache3 = moe_gemm_a8w4( + intermediate_cache1, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_scale, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + + # Restore the per-token weight sum (e.g. routed_scaling_factor) that the + # log+softmax normalization divided out. Output weighting (gammas on GEMM2) is + # linear in the gate, so a single post-scale is exact. + logger.info( + "[aiter-w4a8] gemm2 OK ic3=%s %s; rescaling + writing output=%s", + tuple(intermediate_cache3.shape), + intermediate_cache3.dtype, + tuple(output.shape), + ) + out = intermediate_cache3.to(output.dtype) * tw.sum(dim=1, keepdim=True) + output.copy_(out.reshape(output.shape)) + logger.info("[aiter-w4a8] DONE") + return output + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + # Use local variable to help mypy narrow the type after None check + quant_config = self.quant_config + if quant_config is None: + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG + + global_topk_ids = topk_ids + if expert_map is not None: + topk_ids = expert_map[topk_ids] + + local_num_experts = w1.shape[0] + if global_num_experts == -1: + global_num_experts = local_num_experts + + routing_data, gather_indx, scatter_indx = self._make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + + topk = topk_ids.size(1) + + # 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, only check non-mxfp4 + assert hidden_states.ndim == 2 + assert hidden_states.shape[-1] == w1.shape[-2] + assert w2.shape[-1] == w1.shape[1] + + batch_dim = 1 + M, K = hidden_states.shape + E, _, N = w1.shape + + if global_num_experts == -1: + global_num_experts = E + + # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from + # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs + # unswizzle. Output-side router weighting only (the post-scale below assumes + # it); LoRA still uses the matmul_ogs path further down. + if ( + (quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16) + and not apply_router_weight_on_input + and self._lora_context is None + and rocm_aiter_ops.is_enabled() + ): + if ( + self._try_apply_aiter_w4a8( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + apply_router_weight_on_input, + ) + is not None + ): + return + + # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids + # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for + # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). + # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not + # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA + # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below + # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA + # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. + + # Note that the output tensor might be in workspace13 + intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) + intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) + activation_out_dim = self.adjust_N_for_activation(N, activation) + intermediate_cache2 = _resize_cache(workspace13, (M * topk, activation_out_dim)) + + gammas = routing_data.gate_scal if routing_data else None + + hidden_bf16_for_lora = hidden_states + + _moe_gemm_a8w4_fn = None + _downcast_static_fp8_fn = None + # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's + # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also + # serve the w4a16 case by dynamically quantizing the bf16 activations to + # FP8 here and passing the scale. This routes the MoE through aiter's + # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. + _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 + if _mxfp4_weights and rocm_aiter_ops.is_enabled(): + try: + from aiter.ops.triton.moe_op_gemm_a8w4 import ( + moe_gemm_a8w4 as _moe_gemm_a8w4_fn, + ) + from aiter.ops.triton.quant_moe import ( + downcast_to_static_fp8 as _downcast_static_fp8_fn, + ) + except ImportError: + pass + use_aiter_unfused_a8w4 = ( + _mxfp4_weights + and _moe_gemm_a8w4_fn is not None + and _downcast_static_fp8_fn is not None + ) + + if use_aiter_unfused_a8w4: + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(self.moe_config) + ) + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.702 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=1, + activations=hidden_bf16_for_lora, + ) + # GEMM1 passes W2's static scale through to the kernel; resolve it before + # GEMM2 activations exist (fallback uses token hidden states for amax). + gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=hidden_bf16_for_lora, + ) + hidden_fp8 = _downcast_static_fp8_fn( + hidden_bf16_for_lora, + gemm1_lhs_scale, + ) + gemm1_out = _moe_gemm_a8w4_fn( + hidden_fp8, + w1.storage.data, + None, + quant_config.w1_precision.weight_scale.storage.data, + gemm1_lhs_scale, + gemm2_lhs_scale_proxy, + 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.bfloat16, + apply_swiglu=False, + alpha=swiglu_alpha, + limit=swiglu_limit, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) + else: + hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + hidden_bf16_for_lora, + quant_config, + ) + ) + w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w1_precision, + gemm1_lhs_scale_fp8, + ) + + matmul_ogs( + hidden_for_gemm1, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w1_precision_for_gemm, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=None, + y=intermediate_cache1, + ) + + # w13 LoRA: gather the activation input from expert-sorted + # intermediate_cache1, then add the LoRA delta in-place on that copy + # before passing it to activation — exactly mirroring the old + # decorator approach which modified the gathered tensor in-place. + act_input = intermediate_cache1.view(-1, N)[gather_indx.dst_indx] + + sorted_token_ids_lora = None + expert_ids_lora = None + num_tokens_post_padded_lora = None + token_lora_mapping = None + lora_context = self._lora_context + if lora_context is not None: + ( + sorted_token_ids_lora, + expert_ids_lora, + num_tokens_post_padded_lora, + token_lora_mapping, + ) = self.apply_w13_lora( + lora_context, + y=act_input, + x=hidden_bf16_for_lora, + topk_ids=global_topk_ids, + topk_weights=topk_weights, + expert_map=expert_map, + w1=w1, + w2=w2, + num_tokens=M, + top_k_num=topk, + ) + + self.activation( + activation, + intermediate_cache2, + act_input, + ) + + # matmul_ogs grouped reduction fuses sum across multiple experts: + # y[dst_indx // n_expts_act, :] += x + # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. + routing_data.n_expts_act = 1 + + gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] + if use_aiter_unfused_a8w4: + gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, + gemm_num=2, + activations=gemm2_input_bf16, + ) + gemm2_fp8 = _downcast_static_fp8_fn( + gemm2_input_bf16, + gemm2_lhs_scale, + ) + gemm2_out = _moe_gemm_a8w4_fn( + gemm2_fp8, + w2.storage.data, + None, + quant_config.w2_precision.weight_scale.storage.data, + gemm2_lhs_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, + ) + intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) + else: + gemm2_input, gemm2_lhs_scale_fp8 = ( + _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( + gemm2_input_bf16, + quant_config, + gemm_num=2, + ) + ) + w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( + quant_config.w2_precision, + gemm2_lhs_scale_fp8, + ) + + matmul_ogs( + gemm2_input, + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=w2_precision_for_gemm, + gammas=None if apply_router_weight_on_input else gammas, + y=intermediate_cache3, + ) + + # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is + # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. + if lora_context is not None: + self.apply_w2_lora( + lora_context, + y=intermediate_cache3.view(-1, topk, K), + x=intermediate_cache2, + topk_weights=topk_weights, + sorted_token_ids_lora=sorted_token_ids_lora, + expert_ids_lora=expert_ids_lora, + num_tokens_post_padded_lora=num_tokens_post_padded_lora, + token_lora_mapping=token_lora_mapping, + num_tokens=M, + w1=w1, + w2=w2, + top_k_num=topk, + ) + + self.moe_sum(intermediate_cache3.view(-1, topk, K), output) + + +class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + """Monolithic Triton MXFP4 expert. Wraps triton_kernel_moe_forward().""" + + 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: + return _triton_kernel_moe_supports_current_device() and has_triton_kernels() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + 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 + + @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: + return triton_kernel_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, + ) diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 1efa987fe7bd..f129f19a5716 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -13,6 +13,9 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( + _fused_kv_compress_norm_rope_insert_indexer_attn, + _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + _fused_kv_compress_norm_rope_insert_sparse_attn, compress_norm_rope_store_triton, ) from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 69b35830add2..004d4a426780 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -116,7 +116,7 @@ def __init__( logger.info_once( "Using aiter sampler on ROCm (lazy import, sampling-only)." ) - self.forward = self.forward_hip + self.forward = self.forward_native else: self.forward = self.forward_native From 4b0135b4b197c64605202a03f13ed080a28d2688 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Tue, 23 Jun 2026 16:39:53 -0500 Subject: [PATCH 31/68] Small import error Signed-off-by: jpvillam --- .../layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 index adae262888d3..cc616f13cbeb 100644 --- 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 @@ -139,8 +139,11 @@ def triton_kernel_fused_mxfp4_w4a8_experts( from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( should_use_cdna4_mx_scale_swizzle, ) + from vllm.platforms.rocm import on_gfx1250 _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + # TODO (JPVILLAM): merge conflict resolve later if _swizzle_mx_scale is enough + mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" assert quant_config.w1_precision is not None, ( "w1_precision in quant config can't be None" @@ -158,9 +161,6 @@ def triton_kernel_fused_mxfp4_w4a8_experts( # CDNA4-swizzled scale as garbage (validated on the FFM sim: CDNA4_SCALE -> # maxrel ~7e4, plain/None -> ~6e-3); pass swizzle_mx_scale=None there. # gfx950 uses the CDNA4 swizzle layout. - from vllm.platforms.rocm import on_gfx1250 - - mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" intermediate_cache1 = moe_gemm_a8w4( hidden_states, From dd10b54d802ad1fb5bbe8fd2e3e436eda501853a Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 29 Jun 2026 14:21:16 -0500 Subject: [PATCH 32/68] Swap back to whl and apply patches Signed-off-by: jpvillam --- cmake/external_projects/triton_kernels.cmake | 7 ++- docker/Dockerfile.rocm_base | 64 ++++++++++++-------- vllm/utils/import_utils.py | 4 +- 3 files changed, 46 insertions(+), 29 deletions(-) diff --git a/cmake/external_projects/triton_kernels.cmake b/cmake/external_projects/triton_kernels.cmake index 2966c78030bd..dba2e36608d3 100644 --- a/cmake/external_projects/triton_kernels.cmake +++ b/cmake/external_projects/triton_kernels.cmake @@ -1,6 +1,6 @@ # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels -set(DEFAULT_TRITON_KERNELS_TAG "v3.5.1") +set(DEFAULT_TRITON_KERNELS_TAG "padroute") # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to # be directly set to the triton_kernels python directory. @@ -11,13 +11,14 @@ if (DEFINED ENV{TRITON_KERNELS_SRC_DIR}) SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR} ) +#TODO (JPVILLAM): Test new release of triton kernels and put back to normal else() - set(TRITON_GIT "https://github.com/triton-lang/triton.git") + set(TRITON_GIT "https://github.com/jpvillam-amd/triton.git") message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}") FetchContent_Declare( triton_kernels # TODO (varun) : Fetch just the triton_kernels directory from Triton - GIT_REPOSITORY https://github.com/triton-lang/triton.git + GIT_REPOSITORY https://github.com/jpvillam-amd/triton.git GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG} GIT_PROGRESS TRUE SOURCE_SUBDIR python/triton_kernels/triton_kernels diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index f93c71611bb9..4ead2b19c075 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,9 +1,13 @@ -ARG BASE_IMAGE=registry-sc-harbor.amd.com/framework/therock-npi:pytorch-2.10.0-rocm7.14.0a20260605.a0-7.14.0a20260605-a0-nightly-ubuntu24.04-gfx1250 -ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250 -ARG TORCH_VERSION=2.10.0+rocm7.14.0a20260605.a0 -ARG TORCHVISION_VERSION=0.28.0a0+rocm7.14.0a20260605 -ARG TORCHAUDIO_VERSION=2.11.0a0+rocm7.14.0a20260605 -ARG ROCM_SDK_VERSION=7.14.0a20260605 +ARG BASE_IMAGE=ubuntu:24.04 +ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ +ARG ROCM_SDK_VERSION=7.14.0a20260623 +ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 +ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 +ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 + + ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" ARG AITER_BRANCH="jpvillam/gfx1250_0604" @@ -83,25 +87,35 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ && sccache --version; \ fi -### -### Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index -### -#ARG ROCM_WHEEL_INDEX -#ARG TORCH_VERSION -#ARG TORCHVISION_VERSION -#ARG TORCHAUDIO_VERSION -#ARG ROCM_SDK_VERSION -## torch/torchvision/torchaudio must be pinned to mutually-consistent builds -## (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm -## sdk version is derived from torch's own dependency pin unless overridden, -## which keeps the set consistent and avoids pip backtracking. -#RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ -# "torch==${TORCH_VERSION}" \ -# "torchvision==${TORCHVISION_VERSION}" \ -# "torchaudio==${TORCHAUDIO_VERSION}" \ -# "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ -# rocm-sdk init - +## +## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +## +ARG ROCM_WHEEL_INDEX +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG ROCM_SDK_VERSION +ARG APEX_VERSION +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + "torch==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "apex==${APEX_VERSION}" \ + "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + + # Torch runtime deps that may not be published on the ROCm wheel index; # install them from PyPI afterwards. RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 043798a584b8..7cbf63e5040f 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -500,7 +500,9 @@ def has_triton_kernels() -> bool: @cache def has_tilelang() -> bool: """Whether the optional `tilelang` package is available.""" - return _has_module("tilelang") + from vllm.platforms.rocm import on_gfx1250 + + return _has_module("tilelang") and not on_gfx1250() def has_arctic_inference() -> bool: From ac50ce2543773d0adc50796873f83b4073662225 Mon Sep 17 00:00:00 2001 From: Daniel Date: Tue, 23 Jun 2026 17:58:56 -0400 Subject: [PATCH 33/68] adding vllm basic functionality smoketest Signed-off-by: Daniel --- benchmarks/vllm_smoketest.sh | 103 +++++++++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 benchmarks/vllm_smoketest.sh diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh new file mode 100644 index 000000000000..959dc77f1a92 --- /dev/null +++ b/benchmarks/vllm_smoketest.sh @@ -0,0 +1,103 @@ +#!/bin/bash +set -euo pipefail + +# add model args here +declare -A MODEL_PATHS=( + [gpt-oss-120b-mxfp4]="/data/amd/gpt-oss-120b-w-mxfp4-a-fp8" + [DeepSeek-R1-0528-MXFP4]="/data/amd/DeepSeek-R1-0528-MXFP4" +) +declare -A MODEL_SERVE_ARGS=( + [gpt-oss-120b-mxfp4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --attention-backend ROCM_AITER_UNIFIED_ATTN --compilation-config {\"mode\":\"None\",\"cudagraph_mode\":\"FULL\",\"cudagraph_capture_sizes\":[1]}" + [DeepSeek-R1-0528-MXFP4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --dtype auto --no-enable-prefix-caching --disable-uvicorn-access-log --trust-remote-code" +) +declare -A MODEL_ENV=( + [gpt-oss-120b-mxfp4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [DeepSeek-R1-0528-MXFP4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" +) + +MODEL_NAME="gpt-oss-120b-mxfp4" # default +MODEL_PATH_OVERRIDE="" # overrided path +MODEL_SET=0 # whether --model was explicitly passed +RUN_LM_EVAL=0 +PORT=8000 + +usage() { + cat <<'EOF' +Usage: ./vllm_smoketest.sh [--model NAME] [--model-path PATH] [--lm-eval] [--port PORT] [--list] + --model NAME Registry key to run (default: gpt-oss-120b-mxfp4) + --model-path PATH Override model path for this run + --lm-eval Run lm_eval after curl check + --port PORT Server port (default: 8000) + --list List available models +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL_NAME="$2"; MODEL_SET=1; shift 2 ;; + --model-path) MODEL_PATH_OVERRIDE="$2"; shift 2 ;; + --lm-eval) RUN_LM_EVAL=1; shift ;; + --port) PORT="$2"; shift 2 ;; + --list) for n in "${!MODEL_PATHS[@]}"; do echo "$n -> ${MODEL_PATHS[$n]}"; done; exit 0 ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown argument: $1" >&2; usage; exit 1 ;; + esac +done + +if [[ -n "$MODEL_PATH_OVERRIDE" && $MODEL_SET -eq 0 ]]; then + echo "ERROR: --model-path requires --model (which model's serve args/env to use)." >&2 + usage; exit 1 +fi + +MODEL="${MODEL_PATHS[$MODEL_NAME]:-$MODEL_NAME}" +[[ -n "$MODEL_PATH_OVERRIDE" ]] && MODEL="$MODEL_PATH_OVERRIDE" +SERVE_ARGS="${MODEL_SERVE_ARGS[$MODEL_NAME]:-}" +MODEL_ENV_ARGS="${MODEL_ENV[$MODEL_NAME]:-}" +echo "Model: $MODEL | Port: $PORT | lm_eval: $RUN_LM_EVAL" + +# temp +pip3 install "fastapi<0.137" + +LOG_DIR="$(pwd)/vllm_smoke_test_logs" +mkdir -p "$LOG_DIR" +echo "Logs: $LOG_DIR" + +server_pid="" +cleanup() { [[ -n "$server_pid" ]] && kill "$server_pid" 2>/dev/null && wait "$server_pid" 2>/dev/null; true; } +trap cleanup EXIT +trap 'exit 130' INT TERM + +# --- start server --- +env $MODEL_ENV_ARGS \ +vllm serve --model "$MODEL" --host localhost --port "$PORT" $SERVE_ARGS \ + > "$LOG_DIR/server.log" 2>&1 & +server_pid=$! +echo "Waiting for server (pid $server_pid)..." + +until curl -s "http://localhost:$PORT/health" &>/dev/null; do + ps -p "$server_pid" >/dev/null || { echo "ERROR: server died"; tail -n 50 "$LOG_DIR/server.log"; exit 1; } + sleep 5 +done +echo "Server ready." + +# --- check if text was generated --- +echo "=== Curl completion check ===" +http_code=$(curl -s -o "$LOG_DIR/curl_completion.log" -w '%{http_code}' \ + "http://localhost:$PORT/v1/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL\",\"prompt\":\"San Francisco is a\",\"max_tokens\":32,\"temperature\":0}") +cat "$LOG_DIR/curl_completion.log"; echo +[[ "$http_code" == "200" ]] || { echo "FAIL: HTTP $http_code" >&2; exit 1; } +grep -q '"text"' "$LOG_DIR/curl_completion.log" || { echo "FAIL: no text in response" >&2; exit 1; } +echo "PASS" + +# --- lm_eval --- +if [[ $RUN_LM_EVAL -eq 1 ]]; then + echo "=== lm_eval ===" + pip install "lm-eval[api]" + lm_eval --model local-completions \ + --model_args "model=$MODEL,base_url=http://localhost:$PORT/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False" \ + --tasks gsm8k --num_fewshot 3 --limit 100 2>&1 | tee "$LOG_DIR/lm_eval.log" +fi + +echo "Done." From a4d5bba516c6251559e2bfa4d877e1faf62350e5 Mon Sep 17 00:00:00 2001 From: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:15:56 -0500 Subject: [PATCH 34/68] 455 wip 625 (#1023) * Initial fixes to get dsr4 working * DSv4 on aiter and gpt working --------- Co-authored-by: root Co-authored-by: jpvillam Signed-off-by: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> --- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 1 + .../experts/gpt_oss_triton_kernels_moe.py | 686 +++++++++++++----- vllm/model_executor/layers/utils.py | 8 +- 3 files changed, 506 insertions(+), 189 deletions(-) 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 index cc616f13cbeb..ed1da9416c40 100644 --- 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 @@ -141,6 +141,7 @@ def triton_kernel_fused_mxfp4_w4a8_experts( ) from vllm.platforms.rocm import on_gfx1250 + from vllm.platforms.rocm import on_gfx1250 _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None # TODO (JPVILLAM): merge conflict resolve later if _swizzle_mx_scale is enough mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" 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 adb24d03dc5d..fce6d01145b2 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 @@ -574,6 +574,44 @@ def triton_kernel_moe_forward( ) -> torch.Tensor: sm_first = not renormalize + # ROCm aiter-native MXFP4 fast path (mirrors ROCm/ATOM triton_kernel_moe_forward): + # native aiter routing on the raw logits + the two-GEMM aiter compute. Falls back + # to the vendored triton_kernels routing + matmul_ogs path below when unavailable + # or out of the ported scope (SiLU is FP8-only). + if ( + expert_map is None + and quant_config is not None + and rocm_aiter_ops.is_enabled() + and _aiter_act_quant_mode(quant_config) is not None + ): + _act_mode = _aiter_act_quant_mode(quant_config) + _aiter_ops = _import_aiter_moe_ops() + if _aiter_ops is not None and ( + activation == MoEActivation.SWIGLUOAI or _act_mode == _AITER_ACT_FP8 + ): + routing_data, gather_idx, scatter_idx = _aiter_ops["routing"]( + gating_output, topk, sm_first=sm_first + ) + swiglu_alpha, swiglu_limit = _aiter_swiglu_params(quant_config) + # Native aiter routing applies router weights in-kernel, so no post-scale. + return _aiter_moe_two_gemm( + _aiter_ops, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk, + quant_config, + _act_mode, + activation, + swiglu_alpha, + swiglu_limit, + apply_router_weight_on_input, + None, + ) + # When no expert map is provided (no EP), call the fused `routing()` # kernel directly. It combines softmax, topk, bitmatrix packing, and # routing-metadata construction in a single launch, instead of the @@ -826,7 +864,7 @@ def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( ) -> tuple[torch.Tensor, torch.Tensor | None]: """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). - Uses ``aiter.ops.triton.quant_moe.downcast_to_static_fp8`` with a scalar scale: + Uses ``aiter.ops.triton.moe.quant_moe.downcast_to_static_fp8`` with a scalar scale: for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static @@ -841,7 +879,7 @@ def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( if not quant_config.use_mxfp4_w4a16: return hidden_states, None try: - from aiter.ops.triton.quant_moe import downcast_to_static_fp8 + from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 except ImportError: return hidden_states, None if not rocm_aiter_ops.is_enabled(): @@ -929,6 +967,392 @@ def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( return replace(precision_config, flex_ctx=new_flex) +# --------------------------------------------------------------------------- +# aiter `.moe.*` MXFP4 MoE path (ported from ROCm/ATOM +# atom/model_ops/fused_moe_triton.py). +# +# Scope (agreed): the reference's *complete* paths only -- +# * SwiGLU (interleaved, gpt-oss): FP8 -> a8w4, FP4/BF16 -> a16w4, with the +# activation fused into the GEMM (apply_swiglu=True, swiglu_add_residual). +# * SiLU (half/half, e.g. DeepSeek-V4): FP8 -> a8w4 via dynamic MXFP8 quant + +# fused_clamp_act_mul. +# The reference's SiLU FP4 (a4w4) / BF16 (a16w4) paths feed an *unactivated* +# buffer to GEMM2 and are therefore intentionally NOT ported -- they raise. +# --------------------------------------------------------------------------- + +# aiter MXFP4 activation-quant modes (mirror ATOM MoEActivationQuant). +_AITER_ACT_FP8 = "fp8" # a8w4: FP8 activations (use_mxfp4_w4a8) +_AITER_ACT_FP4 = "fp4" # a4w4: MXFP4 activations (use_mxfp4_w4a4) +_AITER_ACT_BF16 = "bf16" # a16w4: BF16 activations (use_mxfp4_w4a16) + + +def _aiter_raw(t): + """Unwrap a triton_kernels ``Tensor`` to its backing ``torch.Tensor``. + + aiter's ``.moe.*`` kernels take plain ``torch.Tensor`` weights/scales, while + vLLM stores MXFP4 weights as triton_kernels ``Tensor`` wrappers + (``.storage.data``). Plain tensors pass through unchanged. + """ + return t.storage.data if hasattr(t, "storage") else t + + +def _aiter_act_quant_mode(quant_config: FusedMoEQuantConfig) -> str | None: + """Activation-quant mode for aiter MXFP4 MoE, or ``None`` if not MXFP4.""" + if quant_config.use_mxfp4_w4a8: + return _AITER_ACT_FP8 + if quant_config.use_mxfp4_w4a4: + return _AITER_ACT_FP4 + if quant_config.use_mxfp4_w4a16: + return _AITER_ACT_BF16 + return None + + +def _aiter_swiglu_params(quant_config: FusedMoEQuantConfig) -> tuple[float, float]: + """(alpha, clamp limit) for the SwiGLU/SiLU activation; DeepSeek-V4 defaults.""" + alpha = quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + return alpha, limit + + +def _import_aiter_moe_ops() -> dict | None: + """Import the aiter ``.moe.*`` triton MoE kernels (ROCm/ATOM API), or ``None``. + + Mirrors ROCm/ATOM ``atom/model_ops/fused_moe_triton.py``. Returns a dict of + callables on success; ``None`` if aiter (or this kernel family) is absent, so + the caller can fall back to the vendored ``matmul_ogs`` path. + """ + try: + from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul + from aiter.ops.triton.moe.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe.moe_routing.routing import routing as aiter_routing + from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 + from aiter.ops.triton.quant.quant import dynamic_mxfp8_quant + from aiter.ops.triton.utils._triton.arch_info import get_arch + except ImportError: + return None + return { + "moe_gemm_a8w4": moe_gemm_a8w4, + "moe_gemm_a16w4": moe_gemm_a16w4, + "routing": aiter_routing, + "downcast_to_static_fp8": downcast_to_static_fp8, + "dynamic_mxfp8_quant": dynamic_mxfp8_quant, + "fused_clamp_act_mul": fused_clamp_act_mul, + "get_arch": get_arch, + } + + +def _aiter_routing_from_topk( + aiter_ops: dict, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_experts: int, + device: torch.device, +) -> tuple: + """Build aiter ``RoutingData`` from a precomputed topk selection. + + vLLM selects experts before the expert kernel runs, but aiter's MXFP4 GEMMs + need aiter routing structs (distinct from the vendored triton_kernels ones). + aiter ``routing`` always softmaxes its logits, so we scatter ``log(weight)`` + into a dense ``-inf`` logit matrix: softmax-over-selected reproduces both the + exact expert selection and ``weight / sum(weight)``. The per-token weight sum + (e.g. ``routed_scaling_factor``) divided out here is restored by the caller. + + Returns ``(routing_data, gather_idx, scatter_idx, topk_weights_f32)``. + """ + routing = aiter_ops["routing"] + M = topk_weights.shape[0] + topk = topk_ids.size(1) + tw = topk_weights.to(torch.float32) + logits = torch.full((M, num_experts), -1e30, device=device, dtype=torch.float32) + logits.scatter_( + 1, + topk_ids.long().clamp(min=0, max=num_experts - 1), + torch.log(tw.clamp(min=1e-20)), + ) + routing_data, gather_idx, scatter_idx = routing(logits, topk, sm_first=False) + return routing_data, gather_idx, scatter_idx, tw + + +def _aiter_moe_two_gemm( + aiter_ops: dict, + hidden_states: torch.Tensor, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk: int, + quant_config: FusedMoEQuantConfig, + act_mode: str, + activation: MoEActivation, + swiglu_alpha: float, + swiglu_limit: float, + apply_router_weight_on_input: bool, + moe_config: FusedMoEConfig | None, +) -> torch.Tensor: + """Two-GEMM aiter MXFP4 MoE compute (ported from ROCm/ATOM fused_moe_triton). + + vLLM/gfx1250 adaptations vs the reference: weights/scales are unwrapped from + triton_kernels ``Tensor`` (``.storage.data``); scales are passed unswizzled + with logical ``unpadded_N/K`` (the gfx1250 convention used by the validated + aiter path); and GEMM1's in-kernel gather -- numerically broken on gfx1250 -- + is replaced by an explicit torch gather there. + + Returns the (M, K) bf16 result *before* any router-weight post-scale; the + caller applies that when routing was reconstructed from a topk selection. + """ + moe_gemm_a8w4 = aiter_ops["moe_gemm_a8w4"] + moe_gemm_a16w4 = aiter_ops["moe_gemm_a16w4"] + dynamic_mxfp8_quant = aiter_ops["dynamic_mxfp8_quant"] + fused_clamp_act_mul = aiter_ops["fused_clamp_act_mul"] + downcast_to_static_fp8 = aiter_ops["downcast_to_static_fp8"] + get_arch = aiter_ops["get_arch"] + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + w1_data = _aiter_raw(w1) + w2_data = _aiter_raw(w2) + w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) + w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + w1_bias = quant_config.w1_bias + w2_bias = quant_config.w2_bias + + if moe_config is not None: + unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( + _mxfp4_w4a8_unpadded_dims(moe_config) + ) + else: + unpadded_n_w1 = unpadded_k_w1 = unpadded_n_w2 = unpadded_k_w2 = None + + arch = get_arch() + quant_dtype = torch.float8_e4m3fnuz if arch == "gfx942" else torch.float8_e4m3fn + + # gfx1250's in-kernel gather miscomputes (validated on the FFM sim), so gather + # rows into expert-sorted order in torch and pass gather_indx=None instead. Per + # aiter's moe_gemm_torch, sorted row i reads source token gather_idx[i] // topk. + if arch == "gfx1250": + gather_src = gather_idx.to(torch.long) // topk + gemm1_input = hidden_states[gather_src] + gemm1_gather_indx = None + else: + gemm1_input = hidden_states + gemm1_gather_indx = gather_idx + + gammas = routing_data.gate_scal if routing_data else None + g1_gammas = gammas if apply_router_weight_on_input else None + g2_gammas = None if apply_router_weight_on_input else gammas + + if activation == MoEActivation.SWIGLUOAI: + # Interleaved [gate, up] (gpt-oss): activation fused into the GEMM. + if act_mode == _AITER_ACT_FP8: + a13_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=1, activations=gemm1_input + ) + a2_scale = _mxfp4_w4a8_resolve_lhs_scale( + quant_config, gemm_num=2, activations=gemm1_input + ) + hidden_fp8 = downcast_to_static_fp8(gemm1_input, a13_scale) + interm = moe_gemm_a8w4( + hidden_fp8, + w1_data, + None, + w1_wscale, + a13_scale, + a2_scale, + w1_bias, + routing_data, + gather_indx=gemm1_gather_indx, + gammas=g1_gammas, + swizzle_mx_scale=None, + out_dtype=quant_dtype, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=True, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + out = moe_gemm_a8w4( + interm, + w2_data, + None, + w2_wscale, + a2_scale, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + else: + # FP4 and BF16 activations both route through the a16w4 kernel. + interm = moe_gemm_a16w4( + gemm1_input, + w1_data, + None, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gemm1_gather_indx, + gammas=g1_gammas, + swizzle_mx_scale=None, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=True, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + out = moe_gemm_a16w4( + interm, + w2_data, + None, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + return out + + # SiLU (concatenated [gate | up], half/half): manual activation between GEMMs. + # FP8 (a8w4) and BF16 (w4a16 served as a8w4) both run through moe_gemm_a8w4 by + # dynamically quantizing activations to FP8 (dynamic_mxfp8_quant below). Only + # FP4 (a4w4) SiLU is unsupported (ATOM reference a4w4 SiLU path is incomplete). + if act_mode == _AITER_ACT_FP4: + raise NotImplementedError( + "aiter SiLU MoE is unsupported for FP4 (a4w4) activations because the " + "ATOM reference a4w4 SiLU path is incomplete; use FP8 or BF16 (w4a16)." + ) + + hidden_q, a13_scale = dynamic_mxfp8_quant(gemm1_input, quant_dtype=quant_dtype) + raw_gate_up = moe_gemm_a8w4( + hidden_q, + w1_data, + a13_scale, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gemm1_gather_indx, + gammas=g1_gammas, + swizzle_mx_scale=None, + out_dtype=torch.bfloat16, + apply_swiglu=False, + unpadded_N=unpadded_n_w1, + unpadded_K=unpadded_k_w1, + ) + if unpadded_n_w1 is not None: + raw_gate_up = raw_gate_up[:, :unpadded_n_w1] + # SiLU + clamp + (gate * up), fused with dynamic MXFP8 quant of the result. + interm_fp8, a2_scale = fused_clamp_act_mul( + raw_gate_up, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=quant_dtype, + scale_dtype_fmt="ue8m0", + quant_block_size=32, + ) + out = moe_gemm_a8w4( + interm_fp8, + w2_data, + a2_scale, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_n_w2, + unpadded_K=unpadded_k_w2, + ) + return out + + +def _aiter_mxfp4_fused_experts( + output: torch.Tensor, + hidden_states: torch.Tensor, + w1, + w2, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_config: FusedMoEQuantConfig, + activation: MoEActivation, + apply_router_weight_on_input: bool, + moe_config: FusedMoEConfig | None, +) -> torch.Tensor | None: + """aiter-native MXFP4 MoE for the modular expert path (ROCm fast path). + + Builds aiter routing from ``topk_ids``/``topk_weights``, runs the two-GEMM + aiter compute, restores the per-token router-weight sum that routing + reconstruction normalized away, and writes ``output``. + + Returns ``output`` on success, or ``None`` (caller falls back to ``matmul_ogs``) + when aiter is unavailable, the weights are not MXFP4, or the requested + activation/quant combination is out of the ported scope. + """ + if not rocm_aiter_ops.is_enabled(): + return None + act_mode = _aiter_act_quant_mode(quant_config) + if act_mode is None: + return None + if quant_config.w1_precision is None or quant_config.w2_precision is None: + return None + # SiLU served via a8w4 for FP8 and BF16 (dynamic FP8 quant); only FP4 (a4w4) + # SiLU is out of scope (reference a4w4 SiLU path incomplete). This keeps w4a16 + # SiLU on the aiter-native fast path instead of the vendored-routing fallback. + if activation != MoEActivation.SWIGLUOAI and act_mode == _AITER_ACT_FP4: + return None + aiter_ops = _import_aiter_moe_ops() + if aiter_ops is None: + return None + + num_experts = w1.shape[0] + routing_data, gather_idx, scatter_idx, tw = _aiter_routing_from_topk( + aiter_ops, topk_ids, topk_weights, num_experts, hidden_states.device + ) + swiglu_alpha, swiglu_limit = _aiter_swiglu_params(quant_config) + out = _aiter_moe_two_gemm( + aiter_ops, + hidden_states, + w1, + w2, + routing_data, + gather_idx, + scatter_idx, + topk_ids.size(1), + quant_config, + act_mode, + activation, + swiglu_alpha, + swiglu_limit, + apply_router_weight_on_input, + moe_config, + ) + # Restore the per-token weight sum that routing reconstruction normalized out. + out = out.to(output.dtype) * tw.sum(dim=1, keepdim=True) + output.copy_(out.reshape(output.shape)) + return output + + @triton.jit def _masked_topk_sum_kernel( inp_ptr, # (M, topk, K) contiguous @@ -1139,9 +1563,49 @@ def apply( self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG if expert_map is not None: - # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch): - # make_routing_data treats -1 as the skip sentinel. - topk_ids = remap_topk_to_local(topk_ids, expert_map) + topk_ids = expert_map[topk_ids] + + # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible + # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. + if ( + not apply_router_weight_on_input + and rocm_aiter_ops.is_enabled() + and _aiter_mxfp4_fused_experts( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + self.quant_config, + activation, + apply_router_weight_on_input, + getattr(self, "moe_config", None), + ) + is not None + ): + return + + # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible + # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. + if ( + not apply_router_weight_on_input + and rocm_aiter_ops.is_enabled() + and _aiter_mxfp4_fused_experts( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + self.quant_config, + activation, + apply_router_weight_on_input, + getattr(self, "moe_config", None), + ) + is not None + ): + return local_num_experts = w1.shape[0] if global_num_experts == -1: @@ -1267,193 +1731,34 @@ def _try_apply_aiter_w4a8( topk_weights: torch.Tensor, topk_ids: torch.Tensor, quant_config: FusedMoEQuantConfig, + activation: MoEActivation, apply_router_weight_on_input: bool, ) -> torch.Tensor | None: - """ROCm/gfx1250 MXFP4 MoE via aiter's W4A8 ``moe_gemm_a8w4``. + """ROCm aiter-native MXFP4 MoE fast path; ``None`` -> caller uses matmul_ogs. The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and - fails to compile on gfx1250. aiter's gluon ``moe_gemm_a8w4`` (FP8 act x MXFP4 - weight) does run on gfx1250, so we route the MoE through it. - - Routing is built aiter-native directly from ``topk_ids``/``topk_weights`` so we - never touch the diverged vendored ``triton_kernels`` routing structs (whose - ``RoutingData``/``GatherIndx``/``ExptData`` field names and layout differ from - aiter's). This mirrors ``triton_kernel_fused_mxfp4_w4a8_experts`` - (aiter_mxfp4_w4a8_moe.py) but resolves the FP8 LHS scales dynamically because - DeepSeek-V4 is ``activation_scheme: dynamic`` (no calibrated static scale). - - aiter's ``routing`` always softmaxes the gate weights, so we feed - ``log(topk_weights)``: softmax-over-selected then recovers - ``topk_weights / sum(topk_weights)``. DeepSeek-V4 weights are - norm_topk_prob-normalized then scaled by ``routed_scaling_factor`` (so they sum - to that factor, not 1); we restore the exact weighting with a per-token output - rescale by the weight sum (== 1 for plain renormalize routing, a no-op there). - - Returns ``output`` on success, or ``None`` if aiter is unavailable (caller then - falls back to ``matmul_ogs``). + fails to compile on gfx1250. aiter's ``.moe.*`` gluon kernels run on gfx1250, + so we route the MoE through them. + + Delegates to :func:`_aiter_mxfp4_fused_experts` (ported from ROCm/ATOM + ``fused_moe_triton.py``): a8w4 SiLU via dynamic MXFP8 + ``fused_clamp_act_mul``, + and the fused-SwiGLU modes. Routing is built aiter-native from + ``topk_ids``/``topk_weights`` and the router-weight sum is restored afterward; + see that function for the routing/scaling details. """ - try: - from aiter.ops.triton.moe.moe_routing.routing import ( - routing as aiter_routing, - ) - from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 - from aiter.ops.triton.quant_moe import downcast_to_static_fp8 - except ImportError: - return None - - assert quant_config.w1_precision is not None - assert quant_config.w2_precision is not None - - M = hidden_states.shape[0] - num_experts = w1.shape[0] - topk = topk_ids.size(1) - - # Reconstruct dense gating logits from the sparse topk selection. log(weight) - # padded with a large-negative sentinel makes aiter's softmax reproduce both - # this exact expert selection and weight = topk_weights / per-token-sum. - tw = topk_weights.to(torch.float32) - logits = torch.full( - (M, num_experts), -1e30, device=hidden_states.device, dtype=torch.float32 - ) - logits.scatter_( - 1, - topk_ids.long().clamp(min=0, max=num_experts - 1), - torch.log(tw.clamp(min=1e-20)), - ) - logger.info( - "[aiter-w4a8] ENTER M=%d K=%d num_experts=%d topk=%d " - "w1=%s w2=%s w1_wscale=%s w2_wscale=%s", - M, - hidden_states.shape[1], - num_experts, - topk, - tuple(w1.storage.data.shape), - tuple(w2.storage.data.shape), - tuple(quant_config.w1_precision.weight_scale.storage.data.shape), - tuple(quant_config.w2_precision.weight_scale.storage.data.shape), - ) - routing_data, gather_idx, scatter_idx = aiter_routing( - logits, topk, sm_first=False - ) - gammas = routing_data.gate_scal - logger.info( - "[aiter-w4a8] routing OK block_m=%s gate_scal=%s gather=%s scatter=%s", - getattr(routing_data, "block_m", None), - tuple(gammas.shape), - tuple(gather_idx.shape), - tuple(scatter_idx.shape), - ) - - unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( - _mxfp4_w4a8_unpadded_dims(self.moe_config) - ) - swiglu_alpha = ( - quant_config.gemm1_alpha - if quant_config.gemm1_alpha is not None - else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default - ) - swiglu_limit = ( - quant_config.gemm1_clamp_limit - if quant_config.gemm1_clamp_limit is not None - else 7.0 - ) - - # FP8 LHS scales (static if calibrated, else dynamic amax/448). GEMM1 emits FP8 - # directly (apply_swiglu=True, out_dtype fp8) quantized with GEMM2's LHS scale, - # so resolve GEMM2's scale up front (hidden states as the amax proxy) and reuse - # it as GEMM2's input scale so quantize/dequantize are consistent. - gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, gemm_num=1, activations=hidden_states - ) - gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, gemm_num=2, activations=hidden_states - ) - # aiter's in-kernel gather is numerically broken on gfx1250 (validated on the - # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted - # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, sorted - # row i reads source token gather_idx[i] // n_expts_act, so this reproduces the - # in-kernel gather exactly (validated: manual gather -> maxrel ~5e-3). - gather_src = gather_idx.to(torch.long) // topk - hidden_sorted = hidden_states[gather_src] - hidden_fp8 = downcast_to_static_fp8(hidden_sorted, gemm1_lhs_scale) - logger.info( - "[aiter-w4a8] scales g1_lhs=%s g2_lhs=%s hidden_fp8=%s; calling gemm1 " - "(unpadded N/K w1=%s/%s w2=%s/%s, alpha=%.4f limit=%.2f)", - float(gemm1_lhs_scale), - float(gemm2_lhs_scale), - tuple(hidden_fp8.shape), - unpadded_n_w1, - unpadded_k_w1, - unpadded_n_w2, - unpadded_k_w2, - swiglu_alpha, - swiglu_limit, - ) - - # GEMM1 WITHOUT fused swiglu. aiter's fused _swiglu splits gate/up by - # INTERLEAVING, but DeepSeek-V4 (SILU) stores w13 half/half, so emit the raw - # gate_up (bf16) and apply the half/half SILU+clamp ourselves -- identical to - # the unfused activation() path (swiglu_limit_func) -- then requantize to fp8. - raw_gate_up = moe_gemm_a8w4( - hidden_fp8, - w1.storage.data, - None, - quant_config.w1_precision.weight_scale.storage.data, - gemm1_lhs_scale, - None, - quant_config.w1_bias, - routing_data, - gather_indx=None, - gammas=gammas if apply_router_weight_on_input else None, - swizzle_mx_scale=None, - out_dtype=torch.bfloat16, - apply_swiglu=False, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - gate_up = raw_gate_up[:, :unpadded_n_w1] - _act = torch.empty( - (gate_up.shape[0], unpadded_n_w1 // 2), - dtype=torch.bfloat16, - device=gate_up.device, - ) - swiglu_limit_func(_act, gate_up, swiglu_limit) - intermediate_cache1 = downcast_to_static_fp8(_act, gemm2_lhs_scale) - logger.info( - "[aiter-w4a8] gemm1(raw)+halfsplit-silu OK ic1=%s %s; calling gemm2", - tuple(intermediate_cache1.shape), - intermediate_cache1.dtype, - ) - intermediate_cache3 = moe_gemm_a8w4( - intermediate_cache1, - w2.storage.data, - None, - quant_config.w2_precision.weight_scale.storage.data, - gemm2_lhs_scale, - None, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_idx, - gammas=None if apply_router_weight_on_input else gammas, - swizzle_mx_scale=None, - unpadded_N=unpadded_n_w2, - unpadded_K=unpadded_k_w2, - ) - - # Restore the per-token weight sum (e.g. routed_scaling_factor) that the - # log+softmax normalization divided out. Output weighting (gammas on GEMM2) is - # linear in the gate, so a single post-scale is exact. - logger.info( - "[aiter-w4a8] gemm2 OK ic3=%s %s; rescaling + writing output=%s", - tuple(intermediate_cache3.shape), - intermediate_cache3.dtype, - tuple(output.shape), + return _aiter_mxfp4_fused_experts( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + quant_config, + activation, + apply_router_weight_on_input, + self.moe_config, ) - out = intermediate_cache3.to(output.dtype) * tw.sum(dim=1, keepdim=True) - output.copy_(out.reshape(output.shape)) - logger.info("[aiter-w4a8] DONE") - return output def apply( self, @@ -1520,7 +1825,11 @@ def apply( # unswizzle. Output-side router weighting only (the post-scale below assumes # it); LoRA still uses the matmul_ogs path further down. if ( - (quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16) + ( + quant_config.use_mxfp4_w4a8 + or quant_config.use_mxfp4_w4a16 + or quant_config.use_mxfp4_w4a4 + ) and not apply_router_weight_on_input and self._lora_context is None and rocm_aiter_ops.is_enabled() @@ -1534,6 +1843,7 @@ def apply( topk_weights, topk_ids, quant_config, + activation, apply_router_weight_on_input, ) is not None @@ -1569,10 +1879,10 @@ def apply( _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 if _mxfp4_weights and rocm_aiter_ops.is_enabled(): try: - from aiter.ops.triton.moe_op_gemm_a8w4 import ( + from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( moe_gemm_a8w4 as _moe_gemm_a8w4_fn, ) - from aiter.ops.triton.quant_moe import ( + from aiter.ops.triton.moe.quant_moe import ( downcast_to_static_fp8 as _downcast_static_fp8_fn, ) except ImportError: diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index a01acec2ac2e..a4e66ca06e82 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -164,7 +164,13 @@ def rocm_unquantized_gemm_impl( if use_skinny_reduce_counting: return ops.wvSplitKrc(x, weight, cu_count, bias) - if use_aiter_triton_gemm(n, m, k, x.dtype): + # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires + # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a + # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b + # hidden), so skip aiter there and fall back to the torch GEMM path below. + if use_aiter_triton_gemm(n, m, k, x.dtype) and not ( + on_gfx1250() and k % 256 != 0 + ): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) From 79c79bcefa9ca981f61c7fdda4957bcb1d415e1e Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Mon, 29 Jun 2026 13:11:11 -0400 Subject: [PATCH 35/68] 455 platform enablement (#1031) Signed-off-by: Jaden Mathias --- .../test_transcription_validation_whisper.py | 6 +- .../test_translation_validation.py | 6 +- .../attention/test_rocm_aiter_unified_attn.py | 4 +- tests/kernels/moe/test_routing.py | 6 +- .../layers/test_rocm_unquantized_gemm.py | 3 + .../generation/test_granite_speech.py | 4 +- .../models/quantization/test_bitsandbytes.py | 4 +- .../test_rocm_attention_backends_selection.py | 12 ++-- vllm/_aiter_ops.py | 14 ++--- .../kernels/linear/scaled_mm/pytorch.py | 4 +- .../kernels/linear/scaled_mm/rocm.py | 6 +- .../fused_moe/experts/fused_batched_moe.py | 8 +-- .../experts/gpt_oss_triton_kernels_moe.py | 4 +- .../quark/schemes/quark_w4a8_mxfp4_fp8.py | 4 +- vllm/model_executor/layers/utils.py | 3 +- vllm/platforms/rocm.py | 59 ++++++++++++++----- vllm/v1/attention/backends/rocm_aiter_fa.py | 6 +- 17 files changed, 92 insertions(+), 61 deletions(-) diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index 511179f7fcb1..34247782b14b 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -31,7 +31,7 @@ def _get_attention_backend_params() -> list[str | None]: falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER) - TRITON_ATTN: always available on ROCm - - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950 + - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950/gfx1250 On non-ROCm platforms, we just run with the default backend. """ @@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]: if current_platform.is_rocm(): backends: list[str | None] = [None, "TRITON_ATTN"] - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import _on_mi3or4 - if _ON_MI3XX: + if _on_mi3or4(): backends.append("ROCM_AITER_UNIFIED_ATTN") return backends except Exception: diff --git a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py index ed3cff5f1c22..de82dc3c4fc3 100644 --- a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py +++ b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py @@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name): if "whisper" in model_name.lower(): try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import _on_mi3or4 - if _ON_MI3XX: + if _on_mi3or4(): return {"backend": "ROCM_AITER_UNIFIED_ATTN"} except ImportError: logger.warning( - "Could not import _ON_MI3XX from rocm platform, " + "Could not import _on_mi3or4 from rocm platform, " "falling back to TRITON_ATTN for Whisper." ) return {"backend": "TRITON_ATTN"} diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index c02a457c98a4..410c1d3ebab6 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -17,9 +17,9 @@ _SKIP_NON_MI3XX = True if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_mi3or4 - _SKIP_NON_MI3XX = not on_mi3xx() + _SKIP_NON_MI3XX = not on_mi3or4() pytestmark = [ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 62a4968a0d1f..c5081a8356ed 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -19,13 +19,13 @@ def _is_aiter_capable() -> bool: - """Check if the platform supports AITER (gfx942/gfx950).""" + """Check if the platform supports AITER (gfx942/gfx950/gfx1250).""" if not current_platform.is_rocm(): return False try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import _on_mi3or4 - return _ON_MI3XX + return _on_mi3or4() except ImportError: return False diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index f4de9bc9038f..f20be7cfb618 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -26,6 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -52,6 +53,7 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -76,6 +78,7 @@ def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t()) diff --git a/tests/models/multimodal/generation/test_granite_speech.py b/tests/models/multimodal/generation/test_granite_speech.py index 3019f5f22d4b..48d32531d290 100644 --- a/tests/models/multimodal/generation/test_granite_speech.py +++ b/tests/models/multimodal/generation/test_granite_speech.py @@ -45,9 +45,9 @@ def vllm_to_hf_output( def granite_speech_attention_config(): """Return attention config for Granite Speech tests on ROCm.""" if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_mi3or4 - if on_mi3xx(): + if on_mi3or4(): return {"backend": "ROCM_AITER_FA"} return {"backend": "TRITON_ATTN"} return None diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index 03c19b0bf62a..62ce36a3f980 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -21,10 +21,10 @@ from ..utils import check_embeddings_close, check_logprobs_close if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import on_gfx9, on_gfx1250 pytestmark = pytest.mark.skipif( - on_gfx9(), + on_gfx9() or on_gfx1250(), reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)", ) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 48c6de8f8bd7..85ce15d2fb68 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -35,9 +35,9 @@ def mock_on_gfx9(): @pytest.fixture -def mock_on_mi3xx(): - """Mock mi3xx arch detection to return True.""" - with patch("vllm.platforms.rocm.on_mi3xx", return_value=True): +def mock_on_mi3or4(): + """Mock mi3or4 arch detection to return True.""" + with patch("vllm.platforms.rocm.on_mi3or4", return_value=True): yield @@ -112,7 +112,7 @@ def test_standard_attention_backend_selection( expected_backend_path, mock_vllm_config, mock_on_gfx9, - mock_on_mi3xx, + mock_on_mi3or4, monkeypatch, ): """Test standard attention backend selection with various configurations.""" @@ -316,9 +316,9 @@ def test_aiter_fa_requires_mi3xx(mock_vllm_config): """Test that ROCM_AITER_FA requires mi3xx architecture.""" from vllm.platforms.rocm import RocmPlatform - # Mock on_mi3xx to return False (used by supports_compute_capability) + # Mock on_mi3or4 to return False (used by supports_compute_capability) with ( - patch("vllm.platforms.rocm.on_mi3xx", return_value=False), + patch("vllm.platforms.rocm.on_mi3or4", return_value=False), pytest.raises( ValueError, match="compute capability not supported", diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index f826c1ee3834..afca102f763d 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -92,7 +92,7 @@ def should_custom_ar(self, inp: torch.Tensor) -> bool: ... def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. - Checks: platform (ROCm), device arch (gfx9), and library existence. + Checks: platform (ROCm), device arch (gfx9 or gfx1250), and library existence. Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled(). This function determines if aiter CAN be used, not if it SHOULD be used. @@ -106,9 +106,9 @@ def is_aiter_found_and_supported() -> bool: VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery. """ if current_platform.is_rocm() and IS_AITER_FOUND: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_mi3or4 - return on_mi3xx() + return on_mi3or4() return False @@ -1893,16 +1893,16 @@ def is_fp8bmm_enabled(cls) -> bool: @classmethod @if_aiter_supported def is_fp4bmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() + return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() # TODO GFX1250: (on_gfx950() or on_gfx1250()) @classmethod @if_aiter_supported def is_linear_hipbmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_mi3or4 - return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED + return cls.is_linear_enabled() and on_mi3or4() and cls._LINEAR_HIPBMM_ENABLED @classmethod @if_aiter_supported diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 2b6d3ed73698..dbb77db5077e 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -104,9 +104,9 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_mi3or4 - if not on_mi3xx(): + if not on_mi3or4(): return False, "requires MI3xx." if compute_capability is not None and compute_capability < 94: diff --git a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py index 64bc5b6c8bbe..446eabd5431c 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py @@ -79,10 +79,10 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_gfx12x, on_mi3xx + from vllm.platforms.rocm import on_gfx12x, on_mi3or4 - if not (on_mi3xx() or on_gfx12x()): - return False, "requires MI3xx or gfx12x" + if not (on_mi3or4() or on_gfx12x()): + return False, "requires MI3or4 or gfx12x" if not envs.VLLM_ROCM_USE_SKINNY_GEMM: return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 21bda8e173fd..520ccd295389 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -758,13 +758,13 @@ def _supports_quant_scheme( ) -> bool: p = current_platform if p.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import on_gfx9, on_gfx1250 - is_rocm_on_gfx9 = on_gfx9() + _rocm_support_fp8 = on_gfx9() or on_gfx1250() else: - is_rocm_on_gfx9 = False + _rocm_support_fp8 = False - device_supports_fp8 = is_rocm_on_gfx9 or ( + device_supports_fp8 = _rocm_support_fp8 or ( p.is_cuda() and p.has_device_capability((8, 9)) ) 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 fce6d01145b2..8b688697d167 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 @@ -51,13 +51,13 @@ def _triton_kernel_moe_supports_current_device() -> bool: # range was not validated. return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) if p.is_rocm(): - from vllm.platforms.rocm import on_gfx1x, on_gfx9 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx1250 # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx); # on_gfx9() already excludes gfx906/gfx908. # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). - return on_gfx9() or on_gfx1x() + return on_gfx9() or on_gfx1x() or on_gfx1250() return False diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py index 29283c7bbda4..7da4a76e6073 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py @@ -64,9 +64,9 @@ def __init__( kernel_supported_gpu = False if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - kernel_supported_gpu = on_gfx950() + kernel_supported_gpu = on_gfx950() or on_gfx1250() self.use_aiter_kernel = ( is_aiter_found_and_supported() diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index a4e66ca06e82..814142ce053f 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -178,9 +178,8 @@ def rocm_unquantized_gemm_impl( use_skinny = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) - # gfx1250 excludes the skinny-GEMM kernels (wvSplitK/LLMM1) from the # build (gfx9/gfx11 ISA); fall back to torch GEMM there. - and not on_gfx1250() + and not on_gfx1250() # TODO GFX1250: Remove once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 32d27ab0ebc4..a857502fed78 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -225,18 +225,22 @@ def _get_gcn_arch() -> str: _ON_GFX1100 = "gfx1100" in _GCN_ARCH _ON_GFX1151 = "gfx1151" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) -_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950", "gfx1250"]) +_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) _ON_GFX9 = any( - arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950", "gfx1250"] -) # TODO(JPVILLAM): Bubblegum patch to unlock gptoss + arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"] +) _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH -# any( -# arch in _GCN_ARCH for arch in ["gfx950", "gfx1250"] -# ) # TODO(JPVILLAM): Bubblegum patch to unlock DSR1 _ON_GFX1250 = "gfx1250" in _GCN_ARCH +_ON_MI3OR4 = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950", "gfx1250"]) # TODO GFX1250: Need a correct name for this +_ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9, gfx1250"]) +_ON_RDNA = any( + arch for arch in _GCN_ARCH + if (arch.startswith("gfx11") or arch.startswith("gfx12")) and arch != "gfx1250" +) +#TODO GFX1250: Use CDNA for MI3OR4 def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: """ @@ -310,7 +314,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: def on_gfx1x() -> bool: - return _ON_GFX1X + return _ON_GFX1X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly def on_gfx11() -> bool: @@ -326,13 +330,21 @@ def on_gfx1151() -> bool: def on_gfx12x() -> bool: - return _ON_GFX12X + return _ON_GFX12X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly + + +def on_gfx1250() -> bool: + return _ON_GFX1250 def on_mi3xx() -> bool: return _ON_MI3XX +def on_mi3or4() -> bool: + return _ON_MI3OR4 + + def on_gfx9() -> bool: return _ON_GFX9 @@ -349,8 +361,25 @@ def on_gfx950() -> bool: return _ON_GFX950 -def on_gfx1250() -> bool: - return _ON_GFX1250 +def on_cdna() -> bool: + return _ON_CDNA + + +def on_rdna() -> bool: + return _ON_RDNA + + +def get_cdna_version() -> int: + if on_gfx90a(): + return 2 + if on_gfx942(): + return 3 + if on_gfx950(): + return 4 + if on_gfx1250(): + return 5 + return -1 + # Enable HIP online tuning early, before hipBLASLt initializes. @@ -359,7 +388,7 @@ def on_gfx1250() -> bool: envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - and on_mi3xx() + and on_mi3or4() ): os.environ["HIP_ONLINE_TUNING"] = "1" @@ -378,7 +407,7 @@ def use_rocm_custom_paged_attention( ) -> bool: # custom paged attn always supported on V0. On V1, requires sliding window # disabled due to observed numerical discrepancy. - if _ON_GFX9: + if _ON_GFX9 or _ON_GFX1250: return ( (sliding_window == 0 or sliding_window == (-1, -1)) and (qtype == torch.half or qtype == torch.bfloat16) @@ -665,12 +694,12 @@ def get_vit_attn_backend( from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_enabled() and on_gfx9(): + if rocm_aiter_ops.is_enabled() and (on_gfx9() or on_gfx1250()): logger.info_once("Using AITER Flash Attention backend for ViT model.") return AttentionBackendEnum.ROCM_AITER_FA if ( - on_gfx9() + (on_gfx9() or on_gfx1250()) and find_spec("flash_attn") is not None and (dtype == torch.float16 or dtype == torch.bfloat16) ): @@ -893,7 +922,7 @@ def supports_mx(cls) -> bool: @classmethod def supports_fp8(cls) -> bool: - return on_gfx9() or on_gfx12x() + return on_gfx9() or on_gfx12x() or on_gfx1250() # TODO GFX1250: We know this is redundant ATM @classmethod def is_fp8_fnuz(cls) -> bool: diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 7e850af3e7bc..fa126a80de2e 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -768,12 +768,12 @@ def get_kv_cache_shape( @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import on_mi3or4 # DeviceCapability is currently created using torch.cuda.get_device_capability() - # which is known to be buggy on rocm systems. on_mi3xx uses amd-smi which is + # which is known to be buggy on rocm systems. on_mi3or4 uses amd-smi which is # more reliable. - return on_mi3xx() + return on_mi3or4() @classmethod def supports_non_causal(cls) -> bool: From d545ca8666ec4622d98bd461e68547a7af3e8a21 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Tue, 30 Jun 2026 03:57:39 +0000 Subject: [PATCH 36/68] Small fixes for whl build Signed-off-by: jpvillam --- benchmarks/vllm_smoketest.sh | 7 +++--- docker/Dockerfile.rocm | 8 +++++++ docker/Dockerfile.rocm_base | 42 ++++++++++++++++-------------------- requirements/rocm.txt | 3 +-- 4 files changed, 30 insertions(+), 30 deletions(-) diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh index 959dc77f1a92..88aaa9107c03 100644 --- a/benchmarks/vllm_smoketest.sh +++ b/benchmarks/vllm_smoketest.sh @@ -7,11 +7,11 @@ declare -A MODEL_PATHS=( [DeepSeek-R1-0528-MXFP4]="/data/amd/DeepSeek-R1-0528-MXFP4" ) declare -A MODEL_SERVE_ARGS=( - [gpt-oss-120b-mxfp4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --attention-backend ROCM_AITER_UNIFIED_ATTN --compilation-config {\"mode\":\"None\",\"cudagraph_mode\":\"FULL\",\"cudagraph_capture_sizes\":[1]}" + [gpt-oss-120b-mxfp4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.7 --attention-backend TRITON_ATTN" [DeepSeek-R1-0528-MXFP4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --dtype auto --no-enable-prefix-caching --disable-uvicorn-access-log --trust-remote-code" ) declare -A MODEL_ENV=( - [gpt-oss-120b-mxfp4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [gpt-oss-120b-mxfp4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=0 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" [DeepSeek-R1-0528-MXFP4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" ) @@ -69,8 +69,7 @@ trap 'exit 130' INT TERM # --- start server --- env $MODEL_ENV_ARGS \ -vllm serve --model "$MODEL" --host localhost --port "$PORT" $SERVE_ARGS \ - > "$LOG_DIR/server.log" 2>&1 & +vllm serve --model "$MODEL" --host localhost --port "$PORT" $SERVE_ARGS & server_pid=$! echo "Waiting for server (pid $server_pid)..." diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 789c4446a29d..f01786b04602 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -803,6 +803,14 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall -y triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + CMD ["/bin/bash"] #Set entrypoint for vllm-openai official images diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 4ead2b19c075..9e5daf66a9f7 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -10,7 +10,7 @@ ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="jpvillam/gfx1250_0604" +ARG AITER_BRANCH="main" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -46,7 +46,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++\ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ @@ -96,42 +96,36 @@ ARG TORCHVISION_VERSION ARG TORCHAUDIO_VERSION ARG ROCM_SDK_VERSION ARG APEX_VERSION +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + # torch/torchvision/torchaudio must be pinned to mutually-consistent builds # (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm # sdk version is derived from torch's own dependency pin unless overridden, # which keeps the set consistent and avoids pip backtracking. RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ - "torch==${TORCH_VERSION}" \ + --extra-index-url https://pypi.org/simple \ + "torch[device-all]==${TORCH_VERSION}" \ "torchvision==${TORCHVISION_VERSION}" \ "torchaudio==${TORCHAUDIO_VERSION}" \ - "apex==${APEX_VERSION}" \ - "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ + "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ rocm-sdk init - -### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better -RUN pip3 uninstall triton && \ - git clone https://github.com/triton-lang/triton.git && \ - cd triton && \ - git checkout c517f38c && \ - TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . - # Torch runtime deps that may not be published on the ROCm wheel index; # install them from PyPI afterwards. RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" -ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages -ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel -ENV ROCM_HOME=${ROCM_PATH} -ENV ROCM_SOURCE_DIR=${ROCM_PATH} -ENV ROCM_BIN=${ROCM_PATH}/bin -ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake -ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode -ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH -ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib -ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake -ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + # Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream # builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 5179f6ee8d74..259016458866 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -6,7 +6,6 @@ grpcio==1.78.0 grpcio-reflection==1.78.0 numba == 0.65.0 # Required for N-gram speculative decoding - # Dependencies for AMD GPUs datasets peft @@ -22,7 +21,7 @@ timm>=1.0.17 # amd-quark: required for Quark quantization on ROCm # To be consistent with test_quark.py amd-quark>=0.8.99 -tilelang==0.1.10 +#tilelang==0.1.10 # Required apache-tvm-ffi matching tilelang version apache-tvm-ffi==0.1.10 # Required for faster safetensors model loading From 038649fe115f83a1df146bbee075a011bb591639 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Mon, 29 Jun 2026 17:47:06 -0400 Subject: [PATCH 37/68] Add PyPi index and move triton patch in dockerfile Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm_base | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 9e5daf66a9f7..e3b4301c2b8f 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -46,7 +46,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++\ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ \ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ @@ -96,6 +96,32 @@ ARG TORCHVISION_VERSION ARG TORCHAUDIO_VERSION ARG ROCM_SDK_VERSION ARG APEX_VERSION +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + --extra-index-url https://pypi.org/simple \ + "torch==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "apex==${APEX_VERSION}" \ + "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel ENV ROCM_HOME=${ROCM_PATH} @@ -119,7 +145,7 @@ RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ "torchaudio==${TORCHAUDIO_VERSION}" \ "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ rocm-sdk init - + # Torch runtime deps that may not be published on the ROCm wheel index; # install them from PyPI afterwards. RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ @@ -159,6 +185,12 @@ 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} +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall -y triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . ### ### AMD SMI Build From 6ccba6f2f43e8fb16d2268022020bbb384e74cae Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Mon, 29 Jun 2026 17:57:56 -0400 Subject: [PATCH 38/68] Remove triton installation Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm_base | 8 -------- 1 file changed, 8 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index e3b4301c2b8f..a239bb9c986f 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -109,14 +109,6 @@ RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ rocm-sdk init -### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better -RUN pip3 uninstall triton && \ - git clone https://github.com/triton-lang/triton.git && \ - cd triton && \ - git checkout c517f38c && \ - TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . - - # Torch runtime deps that may not be published on the ROCm wheel index; # install them from PyPI afterwards. RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ From 05c638f3bd030a6bbffc010002f63b562c304f65 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Tue, 30 Jun 2026 10:33:02 -0400 Subject: [PATCH 39/68] Remove triton patch from rocm_base Signed-off-by: Jaden Mathias --- docker/Dockerfile.rocm_base | 6 ------ 1 file changed, 6 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index a239bb9c986f..50ee792381a3 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -177,12 +177,6 @@ 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} -### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better -RUN pip3 uninstall -y triton && \ - git clone https://github.com/triton-lang/triton.git && \ - cd triton && \ - git checkout c517f38c && \ - TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . ### ### AMD SMI Build From 6dd001ba51b2264ca14899e8e555e54d74960b8f Mon Sep 17 00:00:00 2001 From: jpvillam Date: Tue, 30 Jun 2026 12:11:24 -0500 Subject: [PATCH 40/68] Swap to cdna version check Signed-off-by: jpvillam --- .../test_transcription_validation_whisper.py | 4 +-- .../test_translation_validation.py | 6 ++-- .../attention/test_rocm_aiter_unified_attn.py | 8 ++--- tests/kernels/moe/test_routing.py | 4 +-- .../generation/test_granite_speech.py | 4 +-- .../models/quantization/test_bitsandbytes.py | 6 ++-- .../test_rocm_attention_backends_selection.py | 16 +++------- vllm/_aiter_ops.py | 24 ++++++++------ .../kernels/linear/scaled_mm/pytorch.py | 6 ++-- .../kernels/linear/scaled_mm/rocm.py | 6 ++-- .../fused_moe/experts/fused_batched_moe.py | 4 +-- .../quark/schemes/quark_w4a8_mxfp4_fp8.py | 4 +-- vllm/platforms/rocm.py | 31 +++++++++---------- vllm/v1/attention/backends/rocm_aiter_fa.py | 6 ++-- 14 files changed, 61 insertions(+), 68 deletions(-) diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index 34247782b14b..c50980de2543 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]: if current_platform.is_rocm(): backends: list[str | None] = [None, "TRITON_ATTN"] - from vllm.platforms.rocm import _on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - if _on_mi3or4(): + if get_cdna_version() > 2: backends.append("ROCM_AITER_UNIFIED_ATTN") return backends except Exception: diff --git a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py index de82dc3c4fc3..25f26404d0d9 100644 --- a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py +++ b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py @@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name): if "whisper" in model_name.lower(): try: - from vllm.platforms.rocm import _on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - if _on_mi3or4(): + if get_cdna_version() > 2: return {"backend": "ROCM_AITER_UNIFIED_ATTN"} except ImportError: logger.warning( - "Could not import _on_mi3or4 from rocm platform, " + "Could not check cdna version from rocm platform, " "falling back to TRITON_ATTN for Whisper." ) return {"backend": "TRITON_ATTN"} diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index 410c1d3ebab6..9d538dba43f4 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -15,15 +15,15 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -_SKIP_NON_MI3XX = True +_SKIP_NON_CDNA_2_PLUS = True if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - _SKIP_NON_MI3XX = not on_mi3or4() + _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2 pytestmark = [ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), - pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"), + pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"), ] NUM_Q_HEADS = 8 diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index c5081a8356ed..59e93092281e 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -23,9 +23,9 @@ def _is_aiter_capable() -> bool: if not current_platform.is_rocm(): return False try: - from vllm.platforms.rocm import _on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - return _on_mi3or4() + return get_cdna_version() > 2 except ImportError: return False diff --git a/tests/models/multimodal/generation/test_granite_speech.py b/tests/models/multimodal/generation/test_granite_speech.py index 48d32531d290..de4c0ad2327c 100644 --- a/tests/models/multimodal/generation/test_granite_speech.py +++ b/tests/models/multimodal/generation/test_granite_speech.py @@ -45,9 +45,9 @@ def vllm_to_hf_output( def granite_speech_attention_config(): """Return attention config for Granite Speech tests on ROCm.""" if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - if on_mi3or4(): + if get_cdna_version(): return {"backend": "ROCM_AITER_FA"} return {"backend": "TRITON_ATTN"} return None diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index 62ce36a3f980..06745ae697e9 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -21,11 +21,11 @@ from ..utils import check_embeddings_close, check_logprobs_close if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx9, on_gfx1250 + from vllm.platforms.rocm import on_cdna pytestmark = pytest.mark.skipif( - on_gfx9() or on_gfx1250(), - reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)", + on_cdna(), + reason="bitsandbytes not supported on CDNA (warp size 64 limitation)", ) models_4bit_to_test = [ diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 85ce15d2fb68..67bbdb0a99ea 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -28,16 +28,9 @@ def mock_vllm_config(): @pytest.fixture -def mock_on_gfx9(): - """Mock gfx9 arch detection to return True.""" - with patch("vllm.platforms.rocm.on_gfx9", return_value=True): - yield - - -@pytest.fixture -def mock_on_mi3or4(): - """Mock mi3or4 arch detection to return True.""" - with patch("vllm.platforms.rocm.on_mi3or4", return_value=True): +def mock_get_cdna_version(): + """Mock cdna version arch detection to return True.""" + with patch("vllm.platforms.rocm.get_cdna_version", return_value=3): yield @@ -111,8 +104,7 @@ def test_standard_attention_backend_selection( selected_backend, expected_backend_path, mock_vllm_config, - mock_on_gfx9, - mock_on_mi3or4, + mock_get_cdna_version, monkeypatch, ): """Test standard attention backend selection with various configurations.""" diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index afca102f763d..16142c4eb738 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -92,7 +92,7 @@ def should_custom_ar(self, inp: torch.Tensor) -> bool: ... def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. - Checks: platform (ROCm), device arch (gfx9 or gfx1250), and library existence. + Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence. Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled(). This function determines if aiter CAN be used, not if it SHOULD be used. @@ -106,9 +106,9 @@ def is_aiter_found_and_supported() -> bool: VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery. """ if current_platform.is_rocm() and IS_AITER_FOUND: - from vllm.platforms.rocm import on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - return on_mi3or4() + return get_cdna_version() > 2 return False @@ -224,9 +224,7 @@ def _rocm_aiter_fused_moe_triton_gemm_a4w4( m, _ = topk_ids.shape num_experts = w1.shape[0] - logits = torch.full( - (m, num_experts), -1e9, device=device, dtype=torch.float32 - ) + logits = torch.full((m, num_experts), -1e9, device=device, dtype=torch.float32) tid = topk_ids.long().clamp(min=0, max=num_experts - 1) logits.scatter_( 1, @@ -1893,16 +1891,22 @@ def is_fp8bmm_enabled(cls) -> bool: @classmethod @if_aiter_supported def is_fp4bmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_gfx950, on_gfx1250 + from vllm.platforms.rocm import get_cdna_version - return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() # TODO GFX1250: (on_gfx950() or on_gfx1250()) + return ( + cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4 + ) # TODO GFX1250: Swap to > 3 @classmethod @if_aiter_supported def is_linear_hipbmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - return cls.is_linear_enabled() and on_mi3or4() and cls._LINEAR_HIPBMM_ENABLED + return ( + cls.is_linear_enabled() + and (get_cdna_version() > 2) + and cls._LINEAR_HIPBMM_ENABLED + ) @classmethod @if_aiter_supported diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index dbb77db5077e..a04be2e097c0 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -104,10 +104,10 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - if not on_mi3or4(): - return False, "requires MI3xx." + if get_cdna_version() <= 2: + return False, "requires CDNA3+" if compute_capability is not None and compute_capability < 94: return False, "requires compute capability 94 and above." diff --git a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py index 446eabd5431c..688338b31188 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py @@ -79,10 +79,10 @@ def is_supported( if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_gfx12x, on_mi3or4 + from vllm.platforms.rocm import get_cdna_version - if not (on_mi3or4() or on_gfx12x()): - return False, "requires MI3or4 or gfx12x" + if get_cdna_version() <= 2: + return False, "requires CDNA capabilities greater than 2" if not envs.VLLM_ROCM_USE_SKINNY_GEMM: return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 520ccd295389..c4863c7d623f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -758,9 +758,9 @@ def _supports_quant_scheme( ) -> bool: p = current_platform if p.is_rocm(): - from vllm.platforms.rocm import on_gfx9, on_gfx1250 + from vllm.platforms.rocm import get_cdna_version - _rocm_support_fp8 = on_gfx9() or on_gfx1250() + _rocm_support_fp8 = get_cdna_version() > 2 else: _rocm_support_fp8 = False diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py index 7da4a76e6073..1ca7cbfb2fe5 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py @@ -64,9 +64,9 @@ def __init__( kernel_supported_gpu = False if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx950, on_gfx1250 + from vllm.platforms.rocm import get_cdna_version - kernel_supported_gpu = on_gfx950() or on_gfx1250() + kernel_supported_gpu = get_cdna_version() > 3 self.use_aiter_kernel = ( is_aiter_found_and_supported() diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index a857502fed78..062af2993933 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -226,21 +226,20 @@ def _get_gcn_arch() -> str: _ON_GFX1151 = "gfx1151" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) -_ON_GFX9 = any( - arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"] -) +_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH _ON_GFX1250 = "gfx1250" in _GCN_ARCH -_ON_MI3OR4 = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950", "gfx1250"]) # TODO GFX1250: Need a correct name for this _ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9, gfx1250"]) _ON_RDNA = any( - arch for arch in _GCN_ARCH + arch + for arch in _GCN_ARCH if (arch.startswith("gfx11") or arch.startswith("gfx12")) and arch != "gfx1250" ) -#TODO GFX1250: Use CDNA for MI3OR4 +# TODO GFX1250: Use CDNA for MI3OR4 + def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: """ @@ -314,7 +313,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: def on_gfx1x() -> bool: - return _ON_GFX1X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly + return _ON_GFX1X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly def on_gfx11() -> bool: @@ -330,7 +329,9 @@ def on_gfx1151() -> bool: def on_gfx12x() -> bool: - return _ON_GFX12X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly + return ( + _ON_GFX12X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly + ) def on_gfx1250() -> bool: @@ -341,10 +342,6 @@ def on_mi3xx() -> bool: return _ON_MI3XX -def on_mi3or4() -> bool: - return _ON_MI3OR4 - - def on_gfx9() -> bool: return _ON_GFX9 @@ -388,7 +385,7 @@ def get_cdna_version() -> int: envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - and on_mi3or4() + and get_cdna_version() > 2 ): os.environ["HIP_ONLINE_TUNING"] = "1" @@ -407,7 +404,7 @@ def use_rocm_custom_paged_attention( ) -> bool: # custom paged attn always supported on V0. On V1, requires sliding window # disabled due to observed numerical discrepancy. - if _ON_GFX9 or _ON_GFX1250: + if on_cdna(): return ( (sliding_window == 0 or sliding_window == (-1, -1)) and (qtype == torch.half or qtype == torch.bfloat16) @@ -694,12 +691,12 @@ def get_vit_attn_backend( from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_enabled() and (on_gfx9() or on_gfx1250()): + if rocm_aiter_ops.is_enabled() and on_cdna(): logger.info_once("Using AITER Flash Attention backend for ViT model.") return AttentionBackendEnum.ROCM_AITER_FA if ( - (on_gfx9() or on_gfx1250()) + on_cdna() and find_spec("flash_attn") is not None and (dtype == torch.float16 or dtype == torch.bfloat16) ): @@ -922,7 +919,7 @@ def supports_mx(cls) -> bool: @classmethod def supports_fp8(cls) -> bool: - return on_gfx9() or on_gfx12x() or on_gfx1250() # TODO GFX1250: We know this is redundant ATM + return on_cdna() or on_gfx12x() # TODO GFX1250: We know this is redundant ATM @classmethod def is_fp8_fnuz(cls) -> bool: diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index fa126a80de2e..cea954b47643 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -768,12 +768,12 @@ def get_kv_cache_shape( @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: - from vllm.platforms.rocm import on_mi3or4 + from vllm.platforms.rocm import get_cdna_version # DeviceCapability is currently created using torch.cuda.get_device_capability() - # which is known to be buggy on rocm systems. on_mi3or4 uses amd-smi which is + # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is # more reliable. - return on_mi3or4() + return get_cdna_version() > 2 @classmethod def supports_non_causal(cls) -> bool: From ff34a57d9d2846378a206c805345bc7f409fb6f1 Mon Sep 17 00:00:00 2001 From: Daniel Date: Tue, 30 Jun 2026 15:38:16 -0400 Subject: [PATCH 41/68] cleaning dockerfile Signed-off-by: Daniel --- docker/Dockerfile.rocm | 130 ++++++----------------------------------- 1 file changed, 19 insertions(+), 111 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index f01786b04602..e16770b5c5f1 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,24 +1,7 @@ # default base image ARG REMOTE_VLLM="0" ARG COMMON_WORKDIR=/app -ARG BASE_IMAGE=rocm/vllm-private:juan_455_npi_base -ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base -# NIC backend for MoRI RDMA support. -# By default (all), drivers and userspace libraries for all supported NIC types -# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. -# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: -# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only -# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only -# --build-arg NIC_BACKEND=none # Install nothing. -ARG NIC_BACKEND=all -# AMD AINIC apt repo settings -# Users can specify a custom version compatible with their host drivers. -# The default version has been tested with ioinic-dkms=25.11.1.001 -ARG AINIC_VERSION=1.117.3-hydra -ARG UBUNTU_CODENAME=jammy - -# Sccache configuration. Release builds use this today; CI can opt in when a -# shared S3-compatible cache backend is available. +ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types @@ -48,7 +31,6 @@ FROM ${BASE_IMAGE} AS base ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} -# Install build dependencies and utilities # Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ @@ -160,27 +142,25 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ COPY tools/install_protoc.sh /tmp/install_protoc.sh RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh -# Install rustup; the toolchain itself is pinned by rust-toolchain.toml. -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ - sh -s -- -y --profile minimal --default-toolchain none -ENV PATH="/root/.cargo/bin:${PATH}" - # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 ENV CARGO_NET_RETRY=10 ENV RUSTUP_MAX_RETRIES=10 +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt + # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while -# keeping the cache benefit. Copy the binary out so it persists into the image -# layer for later COPY --from=rust-build. +# keeping the cache benefit. Do not cache target/, because stale target metadata +# can outlive source updates across BuildKit cache reuse. RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ - --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ - && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \ - && test -x /tmp/vllm-rs + && bash build_rust.sh \ + && test -x vllm/vllm-rs # ----------------------- # vLLM native build stages @@ -200,6 +180,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ # pyproject.toml is bind-mounted in the RUN step so metadata-only changes do # not invalidate the expensive native build layer. COPY setup.py CMakeLists.txt ./ +COPY tools/build_rust.py tools/build_rust.py COPY cmake cmake/ COPY csrc csrc/ COPY vllm/envs.py vllm/envs.py @@ -231,9 +212,10 @@ ENV VLLM_TARGET_DEVICE=rocm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd vllm \ @@ -256,7 +238,6 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh / COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -287,14 +268,10 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ ibverbs-providers \ && rm -rf /var/lib/apt/lists/* -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system meson auditwheel patchelf tomlkit RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system meson auditwheel patchelf tomlkit RUN --mount=type=cache,target=/root/.cache/ccache \ - cd /usr/local/src && \ - RUN --mount=type=cache,target=/root/.cache/ccache \ cd /usr/local/src && \ git clone ${UCX_REPO} && \ cd ucx && \ @@ -302,7 +279,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ ./autogen.sh && \ mkdir build && cd build && \ CC="ccache gcc" CXX="ccache g++" \ - CC="ccache gcc" CXX="ccache g++" \ ../configure \ --prefix=/usr/local/ucx \ --enable-shared \ @@ -321,20 +297,16 @@ ENV PATH=/usr/local/ucx/bin:$PATH ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} RUN --mount=type=cache,target=/root/.cache/ccache \ - git clone ${RIXL_REPO} /opt/rixl && \ - RUN --mount=type=cache,target=/root/.cache/ccache \ git clone ${RIXL_REPO} /opt/rixl && \ cd /opt/rixl && \ git checkout ${RIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ - CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ --force-fallback-for=abseil-cpp \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ ninja -j$(nproc) && \ - ninja -j$(nproc) && \ ninja install # Generate RIXL wheel @@ -455,9 +427,10 @@ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR -# Drop the pre-built rust frontend binary into the source tree. setup.py -# detects it and ships it as-is, skipping the local cargo build. -COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ # Create /install directory for custom wheels RUN mkdir -p /install @@ -491,11 +464,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) # This ensures setuptools_scm sees clean repo state for version detection RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/uv \ cd vllm \ && uv pip install --system setuptools_scm regex \ - && uv pip install --system setuptools_scm regex \ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ && echo "Detected vLLM version: ${VLLM_VERSION}" \ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt @@ -532,26 +503,19 @@ RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ # Install dependencies using custom wheels from /install RUN --mount=type=cache,target=/root/.cache/uv \ - cd vllm \ - RUN --mount=type=cache,target=/root/.cache/uv \ cd vllm \ && echo "Building vLLM with custom wheels from /install" \ && uv pip install --system --find-links /install -r requirements/rocm.txt -&& uv pip install --system --find-links /install -r requirements/rocm.txt # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt # (setup.py auto-detects ccache/sccache in PATH) -# (setup.py auto-detects ccache/sccache in PATH) RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ && export CCACHE_BASEDIR="$PWD" \ - && export CCACHE_BASEDIR="$PWD" \ && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist -&& MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist FROM scratch AS export_vllm_wheel_release ARG COMMON_WORKDIR @@ -564,7 +528,6 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchc COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- @@ -588,17 +551,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ uv pip install --system /rixl_install/*.whl /deep_install/*.whl ---mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ - uv pip install --system /rixl_install/*.whl /deep_install/*.whl # Copy ROCShmem runtime libraries. COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem -# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. -RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - # Copy ROCShmem runtime libraries. - COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem - # RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ librdmacm1 \ @@ -607,24 +563,16 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ ibverbs-utils \ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ - pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ - libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* -# Install torchcodec from source for ROCm/torch ABI compatibility. # Install torchcodec from source for ROCm/torch ABI compatibility. COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=cache,target=/root/.cache/pip \ - --mount=type=cache,target=/root/.cache/torchcodec-wheels \ - bash /tmp/install_torchcodec.sh \ - RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/pip \ --mount=type=cache,target=/root/.cache/torchcodec-wheels \ bash /tmp/install_torchcodec.sh \ && rm /tmp/install_torchcodec.sh \ && apt-get clean && rm -rf /var/lib/apt/lists/* -&& apt-get clean && rm -rf /var/lib/apt/lists/* # Pre-install shared ROCm runtime dependencies. COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ @@ -650,7 +598,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \ --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ && rm /tmp/rocm-test-reqs.txt -# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # See: https://github.com/pytorch/pytorch/issues/169857 ENV MIOPEN_DEBUG_CONV_DIRECT=0 @@ -690,41 +637,6 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace # Copy in the v1 package (for python-only install test group). COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 -# Hide source under src/ so it won't shadow the installed package in tests. -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. -# See: https://github.com/ROCm/rocm-libraries/issues/6266 -ENV HSA_ENABLE_IPC_MODE_LEGACY=1 - -# ROCm profiler limits workaround. -RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf -ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" - -# Install vllm_test_utils in ci_base for ci_base + wheel parity. -COPY tests/vllm_test_utils /tmp/vllm_test_utils -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system /tmp/vllm_test_utils \ - && rm -rf /tmp/vllm_test_utils - -# ----------------------- -# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. -FROM ${CI_BASE_IMAGE} AS test -ARG COMMON_WORKDIR - -# Install the vLLM wheel (--no-deps: all deps already in ci_base). -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && uv pip install --system --no-deps *.whl - -# Store the vLLM wheel in the image for python-only install tests. -COPY --from=export_vllm /*.whl /opt/vllm-wheels/ - -WORKDIR /vllm-workspace -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace - -# Copy in the v1 package (for python-only install test group). -COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 - # Hide source under src/ so it won't shadow the installed package in tests. RUN mkdir src && mv vllm src/vllm @@ -744,7 +656,6 @@ RUN rm -f /usr/bin/sccache || true \ ENV SCCACHE_BUCKET= ENV SCCACHE_REGION= ENV SCCACHE_ENDPOINT= -ENV SCCACHE_ENDPOINT= ENV SCCACHE_S3_NO_CREDENTIALS= ENV SCCACHE_IDLE_TIMEOUT= @@ -793,9 +704,6 @@ ENV SAFETENSORS_FAST_GPU=1 # Performance environment variable. ENV HIP_FORCE_DEV_KERNARG=1 -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" @@ -815,4 +723,4 @@ CMD ["/bin/bash"] #Set entrypoint for vllm-openai official images FROM final AS vllm-openai -ENTRYPOINT ["vllm", "serve"] +ENTRYPOINT ["vllm", "serve"] \ No newline at end of file From 7ff9722b98ceed30345c02751f004bae8af0ffdc Mon Sep 17 00:00:00 2001 From: Daniel Date: Tue, 30 Jun 2026 16:47:51 -0400 Subject: [PATCH 42/68] fixing dockerfile linting Signed-off-by: Daniel --- docker/Dockerfile.rocm | 262 ++++++++++++++++++------------------ docker/Dockerfile.rocm_base | 22 +-- 2 files changed, 133 insertions(+), 151 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index e16770b5c5f1..9cc12e1142e5 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -44,9 +44,9 @@ RUN --mount=type=cache,target=/root/.cache/pip \ # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE RUN if [ "$USE_SCCACHE" != "1" ]; then \ - apt-get purge -y sccache || true; \ - python3 -m pip uninstall -y sccache || true; \ - rm -f "$(which sccache)" || true; \ + apt-get purge -y sccache || true; \ + python3 -m pip uninstall -y sccache || true; \ + rm -f "$(which sccache)" || true; \ fi # Install UV — download first, then run, so a curl failure is not masked by the pipe @@ -80,21 +80,21 @@ ARG SCCACHE_BUCKET_NAME ARG SCCACHE_REGION_NAME ARG SCCACHE_S3_NO_CREDENTIALS RUN if [ "$USE_SCCACHE" = "1" ]; then \ - if command -v sccache >/dev/null 2>&1; then \ - echo "sccache already installed, skipping installation"; \ - sccache --version; \ - else \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi; \ + if command -v sccache >/dev/null 2>&1; then \ + echo "sccache already installed, skipping installation"; \ + sccache --version; \ + else \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi; \ fi # Set sccache environment variables only when USE_SCCACHE=1 @@ -119,12 +119,12 @@ ARG VLLM_BRANCH="455_wip" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ - && cd vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git checkout FETCH_HEAD \ - && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ - git remote add upstream "https://github.com/vllm-project/vllm.git" \ - && git fetch upstream ; fi + && cd vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git checkout FETCH_HEAD \ + && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ + git remote add upstream "https://github.com/vllm-project/vllm.git" \ + && git fetch upstream ; fi FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm # ----------------------- @@ -136,7 +136,7 @@ ARG COMMON_WORKDIR # protoc is used by tonic-build/prost-build. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - ca-certificates curl unzip \ + ca-certificates curl unzip \ && rm -rf /var/lib/apt/lists/* COPY tools/install_protoc.sh /tmp/install_protoc.sh @@ -280,16 +280,16 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ mkdir build && cd build && \ CC="ccache gcc" CXX="ccache g++" \ ../configure \ - --prefix=/usr/local/ucx \ - --enable-shared \ - --disable-static \ - --disable-doxygen-doc \ - --enable-optimizations \ - --enable-devel-headers \ - --with-rocm=${ROCM_PATH} \ - --with-verbs \ - --with-dm \ - --enable-mt && \ + --prefix=/usr/local/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-devel-headers \ + --with-rocm=${ROCM_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt && \ make -j$(nproc) && \ make install @@ -302,9 +302,9 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ git checkout ${RIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ - --force-fallback-for=abseil-cpp \ - -Ducx_path=${UCX_HOME} \ - -Drocm_path=${ROCM_PATH} && \ + --force-fallback-for=abseil-cpp \ + -Ducx_path=${UCX_HOME} \ + -Drocm_path=${ROCM_PATH} && \ cd build && \ ninja -j$(nproc) && \ ninja install @@ -314,18 +314,18 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ # that are not shipped in the wheel and vary across base images. RUN cd /opt/rixl && \ sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ - contrib/build-wheel.sh && \ + contrib/build-wheel.sh && \ # The wheel build re-runs meson via meson-python; force the bundled abseil sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ - pyproject.toml && \ + pyproject.toml && \ grep -q 'force-fallback-for' pyproject.toml && \ mkdir -p /app/install && \ _ucx_install_dir=${UCX_HOME} \ ./contrib/build-wheel.sh \ - --output-dir /app/install \ - --rocm-dir ${ROCM_PATH} \ - --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ - --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + --output-dir /app/install \ + --rocm-dir ${ROCM_PATH} \ + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ + --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins # ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not # invalidate the slow ROCShmem build. @@ -340,16 +340,16 @@ ENV ROCSHMEM_DIR=/opt/rocshmem RUN --mount=type=cache,target=/root/.cache/ccache \ git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ - && cd rocm-systems \ - && git sparse-checkout set --cone projects/rocshmem \ - && git checkout ${ROCSHMEM_BRANCH} \ - && mkdir -p projects/rocshmem/build \ - && cd projects/rocshmem/build \ - && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ + && git checkout ${ROCSHMEM_BRANCH} \ + && mkdir -p projects/rocshmem/build \ + && cd projects/rocshmem/build \ + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ bash ../scripts/build_configs/all_backends \ - -DROCM_PATH=${ROCM_PATH} \ - -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ - -DUSE_EXTERNAL_MPI=OFF + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF # DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. FROM build_rocshmem AS build_deepep @@ -361,10 +361,10 @@ ARG DEEPEP_NIC="cx7" # DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. RUN --mount=type=cache,target=/root/.cache/ccache \ export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ - && git clone ${DEEPEP_REPO} \ - && cd DeepEP \ - && git checkout ${DEEPEP_BRANCH} \ - && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + && git clone ${DEEPEP_REPO} \ + && cd DeepEP \ + && git checkout ${DEEPEP_BRANCH} \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do # not force users to rebuild the long-lived Dockerfile.rocm_base image. @@ -373,52 +373,52 @@ ARG NIC_BACKEND ARG AINIC_VERSION ARG UBUNTU_CODENAME RUN /bin/bash -lc 'set -euo pipefail; \ - \ - install_ainic() { \ - apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ - rm -rf /var/lib/apt/lists/*; \ - mkdir -p /etc/apt/keyrings; \ - curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ - echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ - > /etc/apt/sources.list.d/amdainic.list; \ - apt-get update && apt-get install -y --no-install-recommends \ - libionic-dev \ - ionic-common \ - ; \ - rm -rf /var/lib/apt/lists/*; \ - }; \ - \ - # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ - # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ - install_bnxt() { \ - install -m 0755 -d /etc/apt/keyrings; \ - curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ - -o /etc/apt/keyrings/broadcom-nic.asc; \ - chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ - echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ - > /etc/apt/sources.list.d/broadcom-nic.list; \ - apt-get update && apt-get install -y --no-install-recommends \ - bnxt-rocelib=235.2.86.0 \ - ; \ - cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ - ldconfig; \ - rm -rf /var/lib/apt/lists/*; \ - }; \ - \ - echo "[MORI] Install MoRI proxy deps"; \ - pip install --quiet --ignore-installed blinker && \ - pip install --quiet quart msgpack aiohttp pyzmq; \ - echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ - \ - # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ - # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ - case "${NIC_BACKEND}" in \ - none) ;; \ - all) install_ainic; install_bnxt ;; \ - ainic) install_ainic ;; \ - bnxt) install_bnxt ;; \ - *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ - esac' + \ + install_ainic() { \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ + rm -rf /var/lib/apt/lists/*; \ + mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ + # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ + install_bnxt() { \ + install -m 0755 -d /etc/apt/keyrings; \ + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ + -o /etc/apt/keyrings/broadcom-nic.asc; \ + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ + > /etc/apt/sources.list.d/broadcom-nic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + bnxt-rocelib=235.2.86.0 \ + ; \ + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ + ldconfig; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + \ + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ + # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ + case "${NIC_BACKEND}" in \ + none) ;; \ + all) install_ainic; install_bnxt ;; \ + ainic) install_ainic ;; \ + bnxt) install_bnxt ;; \ + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ + esac' # ----------------------- # vLLM wheel release build stage (for building distributable wheels) @@ -443,22 +443,22 @@ COPY docker/context/base-wheels/ /tmp/base-wheels/ # If there are not wheels found there, we are not building for a wheel release. # So we exit with an error. To skip this stage. RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ - echo "Found custom wheels - copying to /install"; \ - cp /tmp/base-wheels/*.whl /install/ && \ - echo "Copied custom wheels:"; \ - ls -lh /install/; \ + echo "Found custom wheels - copying to /install"; \ + cp /tmp/base-wheels/*.whl /install/ && \ + echo "Copied custom wheels:"; \ + ls -lh /install/; \ else \ - echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ - echo "Wheel releases require pre-built ROCm wheels."; \ - exit 1; \ + echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ + echo "Wheel releases require pre-built ROCm wheels."; \ + exit 1; \ fi # GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) # This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm ARG GIT_REPO_CHECK=0 RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ - echo "Running repository checks..."; \ - cd vllm && bash tools/check_repo.sh; \ + echo "Running repository checks..."; \ + cd vllm && bash tools/check_repo.sh; \ fi # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) @@ -477,22 +477,22 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \ RUN echo "Checking for git-based packages in requirements files..." \ && echo "Checking common.txt for git-based packages:" \ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ - echo "ERROR: Git-based packages found in common.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in common.txt"; \ - fi \ + echo "ERROR: Git-based packages found in common.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in common.txt"; \ + fi \ && echo "Checking rocm.txt for git-based packages:" \ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ - echo "ERROR: Git-based packages found in rocm.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in rocm.txt"; \ - fi \ + echo "ERROR: Git-based packages found in rocm.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in rocm.txt"; \ + fi \ && echo "All requirements files are clean - no git-based packages found" # Pin vLLM dependencies to exact versions of custom ROCm wheels @@ -545,7 +545,7 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ mkdir -p build && cd build && \ cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ - fi +fi # Install RIXL + DeepEP wheels. RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ @@ -595,7 +595,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ && test -n "${FASTSAFETENSORS_REQ}" \ && python3 -m pip install --force-reinstall --no-deps \ - --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ && rm /tmp/rocm-test-reqs.txt # Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. @@ -662,9 +662,9 @@ ENV SCCACHE_IDLE_TIMEOUT= # Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. # Manually remove it so that later steps of numpy upgrade can continue RUN case "$(which python3)" in \ - *"/opt/conda/envs/py_3.9"*) \ - rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ - *) ;; esac + *"/opt/conda/envs/py_3.9"*) \ + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ + *) ;; esac RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system --upgrade huggingface-hub[cli] diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 50ee792381a3..9e5daf66a9f7 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -46,7 +46,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++\ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ @@ -96,24 +96,6 @@ ARG TORCHVISION_VERSION ARG TORCHAUDIO_VERSION ARG ROCM_SDK_VERSION ARG APEX_VERSION -# torch/torchvision/torchaudio must be pinned to mutually-consistent builds -# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm -# sdk version is derived from torch's own dependency pin unless overridden, -# which keeps the set consistent and avoids pip backtracking. -RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ - --extra-index-url https://pypi.org/simple \ - "torch==${TORCH_VERSION}" \ - "torchvision==${TORCHVISION_VERSION}" \ - "torchaudio==${TORCHAUDIO_VERSION}" \ - "apex==${APEX_VERSION}" \ - "rocm[libraries,devel]==${ROCM_SDK_VERSION}" && \ - rocm-sdk init - -# Torch runtime deps that may not be published on the ROCm wheel index; -# install them from PyPI afterwards. -RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ - "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" - ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel ENV ROCM_HOME=${ROCM_PATH} @@ -137,7 +119,7 @@ RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ "torchaudio==${TORCHAUDIO_VERSION}" \ "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ rocm-sdk init - + # Torch runtime deps that may not be published on the ROCm wheel index; # install them from PyPI afterwards. RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ From 7cc3c73eed5bf35416d75cb3bcaa7ac6d3fd127c Mon Sep 17 00:00:00 2001 From: jpvillam Date: Tue, 30 Jun 2026 16:01:38 -0500 Subject: [PATCH 43/68] Clean up Signed-off-by: jpvillam --- .buildkite/release-pipeline.yaml | 693 ++++++++++-------- .buildkite/scripts/annotate-rocm-release.sh | 15 +- benchmarks/curl_gpt.sh | 7 - benchmarks/gpt_fp4_serve.sh | 7 - docker/Dockerfile.rocm_1250 | 114 --- docker/Dockerfile.rocm_base | 20 +- .../test_rocm_attention_backends_selection.py | 6 +- vllm/_aiter_ops.py | 4 +- .../kernels/linear/scaled_mm/aiter.py | 2 +- vllm/v1/sample/ops/topk_topp_sampler.py | 4 +- 10 files changed, 412 insertions(+), 460 deletions(-) delete mode 100644 benchmarks/curl_gpt.sh delete mode 100644 benchmarks/gpt_fp4_serve.sh delete mode 100644 docker/Dockerfile.rocm_1250 diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 7029c39f994b..897c98145341 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -38,14 +38,13 @@ steps: depends_on: ~ id: build-wheel-arm64-cuda-12-9 agents: - queue: arm64_cpu_queue_postmerge + queue: arm64_cpu_queue_release commands: - # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: - # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 - - "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=12.9.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --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" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -53,14 +52,13 @@ steps: depends_on: ~ id: build-wheel-arm64-cuda-13-0 agents: - queue: arm64_cpu_queue_postmerge + queue: arm64_cpu_queue_release commands: - # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here: - # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7 - - "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.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-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 manylinux_2_35" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -68,12 +66,13 @@ steps: depends_on: ~ id: build-wheel-arm64-cpu agents: - queue: arm64_cpu_queue_postmerge + queue: arm64_cpu_queue_release commands: - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." - "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 manylinux_2_35" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -81,12 +80,13 @@ steps: depends_on: ~ id: build-wheel-x86-cuda-12-9 agents: - queue: cpu_queue_postmerge + 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=12.9.1 --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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --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 manylinux_2_31" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -94,12 +94,13 @@ steps: depends_on: ~ id: build-wheel-x86-cuda-13-0 agents: - queue: cpu_queue_postmerge + 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.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-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 manylinux_2_35" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" @@ -107,80 +108,265 @@ steps: depends_on: ~ id: build-wheel-x86-cpu agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ." - "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 manylinux_2_35" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"' env: DOCKER_BUILDKIT: "1" + - label: "Generate and upload wheel indices" + depends_on: "build-wheels" + allow_dependency_failure: true + agents: + queue: cpu_queue_release + commands: + - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + + - block: "Unblock to build release Docker images" + depends_on: ~ + key: block-build-release-images + if: build.env("NIGHTLY") != "1" + - group: "Build release Docker images" key: "build-release-images" + depends_on: block-build-release-images + allow_dependency_failure: true steps: - - label: "Build release image - x86_64 - CUDA 12.9" + - label: "Build release image - x86_64 - CUDA 13.0" depends_on: ~ id: build-release-image-x86 agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "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=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh) \ + --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 INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" # re-tag to default image tag and push, just in case arm64 build fails - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"' - - label: "Build release image - aarch64 - CUDA 12.9" + - label: "Build release image - aarch64 - CUDA 13.0" depends_on: ~ id: build-release-image-arm64 agents: - queue: arm64_cpu_queue_postmerge + queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "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=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ." + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh) \ + --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 INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"' - - label: "Build release image - x86_64 - CUDA 13.0" + - label: "Build release image - x86_64 - CUDA 12.9" depends_on: ~ - id: build-release-image-x86-cuda-13-0 + id: build-release-image-x86-cuda-12-9 agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "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.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \ + --build-arg max_jobs=16 \ + --build-arg USE_SCCACHE=1 \ + --build-arg GIT_REPO_CHECK=1 \ + --build-arg CUDA_VERSION=12.9.1 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" # re-tag to default image tag and push, just in case arm64 build fails - - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"' - - label: "Build release image - aarch64 - CUDA 13.0" + - label: "Build release image - aarch64 - CUDA 12.9" + depends_on: ~ + id: build-release-image-arm64-cuda-12-9 + agents: + queue: arm64_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \ + --build-arg max_jobs=16 \ + --build-arg USE_SCCACHE=1 \ + --build-arg GIT_REPO_CHECK=1 \ + --build-arg CUDA_VERSION=12.9.1 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"' + + - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04" + depends_on: ~ + id: build-release-image-x86-ubuntu2404 + agents: + queue: cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \ + --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 UBUNTU_VERSION=24.04 \ + --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"' + + - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04" + depends_on: ~ + id: build-release-image-arm64-ubuntu2404 + agents: + queue: arm64_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \ + --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 UBUNTU_VERSION=24.04 \ + --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"' + + - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04" depends_on: ~ - id: build-release-image-arm64-cuda-13-0 + id: build-release-image-x86-cuda-12-9-ubuntu2404 agents: - queue: arm64_cpu_queue_postmerge + queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark - - "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.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ." - - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \ + --build-arg max_jobs=16 \ + --build-arg USE_SCCACHE=1 \ + --build-arg GIT_REPO_CHECK=1 \ + --build-arg CUDA_VERSION=12.9.1 \ + --build-arg UBUNTU_VERSION=24.04 \ + --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' + + - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04" + depends_on: ~ + id: build-release-image-arm64-cuda-12-9-ubuntu2404 + agents: + queue: arm64_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \ + --build-arg max_jobs=16 \ + --build-arg USE_SCCACHE=1 \ + --build-arg GIT_REPO_CHECK=1 \ + --build-arg CUDA_VERSION=12.9.1 \ + --build-arg UBUNTU_VERSION=24.04 \ + --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \ + --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \ + --build-arg INSTALL_KV_CONNECTORS=true \ + --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \ + --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \ + --target vllm-openai \ + --progress plain \ + -f docker/Dockerfile . + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"' - block: "Build release image for x86_64 CPU" key: block-cpu-release-image-build depends_on: ~ - label: "Build release image - x86_64 - CPU" + key: build-cpu-release-image-x86 depends_on: - block-cpu-release-image-build - input-release-version agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest" - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"' env: DOCKER_BUILDKIT: "1" @@ -189,61 +375,82 @@ steps: depends_on: ~ - label: "Build release image - arm64 - CPU" - depends_on: + key: build-cpu-release-image-arm64 + depends_on: - block-arm64-cpu-release-image-build - input-release-version agents: - queue: arm64_cpu_queue_postmerge + queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ." - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest" - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"' env: DOCKER_BUILDKIT: "1" - group: "Publish release images" key: "publish-release-images" steps: - - label: "Create multi-arch manifest - CUDA 12.9" + - label: "Create multi-arch manifest - CUDA 13.0" depends_on: - build-release-image-x86 - build-release-image-arm64 id: create-multi-arch-manifest agents: - queue: small_cpu_queue_postmerge + queue: small_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend" - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"' - - label: "Annotate release workflow - CUDA 12.9" + - label: "Create multi-arch manifest - CUDA 12.9" depends_on: - - create-multi-arch-manifest - id: annotate-release-workflow + - build-release-image-x86-cuda-12-9 + - build-release-image-arm64-cuda-12-9 + id: create-multi-arch-manifest-cuda-12-9 agents: - queue: small_cpu_queue_postmerge + queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/annotate-release.sh" + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"' - - label: "Create multi-arch manifest - CUDA 13.0" + - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04" + depends_on: + - build-release-image-x86-ubuntu2404 + - build-release-image-arm64-ubuntu2404 + id: create-multi-arch-manifest-ubuntu2404 + agents: + queue: small_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"' + + - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04" depends_on: - - build-release-image-x86-cuda-13-0 - - build-release-image-arm64-cuda-13-0 - id: create-multi-arch-manifest-cuda-13-0 + - build-release-image-x86-cuda-12-9-ubuntu2404 + - build-release-image-arm64-cuda-12-9-ubuntu2404 + id: create-multi-arch-manifest-cuda-12-9-ubuntu2404 agents: - queue: small_cpu_queue_postmerge + queue: small_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend" - - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"' - label: "Publish nightly multi-arch image to DockerHub" depends_on: - create-multi-arch-manifest if: build.env("NIGHTLY") == "1" agents: - queue: small_cpu_queue_postmerge + queue: small_cpu_queue_release commands: - "bash .buildkite/scripts/push-nightly-builds.sh" # Clean up old nightly builds (keep only last 14) @@ -256,16 +463,16 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" - - label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0" + - label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9" depends_on: - - create-multi-arch-manifest-cuda-13-0 + - create-multi-arch-manifest-cuda-12-9 if: build.env("NIGHTLY") == "1" agents: - queue: small_cpu_queue_postmerge + queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/push-nightly-builds.sh cu130" + - "bash .buildkite/scripts/push-nightly-builds.sh cu129" # Clean up old nightly builds (keep only last 14) - - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-" plugins: - docker-login#v3.0.0: username: vllmbot @@ -274,24 +481,6 @@ steps: DOCKER_BUILDKIT: "1" DOCKERHUB_USERNAME: "vllmbot" - - group: "Publish wheels" - key: "publish-wheels" - steps: - - block: "Confirm update release wheels to PyPI (experimental, use with caution)?" - key: block-upload-release-wheels - depends_on: - - input-release-version - - build-wheels - - - label: "Upload release wheels to PyPI" - depends_on: - - block-upload-release-wheels - id: upload-release-wheels - agents: - queue: small_cpu_queue_postmerge - commands: - - "bash .buildkite/scripts/upload-release-wheels-pypi.sh" - # ============================================================================= # ROCm Release Pipeline (x86_64 only) # ============================================================================= @@ -300,184 +489,112 @@ steps: # To build a specific version, trigger the build from that branch/tag. # # Environment variables for ROCm builds (set via Buildkite UI or schedule): - # ROCM_PYTHON_VERSION: Python version (default: 3.12) - # PYTORCH_ROCM_ARCH: GPU architectures (default: gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151) - # ROCM_UPLOAD_WHEELS: Upload to S3 (default: false for nightly, true for releases) - # ROCM_FORCE_REBUILD: Force rebuild base wheels, ignore S3 cache (default: false) # # Note: ROCm version is determined by BASE_IMAGE in docker/Dockerfile.rocm_base - # (currently rocm/dev-ubuntu-22.04:7.1-complete) # # ============================================================================= - # ROCm Input Step - Collect build configuration (manual trigger only) - - input: "ROCm Wheel Release Build Configuration" - key: input-rocm-config - depends_on: ~ - if: build.source == "ui" - fields: - - text: "Python Version" - key: "rocm-python-version" - default: "3.12" - hint: "Python version (e.g., 3.12)" - - text: "GPU Architectures" - key: "rocm-pytorch-rocm-arch" - default: "gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" - hint: "Semicolon-separated GPU architectures" - - select: "Upload Wheels to S3" - key: "rocm-upload-wheels" - default: "true" - options: - - label: "No - Build only (nightly/dev)" - value: "false" - - label: "Yes - Upload to S3 (release)" - value: "true" - - select: "Force Rebuild Base Wheels" - key: "rocm-force-rebuild" - default: "false" - hint: "Ignore S3 cache and rebuild base wheels from scratch" - options: - - label: "No - Use cached wheels if available" - value: "false" - - label: "Yes - Rebuild even if cache exists" - value: "true" - # ROCm Job 1: Build ROCm Base Wheels (with S3 caching) - - label: ":rocm: Build ROCm Base Wheels" + - label: ":rocm: Build ROCm Base Image & Wheels" id: build-rocm-base-wheels - depends_on: - - step: input-rocm-config - allow_failure: true # Allow failure so non-UI builds can proceed (input step is skipped) + depends_on: ~ agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - # Set configuration and check cache - | set -euo pipefail - # Get values from meta-data (set by input step) or use defaults - PYTHON_VERSION="$$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo '')" - export PYTHON_VERSION="$${PYTHON_VERSION:-3.12}" - - PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" - export PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" - - # Check for force rebuild flag - ROCM_FORCE_REBUILD="$${ROCM_FORCE_REBUILD:-}" - if [ -z "$${ROCM_FORCE_REBUILD}" ]; then - ROCM_FORCE_REBUILD="$$(buildkite-agent meta-data get rocm-force-rebuild 2>/dev/null || echo '')" - fi + # Generate cache key + CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base" echo "========================================" - echo "ROCm Base Wheels Build Configuration" + echo "ROCm Base Build Configuration" echo "========================================" - echo " PYTHON_VERSION: $${PYTHON_VERSION}" - echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}" - echo " ROCM_FORCE_REBUILD: $${ROCM_FORCE_REBUILD:-false}" + echo " CACHE_KEY: $${CACHE_KEY}" + echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}" echo "========================================" + + # Login to ECR + aws ecr-public get-login-password --region us-east-1 | \ + docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 + + IMAGE_EXISTS=false + WHEELS_EXIST=false + + # Check ECR for Docker image - # Save resolved config for later jobs - buildkite-agent meta-data set "rocm-python-version" "$${PYTHON_VERSION}" - buildkite-agent meta-data set "rocm-pytorch-rocm-arch" "$${PYTORCH_ROCM_ARCH}" - - # Check S3 cache for pre-built wheels - CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key) - CACHE_PATH=$$(.buildkite/scripts/cache-rocm-base-wheels.sh path) - echo "" - echo "Cache key: $${CACHE_KEY}" - echo "Cache path: $${CACHE_PATH}" - - # Save cache key for downstream jobs - buildkite-agent meta-data set "rocm-cache-key" "$${CACHE_KEY}" - - CACHE_STATUS="miss" - if [ "$${ROCM_FORCE_REBUILD}" != "true" ]; then - CACHE_STATUS=$$(.buildkite/scripts/cache-rocm-base-wheels.sh check) - else - echo "Force rebuild requested, skipping cache check" + if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then + IMAGE_EXISTS=true + echo "ECR image cache HIT" + fi + + # Check S3 for wheels + WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check) + if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then + WHEELS_EXIST=true + echo "S3 wheels cache HIT" fi - if [ "$${CACHE_STATUS}" = "hit" ]; then + + # Scenario 1: Both cached (best case) + if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then echo "" - echo "CACHE HIT! Downloading pre-built wheels..." + echo "FULL CACHE HIT - Reusing both image and wheels" echo "" - .buildkite/scripts/cache-rocm-base-wheels.sh download - - # Set the S3 path for the cached Docker image (for Job 2 to download) - S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}" - buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" - # Mark that we used cache (for Docker image handling) - buildkite-agent meta-data set "rocm-used-cache" "true" - - echo "" - echo "Cache download complete. Skipping Docker build." - echo "Docker image will be downloaded from: $${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" + # Download wheels + .buildkite/scripts/cache-rocm-base-wheels.sh download + + # Save ECR tag for downstream jobs + buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" + + # Scenario 2: Full rebuild needed else echo "" - echo "CACHE MISS. Building from scratch..." + echo " CACHE MISS - Building from scratch..." echo "" - - # Build full base image (for later vLLM build) + + # Build full base image and push to ECR DOCKER_BUILDKIT=1 docker buildx build \ --file docker/Dockerfile.rocm_base \ - --tag rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} \ - --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ - --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \ + --tag "$${ECR_CACHE_TAG}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ - --load \ + --push \ . - - # Build debs_wheel_release stage for wheel extraction + + # Build wheel extraction stage DOCKER_BUILDKIT=1 docker buildx build \ --file docker/Dockerfile.rocm_base \ --tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \ --target debs_wheel_release \ - --build-arg PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ - --build-arg PYTHON_VERSION="$${PYTHON_VERSION}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ --load \ . - - # Extract wheels from Docker image + + # Extract and upload wheels mkdir -p artifacts/rocm-base-wheels - container_id=$$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) - docker cp $${container_id}:/app/debs/. artifacts/rocm-base-wheels/ - docker rm $${container_id} - echo "Extracted base wheels:" - ls -lh artifacts/rocm-base-wheels/ - - # Upload wheels to S3 cache for future builds - echo "" - echo "Uploading wheels to S3 cache..." + cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}) + docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/ + docker rm $${cid} + .buildkite/scripts/cache-rocm-base-wheels.sh upload - # Export base Docker image for reuse in vLLM build - mkdir -p artifacts/rocm-docker-image - docker save rocm/vllm-dev:base-$${BUILDKITE_BUILD_NUMBER} | gzip > artifacts/rocm-docker-image/rocm-base-image.tar.gz - echo "Docker image size:" - ls -lh artifacts/rocm-docker-image/ - - # Upload large Docker image to S3 (also cached by cache key) - S3_ARTIFACT_PATH="s3://$${S3_BUCKET}/rocm/cache/$${CACHE_KEY}" - echo "Uploading Docker image to $${S3_ARTIFACT_PATH}/" - aws s3 cp artifacts/rocm-docker-image/rocm-base-image.tar.gz "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" - - # Save the S3 path for downstream jobs - buildkite-agent meta-data set "rocm-docker-image-s3-path" "$${S3_ARTIFACT_PATH}/rocm-base-image.tar.gz" - - # Mark that we did NOT use cache - buildkite-agent meta-data set "rocm-used-cache" "false" - + # Cache base docker image to ECR + docker push "$${ECR_CACHE_TAG}" + + buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}" + echo "" - echo "Build complete. Wheels cached for future builds." + echo " Build complete - Image and wheels cached" fi + artifact_paths: - "artifacts/rocm-base-wheels/*.whl" env: @@ -491,7 +608,7 @@ steps: - step: build-rocm-base-wheels allow_failure: false agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release timeout_in_minutes: 180 commands: # Download artifacts and prepare Docker image @@ -521,31 +638,25 @@ steps: echo "Downloading wheel artifacts from current build" buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . - # Download Docker image from S3 (too large for Buildkite artifacts) - DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')" - if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then - echo "ERROR: rocm-docker-image-s3-path metadata not found" + # Get ECR image tag from metadata (set by build-rocm-base-wheels) + ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" + if [ -z "$${ECR_IMAGE_TAG}" ]; then + echo "ERROR: rocm-base-image-tag metadata not found" echo "This should have been set by the build-rocm-base-wheels job" exit 1 fi - echo "Downloading Docker image from $${DOCKER_IMAGE_S3_PATH}" - mkdir -p artifacts/rocm-docker-image - aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz - - # Load base Docker image and capture the tag - echo "Loading base Docker image..." - LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load) - echo "$${LOAD_OUTPUT}" - # Extract the actual loaded image tag from "Loaded image: " output - # This avoids picking up stale images (like rocm/vllm-dev:nightly) already on the agent - BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //') - if [ -z "$${BASE_IMAGE_TAG}" ]; then - echo "ERROR: Failed to extract image tag from docker load output" - echo "Load output was: $${LOAD_OUTPUT}" - exit 1 - fi - echo "Loaded base image: $${BASE_IMAGE_TAG}" - + + echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" + + # Login to ECR + aws ecr-public get-login-password --region us-east-1 | \ + docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 + + # Pull base Docker image from ECR + docker pull "$${ECR_IMAGE_TAG}" + + echo "Loaded base image: $${ECR_IMAGE_TAG}" + # Prepare base wheels for Docker build context mkdir -p docker/context/base-wheels touch docker/context/base-wheels/.keep @@ -553,16 +664,11 @@ steps: echo "Base wheels for vLLM build:" ls -lh docker/context/base-wheels/ - # Get GPU architectures from meta-data - PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" - PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" - echo "========================================" echo "Building vLLM wheel with:" echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" - echo " PYTORCH_ROCM_ARCH: $${PYTORCH_ROCM_ARCH}" - echo " BASE_IMAGE: $${BASE_IMAGE_TAG}" + echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" echo "========================================" # Build vLLM wheel using local checkout (REMOTE_VLLM=0) @@ -570,8 +676,7 @@ steps: --file docker/Dockerfile.rocm \ --target export_vllm_wheel_release \ --output type=local,dest=rocm-dist \ - --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \ - --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ + --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ --build-arg REMOTE_VLLM=0 \ --build-arg GIT_REPO_CHECK=1 \ --build-arg USE_SCCACHE=1 \ @@ -579,10 +684,8 @@ steps: --build-arg SCCACHE_REGION_NAME=us-west-2 \ --build-arg SCCACHE_S3_NO_CREDENTIALS=0 \ . - echo "Built vLLM wheel:" ls -lh rocm-dist/*.whl - # Copy wheel to artifacts directory mkdir -p artifacts/rocm-vllm-wheel cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/ @@ -601,35 +704,13 @@ steps: - step: build-rocm-vllm-wheel allow_failure: false agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release timeout_in_minutes: 60 commands: # Download all wheel artifacts and run upload - | set -euo pipefail - # Check if upload is enabled (from env var, meta-data, or release branch) - ROCM_UPLOAD_WHEELS="$${ROCM_UPLOAD_WHEELS:-}" - if [ -z "$${ROCM_UPLOAD_WHEELS}" ]; then - # Try to get from meta-data (input form) - ROCM_UPLOAD_WHEELS="$$(buildkite-agent meta-data get rocm-upload-wheels 2>/dev/null || echo '')" - fi - - echo "========================================" - echo "Upload check:" - echo " ROCM_UPLOAD_WHEELS: $${ROCM_UPLOAD_WHEELS}" - echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}" - echo "========================================" - - # Skip upload if not enabled - if [ "$${ROCM_UPLOAD_WHEELS}" != "true" ]; then - echo "Skipping S3 upload (ROCM_UPLOAD_WHEELS != true, NIGHTLY != 1, not a release branch)" - echo "To enable upload, set 'Upload Wheels to S3' to 'Yes' in the build configuration" - exit 0 - fi - - echo "Upload enabled, proceeding..." - # Download artifacts from current build echo "Downloading artifacts from current build" buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . @@ -645,12 +726,9 @@ steps: - label: ":memo: Annotate ROCm wheel release" id: annotate-rocm-release depends_on: - - step: upload-rocm-wheels - allow_failure: true - - step: input-release-version - allow_failure: true + - upload-rocm-wheels agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - "bash .buildkite/scripts/annotate-rocm-release.sh" env: @@ -667,61 +745,60 @@ steps: depends_on: block-generate-root-index-rocm-wheels id: generate-root-index-rocm-wheels agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release commands: - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" env: S3_BUCKET: "vllm-wheels" VARIANT: "rocm723" - # ROCm Job 5: Build ROCm Release Docker Image + # ROCm Job 6: Build ROCm Release Docker Image - label: ":docker: Build release image - x86_64 - ROCm" id: build-rocm-release-image depends_on: + - step: block-build-release-images + allow_failure: true - step: build-rocm-base-wheels allow_failure: false agents: - queue: cpu_queue_postmerge + queue: cpu_queue_release timeout_in_minutes: 60 commands: - | set -euo pipefail - + # Login to ECR aws ecr-public get-login-password --region us-east-1 | \ docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Download Docker image from S3 (set by build-rocm-base-wheels) - DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')" - if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then - echo "ERROR: rocm-docker-image-s3-path metadata not found" + + # Get ECR image tag from metadata (set by build-rocm-base-wheels) + ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')" + if [ -z "$${ECR_IMAGE_TAG}" ]; then + echo "ERROR: rocm-base-image-tag metadata not found" + echo "This should have been set by the build-rocm-base-wheels job" exit 1 fi - - echo "Downloading base image from $${DOCKER_IMAGE_S3_PATH}" - mkdir -p artifacts/rocm-docker-image - aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz - - # Load base Docker image - echo "Loading base Docker image..." - LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load) - BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //') - echo "Loaded base image: $${BASE_IMAGE_TAG}" - - # Tag and push the base image to ECR - docker tag "$${BASE_IMAGE_TAG}" public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base - docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base - echo "Pushed base image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base" - - # Get GPU architectures from meta-data - PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')" - PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}" - + + echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}" + + # Pull base Docker image from ECR + docker pull "$${ECR_IMAGE_TAG}" + + echo "Loaded base image: $${ECR_IMAGE_TAG}" + + # Pass the base image ECR tag to downstream steps (nightly publish) + buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}" + + echo "========================================" + echo "Building vLLM ROCm release image with:" + echo " BASE_IMAGE: $${ECR_IMAGE_TAG}" + echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}" + echo "========================================" + # Build vLLM ROCm release image using cached base DOCKER_BUILDKIT=1 docker build \ --build-arg max_jobs=16 \ - --build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \ - --build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \ + --build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \ --build-arg USE_SCCACHE=1 \ --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ --build-arg SCCACHE_REGION_NAME=us-west-2 \ @@ -730,10 +807,14 @@ steps: --target vllm-openai \ --progress plain \ -f docker/Dockerfile.rocm . - + # Push to ECR docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm - echo "Pushed: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" + + echo "" + echo " Successfully built and pushed ROCm release image" + echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm" + echo "" env: DOCKER_BUILDKIT: "1" S3_BUCKET: "vllm-wheels" diff --git a/.buildkite/scripts/annotate-rocm-release.sh b/.buildkite/scripts/annotate-rocm-release.sh index f42f77fc9a85..d66129722748 100755 --- a/.buildkite/scripts/annotate-rocm-release.sh +++ b/.buildkite/scripts/annotate-rocm-release.sh @@ -5,20 +5,21 @@ # Generate Buildkite annotation for ROCm wheel release set -ex -# Get build configuration from meta-data +# Extract build configuration from Dockerfile.rocm_base (single source of truth) # Extract ROCm version dynamically from Dockerfile.rocm_base # BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0" ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown") -PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12") -PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1250;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151") +PYTHON_VERSION=$(grep '^ARG PYTHON_VERSION=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTHON_VERSION=//') +PYTORCH_ROCM_ARCH=$(grep '^ARG PYTORCH_ROCM_ARCH=' docker/Dockerfile.rocm_base | sed 's/^ARG PYTORCH_ROCM_ARCH=//') -# TODO: Enable the nightly build for ROCm # Get release version, default to 1.0.0.dev for nightly/per-commit builds RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "") if [ -z "${RELEASE_VERSION}" ]; then RELEASE_VERSION="1.0.0.dev" fi +ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + # S3 URLs S3_BUCKET="${S3_BUCKET:-vllm-wheels}" S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}" @@ -68,7 +69,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl . -aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl . +aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl . aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl . \`\`\` @@ -80,7 +81,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash- - **torchvision**: TorchVision for ROCm PyTorch - **torchaudio**: Torchaudio for ROCm PyTorch - **amdsmi**: AMD SMI Python bindings -- **aiter**: Aiter for ROCm +- **amd_aiter**: Aiter for ROCm - **flash-attn**: Flash Attention for ROCm ### :warning: Notes @@ -96,7 +97,7 @@ To download and upload the image: docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base docker push vllm/vllm-openai-rocm:latest-base diff --git a/benchmarks/curl_gpt.sh b/benchmarks/curl_gpt.sh deleted file mode 100644 index 51a467bfbc2f..000000000000 --- a/benchmarks/curl_gpt.sh +++ /dev/null @@ -1,7 +0,0 @@ -MODEL="${MODEL:-/data/models/gpt-oss-120b-w-mxfp4-a-fp8}" \ -curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{ - "model": ${MODEL}, - "prompt": "What is the capital of France, and what is it known for?", - "temperature": 0.0, - "max_tokens": 100 -}' diff --git a/benchmarks/gpt_fp4_serve.sh b/benchmarks/gpt_fp4_serve.sh deleted file mode 100644 index ef23808e95ab..000000000000 --- a/benchmarks/gpt_fp4_serve.sh +++ /dev/null @@ -1,7 +0,0 @@ -#ROCR_VISIBLE_DEVICE=0 \ -MODEL="${MODEL:-/data/models/gpt-oss-120b-w-mxfp4-a-fp8}" \ -HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 \ -VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 \ -VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 \ -VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 \ -vllm serve --model ${MODEL} --host localhost --port 8000 --tensor-parallel-size 1 --gpu_memory_utilization 0.7 #--compilation-config '{"mode":"None","cudagraph_mode": "FULL", "cudagraph_capture_sizes": [1]}' diff --git a/docker/Dockerfile.rocm_1250 b/docker/Dockerfile.rocm_1250 deleted file mode 100644 index 697d92dda512..000000000000 --- a/docker/Dockerfile.rocm_1250 +++ /dev/null @@ -1,114 +0,0 @@ -FROM ubuntu:24.04 -ENV DEBIAN_FRONTEND=noninteractive - -# --- ROCm + PyTorch versioning (override any of these at build time) ------- -ARG GFX_ARCH=gfx1250 -ARG ROCM_VERSION=7.14.0a20260604 -ARG ROCM_BUILD=a0 -ARG TORCH_VERSION=2.10.0 -ARG ROCM_WHEEL_INDEX=https://rocm.genesis.amd.com/whl/gfx1250/ -ARG VLLM_REPO="https://github.com/ROCm/vllm.git" -ARG VLLM_BRANCH="455_wip" -ARG AITER_BRANCH="jpvillam/gfx1250_0604" - -ENV VLLM_TARGET_DEVICE=rocm -ENV PYTORCH_ROCM_ARCH=${GFX_ARCH} -ENV CMAKE_BUILD_TYPE=Release -ENV GPU_TARGET=${GFX_ARCH} -ENV VIRTUAL_ENV=/opt/venv - -RUN apt-get update && \ - apt-get install -y --no-install-recommends \ - build-essential clang-19 lld ccache ninja-build cmake \ - git curl libcurl4-openssl-dev rsync ssh wget ca-certificates \ - python3 python3-pip python3-dev python3-venv \ - less vim libzstd-dev numactl libelf1 m4 && \ - apt-get clean && rm -rf /var/lib/apt/lists/* - -SHELL ["/bin/bash", "-e", "-u", "-o", "pipefail", "-c"] - -RUN python3 -m venv "${VIRTUAL_ENV}" && \ - "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML -ENV PATH=${VIRTUAL_ENV}/bin:$PATH - -# COPY ffm/ /root/x/ffm/ -# ENV PATH=${ROCM_PATH}/llvm/bin:${ROCM_PATH}/bin:$PATH -# ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${ROCM_PATH}/lib/rocm_sysdeps/lib:${ROCM_PATH}/llvm/lib -# ENV HSA_MODEL_LIB=/root/x/ffm/libhsakmtmodel.so -# ENV HSA_MODEL_TOML="/root/x/ffm/ffm_config.toml" -# ENV HSA_MODEL_ARGS=ffm_enable_time_slicing -# ENV HSA_MODEL_TOPOLOGY=/root/x/ffm/topology/mi450 -# ENV HSA_MODEL_NUM_THREADS=64 - -ENV HSA_ENABLE_SDMA=0 -ENV HSA_ENABLE_INTERRUPT=0 - -# Install the TheRock PyTorch wheel and the matching rocm-sdk wheels into the venv -RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ - "torch==${TORCH_VERSION}+rocm${ROCM_VERSION}${ROCM_BUILD:+.${ROCM_BUILD}}" \ - "rocm[libraries,devel]==${ROCM_VERSION}${ROCM_BUILD:++${ROCM_BUILD}}" && \ - rocm-sdk init - -RUN pip install \ - filelock \ - "typing-extensions>=4.10.0" \ - "sympy>=1.13.3" \ - "networkx>=2.5.1" \ - jinja2 \ - "fsspec>=0.8.5" - -# Set ROCm, Python, Paths, and Environment Variables -ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python3.12/site-packages -ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel -ENV ROCM_HOME=${ROCM_PATH} -ENV ROCM_SOURCE_DIR=${ROCM_PATH} -ENV ROCM_BIN=${ROCM_PATH}/bin -ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake -ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode -ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH -ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib:${SITE_PACKAGES}/_rocm_sdk_libraries_${GFX_ARCH}/lib -ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake -ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi - -# NPI base-image fix-ups so the wheel ROCm behaves like a normal install: -# - build amd_smi from the bundled source (provides the `amdsmi` module), -# - flip rocm_sdk's library preload off RTLD_GLOBAL (avoids symbol clashes), -RUN if [ -d "${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi" ]; then \ - cd "${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi" && pip install .; \ - fi && \ - if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ - sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ - "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ - fi - -# Install AITER (https://github.com/ROCm/aiter): plain clone WITHOUT submodules -# (so 3rdparty/composable_kernel is absent), built with Composable Kernel disabled. -RUN git clone -b ${AITER_BRANCH} --depth 1 https://github.com/ROCm/aiter.git /app/aiter && \ - cd /app/aiter && \ - pip install packaging "cmake<4" ninja wheel setuptools_scm vcs_versioning pybind11 numpy && \ - ENABLE_CK=0 AITER_USE_SYSTEM_TRITON=1 pip install --no-build-isolation -e . - - -RUN apt-get update && \ - apt-get install -y --no-install-recommends pkg-config libnuma-dev libdrm-dev && \ - apt-get clean && rm -rf /var/lib/apt/lists/* - -RUN pip install "setuptools>=77.0.3,<81.0.0" setuptools_scm setuptools_rust wheel packaging jinja2 - -# Build + install vLLM for gfx1250 (cloned from VLLM_REPO @ VLLM_BRANCH) -RUN git clone "${VLLM_REPO}" /app/vllm \ - && cd /app/vllm \ - && git fetch -v --prune -- origin "${VLLM_BRANCH}" \ - && git checkout FETCH_HEAD - -# Compile the C++/HIP extensions (_C, _rocm_C) for gfx1250 and editable-install. -# MAX_JOBS is a build-arg so it can be tuned per machine (docker build --build-arg). -RUN cd /app/vllm && \ - pip install -e . --no-build-isolation --no-deps -v - -# vLLM Python runtime deps. common.txt pins no torch, so the NPI ROCm torch -# installed above is left untouched. -RUN pip install -r /app/vllm/requirements/common.txt - -WORKDIR /app/vllm -ENTRYPOINT ["/bin/bash"] \ No newline at end of file diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 9e5daf66a9f7..ea0697897dd5 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -25,10 +25,11 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG PYTORCH_ROCM_ARCH=gfx1250 +ARG PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx1250 +ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV MORI_GPU_ARCHS=gfx942;gfx950 +ENV FA_GPU_ARCHS=gfx942;gfx950 # TODO: Unset these when support is available for gfx1250 ENV ENABLE_CK=0 @@ -175,17 +176,14 @@ RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/di FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO +ARG MORI_GPU_ARCHS RUN mkdir -p /app/install; \ - if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ - echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ - else \ git clone ${MORI_REPO} \ && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ - && cp /app/mori/dist/*.whl /app/install; \ - fi + && cp /app/mori/dist/*.whl /app/install; ### @@ -196,9 +194,6 @@ ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE RUN mkdir -p /app/install; \ - if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ - echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping FlashAttention build"; \ - else \ git clone ${FA_REPO} \ && cd flash-attention \ && git checkout ${FA_BRANCH} \ @@ -207,10 +202,9 @@ RUN mkdir -p /app/install; \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ && sccache --show-stats; \ fi \ - && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && cp dist/*.whl /app/install; \ - fi + && cp dist/*.whl /app/install; ### diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 67bbdb0a99ea..f1bd319215ec 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -305,12 +305,12 @@ def test_mla_backend_selection( def test_aiter_fa_requires_mi3xx(mock_vllm_config): - """Test that ROCM_AITER_FA requires mi3xx architecture.""" + """Test that ROCM_AITER_FA requires CDNA3+ architecture.""" from vllm.platforms.rocm import RocmPlatform - # Mock on_mi3or4 to return False (used by supports_compute_capability) + # Mock cdna version to return 1 (used by supports_compute_capability) with ( - patch("vllm.platforms.rocm.on_mi3or4", return_value=False), + patch("vllm.platforms.rocm.get_cdna_version", return_value=1), pytest.raises( ValueError, match="compute capability not supported", diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 16142c4eb738..d439acf34abe 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2720,7 +2720,9 @@ def group_fp8_quant( @staticmethod def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool: - return (n, k) in [ + from vllm.platforms.rocm import on_gfx1250 + + return on_gfx1250() or (n, k) in [ (1024, 8192), (2112, 7168), (3072, 1536), diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index cb175acf45c9..1b39491ab346 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -373,7 +373,7 @@ def __init__(self, config: FP8ScaledMMLinearLayerConfig): self.use_triton = ( not current_platform.is_fp8_fnuz() and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) - ) or True + ) @classmethod def is_supported(cls, compute_capability=None): diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 004d4a426780..344c53b04b3a 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -10,6 +10,7 @@ from vllm.config.model import LogprobsMode from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform +from vllm.platforms.rocm import on_gfx1250 from vllm.triton_utils import HAS_TRITON if HAS_TRITON: @@ -110,13 +111,14 @@ def __init__( elif ( logprobs_mode not in ("processed_logits", "processed_logprobs") and rocm_aiter_ops.is_enabled() + and not on_gfx1250() # TODO (JPVILLAM): Enable this path ): self.aiter_ops = None self._aiter_ops_import_failed = False logger.info_once( "Using aiter sampler on ROCm (lazy import, sampling-only)." ) - self.forward = self.forward_native + self.forward = self.forward_hip else: self.forward = self.forward_native From f0918405b5cf5b53fbcd04f4319779789e687874 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Wed, 1 Jul 2026 13:07:31 -0500 Subject: [PATCH 44/68] Add numa dep Signed-off-by: jpvillam --- docker/Dockerfile.rocm_base | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index ea0697897dd5..ca460f8079dd 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -47,7 +47,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++\ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ libnuma-dev\ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ From 5c0010ce356c83f9b1425275e46b98e836542cf7 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Wed, 1 Jul 2026 14:15:09 -0500 Subject: [PATCH 45/68] Remove mori build again for now Signed-off-by: jpvillam --- docker/Dockerfile.rocm_base | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index ca460f8079dd..ec5e85cd943a 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -47,7 +47,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ libnuma-dev\ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ \ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ @@ -171,19 +171,23 @@ RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/di ### -### MORI Build +### MORI Build TODO(Build needs fixing) ### FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO ARG MORI_GPU_ARCHS RUN mkdir -p /app/install; \ + if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ + echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ + else \ git clone ${MORI_REPO} \ && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ - && cp /app/mori/dist/*.whl /app/install; + && cp /app/mori/dist/*.whl /app/install; \ + fi ### From 0fa6bd392a0c2dbd7666300458c3c05069856b89 Mon Sep 17 00:00:00 2001 From: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Date: Mon, 6 Jul 2026 09:49:27 -0500 Subject: [PATCH 46/68] dllehr_gpt: gpt-oss gfx1250 ATOM-parity perf patches (#1035) (cherry picked from commit c16f522a089f11578cd39e1403e40fffea20c9fb) Co-authored-by: dllehr Signed-off-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- .../layers/attention/attention.py | 9 +- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 17 +-- .../experts/gpt_oss_triton_kernels_moe.py | 21 +-- .../layers/quantization/utils/mxfp4_utils.py | 20 +++ vllm/model_executor/layers/utils.py | 25 +++- vllm/model_executor/models/gpt_oss.py | 67 ++++++++- vllm/triton_utils/jit_monitor.py | 141 ++++++++++++++++++ .../backends/rocm_aiter_unified_attn.py | 88 ++++++++++- vllm/v1/sample/ops/topk_topp_sampler.py | 53 +++++++ vllm/v1/sample/sampler.py | 93 ++++++++++++ 10 files changed, 497 insertions(+), 37 deletions(-) create mode 100644 vllm/triton_utils/jit_monitor.py diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 32562f0d9d9f..69f406e07264 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -431,8 +431,15 @@ def __init__( # for attn backends supporting query quantization self.query_quant = None + # ATOM keeps the query in bf16 (fp8 KV cache, q_descale=None). Skip the + # per-layer fp8 query quant (the aiter static scaled_quant kernel) when + # VLLM_DISABLE_QUERY_QUANT=1 to match ATOM and drop ~2.3us/marker. + import os as _os + + _disable_qq = _os.environ.get("VLLM_DISABLE_QUERY_QUANT", "0") == "1" if ( - self.impl.supports_quant_query_input + not _disable_qq + and self.impl.supports_quant_query_input and ( self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4" ) 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 index ed1da9416c40..c447073298d6 100644 --- 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 @@ -63,16 +63,10 @@ def aiter_triton_kernel_w4a8_moe_forward( gating_output, topk, sm_first=not renormalize ) - # gfx1250: aiter's in-kernel gather is numerically broken (validated on the - # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted - # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, - # sorted row i reads source token gather_idx[i] // n_expts_act, so this - # reproduces the in-kernel gather exactly (manual gather -> maxrel ~5e-3). - # gfx950 keeps the (working) in-kernel gather. - if on_gfx1250(): - gather_src = gather_idx.to(torch.long) // topk - hidden_states = hidden_states[gather_src] - gather_idx = None + # PATCH_GFX1250_INKERNEL_GATHER: ATOM passes gather_idx in-kernel on gfx1250 + # with GFX1250_SCALE-swizzled weight scales (see patch_swizzle.py). The old + # "in-kernel gather broken" was with UNSWIZZLED scales; with the correct + # GFX1250_SCALE swizzle the in-kernel gather is correct, so keep gather_idx. return triton_kernel_fused_mxfp4_w4a8_experts( None, @@ -141,10 +135,9 @@ def triton_kernel_fused_mxfp4_w4a8_experts( ) from vllm.platforms.rocm import on_gfx1250 - from vllm.platforms.rocm import on_gfx1250 _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None # TODO (JPVILLAM): merge conflict resolve later if _swizzle_mx_scale is enough - mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" + mx_scale_swizzle = "GFX1250_SCALE" if on_gfx1250() else "CDNA4_SCALE" assert quant_config.w1_precision is not None, ( "w1_precision in quant config can't be None" 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 8b688697d167..e50d10201cac 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 @@ -6,7 +6,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -57,7 +56,7 @@ def _triton_kernel_moe_supports_current_device() -> bool: # on_gfx9() already excludes gfx906/gfx908. # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). - return on_gfx9() or on_gfx1x() or on_gfx1250() + return on_gfx9() or on_gfx1x() or on_gfx1250() return False @@ -1132,16 +1131,12 @@ def _aiter_moe_two_gemm( arch = get_arch() quant_dtype = torch.float8_e4m3fnuz if arch == "gfx942" else torch.float8_e4m3fn - # gfx1250's in-kernel gather miscomputes (validated on the FFM sim), so gather - # rows into expert-sorted order in torch and pass gather_indx=None instead. Per - # aiter's moe_gemm_torch, sorted row i reads source token gather_idx[i] // topk. - if arch == "gfx1250": - gather_src = gather_idx.to(torch.long) // topk - gemm1_input = hidden_states[gather_src] - gemm1_gather_indx = None - else: - gemm1_input = hidden_states - gemm1_gather_indx = gather_idx + # PATCH_GATHER: in-kernel gather on gfx1250 (ATOM behavior). ATOM (same aiter + # commit) passes gather_indx straight to moe_gemm_a8w4 on gfx1250; do the same + # to drop the eager vectorized_gather + floordiv. (If this miscomputes, the + # gfx1250 in-kernel gather needs GFX1250_SCALE-swizzled scales -> add swizzle.) + gemm1_input = hidden_states + gemm1_gather_indx = gather_idx gammas = routing_data.gate_scal if routing_data else None g1_gammas = gammas if apply_router_weight_on_input else None @@ -1563,7 +1558,7 @@ def apply( self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG if expert_map is not None: - topk_ids = expert_map[topk_ids] + topk_ids = expert_map[topk_ids] # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 5110bcb7be70..fc26d76b2e22 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -102,6 +102,26 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # transpose the tensor so that the quantization axis is on dim1 quant_tensor = quant_tensor.transpose(-2, -1) scale = scale.transpose(-2, -1) + # GFX1250: apply ATOM's gfx1250 MX-scale swizzle so the in-kernel gather and + # moe_gemm_a8w4 read the scale correctly (paired with swizzle_mx_scale= + # "GFX1250_SCALE" in aiter_mxfp4_w4a8_moe). scale is (E, K_SCALE, N) here. + try: + from vllm.platforms.rocm import on_gfx1250 as _on_gfx1250 + except Exception: + _on_gfx1250 = lambda: False + if current_platform.is_rocm() and _on_gfx1250(): + from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( + swizzle_scales_gfx1250 as _swz_gfx1250, + ) + + assert ( + scale.dim() == 3 and scale.shape[-1] % 32 == 0 and scale.shape[-2] % 8 == 0 + ), f"GFX1250 scale swizzle needs (E,K_SCALE%8,N%32); got {tuple(scale.shape)}" + # ATOM passes the raw swizzle_scales_gfx1250 output directly (no + # .contiguous(), no re-layout) -- the kernel's GFX1250_SCALE TMA + # descriptor expects exactly these strides. Adding .contiguous() + # changes the strides and breaks tt.make_tensor_descriptor. + scale = _swz_gfx1250(scale) quant_tensor = convert_layout( wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts ) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 814142ce053f..fba82d762086 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -99,6 +99,8 @@ def default_unquantized_gemm( def use_aiter_triton_gemm(n, m, k, dtype): + from vllm.platforms.rocm import on_gfx1250 + if ( not rocm_aiter_ops.is_triton_gemm_enabled() # MI300's - fp8nuz=True @@ -107,9 +109,17 @@ def use_aiter_triton_gemm(n, m, k, dtype): ): return False - # use hipblaslt for the larger GEMMs - if n > 2048 and m > 512: + # PATCH: gfx1250 branch-aiter gemm_a16w16 is 8-16x FASTER than rocBLAS even + # for large prefill batches (validated: 16384x5120x2880 -> 685us vs 5917us). + # The original "use hipblaslt for larger GEMMs" guard sent big prefill qkv/o_proj + # to rocBLAS (Cijk). Disable it on gfx1250 so large prefill GEMMs use aiter too. + if n > 2048 and m > 512 and not on_gfx1250(): return False + # PATCH: gfx1250 aiter gluon gemm_a16w16 handles arbitrary shapes (incl. + # lm_head m=201088,k=2880 and K%256!=0). Route ALL bf16 dense GEMMs through + # it so lm_head (and any other shape) stops falling back to rocBLAS (Cijk). + if on_gfx1250(): + return True return ( (m == 5120 and k == 2880) or (m == 2880 and k == 4096) @@ -122,7 +132,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1250, on_gfx1x, on_gfx9, on_gfx950 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -168,9 +178,10 @@ def rocm_unquantized_gemm_impl( # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b # hidden), so skip aiter there and fall back to the torch GEMM path below. - if use_aiter_triton_gemm(n, m, k, x.dtype) and not ( - on_gfx1250() and k % 256 != 0 - ): + # PATCH: branch aiter gemm_a16w16 handles K%256!=0 (e.g. gpt-oss-120b K=2880), + # validated numerically. Drop the stale gfx1250 guard so these bf16 projections + # use aiter gluon GEMM instead of falling back to slow rocBLAS (Cijk) torch.linear. + if use_aiter_triton_gemm(n, m, k, x.dtype): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) @@ -179,7 +190,7 @@ def rocm_unquantized_gemm_impl( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) # build (gfx9/gfx11 ISA); fall back to torch GEMM there. - and not on_gfx1250() # TODO GFX1250: Remove once skinny GEMM is supported on gfx1250 + and not on_gfx1250() # TODO GFX1250: Remove once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 863d8ea3ba7f..8c054c9ff6e3 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -36,6 +36,59 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.utils import rocm_unquantized_gemm +from vllm.utils.torch_utils import direct_register_custom_op + +_GPT_OSS_NORMPAD_LOGGED = False + + +def _gpt_oss_rmsnorm_pad_add( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + pad_to: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # ATOM-equivalent: add + rmsnorm (over the original N) + zero-pad to `pad_to`. + from aiter.ops.triton.normalization.fused_add_rmsnorm_pad import ( + fused_add_rmsnorm_pad, + ) + + out, res_out = fused_add_rmsnorm_pad(x, weight, eps, residual, pad_to) + global _GPT_OSS_NORMPAD_LOGGED + if not _GPT_OSS_NORMPAD_LOGGED: + import sys + + print( + f"[normpad] x={tuple(x.shape)} pad_to={pad_to} out={tuple(out.shape)} " + f"res_out={tuple(res_out.shape)}", + file=sys.stderr, + flush=True, + ) + _GPT_OSS_NORMPAD_LOGGED = True + return out, res_out + + +def _gpt_oss_rmsnorm_pad_add_fake( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + pad_to: int, +) -> tuple[torch.Tensor, torch.Tensor]: + M, N = x.shape + n_out = ((N + pad_to - 1) // pad_to) * pad_to if pad_to > 0 else N + return ( + torch.empty((M, n_out), dtype=x.dtype, device=x.device), + torch.empty_like(residual), + ) + + +direct_register_custom_op( + op_name="gpt_oss_rmsnorm_pad_add", + op_func=_gpt_oss_rmsnorm_pad_add, + mutates_args=[], + fake_impl=_gpt_oss_rmsnorm_pad_add_fake, +) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -283,7 +336,19 @@ def forward( hidden_states = self.attn(hidden_states, positions) # Fully Connected - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + _pad_to = self.mlp.experts.moe_config.hidden_dim + if residual is not None and _pad_to > hidden_states.shape[-1]: + hidden_states, residual = torch.ops.vllm.gpt_oss_rmsnorm_pad_add( + hidden_states, + residual, + self.post_attention_layernorm.weight, + self.post_attention_layernorm.variance_epsilon, + _pad_to, + ) + else: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) output = self.mlp(hidden_states) return output, residual diff --git a/vllm/triton_utils/jit_monitor.py b/vllm/triton_utils/jit_monitor.py new file mode 100644 index 000000000000..e7d303440281 --- /dev/null +++ b/vllm/triton_utils/jit_monitor.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Monitor unexpected Triton kernel JIT compilation during inference. + +After server warmup completes, any Triton JIT compilation or autotuning +event indicates a cache miss or unexpected input shape that causes a +latency spike. This module registers hooks in the Triton runtime to +detect and log such events so they can be investigated. + +Set ``--jit-monitor-verbose`` to log every Triton JIT compile with its +dispatch key. This is intentionally opt-in because it can emit many logs and +add overhead. + +Currently monitors: +- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``) +- Triton ``@triton.jit`` first-time compilations + (via ``knobs.runtime.jit_post_compile_hook``) +""" + +import os + +from vllm.logger import init_logger +from vllm.triton_utils.importing import HAS_TRITON + +logger = init_logger(__name__) + +_active: bool = False +_verbose: bool = False + + +def is_active() -> bool: + """Return whether the JIT compilation monitor is currently active.""" + return _active + + +def activate(*, verbose: bool = False) -> None: + """Enable JIT compilation monitoring after warmup. + + Call once per worker process at the end of + :func:`compile_or_warm_up_model`. After activation every Triton + kernel compilation or autotuning benchmark that happens during + inference will be logged as a warning. + + Safe to call multiple times — subsequent calls are no-ops. + + If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in + their environment, autotuning printing is left disabled; the JIT + compilation hook is still registered regardless. + """ + global _active, _verbose + if _active: + return + _active = True + _verbose = verbose + + _setup_triton_autotuning_print() + _setup_triton_jit_hook() + + logger.info( + "Kernel JIT monitor activated — Triton JIT compilations " + "during inference will be logged as warnings." + ) + + +# ------------------------------------------------------------------ +# Triton autotuning print +# ------------------------------------------------------------------ + + +def _setup_triton_autotuning_print() -> None: + """Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out.""" + if not HAS_TRITON: + return + from triton import knobs # type: ignore[import-untyped] + + user_val = os.environ.get("TRITON_PRINT_AUTOTUNING") + if user_val == "0": + logger.debug( + "TRITON_PRINT_AUTOTUNING=0 set by user — " + "autotuning messages will stay suppressed." + ) + return + + knobs.autotuning.print = True + + +# ------------------------------------------------------------------ +# Triton JIT compilation hook +# ------------------------------------------------------------------ + + +def _log_jit_compile(fn_name: str, kwargs) -> None: + if _verbose: + compile_info = kwargs.get("compile") + if not isinstance(compile_info, dict): + compile_info = {} + logger.warning( + "Triton %sJIT compilation during inference: %s (key=%s).", + "autotune/warmup candidate " if kwargs.get("warmup") else "kernel ", + fn_name, + compile_info.get("key") or kwargs.get("key"), + ) + return + + logger.warning_once( + "Triton kernel JIT compilation during inference: %s. " + "This causes a latency spike; consider extending warmup " + "to cover this shape/config.", + fn_name, + ) + + +def _setup_triton_jit_hook() -> None: + """Register a ``jit_post_compile_hook`` that warns on compilation.""" + # PATCH: upstream triton efbae9a's serialize_specialization_data cannot + # JSON-encode Gluon PaddedSharedLayout constexprs (passed by e.g. aiter + # gemm_a16w16), and triton builds spec-data for ANY registered hook. + # This monitor is purely diagnostic, so skip registering to avoid + # crashing live kernel compilation. + return + if not HAS_TRITON: + return + from triton import knobs # type: ignore[import-untyped] + + existing_hook = knobs.runtime.jit_post_compile_hook + + def _on_jit_compile(**kwargs): + # `jit_post_compile_hook` is Triton internal API and its + # signature has changed across releases (kwargs added/renamed). + # Accept **kwargs so an upstream change cannot crash this hook + # with TypeError, and forward the full kwarg set to any + # pre-existing hook unchanged. + fn = kwargs.get("fn") + fn_name = getattr(fn, "name", "") + _log_jit_compile(fn_name, kwargs) + if existing_hook is not None: + return existing_hook(**kwargs) + return None + + knobs.runtime.jit_post_compile_hook = _on_jit_compile diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index d8363169a8c8..abc7bc0b3025 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -85,6 +85,13 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") + if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): + # ATOM-style K/V-outermost layout: K=cache[0], V=cache[1] each + # standalone-contiguous (block stride == block content). Required by + # the gfx1250 gluon unified-attn-3d kernel (TDM descriptors assume + # contiguous per-block K/V). The default (num_blocks, 2, ...) + # interleaves K/V per block (2x gap) -> gluon reads garbage. + return (2, num_blocks, block_size, num_kv_heads, head_size) return (num_blocks, 2, block_size, num_kv_heads, head_size) @staticmethod @@ -148,6 +155,12 @@ def __init__( def _split_kv_cache( self, kv_cache: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: + if ( + rocm_aiter_ops.is_shuffle_kv_cache_enabled() + and self.attn_type != AttentionType.ENCODER_DECODER + ): + # ATOM layout (2, num_blocks, block_size, num_kv_heads, head_size). + return kv_cache.unbind(0) if self.attn_type != AttentionType.ENCODER_DECODER: return kv_cache.unbind(1) @@ -169,6 +182,18 @@ def _split_kv_cache( ) return kv_cache.unbind(0) + def _shuffle_kv_views( + self, key_cache: torch.Tensor, value_cache: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + # No-copy reinterpret of standalone-contiguous K/V (num_blocks, + # block_size, num_kv_heads, head_size) as the 5D shuffle layout the + # gluon kernel reads (the shuffle WRITER laid bytes in this order). + nb, bs, nkv, hd = key_cache.shape + x = 16 // key_cache.element_size() + kc = key_cache.reshape(nb, nkv, hd // x, bs, x) + vc = value_cache.reshape(nb, nkv, bs // x, hd, x) + return kc, vc + def forward( self, layer: torch.nn.Module, @@ -235,6 +260,16 @@ def forward( if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) + _shuffled = ( + rocm_aiter_ops.is_shuffle_kv_cache_enabled() + and self.attn_type != AttentionType.ENCODER_DECODER + ) + if _shuffled: + key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache) + # Shuffle path pre-divides K/V by k/v_scale before the (scale=1.0) + # writer, so the stored fp8 == reshape_and_cache_flash's; descale with + # the same layer scales as the vanilla path. + _kd, _vd = layer._k_scale, layer._v_scale cu_seqlens_q = attn_metadata.query_start_loc seqused_k = attn_metadata.seq_lens @@ -258,10 +293,11 @@ def forward( block_table=block_table, softcap=self.logits_soft_cap, q_descale=layer._q_scale if query.dtype == self.fp8_dtype else None, - k_descale=layer._k_scale, - v_descale=layer._v_scale, + k_descale=_kd, + v_descale=_vd, sinks=self.sinks, output_scale=output_scale, + shuffled_kv_cache=_shuffled, ) return output @@ -280,6 +316,39 @@ def do_kv_cache_update( return key_cache, value_cache = self._split_kv_cache(kv_cache) + if ( + rocm_aiter_ops.is_shuffle_kv_cache_enabled() + and self.attn_type != AttentionType.ENCODER_DECODER + ): + # Write new tokens directly into the 5D shuffle byte layout the + # gluon kernel reads (ATOM-equivalent). Quantizes with scale=1.0. + from vllm.v1.attention.backends.rocm_aiter_fa import ( + reshape_and_cache_shuffle_triton, + ) + + if is_quantized_kv_cache(self.kv_cache_dtype): + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + # PATCH3: the shuffle writer stores at scale=1.0 and the reader + # descales by layer._k/v_scale. For gpt-oss these scales are 1.0, + # so NO pre-division is needed (key/1.0 == key) and the earlier + # `(key/scale).to(dtype)` was a wasted full-K/V fp32 elementwise + # pass every layer (+2.7ms/prefill step). Drop it. + # (If a model ships k/v scales != 1.0, the shuffle writer kernel + # must be taught to apply them; gpt-oss does not.) + _dummy = torch.empty(1, dtype=torch.float32, device=key.device) + reshape_and_cache_shuffle_triton( + key, + value, + key_cache, + value_cache, + slot_mapping, + self.kv_cache_dtype, + _dummy, + _dummy, + ) + return + # Reshape the input keys and values and store them in the cache. ops.reshape_and_cache_flash( key, @@ -312,13 +381,26 @@ def do_rope_and_kv_cache_update( # we use direct Q, K, V tensors without caching return key_cache, value_cache = self._split_kv_cache(kv_cache) - flash_layout = True is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) if is_fp8_kv_cache: key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) + # PATCH4_FUSED_ROPE_SHUFFLE: when the shuffle KV-cache layout is active, + # the fused RoPE+cache kernel must write the SAME 5D shuffle byte layout + # the gluon unified-attn-3d reader expects (ATOM parity). Reinterpret the + # standalone-contiguous K/V as the 5D shuffle views (no copy) and use the + # non-flash (shuffle) layout. Stores at layer k/v scales (=1.0 for + # gpt-oss) -> byte-identical to the non-fused shuffle writer (PATCH3). + _shuffled = ( + rocm_aiter_ops.is_shuffle_kv_cache_enabled() + and self.attn_type != AttentionType.ENCODER_DECODER + ) + if _shuffled: + key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache) + flash_layout = not _shuffled + rocm_aiter_ops.triton_rope_and_cache( query, key, diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 344c53b04b3a..5037054dd533 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -18,6 +18,48 @@ logger = init_logger(__name__) +import os as _os # PATCH5_AITER_TEMP_SAMPLE + +_PATCH5_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" +_PATCH5_DEBUG_DONE = False + + +def _aiter_temp_gumbel_sample(logits, generators, use_fp64_gumbel): + """Fused temperature Gumbel-max sampling via aiter (ATOM parity). + logits are ALREADY temperature-scaled by the parent Sampler, so temps=1. + Returns int64 token ids (num_tokens,).""" + global _PATCH5_DEBUG_DONE + import torch as _torch + from aiter import mixed_sample_outer_exponential + + n, vocab = logits.shape + out = _torch.empty(n, dtype=_torch.int32, device=logits.device) + nseed = len(generators) + if nseed == 0: + # common case: ATOM-style shared (1,vocab) exponential noise (cheap RNG) + exp = ( + _torch.empty((1, vocab), dtype=_torch.float32, device=logits.device) + .exponential_() + .expand(n, vocab) + ) + else: + exp = _torch.empty((n, vocab), dtype=_torch.float32, device=logits.device) + if nseed != n: + exp.exponential_() + for i, g in generators.items(): + exp[i].exponential_(generator=g) + temps = _torch.ones(n, dtype=_torch.float32, device=logits.device) + mixed_sample_outer_exponential(out, logits, exp, temps, eps=1e-10) + if not _PATCH5_DEBUG_DONE: + logger.info( + "PATCH5 aiter fused temp-sample ACTIVE (n=%d vocab=%d seeded=%d)", + n, + vocab, + nseed, + ) + _PATCH5_DEBUG_DONE = True + return out.to(_torch.int64) + def flashinfer_sampler_supported() -> bool: """Decide whether FlashInfer's top-p/top-k sampler can be used. @@ -134,6 +176,17 @@ def forward_native( The logits tensor may be updated in-place. """ + # PATCH5_AITER_TEMP_SAMPLE: fused no-filter temperature sampling (ATOM parity) + if ( + _PATCH5_ON + and k is None + and p is None + and not self.use_fp64_gumbel + and self.logprobs_mode not in ("processed_logits", "processed_logprobs") + ): + return _aiter_temp_gumbel_sample( + logits, generators, self.use_fp64_gumbel + ), None logits = apply_top_k_top_p(logits, k, p) logits_to_return = None if self.logprobs_mode == "processed_logits": diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index bb20432a0815..d55d0a8481ba 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -16,6 +16,93 @@ _SAMPLING_EPS = 1e-5 +import os as _os # PATCH6_AITER_RAW_SAMPLE + +from vllm.logger import init_logger as _p6_init_logger + +_p6log = _p6_init_logger("vllm.v1.sample.patch6") +_PATCH6_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" +_PATCH6_DEBUG = _os.environ.get("VLLM_PATCH6_DEBUG", "0") == "1" +_PATCH6_DBG = False +_PATCH6_ELIG_DBG = False + + +def _patch6_eligible(sm, predict_bonus_token): + global _PATCH6_ELIG_DBG + if not _PATCH6_ON: + return False + if _PATCH6_DEBUG and not _PATCH6_ELIG_DBG: + try: + lp = sm.logitsprocs + h = sm.thinking_budget_state_holder + _p6log.info( + "PATCH6 ELIG: all_random=%r top_k=%r top_p=%r logprobs=%r lpids=%r " + "nopen=%r allow_None=%r bad=%r argmax=%r nonarg=%r think=%r bonus=%r", + sm.all_random, + sm.top_k, + sm.top_p, + sm.max_num_logprobs, + bool(sm.logprob_token_ids), + sm.no_penalties, + sm.allowed_token_ids_mask is None, + bool(sm.bad_words_token_ids), + list(lp.argmax_invariant), + list(lp.non_argmax_invariant), + (h.has_tracked_requests() if h is not None else None), + predict_bonus_token, + ) + except Exception as e: + _p6log.info("PATCH6 ELIG dbg err: %r", e) + _PATCH6_ELIG_DBG = True + if predict_bonus_token: + return False + if not sm.all_random or sm.temperature is None: + return False + if sm.top_k is not None or sm.top_p is not None: + return False + if sm.max_num_logprobs is not None or sm.logprob_token_ids: + return False + if ( + not sm.no_penalties + or sm.allowed_token_ids_mask is not None + or sm.bad_words_token_ids + ): + return False + lp = sm.logitsprocs + if lp.argmax_invariant or lp.non_argmax_invariant: + return False + h = sm.thinking_budget_state_holder + if h is not None and h.has_tracked_requests(): + return False + return True + + +def _aiter_raw_fused_sample(logits, sm): + """Sample token ids directly from RAW (bf16) logits via fused aiter kernel. int32 [n,1].""" + global _PATCH6_DBG + import torch as _t + from aiter import mixed_sample_outer_exponential + + n, vocab = logits.shape + out = _t.empty(n, dtype=_t.int32, device=logits.device) + gens = sm.generators + if not gens: + exp = ( + _t.empty((1, vocab), dtype=_t.float32, device=logits.device) + .exponential_() + .expand(n, vocab) + ) + else: + exp = _t.empty((n, vocab), dtype=_t.float32, device=logits.device) + exp.exponential_() + for i, g in gens.items(): + exp[i].exponential_(generator=g) + mixed_sample_outer_exponential(out, logits, exp, sm.temperature, eps=1e-10) + if not _PATCH6_DBG: + _p6log.info("PATCH6 raw-logits fused sampler ACTIVE (n=%d vocab=%d)", n, vocab) + _PATCH6_DBG = True + return out.unsqueeze(-1) + class Sampler(nn.Module): """ @@ -76,6 +163,12 @@ def forward( predict_bonus_token: bool = False, logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: + # PATCH6_AITER_RAW_SAMPLE: unconstrained all-random batch -> fused raw-logits sampler + if _patch6_eligible(sampling_metadata, predict_bonus_token): + return SamplerOutput( + sampled_token_ids=_aiter_raw_fused_sample(logits, sampling_metadata), + logprobs_tensors=None, + ) logprobs_mode = logprobs_mode_override or self.logprobs_mode # NOTE(woosuk): Use the original logits (before any penalties or # temperature scaling) for the top-k logprobs. From ba24d1df367f732030a03851915db6d3d2bf4ab1 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 6 Jul 2026 10:11:31 -0500 Subject: [PATCH 47/68] Clean up code a bit from prev commit Signed-off-by: jpvillam --- .../layers/attention/attention.py | 20 +++---- .../layers/quantization/utils/mxfp4_utils.py | 12 ++-- vllm/triton_utils/jit_monitor.py | 6 +- vllm/v1/sample/ops/topk_topp_sampler.py | 17 ++---- vllm/v1/sample/sampler.py | 56 +++++-------------- 5 files changed, 36 insertions(+), 75 deletions(-) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 69f406e07264..b1d0429d59ff 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -436,7 +436,12 @@ def __init__( # VLLM_DISABLE_QUERY_QUANT=1 to match ATOM and drop ~2.3us/marker. import os as _os - _disable_qq = _os.environ.get("VLLM_DISABLE_QUERY_QUANT", "0") == "1" + from vllm.platforms.rocm import on_gfx1250 + + _disable_qq = ( + _os.environ.get("VLLM_DISABLE_QUERY_QUANT", "1" if on_gfx1250() else "0") + == "1" + ) if ( not _disable_qq and self.impl.supports_quant_query_input @@ -465,7 +470,6 @@ def forward( # shape does not match the query shape, so we optionally let the model # definition specify the output tensor shape. output_shape: torch.Size | None = None, - output_dtype: torch.dtype | None = None, ) -> torch.Tensor: """ The KV cache is stored inside this class and is accessed via @@ -480,8 +484,7 @@ def forward( torch.ops.vllm.maybe_calc_kv_scales( query, key, value, _encode_layer_name(self.layer_name) ) - if output_dtype is None: - output_dtype = query.dtype + output_dtype = query.dtype if self.query_quant is not None: # quantizing with a simple torch operation enables # torch.compile to fuse this into previous ops @@ -590,14 +593,7 @@ def get_attn_backend(self) -> type[AttentionBackend]: def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # Block size may get updated after model loading, refresh it block_size = vllm_config.cache_config.block_size - # Encoder-only attention is prefill-only and keeps no autoregressive KV - # cache. In hybrid models (e.g. Qwen3.5 / ColQwen3.5: GatedDeltaNet - # linear_attention interleaved with full_attention) the runner iterates - # every attention module to build the KV-cache spec, so an ENCODER_ONLY - # full_attention layer reaches here; it contributes no KV cache group. - if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): - return None - # Should not be called for enc-dec attention. + # Should not be called for enc-dec or encoder-only attention. assert self.attn_type == AttentionType.DECODER quant_mode = get_kv_quant_mode(self.kv_cache_dtype) if self.sliding_window is not None: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index fc26d76b2e22..0bd6964dd018 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -105,13 +105,11 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # GFX1250: apply ATOM's gfx1250 MX-scale swizzle so the in-kernel gather and # moe_gemm_a8w4 read the scale correctly (paired with swizzle_mx_scale= # "GFX1250_SCALE" in aiter_mxfp4_w4a8_moe). scale is (E, K_SCALE, N) here. - try: - from vllm.platforms.rocm import on_gfx1250 as _on_gfx1250 - except Exception: - _on_gfx1250 = lambda: False - if current_platform.is_rocm() and _on_gfx1250(): + from vllm.platforms.rocm import on_gfx1250 + + if current_platform.is_rocm() and on_gfx1250(): from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( - swizzle_scales_gfx1250 as _swz_gfx1250, + swizzle_scales_gfx1250 as swz_gfx1250, ) assert ( @@ -121,7 +119,7 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # .contiguous(), no re-layout) -- the kernel's GFX1250_SCALE TMA # descriptor expects exactly these strides. Adding .contiguous() # changes the strides and breaks tt.make_tensor_descriptor. - scale = _swz_gfx1250(scale) + scale = swz_gfx1250(scale) quant_tensor = convert_layout( wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts ) diff --git a/vllm/triton_utils/jit_monitor.py b/vllm/triton_utils/jit_monitor.py index e7d303440281..27e4666a88bc 100644 --- a/vllm/triton_utils/jit_monitor.py +++ b/vllm/triton_utils/jit_monitor.py @@ -118,7 +118,11 @@ def _setup_triton_jit_hook() -> None: # gemm_a16w16), and triton builds spec-data for ANY registered hook. # This monitor is purely diagnostic, so skip registering to avoid # crashing live kernel compilation. - return + # TODO: Find a better way around this + from vllm.platforms.rocm import on_gfx1250 + + if on_gfx1250(): + return if not HAS_TRITON: return from triton import knobs # type: ignore[import-untyped] diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 5037054dd533..5430c39b4603 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + import torch import torch.nn as nn @@ -18,17 +20,16 @@ logger = init_logger(__name__) -import os as _os # PATCH5_AITER_TEMP_SAMPLE -_PATCH5_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" -_PATCH5_DEBUG_DONE = False +_PATCH5_ON = ( + os.environ.get("VLLM_AITER_TEMP_SAMPLE", "1" if on_gfx1250() else "0") == "1" +) def _aiter_temp_gumbel_sample(logits, generators, use_fp64_gumbel): """Fused temperature Gumbel-max sampling via aiter (ATOM parity). logits are ALREADY temperature-scaled by the parent Sampler, so temps=1. Returns int64 token ids (num_tokens,).""" - global _PATCH5_DEBUG_DONE import torch as _torch from aiter import mixed_sample_outer_exponential @@ -50,14 +51,6 @@ def _aiter_temp_gumbel_sample(logits, generators, use_fp64_gumbel): exp[i].exponential_(generator=g) temps = _torch.ones(n, dtype=_torch.float32, device=logits.device) mixed_sample_outer_exponential(out, logits, exp, temps, eps=1e-10) - if not _PATCH5_DEBUG_DONE: - logger.info( - "PATCH5 aiter fused temp-sample ACTIVE (n=%d vocab=%d seeded=%d)", - n, - vocab, - nseed, - ) - _PATCH5_DEBUG_DONE = True return out.to(_torch.int64) diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index d55d0a8481ba..5c13558e9f50 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -2,11 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """A layer that samples the next tokens from the model's outputs.""" +import os + import torch import torch.nn as nn from vllm.config.model import LogprobsMode -from vllm.utils.torch_utils import PIN_MEMORY +from vllm.platforms.rocm import on_gfx1250 +from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.ops.bad_words import apply_bad_words @@ -16,44 +19,15 @@ _SAMPLING_EPS = 1e-5 -import os as _os # PATCH6_AITER_RAW_SAMPLE - -from vllm.logger import init_logger as _p6_init_logger -_p6log = _p6_init_logger("vllm.v1.sample.patch6") -_PATCH6_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" -_PATCH6_DEBUG = _os.environ.get("VLLM_PATCH6_DEBUG", "0") == "1" -_PATCH6_DBG = False -_PATCH6_ELIG_DBG = False +_PATCH6_ON = ( + os.environ.get("VLLM_AITER_TEMP_SAMPLE", "1" if on_gfx1250() else "0") == "1" +) def _patch6_eligible(sm, predict_bonus_token): - global _PATCH6_ELIG_DBG if not _PATCH6_ON: return False - if _PATCH6_DEBUG and not _PATCH6_ELIG_DBG: - try: - lp = sm.logitsprocs - h = sm.thinking_budget_state_holder - _p6log.info( - "PATCH6 ELIG: all_random=%r top_k=%r top_p=%r logprobs=%r lpids=%r " - "nopen=%r allow_None=%r bad=%r argmax=%r nonarg=%r think=%r bonus=%r", - sm.all_random, - sm.top_k, - sm.top_p, - sm.max_num_logprobs, - bool(sm.logprob_token_ids), - sm.no_penalties, - sm.allowed_token_ids_mask is None, - bool(sm.bad_words_token_ids), - list(lp.argmax_invariant), - list(lp.non_argmax_invariant), - (h.has_tracked_requests() if h is not None else None), - predict_bonus_token, - ) - except Exception as e: - _p6log.info("PATCH6 ELIG dbg err: %r", e) - _PATCH6_ELIG_DBG = True if predict_bonus_token: return False if not sm.all_random or sm.temperature is None: @@ -72,14 +46,12 @@ def _patch6_eligible(sm, predict_bonus_token): if lp.argmax_invariant or lp.non_argmax_invariant: return False h = sm.thinking_budget_state_holder - if h is not None and h.has_tracked_requests(): - return False - return True + return h is not None and h.has_tracked_requests() def _aiter_raw_fused_sample(logits, sm): - """Sample token ids directly from RAW (bf16) logits via fused aiter kernel. int32 [n,1].""" - global _PATCH6_DBG + """Sample token ids directly from RAW (bf16) + logits via fused aiter kernel. int32 [n,1].""" import torch as _t from aiter import mixed_sample_outer_exponential @@ -98,9 +70,6 @@ def _aiter_raw_fused_sample(logits, sm): for i, g in gens.items(): exp[i].exponential_(generator=g) mixed_sample_outer_exponential(out, logits, exp, sm.temperature, eps=1e-10) - if not _PATCH6_DBG: - _p6log.info("PATCH6 raw-logits fused sampler ACTIVE (n=%d vocab=%d)", n, vocab) - _PATCH6_DBG = True return out.unsqueeze(-1) @@ -152,7 +121,7 @@ def __init__( ): super().__init__() self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel) - self.pin_memory = PIN_MEMORY + self.pin_memory = is_pin_memory_available() self.logprobs_mode = logprobs_mode self.use_fp64_gumbel = use_fp64_gumbel @@ -163,7 +132,8 @@ def forward( predict_bonus_token: bool = False, logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: - # PATCH6_AITER_RAW_SAMPLE: unconstrained all-random batch -> fused raw-logits sampler + # PATCH6_AITER_RAW_SAMPLE: + # unconstrained all-random batch -> fused raw-logits sampler if _patch6_eligible(sampling_metadata, predict_bonus_token): return SamplerOutput( sampled_token_ids=_aiter_raw_fused_sample(logits, sampling_metadata), From 3272ef88b972f67b01bd74c34e9ddfeaffe09dff Mon Sep 17 00:00:00 2001 From: Daniel Date: Mon, 6 Jul 2026 16:37:06 -0400 Subject: [PATCH 48/68] adding gfx942 and gfx950 to pytorch archs to enable vllm C++ custom ops Signed-off-by: Daniel --- docker/Dockerfile.rocm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 9cc12e1142e5..23506e259b25 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -28,7 +28,7 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 +ARG ARG_PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install build dependencies and utilities From 672803ce9516a7cad599bf6cdd4430edebb59914 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 6 Jul 2026 21:14:05 -0500 Subject: [PATCH 49/68] Revert "Clean up code a bit from prev commit" until AITER is ready This reverts commit 36e2c1a93aa2f9531f1f81cd4fb97d7bf7837c5d. Signed-off-by: jpvillam --- .../layers/attention/attention.py | 20 ++++--- .../layers/quantization/utils/mxfp4_utils.py | 12 ++-- vllm/triton_utils/jit_monitor.py | 6 +- vllm/v1/sample/ops/topk_topp_sampler.py | 17 ++++-- vllm/v1/sample/sampler.py | 56 ++++++++++++++----- 5 files changed, 75 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index b1d0429d59ff..69f406e07264 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -436,12 +436,7 @@ def __init__( # VLLM_DISABLE_QUERY_QUANT=1 to match ATOM and drop ~2.3us/marker. import os as _os - from vllm.platforms.rocm import on_gfx1250 - - _disable_qq = ( - _os.environ.get("VLLM_DISABLE_QUERY_QUANT", "1" if on_gfx1250() else "0") - == "1" - ) + _disable_qq = _os.environ.get("VLLM_DISABLE_QUERY_QUANT", "0") == "1" if ( not _disable_qq and self.impl.supports_quant_query_input @@ -470,6 +465,7 @@ def forward( # shape does not match the query shape, so we optionally let the model # definition specify the output tensor shape. output_shape: torch.Size | None = None, + output_dtype: torch.dtype | None = None, ) -> torch.Tensor: """ The KV cache is stored inside this class and is accessed via @@ -484,7 +480,8 @@ def forward( torch.ops.vllm.maybe_calc_kv_scales( query, key, value, _encode_layer_name(self.layer_name) ) - output_dtype = query.dtype + if output_dtype is None: + output_dtype = query.dtype if self.query_quant is not None: # quantizing with a simple torch operation enables # torch.compile to fuse this into previous ops @@ -593,7 +590,14 @@ def get_attn_backend(self) -> type[AttentionBackend]: def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # Block size may get updated after model loading, refresh it block_size = vllm_config.cache_config.block_size - # Should not be called for enc-dec or encoder-only attention. + # Encoder-only attention is prefill-only and keeps no autoregressive KV + # cache. In hybrid models (e.g. Qwen3.5 / ColQwen3.5: GatedDeltaNet + # linear_attention interleaved with full_attention) the runner iterates + # every attention module to build the KV-cache spec, so an ENCODER_ONLY + # full_attention layer reaches here; it contributes no KV cache group. + if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER): + return None + # Should not be called for enc-dec attention. assert self.attn_type == AttentionType.DECODER quant_mode = get_kv_quant_mode(self.kv_cache_dtype) if self.sliding_window is not None: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 0bd6964dd018..fc26d76b2e22 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -105,11 +105,13 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # GFX1250: apply ATOM's gfx1250 MX-scale swizzle so the in-kernel gather and # moe_gemm_a8w4 read the scale correctly (paired with swizzle_mx_scale= # "GFX1250_SCALE" in aiter_mxfp4_w4a8_moe). scale is (E, K_SCALE, N) here. - from vllm.platforms.rocm import on_gfx1250 - - if current_platform.is_rocm() and on_gfx1250(): + try: + from vllm.platforms.rocm import on_gfx1250 as _on_gfx1250 + except Exception: + _on_gfx1250 = lambda: False + if current_platform.is_rocm() and _on_gfx1250(): from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( - swizzle_scales_gfx1250 as swz_gfx1250, + swizzle_scales_gfx1250 as _swz_gfx1250, ) assert ( @@ -119,7 +121,7 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # .contiguous(), no re-layout) -- the kernel's GFX1250_SCALE TMA # descriptor expects exactly these strides. Adding .contiguous() # changes the strides and breaks tt.make_tensor_descriptor. - scale = swz_gfx1250(scale) + scale = _swz_gfx1250(scale) quant_tensor = convert_layout( wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts ) diff --git a/vllm/triton_utils/jit_monitor.py b/vllm/triton_utils/jit_monitor.py index 27e4666a88bc..e7d303440281 100644 --- a/vllm/triton_utils/jit_monitor.py +++ b/vllm/triton_utils/jit_monitor.py @@ -118,11 +118,7 @@ def _setup_triton_jit_hook() -> None: # gemm_a16w16), and triton builds spec-data for ANY registered hook. # This monitor is purely diagnostic, so skip registering to avoid # crashing live kernel compilation. - # TODO: Find a better way around this - from vllm.platforms.rocm import on_gfx1250 - - if on_gfx1250(): - return + return if not HAS_TRITON: return from triton import knobs # type: ignore[import-untyped] diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 5430c39b4603..5037054dd533 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os - import torch import torch.nn as nn @@ -20,16 +18,17 @@ logger = init_logger(__name__) +import os as _os # PATCH5_AITER_TEMP_SAMPLE -_PATCH5_ON = ( - os.environ.get("VLLM_AITER_TEMP_SAMPLE", "1" if on_gfx1250() else "0") == "1" -) +_PATCH5_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" +_PATCH5_DEBUG_DONE = False def _aiter_temp_gumbel_sample(logits, generators, use_fp64_gumbel): """Fused temperature Gumbel-max sampling via aiter (ATOM parity). logits are ALREADY temperature-scaled by the parent Sampler, so temps=1. Returns int64 token ids (num_tokens,).""" + global _PATCH5_DEBUG_DONE import torch as _torch from aiter import mixed_sample_outer_exponential @@ -51,6 +50,14 @@ def _aiter_temp_gumbel_sample(logits, generators, use_fp64_gumbel): exp[i].exponential_(generator=g) temps = _torch.ones(n, dtype=_torch.float32, device=logits.device) mixed_sample_outer_exponential(out, logits, exp, temps, eps=1e-10) + if not _PATCH5_DEBUG_DONE: + logger.info( + "PATCH5 aiter fused temp-sample ACTIVE (n=%d vocab=%d seeded=%d)", + n, + vocab, + nseed, + ) + _PATCH5_DEBUG_DONE = True return out.to(_torch.int64) diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index 5c13558e9f50..d55d0a8481ba 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -2,14 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """A layer that samples the next tokens from the model's outputs.""" -import os - import torch import torch.nn as nn from vllm.config.model import LogprobsMode -from vllm.platforms.rocm import on_gfx1250 -from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.ops.bad_words import apply_bad_words @@ -19,15 +16,44 @@ _SAMPLING_EPS = 1e-5 +import os as _os # PATCH6_AITER_RAW_SAMPLE + +from vllm.logger import init_logger as _p6_init_logger -_PATCH6_ON = ( - os.environ.get("VLLM_AITER_TEMP_SAMPLE", "1" if on_gfx1250() else "0") == "1" -) +_p6log = _p6_init_logger("vllm.v1.sample.patch6") +_PATCH6_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" +_PATCH6_DEBUG = _os.environ.get("VLLM_PATCH6_DEBUG", "0") == "1" +_PATCH6_DBG = False +_PATCH6_ELIG_DBG = False def _patch6_eligible(sm, predict_bonus_token): + global _PATCH6_ELIG_DBG if not _PATCH6_ON: return False + if _PATCH6_DEBUG and not _PATCH6_ELIG_DBG: + try: + lp = sm.logitsprocs + h = sm.thinking_budget_state_holder + _p6log.info( + "PATCH6 ELIG: all_random=%r top_k=%r top_p=%r logprobs=%r lpids=%r " + "nopen=%r allow_None=%r bad=%r argmax=%r nonarg=%r think=%r bonus=%r", + sm.all_random, + sm.top_k, + sm.top_p, + sm.max_num_logprobs, + bool(sm.logprob_token_ids), + sm.no_penalties, + sm.allowed_token_ids_mask is None, + bool(sm.bad_words_token_ids), + list(lp.argmax_invariant), + list(lp.non_argmax_invariant), + (h.has_tracked_requests() if h is not None else None), + predict_bonus_token, + ) + except Exception as e: + _p6log.info("PATCH6 ELIG dbg err: %r", e) + _PATCH6_ELIG_DBG = True if predict_bonus_token: return False if not sm.all_random or sm.temperature is None: @@ -46,12 +72,14 @@ def _patch6_eligible(sm, predict_bonus_token): if lp.argmax_invariant or lp.non_argmax_invariant: return False h = sm.thinking_budget_state_holder - return h is not None and h.has_tracked_requests() + if h is not None and h.has_tracked_requests(): + return False + return True def _aiter_raw_fused_sample(logits, sm): - """Sample token ids directly from RAW (bf16) - logits via fused aiter kernel. int32 [n,1].""" + """Sample token ids directly from RAW (bf16) logits via fused aiter kernel. int32 [n,1].""" + global _PATCH6_DBG import torch as _t from aiter import mixed_sample_outer_exponential @@ -70,6 +98,9 @@ def _aiter_raw_fused_sample(logits, sm): for i, g in gens.items(): exp[i].exponential_(generator=g) mixed_sample_outer_exponential(out, logits, exp, sm.temperature, eps=1e-10) + if not _PATCH6_DBG: + _p6log.info("PATCH6 raw-logits fused sampler ACTIVE (n=%d vocab=%d)", n, vocab) + _PATCH6_DBG = True return out.unsqueeze(-1) @@ -121,7 +152,7 @@ def __init__( ): super().__init__() self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel) - self.pin_memory = is_pin_memory_available() + self.pin_memory = PIN_MEMORY self.logprobs_mode = logprobs_mode self.use_fp64_gumbel = use_fp64_gumbel @@ -132,8 +163,7 @@ def forward( predict_bonus_token: bool = False, logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: - # PATCH6_AITER_RAW_SAMPLE: - # unconstrained all-random batch -> fused raw-logits sampler + # PATCH6_AITER_RAW_SAMPLE: unconstrained all-random batch -> fused raw-logits sampler if _patch6_eligible(sampling_metadata, predict_bonus_token): return SamplerOutput( sampled_token_ids=_aiter_raw_fused_sample(logits, sampling_metadata), From 712db32e199b6ba1860579f5ad23a14c11fb26e5 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 6 Jul 2026 21:14:19 -0500 Subject: [PATCH 50/68] Revert "dllehr_gpt: gpt-oss gfx1250 ATOM-parity perf patches (#1035)" until aitre is ready This reverts commit bf0eca886ed424e86c066c9e2cf76b27617475e2. Signed-off-by: jpvillam --- .../layers/attention/attention.py | 9 +- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 17 ++- .../experts/gpt_oss_triton_kernels_moe.py | 21 ++- .../layers/quantization/utils/mxfp4_utils.py | 20 --- vllm/model_executor/layers/utils.py | 25 +--- vllm/model_executor/models/gpt_oss.py | 67 +-------- vllm/triton_utils/jit_monitor.py | 141 ------------------ .../backends/rocm_aiter_unified_attn.py | 88 +---------- vllm/v1/sample/ops/topk_topp_sampler.py | 53 ------- vllm/v1/sample/sampler.py | 93 ------------ 10 files changed, 37 insertions(+), 497 deletions(-) delete mode 100644 vllm/triton_utils/jit_monitor.py diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 69f406e07264..32562f0d9d9f 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -431,15 +431,8 @@ def __init__( # for attn backends supporting query quantization self.query_quant = None - # ATOM keeps the query in bf16 (fp8 KV cache, q_descale=None). Skip the - # per-layer fp8 query quant (the aiter static scaled_quant kernel) when - # VLLM_DISABLE_QUERY_QUANT=1 to match ATOM and drop ~2.3us/marker. - import os as _os - - _disable_qq = _os.environ.get("VLLM_DISABLE_QUERY_QUANT", "0") == "1" if ( - not _disable_qq - and self.impl.supports_quant_query_input + self.impl.supports_quant_query_input and ( self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4" ) 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 index c447073298d6..ed1da9416c40 100644 --- 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 @@ -63,10 +63,16 @@ def aiter_triton_kernel_w4a8_moe_forward( gating_output, topk, sm_first=not renormalize ) - # PATCH_GFX1250_INKERNEL_GATHER: ATOM passes gather_idx in-kernel on gfx1250 - # with GFX1250_SCALE-swizzled weight scales (see patch_swizzle.py). The old - # "in-kernel gather broken" was with UNSWIZZLED scales; with the correct - # GFX1250_SCALE swizzle the in-kernel gather is correct, so keep gather_idx. + # gfx1250: aiter's in-kernel gather is numerically broken (validated on the + # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted + # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, + # sorted row i reads source token gather_idx[i] // n_expts_act, so this + # reproduces the in-kernel gather exactly (manual gather -> maxrel ~5e-3). + # gfx950 keeps the (working) in-kernel gather. + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None return triton_kernel_fused_mxfp4_w4a8_experts( None, @@ -135,9 +141,10 @@ def triton_kernel_fused_mxfp4_w4a8_experts( ) from vllm.platforms.rocm import on_gfx1250 + from vllm.platforms.rocm import on_gfx1250 _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None # TODO (JPVILLAM): merge conflict resolve later if _swizzle_mx_scale is enough - mx_scale_swizzle = "GFX1250_SCALE" if on_gfx1250() else "CDNA4_SCALE" + mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" assert quant_config.w1_precision is not None, ( "w1_precision in quant config can't be None" 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 e50d10201cac..8b688697d167 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 @@ -6,6 +6,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -56,7 +57,7 @@ def _triton_kernel_moe_supports_current_device() -> bool: # on_gfx9() already excludes gfx906/gfx908. # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). - return on_gfx9() or on_gfx1x() or on_gfx1250() + return on_gfx9() or on_gfx1x() or on_gfx1250() return False @@ -1131,12 +1132,16 @@ def _aiter_moe_two_gemm( arch = get_arch() quant_dtype = torch.float8_e4m3fnuz if arch == "gfx942" else torch.float8_e4m3fn - # PATCH_GATHER: in-kernel gather on gfx1250 (ATOM behavior). ATOM (same aiter - # commit) passes gather_indx straight to moe_gemm_a8w4 on gfx1250; do the same - # to drop the eager vectorized_gather + floordiv. (If this miscomputes, the - # gfx1250 in-kernel gather needs GFX1250_SCALE-swizzled scales -> add swizzle.) - gemm1_input = hidden_states - gemm1_gather_indx = gather_idx + # gfx1250's in-kernel gather miscomputes (validated on the FFM sim), so gather + # rows into expert-sorted order in torch and pass gather_indx=None instead. Per + # aiter's moe_gemm_torch, sorted row i reads source token gather_idx[i] // topk. + if arch == "gfx1250": + gather_src = gather_idx.to(torch.long) // topk + gemm1_input = hidden_states[gather_src] + gemm1_gather_indx = None + else: + gemm1_input = hidden_states + gemm1_gather_indx = gather_idx gammas = routing_data.gate_scal if routing_data else None g1_gammas = gammas if apply_router_weight_on_input else None @@ -1558,7 +1563,7 @@ def apply( self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG if expert_map is not None: - topk_ids = expert_map[topk_ids] + topk_ids = expert_map[topk_ids] # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index fc26d76b2e22..5110bcb7be70 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -102,26 +102,6 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): # transpose the tensor so that the quantization axis is on dim1 quant_tensor = quant_tensor.transpose(-2, -1) scale = scale.transpose(-2, -1) - # GFX1250: apply ATOM's gfx1250 MX-scale swizzle so the in-kernel gather and - # moe_gemm_a8w4 read the scale correctly (paired with swizzle_mx_scale= - # "GFX1250_SCALE" in aiter_mxfp4_w4a8_moe). scale is (E, K_SCALE, N) here. - try: - from vllm.platforms.rocm import on_gfx1250 as _on_gfx1250 - except Exception: - _on_gfx1250 = lambda: False - if current_platform.is_rocm() and _on_gfx1250(): - from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( - swizzle_scales_gfx1250 as _swz_gfx1250, - ) - - assert ( - scale.dim() == 3 and scale.shape[-1] % 32 == 0 and scale.shape[-2] % 8 == 0 - ), f"GFX1250 scale swizzle needs (E,K_SCALE%8,N%32); got {tuple(scale.shape)}" - # ATOM passes the raw swizzle_scales_gfx1250 output directly (no - # .contiguous(), no re-layout) -- the kernel's GFX1250_SCALE TMA - # descriptor expects exactly these strides. Adding .contiguous() - # changes the strides and breaks tt.make_tensor_descriptor. - scale = _swz_gfx1250(scale) quant_tensor = convert_layout( wrap_torch_tensor(quant_tensor, dtype=FP4), value_layout, **value_layout_opts ) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index fba82d762086..814142ce053f 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -99,8 +99,6 @@ def default_unquantized_gemm( def use_aiter_triton_gemm(n, m, k, dtype): - from vllm.platforms.rocm import on_gfx1250 - if ( not rocm_aiter_ops.is_triton_gemm_enabled() # MI300's - fp8nuz=True @@ -109,17 +107,9 @@ def use_aiter_triton_gemm(n, m, k, dtype): ): return False - # PATCH: gfx1250 branch-aiter gemm_a16w16 is 8-16x FASTER than rocBLAS even - # for large prefill batches (validated: 16384x5120x2880 -> 685us vs 5917us). - # The original "use hipblaslt for larger GEMMs" guard sent big prefill qkv/o_proj - # to rocBLAS (Cijk). Disable it on gfx1250 so large prefill GEMMs use aiter too. - if n > 2048 and m > 512 and not on_gfx1250(): + # use hipblaslt for the larger GEMMs + if n > 2048 and m > 512: return False - # PATCH: gfx1250 aiter gluon gemm_a16w16 handles arbitrary shapes (incl. - # lm_head m=201088,k=2880 and K%256!=0). Route ALL bf16 dense GEMMs through - # it so lm_head (and any other shape) stops falling back to rocBLAS (Cijk). - if on_gfx1250(): - return True return ( (m == 5120 and k == 2880) or (m == 2880 and k == 4096) @@ -132,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 + from vllm.platforms.rocm import on_gfx1250, on_gfx1x, on_gfx9, on_gfx950 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -178,10 +168,9 @@ def rocm_unquantized_gemm_impl( # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b # hidden), so skip aiter there and fall back to the torch GEMM path below. - # PATCH: branch aiter gemm_a16w16 handles K%256!=0 (e.g. gpt-oss-120b K=2880), - # validated numerically. Drop the stale gfx1250 guard so these bf16 projections - # use aiter gluon GEMM instead of falling back to slow rocBLAS (Cijk) torch.linear. - if use_aiter_triton_gemm(n, m, k, x.dtype): + if use_aiter_triton_gemm(n, m, k, x.dtype) and not ( + on_gfx1250() and k % 256 != 0 + ): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) @@ -190,7 +179,7 @@ def rocm_unquantized_gemm_impl( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) # build (gfx9/gfx11 ISA); fall back to torch GEMM there. - and not on_gfx1250() # TODO GFX1250: Remove once skinny GEMM is supported on gfx1250 + and not on_gfx1250() # TODO GFX1250: Remove once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 8c054c9ff6e3..863d8ea3ba7f 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -36,59 +36,6 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.utils import rocm_unquantized_gemm -from vllm.utils.torch_utils import direct_register_custom_op - -_GPT_OSS_NORMPAD_LOGGED = False - - -def _gpt_oss_rmsnorm_pad_add( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float, - pad_to: int, -) -> tuple[torch.Tensor, torch.Tensor]: - # ATOM-equivalent: add + rmsnorm (over the original N) + zero-pad to `pad_to`. - from aiter.ops.triton.normalization.fused_add_rmsnorm_pad import ( - fused_add_rmsnorm_pad, - ) - - out, res_out = fused_add_rmsnorm_pad(x, weight, eps, residual, pad_to) - global _GPT_OSS_NORMPAD_LOGGED - if not _GPT_OSS_NORMPAD_LOGGED: - import sys - - print( - f"[normpad] x={tuple(x.shape)} pad_to={pad_to} out={tuple(out.shape)} " - f"res_out={tuple(res_out.shape)}", - file=sys.stderr, - flush=True, - ) - _GPT_OSS_NORMPAD_LOGGED = True - return out, res_out - - -def _gpt_oss_rmsnorm_pad_add_fake( - x: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float, - pad_to: int, -) -> tuple[torch.Tensor, torch.Tensor]: - M, N = x.shape - n_out = ((N + pad_to - 1) // pad_to) * pad_to if pad_to > 0 else N - return ( - torch.empty((M, n_out), dtype=x.dtype, device=x.device), - torch.empty_like(residual), - ) - - -direct_register_custom_op( - op_name="gpt_oss_rmsnorm_pad_add", - op_func=_gpt_oss_rmsnorm_pad_add, - mutates_args=[], - fake_impl=_gpt_oss_rmsnorm_pad_add_fake, -) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -336,19 +283,7 @@ def forward( hidden_states = self.attn(hidden_states, positions) # Fully Connected - _pad_to = self.mlp.experts.moe_config.hidden_dim - if residual is not None and _pad_to > hidden_states.shape[-1]: - hidden_states, residual = torch.ops.vllm.gpt_oss_rmsnorm_pad_add( - hidden_states, - residual, - self.post_attention_layernorm.weight, - self.post_attention_layernorm.variance_epsilon, - _pad_to, - ) - else: - hidden_states, residual = self.post_attention_layernorm( - hidden_states, residual - ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) output = self.mlp(hidden_states) return output, residual diff --git a/vllm/triton_utils/jit_monitor.py b/vllm/triton_utils/jit_monitor.py deleted file mode 100644 index e7d303440281..000000000000 --- a/vllm/triton_utils/jit_monitor.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Monitor unexpected Triton kernel JIT compilation during inference. - -After server warmup completes, any Triton JIT compilation or autotuning -event indicates a cache miss or unexpected input shape that causes a -latency spike. This module registers hooks in the Triton runtime to -detect and log such events so they can be investigated. - -Set ``--jit-monitor-verbose`` to log every Triton JIT compile with its -dispatch key. This is intentionally opt-in because it can emit many logs and -add overhead. - -Currently monitors: -- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``) -- Triton ``@triton.jit`` first-time compilations - (via ``knobs.runtime.jit_post_compile_hook``) -""" - -import os - -from vllm.logger import init_logger -from vllm.triton_utils.importing import HAS_TRITON - -logger = init_logger(__name__) - -_active: bool = False -_verbose: bool = False - - -def is_active() -> bool: - """Return whether the JIT compilation monitor is currently active.""" - return _active - - -def activate(*, verbose: bool = False) -> None: - """Enable JIT compilation monitoring after warmup. - - Call once per worker process at the end of - :func:`compile_or_warm_up_model`. After activation every Triton - kernel compilation or autotuning benchmark that happens during - inference will be logged as a warning. - - Safe to call multiple times — subsequent calls are no-ops. - - If the user has explicitly set ``TRITON_PRINT_AUTOTUNING=0`` in - their environment, autotuning printing is left disabled; the JIT - compilation hook is still registered regardless. - """ - global _active, _verbose - if _active: - return - _active = True - _verbose = verbose - - _setup_triton_autotuning_print() - _setup_triton_jit_hook() - - logger.info( - "Kernel JIT monitor activated — Triton JIT compilations " - "during inference will be logged as warnings." - ) - - -# ------------------------------------------------------------------ -# Triton autotuning print -# ------------------------------------------------------------------ - - -def _setup_triton_autotuning_print() -> None: - """Enable ``TRITON_PRINT_AUTOTUNING`` unless the user opted out.""" - if not HAS_TRITON: - return - from triton import knobs # type: ignore[import-untyped] - - user_val = os.environ.get("TRITON_PRINT_AUTOTUNING") - if user_val == "0": - logger.debug( - "TRITON_PRINT_AUTOTUNING=0 set by user — " - "autotuning messages will stay suppressed." - ) - return - - knobs.autotuning.print = True - - -# ------------------------------------------------------------------ -# Triton JIT compilation hook -# ------------------------------------------------------------------ - - -def _log_jit_compile(fn_name: str, kwargs) -> None: - if _verbose: - compile_info = kwargs.get("compile") - if not isinstance(compile_info, dict): - compile_info = {} - logger.warning( - "Triton %sJIT compilation during inference: %s (key=%s).", - "autotune/warmup candidate " if kwargs.get("warmup") else "kernel ", - fn_name, - compile_info.get("key") or kwargs.get("key"), - ) - return - - logger.warning_once( - "Triton kernel JIT compilation during inference: %s. " - "This causes a latency spike; consider extending warmup " - "to cover this shape/config.", - fn_name, - ) - - -def _setup_triton_jit_hook() -> None: - """Register a ``jit_post_compile_hook`` that warns on compilation.""" - # PATCH: upstream triton efbae9a's serialize_specialization_data cannot - # JSON-encode Gluon PaddedSharedLayout constexprs (passed by e.g. aiter - # gemm_a16w16), and triton builds spec-data for ANY registered hook. - # This monitor is purely diagnostic, so skip registering to avoid - # crashing live kernel compilation. - return - if not HAS_TRITON: - return - from triton import knobs # type: ignore[import-untyped] - - existing_hook = knobs.runtime.jit_post_compile_hook - - def _on_jit_compile(**kwargs): - # `jit_post_compile_hook` is Triton internal API and its - # signature has changed across releases (kwargs added/renamed). - # Accept **kwargs so an upstream change cannot crash this hook - # with TypeError, and forward the full kwarg set to any - # pre-existing hook unchanged. - fn = kwargs.get("fn") - fn_name = getattr(fn, "name", "") - _log_jit_compile(fn_name, kwargs) - if existing_hook is not None: - return existing_hook(**kwargs) - return None - - knobs.runtime.jit_post_compile_hook = _on_jit_compile diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index abc7bc0b3025..d8363169a8c8 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -85,13 +85,6 @@ def get_kv_cache_shape( ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") - if rocm_aiter_ops.is_shuffle_kv_cache_enabled(): - # ATOM-style K/V-outermost layout: K=cache[0], V=cache[1] each - # standalone-contiguous (block stride == block content). Required by - # the gfx1250 gluon unified-attn-3d kernel (TDM descriptors assume - # contiguous per-block K/V). The default (num_blocks, 2, ...) - # interleaves K/V per block (2x gap) -> gluon reads garbage. - return (2, num_blocks, block_size, num_kv_heads, head_size) return (num_blocks, 2, block_size, num_kv_heads, head_size) @staticmethod @@ -155,12 +148,6 @@ def __init__( def _split_kv_cache( self, kv_cache: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - if ( - rocm_aiter_ops.is_shuffle_kv_cache_enabled() - and self.attn_type != AttentionType.ENCODER_DECODER - ): - # ATOM layout (2, num_blocks, block_size, num_kv_heads, head_size). - return kv_cache.unbind(0) if self.attn_type != AttentionType.ENCODER_DECODER: return kv_cache.unbind(1) @@ -182,18 +169,6 @@ def _split_kv_cache( ) return kv_cache.unbind(0) - def _shuffle_kv_views( - self, key_cache: torch.Tensor, value_cache: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - # No-copy reinterpret of standalone-contiguous K/V (num_blocks, - # block_size, num_kv_heads, head_size) as the 5D shuffle layout the - # gluon kernel reads (the shuffle WRITER laid bytes in this order). - nb, bs, nkv, hd = key_cache.shape - x = 16 // key_cache.element_size() - kc = key_cache.reshape(nb, nkv, hd // x, bs, x) - vc = value_cache.reshape(nb, nkv, bs // x, hd, x) - return kc, vc - def forward( self, layer: torch.nn.Module, @@ -260,16 +235,6 @@ def forward( if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) - _shuffled = ( - rocm_aiter_ops.is_shuffle_kv_cache_enabled() - and self.attn_type != AttentionType.ENCODER_DECODER - ) - if _shuffled: - key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache) - # Shuffle path pre-divides K/V by k/v_scale before the (scale=1.0) - # writer, so the stored fp8 == reshape_and_cache_flash's; descale with - # the same layer scales as the vanilla path. - _kd, _vd = layer._k_scale, layer._v_scale cu_seqlens_q = attn_metadata.query_start_loc seqused_k = attn_metadata.seq_lens @@ -293,11 +258,10 @@ def forward( block_table=block_table, softcap=self.logits_soft_cap, q_descale=layer._q_scale if query.dtype == self.fp8_dtype else None, - k_descale=_kd, - v_descale=_vd, + k_descale=layer._k_scale, + v_descale=layer._v_scale, sinks=self.sinks, output_scale=output_scale, - shuffled_kv_cache=_shuffled, ) return output @@ -316,39 +280,6 @@ def do_kv_cache_update( return key_cache, value_cache = self._split_kv_cache(kv_cache) - if ( - rocm_aiter_ops.is_shuffle_kv_cache_enabled() - and self.attn_type != AttentionType.ENCODER_DECODER - ): - # Write new tokens directly into the 5D shuffle byte layout the - # gluon kernel reads (ATOM-equivalent). Quantizes with scale=1.0. - from vllm.v1.attention.backends.rocm_aiter_fa import ( - reshape_and_cache_shuffle_triton, - ) - - if is_quantized_kv_cache(self.kv_cache_dtype): - key_cache = key_cache.view(self.fp8_dtype) - value_cache = value_cache.view(self.fp8_dtype) - # PATCH3: the shuffle writer stores at scale=1.0 and the reader - # descales by layer._k/v_scale. For gpt-oss these scales are 1.0, - # so NO pre-division is needed (key/1.0 == key) and the earlier - # `(key/scale).to(dtype)` was a wasted full-K/V fp32 elementwise - # pass every layer (+2.7ms/prefill step). Drop it. - # (If a model ships k/v scales != 1.0, the shuffle writer kernel - # must be taught to apply them; gpt-oss does not.) - _dummy = torch.empty(1, dtype=torch.float32, device=key.device) - reshape_and_cache_shuffle_triton( - key, - value, - key_cache, - value_cache, - slot_mapping, - self.kv_cache_dtype, - _dummy, - _dummy, - ) - return - # Reshape the input keys and values and store them in the cache. ops.reshape_and_cache_flash( key, @@ -381,26 +312,13 @@ def do_rope_and_kv_cache_update( # we use direct Q, K, V tensors without caching return key_cache, value_cache = self._split_kv_cache(kv_cache) + flash_layout = True is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) if is_fp8_kv_cache: key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) - # PATCH4_FUSED_ROPE_SHUFFLE: when the shuffle KV-cache layout is active, - # the fused RoPE+cache kernel must write the SAME 5D shuffle byte layout - # the gluon unified-attn-3d reader expects (ATOM parity). Reinterpret the - # standalone-contiguous K/V as the 5D shuffle views (no copy) and use the - # non-flash (shuffle) layout. Stores at layer k/v scales (=1.0 for - # gpt-oss) -> byte-identical to the non-fused shuffle writer (PATCH3). - _shuffled = ( - rocm_aiter_ops.is_shuffle_kv_cache_enabled() - and self.attn_type != AttentionType.ENCODER_DECODER - ) - if _shuffled: - key_cache, value_cache = self._shuffle_kv_views(key_cache, value_cache) - flash_layout = not _shuffled - rocm_aiter_ops.triton_rope_and_cache( query, key, diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 5037054dd533..344c53b04b3a 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -18,48 +18,6 @@ logger = init_logger(__name__) -import os as _os # PATCH5_AITER_TEMP_SAMPLE - -_PATCH5_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" -_PATCH5_DEBUG_DONE = False - - -def _aiter_temp_gumbel_sample(logits, generators, use_fp64_gumbel): - """Fused temperature Gumbel-max sampling via aiter (ATOM parity). - logits are ALREADY temperature-scaled by the parent Sampler, so temps=1. - Returns int64 token ids (num_tokens,).""" - global _PATCH5_DEBUG_DONE - import torch as _torch - from aiter import mixed_sample_outer_exponential - - n, vocab = logits.shape - out = _torch.empty(n, dtype=_torch.int32, device=logits.device) - nseed = len(generators) - if nseed == 0: - # common case: ATOM-style shared (1,vocab) exponential noise (cheap RNG) - exp = ( - _torch.empty((1, vocab), dtype=_torch.float32, device=logits.device) - .exponential_() - .expand(n, vocab) - ) - else: - exp = _torch.empty((n, vocab), dtype=_torch.float32, device=logits.device) - if nseed != n: - exp.exponential_() - for i, g in generators.items(): - exp[i].exponential_(generator=g) - temps = _torch.ones(n, dtype=_torch.float32, device=logits.device) - mixed_sample_outer_exponential(out, logits, exp, temps, eps=1e-10) - if not _PATCH5_DEBUG_DONE: - logger.info( - "PATCH5 aiter fused temp-sample ACTIVE (n=%d vocab=%d seeded=%d)", - n, - vocab, - nseed, - ) - _PATCH5_DEBUG_DONE = True - return out.to(_torch.int64) - def flashinfer_sampler_supported() -> bool: """Decide whether FlashInfer's top-p/top-k sampler can be used. @@ -176,17 +134,6 @@ def forward_native( The logits tensor may be updated in-place. """ - # PATCH5_AITER_TEMP_SAMPLE: fused no-filter temperature sampling (ATOM parity) - if ( - _PATCH5_ON - and k is None - and p is None - and not self.use_fp64_gumbel - and self.logprobs_mode not in ("processed_logits", "processed_logprobs") - ): - return _aiter_temp_gumbel_sample( - logits, generators, self.use_fp64_gumbel - ), None logits = apply_top_k_top_p(logits, k, p) logits_to_return = None if self.logprobs_mode == "processed_logits": diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index d55d0a8481ba..bb20432a0815 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -16,93 +16,6 @@ _SAMPLING_EPS = 1e-5 -import os as _os # PATCH6_AITER_RAW_SAMPLE - -from vllm.logger import init_logger as _p6_init_logger - -_p6log = _p6_init_logger("vllm.v1.sample.patch6") -_PATCH6_ON = _os.environ.get("VLLM_AITER_TEMP_SAMPLE", "0") == "1" -_PATCH6_DEBUG = _os.environ.get("VLLM_PATCH6_DEBUG", "0") == "1" -_PATCH6_DBG = False -_PATCH6_ELIG_DBG = False - - -def _patch6_eligible(sm, predict_bonus_token): - global _PATCH6_ELIG_DBG - if not _PATCH6_ON: - return False - if _PATCH6_DEBUG and not _PATCH6_ELIG_DBG: - try: - lp = sm.logitsprocs - h = sm.thinking_budget_state_holder - _p6log.info( - "PATCH6 ELIG: all_random=%r top_k=%r top_p=%r logprobs=%r lpids=%r " - "nopen=%r allow_None=%r bad=%r argmax=%r nonarg=%r think=%r bonus=%r", - sm.all_random, - sm.top_k, - sm.top_p, - sm.max_num_logprobs, - bool(sm.logprob_token_ids), - sm.no_penalties, - sm.allowed_token_ids_mask is None, - bool(sm.bad_words_token_ids), - list(lp.argmax_invariant), - list(lp.non_argmax_invariant), - (h.has_tracked_requests() if h is not None else None), - predict_bonus_token, - ) - except Exception as e: - _p6log.info("PATCH6 ELIG dbg err: %r", e) - _PATCH6_ELIG_DBG = True - if predict_bonus_token: - return False - if not sm.all_random or sm.temperature is None: - return False - if sm.top_k is not None or sm.top_p is not None: - return False - if sm.max_num_logprobs is not None or sm.logprob_token_ids: - return False - if ( - not sm.no_penalties - or sm.allowed_token_ids_mask is not None - or sm.bad_words_token_ids - ): - return False - lp = sm.logitsprocs - if lp.argmax_invariant or lp.non_argmax_invariant: - return False - h = sm.thinking_budget_state_holder - if h is not None and h.has_tracked_requests(): - return False - return True - - -def _aiter_raw_fused_sample(logits, sm): - """Sample token ids directly from RAW (bf16) logits via fused aiter kernel. int32 [n,1].""" - global _PATCH6_DBG - import torch as _t - from aiter import mixed_sample_outer_exponential - - n, vocab = logits.shape - out = _t.empty(n, dtype=_t.int32, device=logits.device) - gens = sm.generators - if not gens: - exp = ( - _t.empty((1, vocab), dtype=_t.float32, device=logits.device) - .exponential_() - .expand(n, vocab) - ) - else: - exp = _t.empty((n, vocab), dtype=_t.float32, device=logits.device) - exp.exponential_() - for i, g in gens.items(): - exp[i].exponential_(generator=g) - mixed_sample_outer_exponential(out, logits, exp, sm.temperature, eps=1e-10) - if not _PATCH6_DBG: - _p6log.info("PATCH6 raw-logits fused sampler ACTIVE (n=%d vocab=%d)", n, vocab) - _PATCH6_DBG = True - return out.unsqueeze(-1) - class Sampler(nn.Module): """ @@ -163,12 +76,6 @@ def forward( predict_bonus_token: bool = False, logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: - # PATCH6_AITER_RAW_SAMPLE: unconstrained all-random batch -> fused raw-logits sampler - if _patch6_eligible(sampling_metadata, predict_bonus_token): - return SamplerOutput( - sampled_token_ids=_aiter_raw_fused_sample(logits, sampling_metadata), - logprobs_tensors=None, - ) logprobs_mode = logprobs_mode_override or self.logprobs_mode # NOTE(woosuk): Use the original logits (before any penalties or # temperature scaling) for the top-k logprobs. From d73697e010caf08973d162bc654f5dda7c535e42 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 6 Jul 2026 21:43:41 -0500 Subject: [PATCH 51/68] Small changes to smoketest Signed-off-by: jpvillam --- benchmarks/vllm_smoketest.sh | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh index 88aaa9107c03..facd4f0f98b2 100644 --- a/benchmarks/vllm_smoketest.sh +++ b/benchmarks/vllm_smoketest.sh @@ -11,21 +11,23 @@ declare -A MODEL_SERVE_ARGS=( [DeepSeek-R1-0528-MXFP4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --dtype auto --no-enable-prefix-caching --disable-uvicorn-access-log --trust-remote-code" ) declare -A MODEL_ENV=( - [gpt-oss-120b-mxfp4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=0 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [gpt-oss-120b-mxfp4]="HSA_OVERRIDE_GFX_VERSION=12.5.0 HSA_ENABLE_SDMA=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 " [DeepSeek-R1-0528-MXFP4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" ) MODEL_NAME="gpt-oss-120b-mxfp4" # default -MODEL_PATH_OVERRIDE="" # overrided path +MODEL_PATH_OVERRIDE="" +MODEL_ENV_OVERRIDE="" MODEL_SET=0 # whether --model was explicitly passed RUN_LM_EVAL=0 PORT=8000 usage() { cat <<'EOF' -Usage: ./vllm_smoketest.sh [--model NAME] [--model-path PATH] [--lm-eval] [--port PORT] [--list] +Usage: ./vllm_smoketest.sh [--model NAME] [--model-path PATH] [--env ENV_VARS] [--lm-eval] [--port PORT] [--list] --model NAME Registry key to run (default: gpt-oss-120b-mxfp4) --model-path PATH Override model path for this run + --env ENV_VARS Override environment variables (e.g., "VAR1=val1 VAR2=val2") --lm-eval Run lm_eval after curl check --port PORT Server port (default: 8000) --list List available models @@ -36,6 +38,7 @@ while [[ $# -gt 0 ]]; do case "$1" in --model) MODEL_NAME="$2"; MODEL_SET=1; shift 2 ;; --model-path) MODEL_PATH_OVERRIDE="$2"; shift 2 ;; + --env) MODEL_ENV_OVERRIDE="$2"; shift 2 ;; --lm-eval) RUN_LM_EVAL=1; shift ;; --port) PORT="$2"; shift 2 ;; --list) for n in "${!MODEL_PATHS[@]}"; do echo "$n -> ${MODEL_PATHS[$n]}"; done; exit 0 ;; @@ -53,11 +56,9 @@ MODEL="${MODEL_PATHS[$MODEL_NAME]:-$MODEL_NAME}" [[ -n "$MODEL_PATH_OVERRIDE" ]] && MODEL="$MODEL_PATH_OVERRIDE" SERVE_ARGS="${MODEL_SERVE_ARGS[$MODEL_NAME]:-}" MODEL_ENV_ARGS="${MODEL_ENV[$MODEL_NAME]:-}" +[[ -n "$MODEL_ENV_OVERRIDE" ]] && MODEL_ENV_ARGS="$MODEL_ENV_OVERRIDE" echo "Model: $MODEL | Port: $PORT | lm_eval: $RUN_LM_EVAL" -# temp -pip3 install "fastapi<0.137" - LOG_DIR="$(pwd)/vllm_smoke_test_logs" mkdir -p "$LOG_DIR" echo "Logs: $LOG_DIR" @@ -84,11 +85,12 @@ echo "=== Curl completion check ===" http_code=$(curl -s -o "$LOG_DIR/curl_completion.log" -w '%{http_code}' \ "http://localhost:$PORT/v1/completions" \ -H "Content-Type: application/json" \ - -d "{\"model\":\"$MODEL\",\"prompt\":\"San Francisco is a\",\"max_tokens\":32,\"temperature\":0}") + -d "{\"model\":\"$MODEL\",\"prompt\":\"The capital of France is \",\"max_tokens\":32,\"temperature\":0}") cat "$LOG_DIR/curl_completion.log"; echo [[ "$http_code" == "200" ]] || { echo "FAIL: HTTP $http_code" >&2; exit 1; } grep -q '"text"' "$LOG_DIR/curl_completion.log" || { echo "FAIL: no text in response" >&2; exit 1; } -echo "PASS" +grep -qi "Paris" "$LOG_DIR/curl_completion.log" || { echo "FAIL: 'Paris' not found in response" >&2; exit 1; } +CURL_CHECK_PASSED=1 # --- lm_eval --- if [[ $RUN_LM_EVAL -eq 1 ]]; then @@ -99,4 +101,10 @@ if [[ $RUN_LM_EVAL -eq 1 ]]; then --tasks gsm8k --num_fewshot 3 --limit 100 2>&1 | tee "$LOG_DIR/lm_eval.log" fi +# Shutdown server +cleanup + +# Print PASS only if curl check succeeded +[[ ${CURL_CHECK_PASSED:-0} -eq 1 ]] && echo "PASS" + echo "Done." From f18b3f8c2ed22e4d5105a69d7a81f5359a171b7a Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Wed, 8 Jul 2026 11:43:04 -0400 Subject: [PATCH 52/68] Address TODOs (#1047) * Address TODOs * Revert topk_topp_sampler and Dockerfile.rocm * Revert triton cmake Signed-off-by: Jaden Mathias --- .../layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 6 ------ vllm/model_executor/layers/utils.py | 2 +- vllm/platforms/rocm.py | 9 +++------ 3 files changed, 4 insertions(+), 13 deletions(-) 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 index ed1da9416c40..3c17b8c8b09e 100644 --- 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 @@ -53,8 +53,6 @@ def aiter_triton_kernel_w4a8_moe_forward( except ImportError: from aiter.ops.triton.moe_routing import routing as _routing_mod - # TODO: (JPVILLAM) This causes a tl compile error on 1250. - # Need to figure out why this is a problem and sync with triton team if on_gfx1250(): _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing @@ -139,12 +137,8 @@ def triton_kernel_fused_mxfp4_w4a8_experts( from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( should_use_cdna4_mx_scale_swizzle, ) - from vllm.platforms.rocm import on_gfx1250 - from vllm.platforms.rocm import on_gfx1250 _swizzle_mx_scale = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None - # TODO (JPVILLAM): merge conflict resolve later if _swizzle_mx_scale is enough - mx_scale_swizzle = None if on_gfx1250() else "CDNA4_SCALE" assert quant_config.w1_precision is not None, ( "w1_precision in quant config can't be None" diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 814142ce053f..efe14377324f 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -179,7 +179,7 @@ def rocm_unquantized_gemm_impl( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) # build (gfx9/gfx11 ISA); fall back to torch GEMM there. - and not on_gfx1250() # TODO GFX1250: Remove once skinny GEMM is supported on gfx1250 + # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 062af2993933..0884b874c60f 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -238,7 +238,6 @@ def _get_gcn_arch() -> str: for arch in _GCN_ARCH if (arch.startswith("gfx11") or arch.startswith("gfx12")) and arch != "gfx1250" ) -# TODO GFX1250: Use CDNA for MI3OR4 def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: @@ -313,7 +312,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: def on_gfx1x() -> bool: - return _ON_GFX1X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly + return _ON_GFX1X and not _ON_CDNA def on_gfx11() -> bool: @@ -329,9 +328,7 @@ def on_gfx1151() -> bool: def on_gfx12x() -> bool: - return ( - _ON_GFX12X # and not _ON_GFX1250 TODO GFX1250: should we skip this explicitly - ) + return _ON_GFX12X and not _ON_CDNA def on_gfx1250() -> bool: @@ -919,7 +916,7 @@ def supports_mx(cls) -> bool: @classmethod def supports_fp8(cls) -> bool: - return on_cdna() or on_gfx12x() # TODO GFX1250: We know this is redundant ATM + return on_cdna() or on_gfx12x() @classmethod def is_fp8_fnuz(cls) -> bool: From 952d2069eafb9cf05414aca6851171df3117dbc3 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Wed, 8 Jul 2026 13:22:03 -0400 Subject: [PATCH 53/68] Add gfx1250 gate (#1049) Signed-off-by: Jaden Mathias --- csrc/quickreduce/base.h | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index 86aaf14c64a4..e9af6d7e822c 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -79,6 +79,15 @@ union BufferResource { }; }; +#if !defined(__gfx1250__) +__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( + int32x4_t srsrc, int32_t voffset, int32_t soffset, + int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32"); + +__quickreduce_device_inline__ static void buffer_store_dwordx4( + int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, + int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32"); +#else __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {} @@ -86,6 +95,7 @@ __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( __quickreduce_device_inline__ static void buffer_store_dwordx4( int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {} +#endif __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) { #if defined(__gfx942__) From e63bae591381b2657e26b8a66ce8c2ebdece66f3 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Wed, 8 Jul 2026 14:25:20 -0400 Subject: [PATCH 54/68] Revert FFM Layer passes for hf overrides (#1050) Signed-off-by: Jaden Mathias --- vllm/model_executor/models/gpt_oss.py | 15 --------------- vllm/models/deepseek_v4/amd/model.py | 17 ----------------- 2 files changed, 32 deletions(-) diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 863d8ea3ba7f..01f2752ac54b 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -413,14 +413,6 @@ def _load_weights_mxfp4( if is_pp_missing_parameter(name, self): continue - # Skip checkpoint weights for layers dropped by an hf_overrides - # `num_hidden_layers` prune: the model only builds params for the - # layers it keeps, so extra-layer weights have no destination param. - if "layers." in name: - idx_str = name.split("layers.", 1)[1].split(".", 1)[0] - if idx_str.isdigit() and int(idx_str) >= self.config.num_hidden_layers: - continue - if ".w13_weight_scale" in name: # Handle MLP gate and up projection weights scale if use_ep: @@ -643,13 +635,6 @@ def _get_moe_weight_dtype(layer_id: int = 0) -> str | None: if is_pp_missing_parameter(name, self): continue - # Skip checkpoint weights for layers dropped by an hf_overrides - # `num_hidden_layers` prune (mirrors _load_weights_mxfp4). - if "layers." in name: - idx_str = name.split("layers.", 1)[1].split(".", 1)[0] - if idx_str.isdigit() and int(idx_str) >= self.config.num_hidden_layers: - continue - layer_id, expert_id, fused_name = None, None, None moe_quant_method = None if "experts" in name: diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 2c6501452f87..edb923511501 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -626,16 +626,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: expert_mapping = self.get_expert_mapping() for name, loaded_weight in weights: - # Skip checkpoint weights for transformer layers pruned by a - # num_hidden_layers hf_override (keeps load + run tractable on the - # FFM simulator). Names here are already vLLM-mapped ("layers.N."). - if "layers." in name: - try: - _li = int(name.split("layers.", 1)[1].split(".", 1)[0]) - if _li >= self.config.num_hidden_layers: - continue - except (IndexError, ValueError): - pass for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if ".experts." in name: @@ -692,10 +682,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: elif "attn_sink" in name: if is_pp_missing_parameter(name, self): continue - if name not in params_dict: - # Layer pruned via a num_hidden_layers hf_override; skip - # its checkpoint weights instead of KeyError-ing. - continue narrow_weight = loaded_weight[head_rank_start:head_rank_end] n = narrow_weight.shape[0] params_dict[name][:n].copy_(narrow_weight) @@ -704,9 +690,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: else: if is_pp_missing_parameter(name, self): continue - if name not in params_dict: - # Layer pruned via a num_hidden_layers hf_override; skip. - continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader From b3bcf7155c1ca3cca72ad0f24ed34e7d7230b833 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Wed, 8 Jul 2026 15:37:19 -0400 Subject: [PATCH 55/68] Revert unused imports (#1051) Signed-off-by: Jaden Mathias --- vllm/models/deepseek_v4/compressor.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index f129f19a5716..1efa987fe7bd 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -13,9 +13,6 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( - _fused_kv_compress_norm_rope_insert_indexer_attn, - _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, - _fused_kv_compress_norm_rope_insert_sparse_attn, compress_norm_rope_store_triton, ) from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE From 3be848c1183786f11dcb77ce9a363abf54274b90 Mon Sep 17 00:00:00 2001 From: Daniel Date: Wed, 8 Jul 2026 16:29:24 -0400 Subject: [PATCH 56/68] enabling skinny gemm compilation for non gfx1250 archs Signed-off-by: Daniel --- CMakeLists.txt | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3fada1a6d2c7..b23a307d7d0c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1376,9 +1376,26 @@ if(VLLM_GPU_LANG STREQUAL "HIP") # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude it from the # gfx1250 build and disable its op registrations (VLLM_SKIP_SKINNY_GEMMS); # vLLM falls back to default/Triton GEMM for those ops on gfx1250. - if(VLLM_GPU_ARCHES MATCHES "gfx1250") - message(STATUS "gfx1250 detected: excluding csrc/rocm/skinny_gemms.cu from _rocm_C") - list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES}) + list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250") + if(VLLM_SKINNY_ARCHES) + message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}") + hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "csrc/rocm/skinny_gemms.cu") + unset(_VLLM_LAST_HIPIFY_TARGET) + add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS}) + add_dependencies(_rocm_C_skinny hipify_all) + set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG}) + set_target_properties(_rocm_C_skinny PROPERTIES + ${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}" + POSITION_INDEPENDENT_CODE ON) + target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc) + target_compile_options(_rocm_C_skinny PRIVATE + $<$:${VLLM_ROCM_EXT_FLAGS}>) + target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C") + target_link_libraries(_rocm_C_skinny PRIVATE torch) + else() + message(STATUS "Only gfx1250, skipping skinny_gemms") list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS") endif() @@ -1401,8 +1418,9 @@ if(VLLM_GPU_LANG STREQUAL "HIP") USE_SABI 3 WITH_SOABI) - # Needed to pass skip skinny flag to .cpp targets - if(VLLM_GPU_ARCHES MATCHES "gfx1250") + if(TARGET _rocm_C_skinny) + target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny) + else() target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS) endif() From 9f354bc375671f54c200082c873f9a1e589c4390 Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 9 Jul 2026 11:18:51 -0400 Subject: [PATCH 57/68] temp making a 2 base dockerfiles for cdna and gor gfx1250 builds Signed-off-by: Daniel --- docker/Dockerfile.rocm_base | 2 +- docker/Dockerfile.rocm_base_cdna | 290 +++++++++++++++++++++++++++++++ 2 files changed, 291 insertions(+), 1 deletion(-) create mode 100644 docker/Dockerfile.rocm_base_cdna diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index ec5e85cd943a..ce16a33e5b80 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -25,7 +25,7 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250 +ARG PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV MORI_GPU_ARCHS=gfx942;gfx950 diff --git a/docker/Dockerfile.rocm_base_cdna b/docker/Dockerfile.rocm_base_cdna new file mode 100644 index 000000000000..c51efa060508 --- /dev/null +++ b/docker/Dockerfile.rocm_base_cdna @@ -0,0 +1,290 @@ +ARG BASE_IMAGE=ubuntu:24.04 +ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ +ARG ROCM_SDK_VERSION=7.14.0a20260623 +ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 +ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 +ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 + + +ARG FA_BRANCH="0e60e394" +ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" +ARG AITER_BRANCH="main" +ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG MORI_BRANCH="v1.1.0" +ARG MORI_REPO="https://github.com/ROCm/mori.git" + +# Sccache configuration (only used in release pipeline) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + +FROM ${BASE_IMAGE} AS base + +ARG PYTORCH_ROCM_ARCH=gfx942;gfx950 +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MORI_GPU_ARCHS=gfx942;gfx950 +ENV FA_GPU_ARCHS=gfx942;gfx950 + +# temp dockerfile to enable CK for cdna4 archs +ENV ENABLE_CK=1 +ARG PREBUILD_KERNELS=0 + +# Required for RCCL in ROCm7.1 +ENV HSA_NO_SCRATCH_RECLAIM=1 + +ARG PYTHON_VERSION=3.12 +ENV PYTHON_VERSION=${PYTHON_VERSION} + +RUN mkdir -p /app +WORKDIR /app +ENV DEBIAN_FRONTEND=noninteractive + +# Install Python and other dependencies +RUN apt-get update -y \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ \ + && for i in 1 2 3; do \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + done \ + && apt-get update -y \ + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ + && python3 --version + +ENV VIRTUAL_ENV=/opt/venv +RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH + +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython +RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi + +## +## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +## +ARG ROCM_WHEEL_INDEX +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG ROCM_SDK_VERSION +ARG APEX_VERSION +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + --extra-index-url https://pypi.org/simple \ + "torch[device-all]==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + + + +# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream +# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm; + +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi + +# Setup sccache for HIP compilation via HIP_CLANG_PATH +# This creates wrapper scripts in a separate directory and points HIP to use them +# This avoids modifying the original ROCm binaries which can break detection +# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) +# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +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} + + +### +### AMD SMI Build +### +FROM base AS build_amdsmi +RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ + && pip wheel . --wheel-dir=dist +RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install + + +### +### MORI Build TODO(Build needs fixing) +### +FROM base AS build_mori +ARG MORI_BRANCH +ARG MORI_REPO +ARG MORI_GPU_ARCHS +RUN mkdir -p /app/install; \ + if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ + echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ + else \ + git clone ${MORI_REPO} \ + && cd mori \ + && git checkout ${MORI_BRANCH} \ + && git submodule update --init --recursive \ + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ + && cp /app/mori/dist/*.whl /app/install; \ + fi + + +### +### FlashAttention Build +### +FROM base AS build_fa +ARG FA_BRANCH +ARG FA_REPO +ARG USE_SCCACHE +RUN mkdir -p /app/install; \ + git clone ${FA_REPO} \ + && cd flash-attention \ + && git checkout ${FA_BRANCH} \ + && git submodule update --init \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && cp dist/*.whl /app/install; + + +### +### AITER Build +### +FROM base AS build_aiter +ARG AITER_BRANCH +ARG AITER_REPO +ARG USE_SCCACHE +RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} +RUN cd aiter \ + && git submodule update --init --recursive \ + && pip install -r requirements.txt +RUN pip install pyyaml && cd aiter \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && ls /app/aiter/dist/*.whl +RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install + + +### +### Final Build +### + +# Wheel release stage - +# only includes dependencies used by wheel release pipeline +FROM base AS debs_wheel_release +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs + +# Full debs stage - includes Mori (used by Docker releases) +FROM base AS debs +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi + +FROM base AS final +RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ + pip install /install/*.whl + +ARG BASE_IMAGE +ARG ROCM_WHEEL_INDEX +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG FA_BRANCH +ARG FA_REPO +ARG AITER_BRANCH +ARG AITER_REPO +ARG MORI_BRANCH +ARG MORI_REPO +RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ + && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ + && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ + && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ + && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ + && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ + && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt From a869a25a1aabac7f930d7e4096f67e014c723ee3 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Thu, 9 Jul 2026 13:31:41 -0500 Subject: [PATCH 58/68] Patch for mori build Signed-off-by: Jaden Mathias Signed-off-by: Daniel --- docker/Dockerfile.rocm_base | 33 +++++++++++++++++++++++++++------ 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index ce16a33e5b80..68f3bd55e590 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -47,7 +47,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ @@ -137,6 +137,31 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ fi +# The ROCm SDK wheel ships a broken CMake export for hsakmt +# This patch includes the right paths for the numa build target +RUN <<'EOF' +set -eu +TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" +[ -f "$TARGETS" ] || exit 0 # nothing to patch +grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched + +# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. +sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" + +# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). +sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" + +# 3. Define the numa::numa target the export references but forgot to create. +cat >> "$TARGETS" <<'CMAKE' + +if(NOT TARGET numa::numa) + find_library(NUMA_LIBRARY NAMES numa REQUIRED) + add_library(numa::numa UNKNOWN IMPORTED) + set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") +endif() +CMAKE +EOF + # Setup sccache for HIP compilation via HIP_CLANG_PATH # This creates wrapper scripts in a separate directory and points HIP to use them # This avoids modifying the original ROCm binaries which can break detection @@ -178,16 +203,12 @@ ARG MORI_BRANCH ARG MORI_REPO ARG MORI_GPU_ARCHS RUN mkdir -p /app/install; \ - if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ - echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ - else \ git clone ${MORI_REPO} \ && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ - && cp /app/mori/dist/*.whl /app/install; \ - fi + && cp /app/mori/dist/*.whl /app/install; ### From 9296486daeafa5ab01a7a081f208f27816d0349c Mon Sep 17 00:00:00 2001 From: Daniel Date: Thu, 9 Jul 2026 16:12:11 -0400 Subject: [PATCH 59/68] Apply mori build patch to cdna base Signed-off-by: Daniel --- docker/Dockerfile.rocm_base_cdna | 33 ++++++++++++++++++++++++++------ 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile.rocm_base_cdna b/docker/Dockerfile.rocm_base_cdna index c51efa060508..0354a0fc75f1 100644 --- a/docker/Dockerfile.rocm_base_cdna +++ b/docker/Dockerfile.rocm_base_cdna @@ -47,7 +47,7 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config g++ \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ @@ -137,6 +137,31 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ fi +# The ROCm SDK wheel ships a broken CMake export for hsakmt +# This patch includes the right paths for the numa build target +RUN <<'EOF' +set -eu +TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" +[ -f "$TARGETS" ] || exit 0 # nothing to patch +grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched + +# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. +sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" + +# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). +sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" + +# 3. Define the numa::numa target the export references but forgot to create. +cat >> "$TARGETS" <<'CMAKE' + +if(NOT TARGET numa::numa) + find_library(NUMA_LIBRARY NAMES numa REQUIRED) + add_library(numa::numa UNKNOWN IMPORTED) + set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") +endif() +CMAKE +EOF + # Setup sccache for HIP compilation via HIP_CLANG_PATH # This creates wrapper scripts in a separate directory and points HIP to use them # This avoids modifying the original ROCm binaries which can break detection @@ -178,16 +203,12 @@ ARG MORI_BRANCH ARG MORI_REPO ARG MORI_GPU_ARCHS RUN mkdir -p /app/install; \ - if echo "${PYTORCH_ROCM_ARCH}" | grep -q "gfx1250"; then \ - echo "gfx1250 in PYTORCH_ROCM_ARCH; skipping MORI build"; \ - else \ git clone ${MORI_REPO} \ && cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ - && cp /app/mori/dist/*.whl /app/install; \ - fi + && cp /app/mori/dist/*.whl /app/install; ### From f3e40c2f088eb40a008a07847267c36b9037722f Mon Sep 17 00:00:00 2001 From: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:51:34 -0500 Subject: [PATCH 60/68] Update jit_monitor.py to allow eager mode --- vllm/utils/jit_monitor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 23a8037c5729..c896c275bd4a 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -148,6 +148,7 @@ def _log_triton_jit_compile(fn_name: str, kwargs) -> None: def _setup_triton_jit_hook() -> None: """Register a ``jit_post_compile_hook`` that warns on compilation.""" + return if not HAS_TRITON: return from triton import knobs # type: ignore[import-untyped] From 790a1b5d6466a8abb7980d6dc7428414ca8d7463 Mon Sep 17 00:00:00 2001 From: Daniel Date: Fri, 10 Jul 2026 16:44:16 -0400 Subject: [PATCH 61/68] splitting Dockerfile.rocm and updating defaults to upstream Signed-off-by: Daniel --- docker/Dockerfile.rocm | 66 +- docker/Dockerfile.rocm_base | 320 ++++---- ...base_cdna => Dockerfile.rocm_base_gfx1250} | 6 +- docker/Dockerfile.rocm_gfx1250 | 726 ++++++++++++++++++ 4 files changed, 936 insertions(+), 182 deletions(-) rename docker/{Dockerfile.rocm_base_cdna => Dockerfile.rocm_base_gfx1250} (99%) create mode 100644 docker/Dockerfile.rocm_gfx1250 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 23506e259b25..384b66412da6 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -28,14 +28,14 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG ARG_PYTORCH_ROCM_ARCH=gfx942;gfx950;gfx1250 +ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ - build-essential libnuma-dev ccache mold + libnuma-dev ccache mold RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --upgrade pip # Note: mold is installed but not set as the system default linker because @@ -61,10 +61,6 @@ ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy -# python binary fall back for non venv builds -ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 -# Expose paths from wheel installation -ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} # ccache directory - persisted across layer rebuilds via cache mounts. ENV CCACHE_DIR=/root/.cache/ccache ENV CCACHE_COMPILERCHECK=content @@ -114,8 +110,8 @@ WORKDIR ${COMMON_WORKDIR} FROM base AS fetch_vllm_0 ONBUILD COPY ./ vllm/ FROM base AS fetch_vllm_1 -ARG VLLM_REPO="https://github.com/ROCm/vllm.git" -ARG VLLM_BRANCH="455_wip" +ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" +ARG VLLM_BRANCH="main" ENV VLLM_REPO=${VLLM_REPO} ENV VLLM_BRANCH=${VLLM_BRANCH} ONBUILD RUN git clone ${VLLM_REPO} \ @@ -133,6 +129,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm # don't need the rust toolchain or protoc. FROM fetch_vllm AS rust-build ARG COMMON_WORKDIR +ARG USE_SCCACHE # protoc is used by tonic-build/prost-build. RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ @@ -148,6 +145,10 @@ ENV CARGO_BUILD_JOBS=4 ENV CARGO_NET_RETRY=10 ENV RUSTUP_MAX_RETRIES=10 +# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on +# a separate local sccache daemon while sharing the same remote cache backend. +ENV SCCACHE_SERVER_PORT=4227 + RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd ${COMMON_WORKDIR}/vllm \ && uv pip install --system -r requirements/build/rust.txt @@ -159,8 +160,15 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export RUSTC_WRAPPER=sccache \ + && sccache --show-stats; \ + fi \ && bash build_rust.sh \ - && test -x vllm/vllm-rs + && test -x vllm/vllm-rs \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + sccache --show-stats; \ + fi # ----------------------- # vLLM native build stages @@ -235,9 +243,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -246,7 +262,7 @@ ARG RIXL_BRANCH="39be1de8" ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" ARG UCX_BRANCH="bfb51733" ARG UCX_REPO="https://github.com/openucx/ucx.git" -# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image +ENV ROCM_PATH=/opt/rocm ENV UCX_HOME=/usr/local/ucx ENV RIXL_HOME=/usr/local/rixl ENV RIXL_BENCH_HOME=/usr/local/rixl_bench @@ -302,7 +318,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ git checkout ${RIXL_BRANCH} && \ CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ - --force-fallback-for=abseil-cpp \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ @@ -315,10 +330,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ RUN cd /opt/rixl && \ sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ contrib/build-wheel.sh && \ - # The wheel build re-runs meson via meson-python; force the bundled abseil - sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ - pyproject.toml && \ - grep -q 'force-fallback-for' pyproject.toml && \ mkdir -p /app/install && \ _ucx_install_dir=${UCX_HOME} \ ./contrib/build-wheel.sh \ @@ -335,7 +346,7 @@ ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" # DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so # it can be linked against DeepEP without arch mismatches. ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" -# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image +ENV ROCM_PATH=/opt/rocm ENV ROCSHMEM_DIR=/opt/rocshmem RUN --mount=type=cache,target=/root/.cache/ccache \ @@ -525,9 +536,17 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmar COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- @@ -561,6 +580,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ libibverbs1 \ ibverbs-providers \ ibverbs-utils \ + unzip \ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* @@ -584,6 +604,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ENV HF_XET_HIGH_PERFORMANCE=1 ENV HF_HUB_DOWNLOAD_TIMEOUT=60 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 + # Pre-install vLLM test dependencies. COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt RUN --mount=type=cache,target=/root/.cache/uv \ @@ -704,6 +727,9 @@ ENV SAFETENSORS_FAST_GPU=1 # Performance environment variable. ENV HIP_FORCE_DEV_KERNARG=1 +# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). +ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 + # Workaround for ROCm profiler limits RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" @@ -711,14 +737,6 @@ RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt - -### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better -RUN pip3 uninstall -y triton && \ - git clone https://github.com/triton-lang/triton.git && \ - cd triton && \ - git checkout c517f38c && \ - TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . - CMD ["/bin/bash"] #Set entrypoint for vllm-openai official images diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 68f3bd55e590..8764a410bb80 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,16 +1,15 @@ -ARG BASE_IMAGE=ubuntu:24.04 -ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ -ARG ROCM_SDK_VERSION=7.14.0a20260623 -ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 -ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 -ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 -ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 -ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 - - +ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete +ARG TRITON_BRANCH="0f380657" +ARG TRITON_REPO="https://github.com/ROCm/triton.git" +ARG PYTORCH_BRANCH="d0c8b1f3" # release/2.11 as of 6/09 +ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" +ARG PYTORCH_VISION_BRANCH="v0.24.1" +ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" +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="main" +ARG AITER_BRANCH="v0.1.16.post3" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -25,15 +24,13 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG PYTORCH_ROCM_ARCH=gfx1250 +ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENV ROCM_PATH=/opt/rocm +ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV AITER_ROCM_ARCH=gfx942;gfx950 ENV MORI_GPU_ARCHS=gfx942;gfx950 -ENV FA_GPU_ARCHS=gfx942;gfx950 - -# TODO: Unset these when support is available for gfx1250 -ENV ENABLE_CK=0 -ARG PREBUILD_KERNELS=0 # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 @@ -47,23 +44,19 @@ ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ done \ && apt-get update -y \ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && python3 --version - -ENV VIRTUAL_ENV=/opt/venv -RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ - "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML -ENV PATH=${VIRTUAL_ENV}/bin:$PATH + && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ + && python3 --version && python3 -m pip --version RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* @@ -76,105 +69,31 @@ ARG SCCACHE_BUCKET_NAME ARG SCCACHE_REGION_NAME ARG SCCACHE_S3_NO_CREDENTIALS RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi - -## -## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index -## -ARG ROCM_WHEEL_INDEX -ARG TORCH_VERSION -ARG TORCHVISION_VERSION -ARG TORCHAUDIO_VERSION -ARG ROCM_SDK_VERSION -ARG APEX_VERSION -ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages -ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel -ENV ROCM_HOME=${ROCM_PATH} -ENV ROCM_SOURCE_DIR=${ROCM_PATH} -ENV ROCM_BIN=${ROCM_PATH}/bin -ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake -ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode -ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH -ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib -ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake -ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi - -# torch/torchvision/torchaudio must be pinned to mutually-consistent builds -# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm -# sdk version is derived from torch's own dependency pin unless overridden, -# which keeps the set consistent and avoids pip backtracking. -RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ - --extra-index-url https://pypi.org/simple \ - "torch[device-all]==${TORCH_VERSION}" \ - "torchvision==${TORCHVISION_VERSION}" \ - "torchaudio==${TORCHAUDIO_VERSION}" \ - "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ - rocm-sdk init - -# Torch runtime deps that may not be published on the ROCm wheel index; -# install them from PyPI afterwards. -RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ - "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" - - - -# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream -# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. -RUN ln -sfn "${ROCM_PATH}" /opt/rocm; - -RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ - sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ - "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ fi -# The ROCm SDK wheel ships a broken CMake export for hsakmt -# This patch includes the right paths for the numa build target -RUN <<'EOF' -set -eu -TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" -[ -f "$TARGETS" ] || exit 0 # nothing to patch -grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched - -# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. -sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" - -# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). -sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" - -# 3. Define the numa::numa target the export references but forgot to create. -cat >> "$TARGETS" <<'CMAKE' - -if(NOT TARGET numa::numa) - find_library(NUMA_LIBRARY NAMES numa REQUIRED) - add_library(numa::numa UNKNOWN IMPORTED) - set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") -endif() -CMAKE -EOF - # Setup sccache for HIP compilation via HIP_CLANG_PATH # This creates wrapper scripts in a separate directory and points HIP to use them # This avoids modifying the original ROCm binaries which can break detection # NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) # Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ fi # Set sccache environment variables only when USE_SCCACHE=1 @@ -186,29 +105,103 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} +### +### Triton Build +### +FROM base AS build_triton +ARG TRITON_BRANCH +ARG TRITON_REPO +RUN git clone ${TRITON_REPO} +# Cherry picking the following +# https://github.com/triton-lang/triton/pull/8991 +RUN cd triton \ + && git checkout ${TRITON_BRANCH} \ + && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ + && git cherry-pick 555d04f \ + && if [ ! -f setup.py ]; then cd python; fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && mkdir -p /app/install && cp dist/*.whl /app/install +RUN if [ -d triton/python/triton_kernels ]; then pip install build && cd triton/python/triton_kernels \ + && python3 -m build --wheel && cp dist/*.whl /app/install; fi + + ### ### AMD SMI Build ### FROM base AS build_amdsmi -RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ +RUN cd /opt/rocm/share/amd_smi \ && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install +RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install + + +### +### Pytorch build +### +FROM base AS build_pytorch +ARG PYTORCH_BRANCH +ARG PYTORCH_VISION_BRANCH +ARG PYTORCH_AUDIO_BRANCH +ARG PYTORCH_REPO +ARG PYTORCH_VISION_REPO +ARG PYTORCH_AUDIO_REPO +ARG USE_SCCACHE + +RUN apt-get update && apt-get install -y pkg-config liblzma-dev +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 && python3 tools/amd_build/build_amd.py \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ + && sccache --show-stats; \ + fi \ + && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && pip install dist/*.whl +RUN git clone ${PYTORCH_VISION_REPO} vision +RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && pip install dist/*.whl +RUN git clone ${PYTORCH_AUDIO_REPO} audio +RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ + && git submodule update --init --recursive \ + && pip install -r requirements.txt \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && export CMAKE_C_COMPILER_LAUNCHER=sccache \ + && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ + fi \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && pip install dist/*.whl +RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ + && cp /app/vision/dist/*.whl /app/install \ + && cp /app/audio/dist/*.whl /app/install ### -### MORI Build TODO(Build needs fixing) +### MORI Build ### FROM base AS build_mori ARG MORI_BRANCH ARG MORI_REPO -ARG MORI_GPU_ARCHS -RUN mkdir -p /app/install; \ - git clone ${MORI_REPO} \ - && cd mori \ +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + pip install /install/*.whl +RUN git clone ${MORI_REPO} +RUN cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ - && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ - && cp /app/mori/dist/*.whl /app/install; + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl +RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install ### @@ -218,18 +211,19 @@ FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO ARG USE_SCCACHE -RUN mkdir -p /app/install; \ - git clone ${FA_REPO} \ - && cd flash-attention \ +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + pip install /install/*.whl +RUN git clone ${FA_REPO} +RUN cd flash-attention \ && git checkout ${FA_BRANCH} \ && git submodule update --init \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && cp dist/*.whl /app/install; + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi +RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install ### @@ -239,16 +233,18 @@ FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO ARG USE_SCCACHE +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + pip install /install/*.whl RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} RUN cd aiter \ && git submodule update --init --recursive \ && pip install -r requirements.txt RUN pip install pyyaml && cd aiter \ && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install @@ -262,35 +258,46 @@ RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install # only includes dependencies used by wheel release pipeline FROM base AS debs_wheel_release RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs # Full debs stage - includes Mori (used by Docker releases) FROM base AS debs RUN mkdir /app/debs +RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ cp /install/*.whl /app/debs RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi + cp /install/*.whl /app/debs FROM base AS final RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ pip install /install/*.whl ARG BASE_IMAGE -ARG ROCM_WHEEL_INDEX -ARG ROCM_SDK_VERSION -ARG TORCH_VERSION -ARG TORCHVISION_VERSION -ARG TORCHAUDIO_VERSION +ARG TRITON_BRANCH +ARG TRITON_REPO +ARG PYTORCH_BRANCH +ARG PYTORCH_VISION_BRANCH +ARG PYTORCH_REPO +ARG PYTORCH_VISION_REPO +ARG PYTORCH_AUDIO_BRANCH +ARG PYTORCH_AUDIO_REPO ARG FA_BRANCH ARG FA_REPO ARG AITER_BRANCH @@ -298,14 +305,17 @@ ARG AITER_REPO ARG MORI_BRANCH ARG MORI_REPO RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ - && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ - && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ - && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ - && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ + && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ + && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ + && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ + && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ + && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ + && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ + && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ + && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ - && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \ No newline at end of file diff --git a/docker/Dockerfile.rocm_base_cdna b/docker/Dockerfile.rocm_base_gfx1250 similarity index 99% rename from docker/Dockerfile.rocm_base_cdna rename to docker/Dockerfile.rocm_base_gfx1250 index 0354a0fc75f1..68f3bd55e590 100644 --- a/docker/Dockerfile.rocm_base_cdna +++ b/docker/Dockerfile.rocm_base_gfx1250 @@ -25,14 +25,14 @@ ARG SCCACHE_S3_NO_CREDENTIALS=0 FROM ${BASE_IMAGE} AS base -ARG PYTORCH_ROCM_ARCH=gfx942;gfx950 +ARG PYTORCH_ROCM_ARCH=gfx1250 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV MORI_GPU_ARCHS=gfx942;gfx950 ENV FA_GPU_ARCHS=gfx942;gfx950 -# temp dockerfile to enable CK for cdna4 archs -ENV ENABLE_CK=1 +# TODO: Unset these when support is available for gfx1250 +ENV ENABLE_CK=0 ARG PREBUILD_KERNELS=0 # Required for RCCL in ROCm7.1 diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 new file mode 100644 index 000000000000..9cc12e1142e5 --- /dev/null +++ b/docker/Dockerfile.rocm_gfx1250 @@ -0,0 +1,726 @@ +# default base image +ARG REMOTE_VLLM="0" +ARG COMMON_WORKDIR=/app +ARG BASE_IMAGE=rocm/vllm-dev:base +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base +# NIC backend for MoRI RDMA support. +# By default (all), drivers and userspace libraries for all supported NIC types +# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. +# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: +# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only +# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only +# --build-arg NIC_BACKEND=none # Install nothing. +ARG NIC_BACKEND=all +# AMD AINIC apt repo settings +# Users can specify a custom version compatible with their host drivers. +# The default version has been tested with ioinic-dkms=25.11.1.001 +ARG AINIC_VERSION=1.117.3-hydra +ARG UBUNTU_CODENAME=jammy + +# Sccache configuration. Release builds use this today; CI can opt in when a +# shared S3-compatible cache backend is available. +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + +FROM ${BASE_IMAGE} AS base + +ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 +ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} + +# Install build dependencies and utilities +RUN apt-get update -q -y && apt-get install -q -y \ + sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ + apt-transport-https ca-certificates wget curl \ + build-essential libnuma-dev ccache mold +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --upgrade pip +# Note: mold is installed but not set as the system default linker because +# some packages use JIT compilation at runtime with flags mold does not support. +# Build stages opt in via LDFLAGS="-fuse-ld=mold". +# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) +ARG USE_SCCACHE +RUN if [ "$USE_SCCACHE" != "1" ]; then \ + apt-get purge -y sccache || true; \ + python3 -m pip uninstall -y sccache || true; \ + rm -f "$(which sccache)" || true; \ + fi + +# Install UV — download first, then run, so a curl failure is not masked by the pipe +RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ + && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ + && rm -f /tmp/uv-install.sh \ + && uv --version + +# 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 +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +# Use copy mode to avoid hardlink failures with Docker cache mounts +ENV UV_LINK_MODE=copy +# python binary fall back for non venv builds +ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 +# Expose paths from wheel installation +ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} +# ccache directory - persisted across layer rebuilds via cache mounts. +ENV CCACHE_DIR=/root/.cache/ccache +ENV CCACHE_COMPILERCHECK=content +# Empty by default so build steps fall back to $(nproc); CI can override. +ARG max_jobs +ENV MAX_JOBS=${max_jobs} + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + if command -v sccache >/dev/null 2>&1; then \ + echo "sccache already installed, skipping installation"; \ + sccache --version; \ + else \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +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} + +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR} + + +# ----------------------- +# vLLM fetch stages +FROM base AS fetch_vllm_0 +ONBUILD COPY ./ vllm/ +FROM base AS fetch_vllm_1 +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" +ENV VLLM_REPO=${VLLM_REPO} +ENV VLLM_BRANCH=${VLLM_BRANCH} +ONBUILD RUN git clone ${VLLM_REPO} \ + && cd vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git checkout FETCH_HEAD \ + && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ + git remote add upstream "https://github.com/vllm-project/vllm.git" \ + && git fetch upstream ; fi +FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm + +# ----------------------- +# Rust build stage +# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages +# don't need the rust toolchain or protoc. +FROM fetch_vllm AS rust-build +ARG COMMON_WORKDIR + +# protoc is used by tonic-build/prost-build. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + ca-certificates curl unzip \ + && rm -rf /var/lib/apt/lists/* + +COPY tools/install_protoc.sh /tmp/install_protoc.sh +RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh + +# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit +# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). +ENV CARGO_BUILD_JOBS=4 +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt + +# Build the release binary. Cargo's registry/git caches can be written by +# concurrent BuildKit jobs on shared workers, so lock those cache mounts while +# keeping the cache benefit. Do not cache target/, because stale target metadata +# can outlive source updates across BuildKit cache reuse. +RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + cd ${COMMON_WORKDIR}/vllm \ + && bash build_rust.sh \ + && test -x vllm/vllm-rs + +# ----------------------- +# vLLM native build stages +# +# csrc-build intentionally copies only files that affect ROCm native extension +# compilation. That keeps unrelated CI/test/docs edits from invalidating the +# expensive HIP/C++ build layer. +FROM base AS csrc-build +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR}/vllm + +COPY requirements/rocm.txt requirements/rocm.txt +COPY requirements/common.txt requirements/common.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/rocm.txt + +# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do +# not invalidate the expensive native build layer. +COPY setup.py CMakeLists.txt ./ +COPY tools/build_rust.py tools/build_rust.py +COPY cmake cmake/ +COPY csrc csrc/ +COPY vllm/envs.py vllm/envs.py +COPY vllm/__init__.py vllm/__init__.py + +ENV VLLM_TARGET_DEVICE=rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" + +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + export CCACHE_BASEDIR="$PWD" \ + && echo "=== ccache stats before ROCm native build ===" \ + && (ccache --show-stats || true) \ + && (ccache --zero-stats || true) \ + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null \ + && echo "=== ccache stats after ROCm native build ===" \ + && (ccache --show-stats || true) + +# Build the full vLLM ROCm wheel by reusing the native extension wheel from +# csrc-build. This stage still rebuilds for Python/package changes, but skips +# the expensive HIP/C++ compile when native inputs are unchanged. +FROM fetch_vllm AS build_vllm +ARG COMMON_WORKDIR +ENV VLLM_TARGET_DEVICE=rocm + +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system -r requirements/rocm.txt \ + && export VLLM_USE_PRECOMPILED=1 \ + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ + && export VLLM_DOCKER_BUILD_CONTEXT=1 \ + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null +FROM scratch AS export_vllm +ARG COMMON_WORKDIR +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 + +# RIXL/UCX build stages +FROM base AS build_rixl +ARG RIXL_BRANCH="39be1de8" +ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" +ARG UCX_BRANCH="bfb51733" +ARG UCX_REPO="https://github.com/openucx/ucx.git" +# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image +ENV UCX_HOME=/usr/local/ucx +ENV RIXL_HOME=/usr/local/rixl +ENV RIXL_BENCH_HOME=/usr/local/rixl_bench + +# RIXL build system dependences and RDMA support +RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ + libgrpc-dev \ + libgrpc++-dev \ + libprotobuf-dev \ + protobuf-compiler-grpc \ + libcpprest-dev \ + libaio-dev \ + librdmacm1 \ + librdmacm-dev \ + libibverbs1 \ + libibverbs-dev \ + ibverbs-utils \ + rdmacm-utils \ + ibverbs-providers \ + && rm -rf /var/lib/apt/lists/* + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system meson auditwheel patchelf tomlkit + +RUN --mount=type=cache,target=/root/.cache/ccache \ + cd /usr/local/src && \ + git clone ${UCX_REPO} && \ + cd ucx && \ + git checkout ${UCX_BRANCH} && \ + ./autogen.sh && \ + mkdir build && cd build && \ + CC="ccache gcc" CXX="ccache g++" \ + ../configure \ + --prefix=/usr/local/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-devel-headers \ + --with-rocm=${ROCM_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt && \ + make -j$(nproc) && \ + make install + +ENV PATH=/usr/local/ucx/bin:$PATH +ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone ${RIXL_REPO} /opt/rixl && \ + cd /opt/rixl && \ + git checkout ${RIXL_BRANCH} && \ + CC="ccache gcc" CXX="ccache g++" \ + meson setup build --prefix=${RIXL_HOME} \ + --force-fallback-for=abseil-cpp \ + -Ducx_path=${UCX_HOME} \ + -Drocm_path=${ROCM_PATH} && \ + cd build && \ + ninja -j$(nproc) && \ + ninja install + +# Generate RIXL wheel +# Exclude libcore and libpull from auditwheel: transitive dependencies +# that are not shipped in the wheel and vary across base images. +RUN cd /opt/rixl && \ + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ + contrib/build-wheel.sh && \ + # The wheel build re-runs meson via meson-python; force the bundled abseil + sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ + pyproject.toml && \ + grep -q 'force-fallback-for' pyproject.toml && \ + mkdir -p /app/install && \ + _ucx_install_dir=${UCX_HOME} \ + ./contrib/build-wheel.sh \ + --output-dir /app/install \ + --rocm-dir ${ROCM_PATH} \ + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ + --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not +# invalidate the slow ROCShmem build. +FROM base AS build_rocshmem +ARG ROCSHMEM_BRANCH="f0acb0c6" +ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so +# it can be linked against DeepEP without arch mismatches. +ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" +# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image +ENV ROCSHMEM_DIR=/opt/rocshmem + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ + && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ + && git checkout ${ROCSHMEM_BRANCH} \ + && mkdir -p projects/rocshmem/build \ + && cd projects/rocshmem/build \ + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + bash ../scripts/build_configs/all_backends \ + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF + +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. +FROM build_rocshmem AS build_deepep +ARG DEEPEP_BRANCH="a9ea9774" +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" +ARG DEEPEP_NIC="cx7" + +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. +RUN --mount=type=cache,target=/root/.cache/ccache \ + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ + && git clone ${DEEPEP_REPO} \ + && cd DeepEP \ + && git checkout ${DEEPEP_BRANCH} \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + +# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do +# not force users to rebuild the long-lived Dockerfile.rocm_base image. +FROM base AS mori_base +ARG NIC_BACKEND +ARG AINIC_VERSION +ARG UBUNTU_CODENAME +RUN /bin/bash -lc 'set -euo pipefail; \ + \ + install_ainic() { \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ + rm -rf /var/lib/apt/lists/*; \ + mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ + # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ + install_bnxt() { \ + install -m 0755 -d /etc/apt/keyrings; \ + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ + -o /etc/apt/keyrings/broadcom-nic.asc; \ + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ + > /etc/apt/sources.list.d/broadcom-nic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + bnxt-rocelib=235.2.86.0 \ + ; \ + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ + ldconfig; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + \ + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ + # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ + case "${NIC_BACKEND}" in \ + none) ;; \ + all) install_ainic; install_bnxt ;; \ + ainic) install_ainic ;; \ + bnxt) install_bnxt ;; \ + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ + esac' + +# ----------------------- +# vLLM wheel release build stage (for building distributable wheels) +# This stage pins dependencies to custom ROCm wheel versions and handles version detection +FROM fetch_vllm AS build_vllm_wheel_release + +ARG COMMON_WORKDIR + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ + +# Create /install directory for custom wheels +RUN mkdir -p /install + +# Copy custom ROCm wheels from docker/context if they exist +# COPY ensures Docker cache is invalidated when wheels change +# .keep file ensures directory always exists for COPY to work +COPY docker/context/base-wheels/ /tmp/base-wheels/ +# This is how we know if we are building for a wheel release or not. +# If there are not wheels found there, we are not building for a wheel release. +# So we exit with an error. To skip this stage. +RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ + echo "Found custom wheels - copying to /install"; \ + cp /tmp/base-wheels/*.whl /install/ && \ + echo "Copied custom wheels:"; \ + ls -lh /install/; \ + else \ + echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ + echo "Wheel releases require pre-built ROCm wheels."; \ + exit 1; \ + fi + +# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) +# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm +ARG GIT_REPO_CHECK=0 +RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ + echo "Running repository checks..."; \ + cd vllm && bash tools/check_repo.sh; \ + fi + +# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) +# This ensures setuptools_scm sees clean repo state for version detection +RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system setuptools_scm regex \ + && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ + && echo "Detected vLLM version: ${VLLM_VERSION}" \ + && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + +# Fail if git-based package dependencies are found in requirements files +# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) +# Extra notes: pip install is able to handle git+ URLs, but uv doesn't. +RUN echo "Checking for git-based packages in requirements files..." \ + && echo "Checking common.txt for git-based packages:" \ + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ + echo "ERROR: Git-based packages found in common.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in common.txt"; \ + fi \ + && echo "Checking rocm.txt for git-based packages:" \ + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ + echo "ERROR: Git-based packages found in rocm.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in rocm.txt"; \ + fi \ + && echo "All requirements files are clean - no git-based packages found" + +# Pin vLLM dependencies to exact versions of custom ROCm wheels +# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi +COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py +RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ + && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt + +# Install dependencies using custom wheels from /install +RUN --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + && echo "Building vLLM with custom wheels from /install" \ + && uv pip install --system --find-links /install -r requirements/rocm.txt + +# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# (setup.py auto-detects ccache/sccache in PATH) +RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + cd vllm \ + && export CCACHE_BASEDIR="$PWD" \ + && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist + +FROM scratch AS export_vllm_wheel_release +ARG COMMON_WORKDIR +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 + +# ----------------------- +# CI base image (Tier 1) - stable, rarely changing CI dependencies. +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers +# in the vLLM artifacts for the current commit. +FROM mori_base AS ci_base +ARG COMMON_WORKDIR + +# Update rdma-core to support latest rocshmem. +ARG DEEPEP_NIC +RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ + git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && \ + mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ +fi + +# Install RIXL + DeepEP wheels. +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ + uv pip install --system /rixl_install/*.whl /deep_install/*.whl + +# Copy ROCShmem runtime libraries. +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem + +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + librdmacm1 \ + libibverbs1 \ + ibverbs-providers \ + ibverbs-utils \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ + && rm -rf /var/lib/apt/lists/* + +# Install torchcodec from source for ROCm/torch ABI compatibility. +COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + bash /tmp/install_torchcodec.sh \ + && rm /tmp/install_torchcodec.sh \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Pre-install shared ROCm runtime dependencies. +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ + && rm -rf /tmp/ci-base-requirements + +# Enable fast and less brittle model downloads in tests. +ENV HF_XET_HIGH_PERFORMANCE=1 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60 + +# Pre-install vLLM test dependencies. +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/rocm-test-reqs.txt + +# Rebuild fastsafetensors from source so its C++ extension is compiled with +# USE_ROCM and can detect libamdhip64.so at runtime. +RUN --mount=type=cache,target=/root/.cache/pip \ + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ + && test -n "${FASTSAFETENSORS_REQ}" \ + && python3 -m pip install --force-reinstall --no-deps \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + && rm /tmp/rocm-test-reqs.txt + +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. +# See: https://github.com/pytorch/pytorch/issues/169857 +ENV MIOPEN_DEBUG_CONV_DIRECT=0 +ENV MIOPEN_DEBUG_CONV_GEMM=0 + +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" + +# Install vllm_test_utils in ci_base for ci_base + wheel parity. +COPY tests/vllm_test_utils /tmp/vllm_test_utils +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /tmp/vllm_test_utils \ + && rm -rf /tmp/vllm_test_utils + +# ----------------------- +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. +FROM ${CI_BASE_IMAGE} AS test +ARG COMMON_WORKDIR + +# Install the vLLM wheel (--no-deps: all deps already in ci_base). +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system --no-deps *.whl + +# Store the vLLM wheel in the image for python-only install tests. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ + +WORKDIR /vllm-workspace +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# Copy in the v1 package (for python-only install test group). +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 + +# Hide source under src/ so it won't shadow the installed package in tests. +RUN mkdir src && mv vllm src/vllm + +# ----------------------- +# Final vLLM image +FROM mori_base AS final + +RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* + +# Clean up sccache from release image (not needed at runtime) +# This removes the binary and wrappers that may have been installed during build +RUN rm -f /usr/bin/sccache || true \ + && rm -rf /opt/sccache-wrappers || true + +# Unset sccache environment variables for the release image +# This prevents S3 bucket config from leaking into production images +ENV SCCACHE_BUCKET= +ENV SCCACHE_REGION= +ENV SCCACHE_ENDPOINT= +ENV SCCACHE_S3_NO_CREDENTIALS= +ENV SCCACHE_IDLE_TIMEOUT= + +# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. +# Manually remove it so that later steps of numpy upgrade can continue +RUN case "$(which python3)" in \ + *"/opt/conda/envs/py_3.9"*) \ + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ + *) ;; esac + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system --upgrade huggingface-hub[cli] + +# Install vLLM using uv (inherited from base stage) +# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system -r requirements/rocm.txt \ + && pip uninstall -y vllm \ + && uv pip install --system *.whl + +# Install RIXL wheel +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + uv pip install --system /rixl_install/*.whl + +ARG COMMON_WORKDIR +ARG BASE_IMAGE +ARG NIC_BACKEND +ARG AINIC_VERSION + +# Copy over the benchmark scripts as well +COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks +COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples +COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker + +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +ENV TOKENIZERS_PARALLELISM=false + +# ENV that can improve safe tensor loading, and end-to-end time +ENV SAFETENSORS_FAST_GPU=1 + +# Performance environment variable. +ENV HIP_FORCE_DEV_KERNARG=1 + +# Workaround for ROCm profiler limits +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt + + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall -y triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + +CMD ["/bin/bash"] + +#Set entrypoint for vllm-openai official images +FROM final AS vllm-openai +ENTRYPOINT ["vllm", "serve"] \ No newline at end of file From 9609542f84cd5db2aac9c5f1eec3a60d017ac18c Mon Sep 17 00:00:00 2001 From: Daniel Date: Mon, 13 Jul 2026 11:43:50 -0400 Subject: [PATCH 62/68] guard triton jit when running eager mode Signed-off-by: Daniel --- vllm/utils/jit_monitor.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index c896c275bd4a..a29d222e325a 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -147,11 +147,26 @@ def _log_triton_jit_compile(fn_name: str, kwargs) -> None: def _setup_triton_jit_hook() -> None: - """Register a ``jit_post_compile_hook`` that warns on compilation.""" - return + """Register a jit_post_compile_hook that warns on compilation.""" if not HAS_TRITON: return - from triton import knobs # type: ignore[import-untyped] + from triton import knobs + from triton.runtime import jit as _triton_jit + + # kernels pass non-JSON-serializable constexprs + # make that serialization non-fatal + _orig = _triton_jit.serialize_specialization_data + if not getattr(_orig, "_vllm_guarded", False): + + @functools.wraps(_orig) + def _guarded(*args, **kwargs): + try: + return _orig(*args, **kwargs) + except (TypeError, ValueError): + return None # best-effort metadata; monitor ignores it + + _guarded._vllm_guarded = True + _triton_jit.serialize_specialization_data = _guarded existing_hook = knobs.runtime.jit_post_compile_hook From b0e0bfa3b7db3fcbfc4cdf3a490c6db9d26f21aa Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 13 Jul 2026 11:12:04 -0500 Subject: [PATCH 63/68] Working version of the 4 models Signed-off-by: jpvillam --- benchmarks/vllm_smoketest.sh | 253 ++++++++++++------ vllm/_aiter_ops.py | 2 - .../fused_moe/experts/rocm_aiter_moe.py | 8 +- .../layers/fused_moe/oracle/mxfp4.py | 16 ++ .../layers/quantization/utils/fp8_utils.py | 29 ++ 5 files changed, 214 insertions(+), 94 deletions(-) mode change 100644 => 100755 benchmarks/vllm_smoketest.sh diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh old mode 100644 new mode 100755 index facd4f0f98b2..e45460bbe153 --- a/benchmarks/vllm_smoketest.sh +++ b/benchmarks/vllm_smoketest.sh @@ -1,110 +1,185 @@ #!/bin/bash -set -euo pipefail +# Smoke-test + lm_eval harness for the four gfx1250 models, using the exact +# serve/eval recipes from the per-model scripts in this directory: +# +# gptoss <- gpt.sh + eval_gpt.sh (gpt-oss-120b-w-mxfp4-a-fp8) +# dsr1 <- dsr1.sh + eval_dsr1.sh (DeepSeek-R1-0528-MXFP4) +# dsv4f <- dsr4_accurate.sh+ eval_dsr4.sh (DeepSeek-V4-Flash) +# minimax <- minimax.sh + eval_minimax.sh (MiniMax-M3-MXFP4) +# +# For each selected model it: starts vllm serve with that model's env+args +# (server + lm_eval output stream to the terminal), waits for /health, runs the +# model's gsm8k lm_eval, then shuts the server down (freeing the GPU) before the +# next one. +# +# Usage: +# ./vllm_smoketest.sh # run all four, in order +# ./vllm_smoketest.sh --minimax # run only minimax +# ./vllm_smoketest.sh --gptoss /some/other/path# run gptoss from a custom path +# ./vllm_smoketest.sh --dsr1 --minimax # run a subset +# Each of --gptoss / --dsr1 / --dsv4f / --minimax takes an OPTIONAL model path. +set -uo pipefail -# add model args here +PORT=8000 +CANONICAL_ORDER=(gptoss dsr1 dsv4f minimax) + +# --- default model paths (from the per-model serve scripts) --- declare -A MODEL_PATHS=( - [gpt-oss-120b-mxfp4]="/data/amd/gpt-oss-120b-w-mxfp4-a-fp8" - [DeepSeek-R1-0528-MXFP4]="/data/amd/DeepSeek-R1-0528-MXFP4" -) -declare -A MODEL_SERVE_ARGS=( - [gpt-oss-120b-mxfp4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.7 --attention-backend TRITON_ATTN" - [DeepSeek-R1-0528-MXFP4]="--tensor-parallel-size 1 --gpu_memory_utilization 0.9 --dtype auto --no-enable-prefix-caching --disable-uvicorn-access-log --trust-remote-code" + [gptoss]="/data/models/gpt-oss-120b-w-mxfp4-a-fp8" + [dsr1]="/data/models/DeepSeek-R1-0528-MXFP4" + [dsv4f]="/data/models/DeepSeek-V4-Flash" + [minimax]="/data/models/MiniMax-M3-MXFP4" ) + +# --- serve-time environment (verbatim from each serve script) --- declare -A MODEL_ENV=( - [gpt-oss-120b-mxfp4]="HSA_OVERRIDE_GFX_VERSION=12.5.0 HSA_ENABLE_SDMA=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 " - [DeepSeek-R1-0528-MXFP4]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [gptoss]="HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [dsr1]="HIP_VISIBLE_DEVICES=0 VLLM_DISABLE_COMPILE_CACHE=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MLA=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_FP8BMM=0" + [dsv4f]="ROCR_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_FORCE_TORCH_BLOCK_FP8=1 VLLM_ROCM_USE_AITER_LINEAR=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_TRITON_GEMM=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 ROCR_VISIBLE_DEVICES=2" + [minimax]="HIP_VISIBLE_DEVICES=0 VLLM_DISABLE_COMPILE_CACHE=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MLA=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=0 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_FP8BMM=0" ) -MODEL_NAME="gpt-oss-120b-mxfp4" # default -MODEL_PATH_OVERRIDE="" -MODEL_ENV_OVERRIDE="" -MODEL_SET=0 # whether --model was explicitly passed -RUN_LM_EVAL=0 -PORT=8000 +# --- serve args (everything except --model/--host/--port, which we add). +# NOTE: JSON values must contain NO spaces (they ride through unquoted $VAR +# word-splitting); the dsr1 compilation-config below is space-free. --- +declare -A MODEL_SERVE=( + [gptoss]="--tensor-parallel-size 1 --gpu_memory_utilization 0.7 --attention-backend TRITON_ATTN" + [dsr1]="--trust-remote-code --no-enable-prefix-caching --no-enable-chunked-prefill --max-model-len 8192 --dtype auto --tensor-parallel-size 1 --distributed-executor-backend mp --max-num-batched-tokens 8192 --max-num-seqs 32 --gpu-memory-utilization 0.90 --compilation-config {\"pass_config\":{\"fuse_attn_quant\":true,\"eliminate_noops\":true,\"fuse_norm_quant\":true,\"fuse_mla_dual_rms_norm\":false,\"enable_qk_norm_rope_fusion\":false},\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"custom_ops\":[\"+rms_norm\",\"+silu_and_mul\",\"+quant_fp8\"]}" + [dsv4f]="--tensor-parallel-size 1 --gpu_memory_utilization 0.7 --kv-cache-dtype fp8 --max-model-len 32768 --moe-backend TRITON_UNFUSED" + [minimax]="--trust-remote-code --language-model-only --skip-mm-profiling --block-size 128 --enforce-eager --no-enable-prefix-caching --no-enable-chunked-prefill --max-model-len 32768 --dtype auto --tensor-parallel-size 1 --distributed-executor-backend mp --max-num-batched-tokens 32768 --max-num-seqs 32 --gpu-memory-utilization 0.90 --reasoning-parser minimax_m3 --tool-call-parser minimax_m3 --enable-auto-tool-choice" +) + +# --- lm_eval: kind (chat=chat-completions+template, raw=completions), +# extra model_args, and gen_kwargs (verbatim from each eval_*.sh) --- +declare -A EVAL_KIND=( [gptoss]="chat" [dsr1]="chat" [dsv4f]="raw" [minimax]="chat" ) +declare -A EVAL_MARGS=( + [gptoss]="num_concurrent=64,max_retries=3,tokenized_requests=False,max_length=8192" + [dsr1]="num_concurrent=64,max_retries=3,tokenized_requests=False,max_length=8192,timeout=3600" + [dsv4f]="num_concurrent=8,max_retries=3,tokenized_requests=False,max_length=32768,timeout=1200" + [minimax]="num_concurrent=32,max_retries=3,tokenized_requests=False,max_length=32768,timeout=3600" +) +declare -A EVAL_GEN=( + [gptoss]="max_gen_toks=4096" + [dsr1]="max_gen_toks=4096,do_sample=True,temperature=0.6,top_p=0.95" + [dsv4f]="max_gen_toks=2048,do_sample=True,temperature=1.0,top_p=0.95" + [minimax]="max_gen_toks=2048,do_sample=True,temperature=1.0,top_p=0.95" +) usage() { - cat <<'EOF' -Usage: ./vllm_smoketest.sh [--model NAME] [--model-path PATH] [--env ENV_VARS] [--lm-eval] [--port PORT] [--list] - --model NAME Registry key to run (default: gpt-oss-120b-mxfp4) - --model-path PATH Override model path for this run - --env ENV_VARS Override environment variables (e.g., "VAR1=val1 VAR2=val2") - --lm-eval Run lm_eval after curl check - --port PORT Server port (default: 8000) - --list List available models + cat < gsm8k lm_eval (5-shot, limit 100) -> shutdown. Logs stream to this terminal. EOF } +# --- parse args (each model flag takes an OPTIONAL path) --- +declare -A PATH_OVERRIDE=() +SELECTED=() while [[ $# -gt 0 ]]; do case "$1" in - --model) MODEL_NAME="$2"; MODEL_SET=1; shift 2 ;; - --model-path) MODEL_PATH_OVERRIDE="$2"; shift 2 ;; - --env) MODEL_ENV_OVERRIDE="$2"; shift 2 ;; - --lm-eval) RUN_LM_EVAL=1; shift ;; - --port) PORT="$2"; shift 2 ;; - --list) for n in "${!MODEL_PATHS[@]}"; do echo "$n -> ${MODEL_PATHS[$n]}"; done; exit 0 ;; - -h|--help) usage; exit 0 ;; + --gptoss|--dsr1|--dsv4f|--minimax) + key="${1#--}" + SELECTED+=("$key") + if [[ $# -ge 2 && "$2" != -* ]]; then PATH_OVERRIDE[$key]="$2"; shift 2; else shift; fi + ;; + --port) PORT="$2"; shift 2 ;; + --list) for k in "${CANONICAL_ORDER[@]}"; do printf '%-9s -> %s\n' "$k" "${MODEL_PATHS[$k]}"; done; exit 0 ;; + -h|--help) usage; exit 0 ;; *) echo "Unknown argument: $1" >&2; usage; exit 1 ;; esac done +[[ ${#SELECTED[@]} -eq 0 ]] && SELECTED=("${CANONICAL_ORDER[@]}") + +echo "Models: ${SELECTED[*]} | Port: $PORT (all server + lm_eval logs stream to this terminal)" + +server_pid="" +stop_server() { + [[ -n "$server_pid" ]] && kill "$server_pid" 2>/dev/null + # vLLM renames workers to VLLM::EngineCore / VLLM::Worker; killing the launcher + # alone leaves them holding GPU memory, so sweep them too. + pkill -9 -f "vllm serve" 2>/dev/null + pkill -9 -f "VLLM::" 2>/dev/null + server_pid="" + # wait for GPU0 to actually release before the next model loads + if command -v rocm-smi >/dev/null 2>&1; then + for _ in $(seq 1 40); do + used=$(rocm-smi --showmeminfo vram 2>/dev/null | grep 'GPU\[0\]' | grep -i used | grep -oE '[0-9]+' | tail -1) + [[ -z "$used" ]] && break + (( used < 5000000000 )) && break # < ~5 GiB => free + sleep 3 + done + else + sleep 8 + fi +} +trap 'stop_server; exit 130' INT TERM +trap 'stop_server' EXIT -if [[ -n "$MODEL_PATH_OVERRIDE" && $MODEL_SET -eq 0 ]]; then - echo "ERROR: --model-path requires --model (which model's serve args/env to use)." >&2 - usage; exit 1 -fi +declare -A RESULT +command -v lm_eval >/dev/null 2>&1 || pip install "lm-eval[api]" -MODEL="${MODEL_PATHS[$MODEL_NAME]:-$MODEL_NAME}" -[[ -n "$MODEL_PATH_OVERRIDE" ]] && MODEL="$MODEL_PATH_OVERRIDE" -SERVE_ARGS="${MODEL_SERVE_ARGS[$MODEL_NAME]:-}" -MODEL_ENV_ARGS="${MODEL_ENV[$MODEL_NAME]:-}" -[[ -n "$MODEL_ENV_OVERRIDE" ]] && MODEL_ENV_ARGS="$MODEL_ENV_OVERRIDE" -echo "Model: $MODEL | Port: $PORT | lm_eval: $RUN_LM_EVAL" +run_model() { + local key="$1" + local model="${PATH_OVERRIDE[$key]:-${MODEL_PATHS[$key]}}" + local env="${MODEL_ENV[$key]}" serve="${MODEL_SERVE[$key]}" -LOG_DIR="$(pwd)/vllm_smoke_test_logs" -mkdir -p "$LOG_DIR" -echo "Logs: $LOG_DIR" + echo; echo "==================== $key : $model ====================" + if [[ ! -e "$model" ]]; then + echo "WARNING: skipping '$key' -- model path does not exist:" >&2 + echo " $model" >&2 + if [[ -n "${PATH_OVERRIDE[$key]:-}" ]]; then + echo " (this path was supplied via '--$key $model'; check the path/mount)" >&2 + else + echo " (this is the default path for '$key'; download the model there," >&2 + echo " or point it elsewhere with: --$key /path/to/model)" >&2 + fi + echo " Continuing with the remaining models." >&2 + RESULT[$key]="SKIP (model path not found: $model)" + return + fi -server_pid="" -cleanup() { [[ -n "$server_pid" ]] && kill "$server_pid" 2>/dev/null && wait "$server_pid" 2>/dev/null; true; } -trap cleanup EXIT -trap 'exit 130' INT TERM - -# --- start server --- -env $MODEL_ENV_ARGS \ -vllm serve --model "$MODEL" --host localhost --port "$PORT" $SERVE_ARGS & -server_pid=$! -echo "Waiting for server (pid $server_pid)..." - -until curl -s "http://localhost:$PORT/health" &>/dev/null; do - ps -p "$server_pid" >/dev/null || { echo "ERROR: server died"; tail -n 50 "$LOG_DIR/server.log"; exit 1; } - sleep 5 -done -echo "Server ready." - -# --- check if text was generated --- -echo "=== Curl completion check ===" -http_code=$(curl -s -o "$LOG_DIR/curl_completion.log" -w '%{http_code}' \ - "http://localhost:$PORT/v1/completions" \ - -H "Content-Type: application/json" \ - -d "{\"model\":\"$MODEL\",\"prompt\":\"The capital of France is \",\"max_tokens\":32,\"temperature\":0}") -cat "$LOG_DIR/curl_completion.log"; echo -[[ "$http_code" == "200" ]] || { echo "FAIL: HTTP $http_code" >&2; exit 1; } -grep -q '"text"' "$LOG_DIR/curl_completion.log" || { echo "FAIL: no text in response" >&2; exit 1; } -grep -qi "Paris" "$LOG_DIR/curl_completion.log" || { echo "FAIL: 'Paris' not found in response" >&2; exit 1; } -CURL_CHECK_PASSED=1 - -# --- lm_eval --- -if [[ $RUN_LM_EVAL -eq 1 ]]; then - echo "=== lm_eval ===" - pip install "lm-eval[api]" - lm_eval --model local-completions \ - --model_args "model=$MODEL,base_url=http://localhost:$PORT/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False" \ - --tasks gsm8k --num_fewshot 3 --limit 100 2>&1 | tee "$LOG_DIR/lm_eval.log" -fi - -# Shutdown server -cleanup - -# Print PASS only if curl check succeeded -[[ ${CURL_CHECK_PASSED:-0} -eq 1 ]] && echo "PASS" - -echo "Done." + # --- start server (logs stream to this terminal) --- + echo "[$key] starting vllm serve ..." + env $env vllm serve --model "$model" --host localhost --port "$PORT" $serve & + server_pid=$! + + local ready=0 + for _ in $(seq 1 240); do # up to ~20 min for load + if curl -s "http://localhost:$PORT/health" >/dev/null 2>&1; then ready=1; break; fi + ps -p "$server_pid" >/dev/null 2>&1 || { echo "[$key] ERROR: server died (see output above)"; RESULT[$key]="FAIL (server died at load)"; stop_server; return; } + sleep 5 + done + [[ $ready -eq 1 ]] || { echo "[$key] ERROR: server not ready (timeout)"; RESULT[$key]="FAIL (server timeout)"; stop_server; return; } + echo "[$key] server ready." + + local kind="${EVAL_KIND[$key]}" + + # --- lm_eval (gsm8k, 5-shot, limit 100) --- + echo "[$key] === lm_eval (gsm8k) ===" + local lm_model base_url tmpl=() + if [[ "$kind" == "chat" ]]; then + lm_model="local-chat-completions"; base_url="http://localhost:$PORT/v1/chat/completions"; tmpl=(--apply_chat_template) + else + lm_model="local-completions"; base_url="http://localhost:$PORT/v1/completions" + fi + if lm_eval --model "$lm_model" \ + --model_args "model=$model,base_url=$base_url,${EVAL_MARGS[$key]}" \ + "${tmpl[@]}" \ + --gen_kwargs "${EVAL_GEN[$key]}" \ + --tasks gsm8k --num_fewshot 5 --limit 100 + then RESULT[$key]="PASS (gsm8k table above)"; else RESULT[$key]="FAIL (lm_eval error)"; fi + + stop_server +} + +for key in "${SELECTED[@]}"; do run_model "$key"; done + +echo; echo "==================== SUMMARY ====================" +for key in "${SELECTED[@]}"; do printf '%-9s : %s\n' "$key" "${RESULT[$key]:-UNKNOWN}"; done diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index d439acf34abe..0e9a8fa12f17 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -315,7 +315,6 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) - m_tokens = hidden_states.shape[0] use_triton_a4w4 = ( envs.VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4 and quant_type == QuantType.per_1x32 @@ -324,7 +323,6 @@ def _rocm_aiter_fused_moe_impl( and num_local_tokens is None and a1_scale is None and a2_scale is None - and m_tokens >= envs.VLLM_ROCM_AITER_TRITON_MOE_MIN_TOKENS ) if use_triton_a4w4: return _rocm_aiter_fused_moe_triton_gemm_a4w4( diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 4f191334bb2a..985f540a3a59 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -457,11 +457,13 @@ def _supports_quant_scheme( ] if (weight_key, activation_key) not in SUPPORTED_W_A: return False - # CK MXFP4 MoE kernels are only supported on gfx950. if weight_key == kMxfp4Static: - from vllm.platforms.rocm import on_gfx950 + import vllm.envs as envs + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - if not on_gfx950(): + if not on_gfx950() and not ( + envs.VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4 and on_gfx1250() + ): return False return True diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 55a767f060cb..5dc152be7f9a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -969,6 +969,22 @@ def swap_every_two_rows(x, axis=-1): if w2_bias is not None: w2_bias = w2_bias.data.to(torch.float32) + import vllm.envs as envs + + if envs.VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4: + # Triton moe_gemm_a4w4 path (e.g. gfx1250): the CK e8m0_shuffle / + # shuffle_weights below are the wrong layout for the Triton kernel, + # which consumes native [E, 2*inter, hidden/2] uint8 weights with + # unshuffled scales and maps layout internally. Return untransformed. + return ( + w13_weight, + w2_weight, + w13_weight_scale, + w2_weight_scale, + w13_bias, + w2_bias, + ) + # e8m0_shuffle on weight scales (GFX950 swizzle layout) from aiter.utility.fp4_utils import e8m0_shuffle diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 32a2d86899cc..16bae2158b1c 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -897,6 +897,35 @@ def w8a8_triton_block_scaled_mm( block_size: list[int], output_dtype: torch.dtype = torch.float16, ) -> torch.Tensor: + # TODO (JPVILLAM): Retest and fix, naive impl for now + import os as _os + + if _os.environ.get("VLLM_FORCE_TORCH_BLOCK_FP8") == "1": + # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32. + # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow. + _bn, _bk = block_size[0], block_size[1] + _As = ( + _upcast_e8m0_to_fp32(As) + if As.dtype == torch.float8_e8m0fnu + else As.to(torch.float32) + ) + _Bs = ( + _upcast_e8m0_to_fp32(Bs) + if Bs.dtype == torch.float8_e8m0fnu + else Bs.to(torch.float32) + ) + _K = A.shape[-1] + _N = B.shape[0] + _Af = A.to(torch.float32).reshape(-1, _K) + _Asf = ( + _As.to(torch.float32) + .reshape(-1, _As.shape[-1]) + .repeat_interleave(_bk, dim=1)[:, :_K] + ) + _Bf = B.to(torch.float32) + _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K] + _out = (_Af * _Asf) @ (_Bf * _Bsf).t() + return _out.to(output_dtype).reshape(*A.shape[:-1], _N) """This function performs matrix multiplication with block-wise quantization. It takes two input tensors `A` and `B` with scales `As` and `Bs`. From da553f6b5c1e2fe4dfb7e8cbaa7a453db635a170 Mon Sep 17 00:00:00 2001 From: jpvillam Date: Mon, 13 Jul 2026 13:50:55 -0500 Subject: [PATCH 64/68] Small code review changes Signed-off-by: jpvillam --- benchmarks/vllm_smoketest.sh | 8 ++--- .../layers/quantization/utils/fp8_utils.py | 31 ++++++++++--------- 2 files changed, 20 insertions(+), 19 deletions(-) diff --git a/benchmarks/vllm_smoketest.sh b/benchmarks/vllm_smoketest.sh index e45460bbe153..717d9acdef5b 100755 --- a/benchmarks/vllm_smoketest.sh +++ b/benchmarks/vllm_smoketest.sh @@ -33,10 +33,10 @@ declare -A MODEL_PATHS=( # --- serve-time environment (verbatim from each serve script) --- declare -A MODEL_ENV=( - [gptoss]="HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" - [dsr1]="HIP_VISIBLE_DEVICES=0 VLLM_DISABLE_COMPILE_CACHE=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MLA=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_FP8BMM=0" - [dsv4f]="ROCR_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_FORCE_TORCH_BLOCK_FP8=1 VLLM_ROCM_USE_AITER_LINEAR=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_TRITON_GEMM=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 ROCR_VISIBLE_DEVICES=2" - [minimax]="HIP_VISIBLE_DEVICES=0 VLLM_DISABLE_COMPILE_CACHE=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MLA=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=0 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_FP8BMM=0" + [gptoss]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [dsr1]="VLLM_DISABLE_COMPILE_CACHE=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MLA=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_FP8BMM=0" + [dsv4f]="HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_FORCE_TORCH_BLOCK_FP8=1 VLLM_ROCM_USE_AITER_LINEAR=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_TRITON_GEMM=1 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0" + [minimax]="VLLM_DISABLE_COMPILE_CACHE=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MLA=0 HSA_ENABLE_SDMA=0 USE_SVM=0 HSA_XNACK=0 VLLM_ROCM_AITER_FUSED_MOE_TRITON_GEMM_A4W4=0 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1 VLLM_ROCM_USE_SKINNY_GEMM=0 VLLM_ROCM_USE_AITER_RMSNORM=0 VLLM_ROCM_USE_AITER_FP8BMM=0" ) # --- serve args (everything except --model/--host/--port, which we add). diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 16bae2158b1c..487825171566 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -897,6 +897,21 @@ def w8a8_triton_block_scaled_mm( block_size: list[int], output_dtype: torch.dtype = torch.float16, ) -> torch.Tensor: + """This function performs matrix multiplication with block-wise + quantization. + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should + be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + Returns: + torch.Tensor: The result of matmul. + """ # TODO (JPVILLAM): Retest and fix, naive impl for now import os as _os @@ -926,21 +941,7 @@ def w8a8_triton_block_scaled_mm( _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K] _out = (_Af * _Asf) @ (_Bf * _Bsf).t() return _out.to(output_dtype).reshape(*A.shape[:-1], _N) - """This function performs matrix multiplication with block-wise - quantization. - It takes two input tensors `A` and `B` with scales `As` and `Bs`. - The output is returned in the specified `output_dtype`. - Args: - A: The input tensor, e.g., activation. - B: The input tensor, e.g., weight. - As: The per-token-group quantization scale for `A`. - Bs: The per-block quantization scale for `B`. - block_size: The block size for per-block quantization. It should - be 2-dim, e.g., [128, 128]. - output_dytpe: The dtype of the returned tensor. - Returns: - torch.Tensor: The result of matmul. - """ + assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] From 0e5dd364ee4ca57019b5689fe7bbb2f3ac1437b0 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Thu, 9 Jul 2026 16:23:45 -0400 Subject: [PATCH 65/68] Add new Dsv4 aiterbackend and revert gpt_oss_triton_kernels Signed-off-by: Jaden Mathias --- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 216 +++++ .../experts/gpt_oss_triton_kernels_moe.py | 840 +----------------- .../layers/fused_moe/oracle/mxfp4.py | 5 +- 3 files changed, 247 insertions(+), 814 deletions(-) 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 index 3c17b8c8b09e..dffcb94a1192 100644 --- 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 @@ -20,7 +20,9 @@ __all__ = [ "AiterW4A8ExpertsMonolithic", + "AiterW4A16ExpertsMonolithic", "aiter_triton_kernel_w4a8_moe_forward", + "aiter_triton_kernel_w4a16_moe_forward", ] @@ -324,3 +326,217 @@ def apply( unpadded_N_w2=self.moe_config.hidden_dim_unpadded, unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded, ) + + +def _aiter_raw(t): + return t.storage.data if hasattr(t, "storage") else t + + +def aiter_triton_kernel_w4a16_moe_forward( + hidden_states: torch.Tensor, + w1, + w2, + 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 rocm_aiter_ops.is_enabled() + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe.moe_routing.routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing + + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) + + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + w1_data = _aiter_raw(w1) + w2_data = _aiter_raw(w2) + w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) + w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + + gammas = routing_data.gate_scal if routing_data else None + + swiglu_alpha = ( + quant_config.gemm1_alpha + if quant_config.gemm1_alpha is not None + else 1.0 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + intermediate = moe_gemm_a16w4( + hidden_states, + w1_data, + None, + w1_wscale, + None, + None, + quant_config.w1_bias, + routing_data, + gather_indx=gather_idx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=True, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + + out = moe_gemm_a16w4( + intermediate, + w2_data, + None, + w2_wscale, + None, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + + return out + + +class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + + 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: + if not rocm_aiter_ops.is_enabled(): + return False + from vllm.platforms.rocm import on_gfx950, on_gfx1250 + + return on_gfx950() or on_gfx1250() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kMxfp4Static, None) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + 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 + + @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, + 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_w4a16_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 8b688697d167..4b0a0b8ecad3 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 @@ -1,13 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import replace - import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import _custom_ops as ops -from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -51,13 +47,13 @@ def _triton_kernel_moe_supports_current_device() -> bool: # range was not validated. return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) if p.is_rocm(): - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx1250 + from vllm.platforms.rocm import on_gfx1x, on_gfx9 # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx); # on_gfx9() already excludes gfx906/gfx908. # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). - return on_gfx9() or on_gfx1x() or on_gfx1250() + return on_gfx9() or on_gfx1x() return False @@ -574,44 +570,6 @@ def triton_kernel_moe_forward( ) -> torch.Tensor: sm_first = not renormalize - # ROCm aiter-native MXFP4 fast path (mirrors ROCm/ATOM triton_kernel_moe_forward): - # native aiter routing on the raw logits + the two-GEMM aiter compute. Falls back - # to the vendored triton_kernels routing + matmul_ogs path below when unavailable - # or out of the ported scope (SiLU is FP8-only). - if ( - expert_map is None - and quant_config is not None - and rocm_aiter_ops.is_enabled() - and _aiter_act_quant_mode(quant_config) is not None - ): - _act_mode = _aiter_act_quant_mode(quant_config) - _aiter_ops = _import_aiter_moe_ops() - if _aiter_ops is not None and ( - activation == MoEActivation.SWIGLUOAI or _act_mode == _AITER_ACT_FP8 - ): - routing_data, gather_idx, scatter_idx = _aiter_ops["routing"]( - gating_output, topk, sm_first=sm_first - ) - swiglu_alpha, swiglu_limit = _aiter_swiglu_params(quant_config) - # Native aiter routing applies router weights in-kernel, so no post-scale. - return _aiter_moe_two_gemm( - _aiter_ops, - hidden_states, - w1, - w2, - routing_data, - gather_idx, - scatter_idx, - topk, - quant_config, - _act_mode, - activation, - swiglu_alpha, - swiglu_limit, - apply_router_weight_on_input, - None, - ) - # When no expert map is provided (no EP), call the fused `routing()` # kernel directly. It combines softmax, topk, bitmatrix packing, and # routing-metadata construction in a single launch, instead of the @@ -856,503 +814,6 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx -def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( - hidden_states: torch.Tensor, - quant_config: FusedMoEQuantConfig, - *, - gemm_num: int = 1, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). - - Uses ``aiter.ops.triton.moe.quant_moe.downcast_to_static_fp8`` with a scalar scale: - for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when - single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or - ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static - FP8 path). - - Returns: - (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for - FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). - """ - if gemm_num not in (1, 2): - raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") - if not quant_config.use_mxfp4_w4a16: - return hidden_states, None - try: - from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 - except ImportError: - return hidden_states, None - if not rocm_aiter_ops.is_enabled(): - return hidden_states, None - - qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision - a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale - - lhs_scale = None - if qp is not None and qp.flex_ctx.lhs_data.scale is not None: - s = qp.flex_ctx.lhs_data.scale - if s.numel() == 1: - lhs_scale = s - if lhs_scale is None and a_scale is not None: - s = a_scale - if s.numel() == 1: - lhs_scale = s - if lhs_scale is None: - amax = hidden_states.abs().max().clamp(min=1e-12) - lhs_scale = (amax / 448.0).to(dtype=torch.float32) - - return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale - - -def _mxfp4_w4a8_unpadded_dims( - moe_cfg: FusedMoEConfig, -) -> tuple[int | None, int | None, int | None, int | None]: - """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" - if ( - moe_cfg.intermediate_size_per_partition_unpadded is None - or moe_cfg.hidden_dim_unpadded is None - ): - return None, None, None, None - unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 - unpadded_k_w1 = moe_cfg.hidden_dim_unpadded - unpadded_n_w2 = moe_cfg.hidden_dim_unpadded - unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded - return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 - - -def _mxfp4_w4a8_resolve_lhs_scale( - quant_config: FusedMoEQuantConfig, - *, - gemm_num: int, - activations: torch.Tensor, -) -> torch.Tensor: - """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. - - Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise - ``max(|activations|) / 448`` (e4m3), matching - ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. - """ - if gemm_num not in (1, 2): - raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") - qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision - a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale - lhs_scale = None - if qp is not None and qp.flex_ctx is not None: - ld = qp.flex_ctx.lhs_data - if ld is not None and ld.scale is not None and ld.scale.numel() == 1: - lhs_scale = ld.scale - if lhs_scale is None and a_scale is not None: - s = a_scale - if s.numel() == 1: - lhs_scale = s - if lhs_scale is None: - amax = activations.abs().max().clamp(min=1e-12) - lhs_scale = (amax / 448.0).to(dtype=torch.float32) - return lhs_scale - - -def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( - precision_config, - lhs_scale: torch.Tensor | None, -): - """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" - if precision_config is None or lhs_scale is None: - return precision_config - from triton_kernels.numerics import InFlexData - - new_flex = replace( - precision_config.flex_ctx, - lhs_data=InFlexData(scale=lhs_scale), - ) - return replace(precision_config, flex_ctx=new_flex) - - -# --------------------------------------------------------------------------- -# aiter `.moe.*` MXFP4 MoE path (ported from ROCm/ATOM -# atom/model_ops/fused_moe_triton.py). -# -# Scope (agreed): the reference's *complete* paths only -- -# * SwiGLU (interleaved, gpt-oss): FP8 -> a8w4, FP4/BF16 -> a16w4, with the -# activation fused into the GEMM (apply_swiglu=True, swiglu_add_residual). -# * SiLU (half/half, e.g. DeepSeek-V4): FP8 -> a8w4 via dynamic MXFP8 quant + -# fused_clamp_act_mul. -# The reference's SiLU FP4 (a4w4) / BF16 (a16w4) paths feed an *unactivated* -# buffer to GEMM2 and are therefore intentionally NOT ported -- they raise. -# --------------------------------------------------------------------------- - -# aiter MXFP4 activation-quant modes (mirror ATOM MoEActivationQuant). -_AITER_ACT_FP8 = "fp8" # a8w4: FP8 activations (use_mxfp4_w4a8) -_AITER_ACT_FP4 = "fp4" # a4w4: MXFP4 activations (use_mxfp4_w4a4) -_AITER_ACT_BF16 = "bf16" # a16w4: BF16 activations (use_mxfp4_w4a16) - - -def _aiter_raw(t): - """Unwrap a triton_kernels ``Tensor`` to its backing ``torch.Tensor``. - - aiter's ``.moe.*`` kernels take plain ``torch.Tensor`` weights/scales, while - vLLM stores MXFP4 weights as triton_kernels ``Tensor`` wrappers - (``.storage.data``). Plain tensors pass through unchanged. - """ - return t.storage.data if hasattr(t, "storage") else t - - -def _aiter_act_quant_mode(quant_config: FusedMoEQuantConfig) -> str | None: - """Activation-quant mode for aiter MXFP4 MoE, or ``None`` if not MXFP4.""" - if quant_config.use_mxfp4_w4a8: - return _AITER_ACT_FP8 - if quant_config.use_mxfp4_w4a4: - return _AITER_ACT_FP4 - if quant_config.use_mxfp4_w4a16: - return _AITER_ACT_BF16 - return None - - -def _aiter_swiglu_params(quant_config: FusedMoEQuantConfig) -> tuple[float, float]: - """(alpha, clamp limit) for the SwiGLU/SiLU activation; DeepSeek-V4 defaults.""" - alpha = quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 - limit = ( - quant_config.gemm1_clamp_limit - if quant_config.gemm1_clamp_limit is not None - else 7.0 - ) - return alpha, limit - - -def _import_aiter_moe_ops() -> dict | None: - """Import the aiter ``.moe.*`` triton MoE kernels (ROCm/ATOM API), or ``None``. - - Mirrors ROCm/ATOM ``atom/model_ops/fused_moe_triton.py``. Returns a dict of - callables on success; ``None`` if aiter (or this kernel family) is absent, so - the caller can fall back to the vendored ``matmul_ogs`` path. - """ - try: - from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul - from aiter.ops.triton.moe.moe_op_gemm_a8w4 import moe_gemm_a8w4 - from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 - from aiter.ops.triton.moe.moe_routing.routing import routing as aiter_routing - from aiter.ops.triton.moe.quant_moe import downcast_to_static_fp8 - from aiter.ops.triton.quant.quant import dynamic_mxfp8_quant - from aiter.ops.triton.utils._triton.arch_info import get_arch - except ImportError: - return None - return { - "moe_gemm_a8w4": moe_gemm_a8w4, - "moe_gemm_a16w4": moe_gemm_a16w4, - "routing": aiter_routing, - "downcast_to_static_fp8": downcast_to_static_fp8, - "dynamic_mxfp8_quant": dynamic_mxfp8_quant, - "fused_clamp_act_mul": fused_clamp_act_mul, - "get_arch": get_arch, - } - - -def _aiter_routing_from_topk( - aiter_ops: dict, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - num_experts: int, - device: torch.device, -) -> tuple: - """Build aiter ``RoutingData`` from a precomputed topk selection. - - vLLM selects experts before the expert kernel runs, but aiter's MXFP4 GEMMs - need aiter routing structs (distinct from the vendored triton_kernels ones). - aiter ``routing`` always softmaxes its logits, so we scatter ``log(weight)`` - into a dense ``-inf`` logit matrix: softmax-over-selected reproduces both the - exact expert selection and ``weight / sum(weight)``. The per-token weight sum - (e.g. ``routed_scaling_factor``) divided out here is restored by the caller. - - Returns ``(routing_data, gather_idx, scatter_idx, topk_weights_f32)``. - """ - routing = aiter_ops["routing"] - M = topk_weights.shape[0] - topk = topk_ids.size(1) - tw = topk_weights.to(torch.float32) - logits = torch.full((M, num_experts), -1e30, device=device, dtype=torch.float32) - logits.scatter_( - 1, - topk_ids.long().clamp(min=0, max=num_experts - 1), - torch.log(tw.clamp(min=1e-20)), - ) - routing_data, gather_idx, scatter_idx = routing(logits, topk, sm_first=False) - return routing_data, gather_idx, scatter_idx, tw - - -def _aiter_moe_two_gemm( - aiter_ops: dict, - hidden_states: torch.Tensor, - w1, - w2, - routing_data, - gather_idx, - scatter_idx, - topk: int, - quant_config: FusedMoEQuantConfig, - act_mode: str, - activation: MoEActivation, - swiglu_alpha: float, - swiglu_limit: float, - apply_router_weight_on_input: bool, - moe_config: FusedMoEConfig | None, -) -> torch.Tensor: - """Two-GEMM aiter MXFP4 MoE compute (ported from ROCm/ATOM fused_moe_triton). - - vLLM/gfx1250 adaptations vs the reference: weights/scales are unwrapped from - triton_kernels ``Tensor`` (``.storage.data``); scales are passed unswizzled - with logical ``unpadded_N/K`` (the gfx1250 convention used by the validated - aiter path); and GEMM1's in-kernel gather -- numerically broken on gfx1250 -- - is replaced by an explicit torch gather there. - - Returns the (M, K) bf16 result *before* any router-weight post-scale; the - caller applies that when routing was reconstructed from a topk selection. - """ - moe_gemm_a8w4 = aiter_ops["moe_gemm_a8w4"] - moe_gemm_a16w4 = aiter_ops["moe_gemm_a16w4"] - dynamic_mxfp8_quant = aiter_ops["dynamic_mxfp8_quant"] - fused_clamp_act_mul = aiter_ops["fused_clamp_act_mul"] - downcast_to_static_fp8 = aiter_ops["downcast_to_static_fp8"] - get_arch = aiter_ops["get_arch"] - - assert quant_config.w1_precision is not None - assert quant_config.w2_precision is not None - w1_data = _aiter_raw(w1) - w2_data = _aiter_raw(w2) - w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) - w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) - w1_bias = quant_config.w1_bias - w2_bias = quant_config.w2_bias - - if moe_config is not None: - unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( - _mxfp4_w4a8_unpadded_dims(moe_config) - ) - else: - unpadded_n_w1 = unpadded_k_w1 = unpadded_n_w2 = unpadded_k_w2 = None - - arch = get_arch() - quant_dtype = torch.float8_e4m3fnuz if arch == "gfx942" else torch.float8_e4m3fn - - # gfx1250's in-kernel gather miscomputes (validated on the FFM sim), so gather - # rows into expert-sorted order in torch and pass gather_indx=None instead. Per - # aiter's moe_gemm_torch, sorted row i reads source token gather_idx[i] // topk. - if arch == "gfx1250": - gather_src = gather_idx.to(torch.long) // topk - gemm1_input = hidden_states[gather_src] - gemm1_gather_indx = None - else: - gemm1_input = hidden_states - gemm1_gather_indx = gather_idx - - gammas = routing_data.gate_scal if routing_data else None - g1_gammas = gammas if apply_router_weight_on_input else None - g2_gammas = None if apply_router_weight_on_input else gammas - - if activation == MoEActivation.SWIGLUOAI: - # Interleaved [gate, up] (gpt-oss): activation fused into the GEMM. - if act_mode == _AITER_ACT_FP8: - a13_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, gemm_num=1, activations=gemm1_input - ) - a2_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, gemm_num=2, activations=gemm1_input - ) - hidden_fp8 = downcast_to_static_fp8(gemm1_input, a13_scale) - interm = moe_gemm_a8w4( - hidden_fp8, - w1_data, - None, - w1_wscale, - a13_scale, - a2_scale, - w1_bias, - routing_data, - gather_indx=gemm1_gather_indx, - gammas=g1_gammas, - swizzle_mx_scale=None, - out_dtype=quant_dtype, - apply_swiglu=True, - alpha=swiglu_alpha, - limit=swiglu_limit, - swiglu_add_residual=True, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - out = moe_gemm_a8w4( - interm, - w2_data, - None, - w2_wscale, - a2_scale, - None, - w2_bias, - routing_data, - scatter_indx=scatter_idx, - gammas=g2_gammas, - swizzle_mx_scale=None, - unpadded_N=unpadded_n_w2, - unpadded_K=unpadded_k_w2, - ) - else: - # FP4 and BF16 activations both route through the a16w4 kernel. - interm = moe_gemm_a16w4( - gemm1_input, - w1_data, - None, - w1_wscale, - None, - None, - w1_bias, - routing_data, - gather_indx=gemm1_gather_indx, - gammas=g1_gammas, - swizzle_mx_scale=None, - apply_swiglu=True, - alpha=swiglu_alpha, - limit=swiglu_limit, - swiglu_add_residual=True, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - out = moe_gemm_a16w4( - interm, - w2_data, - None, - w2_wscale, - None, - None, - w2_bias, - routing_data, - scatter_indx=scatter_idx, - gammas=g2_gammas, - swizzle_mx_scale=None, - unpadded_N=unpadded_n_w2, - unpadded_K=unpadded_k_w2, - ) - return out - - # SiLU (concatenated [gate | up], half/half): manual activation between GEMMs. - # FP8 (a8w4) and BF16 (w4a16 served as a8w4) both run through moe_gemm_a8w4 by - # dynamically quantizing activations to FP8 (dynamic_mxfp8_quant below). Only - # FP4 (a4w4) SiLU is unsupported (ATOM reference a4w4 SiLU path is incomplete). - if act_mode == _AITER_ACT_FP4: - raise NotImplementedError( - "aiter SiLU MoE is unsupported for FP4 (a4w4) activations because the " - "ATOM reference a4w4 SiLU path is incomplete; use FP8 or BF16 (w4a16)." - ) - - hidden_q, a13_scale = dynamic_mxfp8_quant(gemm1_input, quant_dtype=quant_dtype) - raw_gate_up = moe_gemm_a8w4( - hidden_q, - w1_data, - a13_scale, - w1_wscale, - None, - None, - w1_bias, - routing_data, - gather_indx=gemm1_gather_indx, - gammas=g1_gammas, - swizzle_mx_scale=None, - out_dtype=torch.bfloat16, - apply_swiglu=False, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - if unpadded_n_w1 is not None: - raw_gate_up = raw_gate_up[:, :unpadded_n_w1] - # SiLU + clamp + (gate * up), fused with dynamic MXFP8 quant of the result. - interm_fp8, a2_scale = fused_clamp_act_mul( - raw_gate_up, - swiglu_limit=swiglu_limit, - activation="silu", - dtype_quant=quant_dtype, - scale_dtype_fmt="ue8m0", - quant_block_size=32, - ) - out = moe_gemm_a8w4( - interm_fp8, - w2_data, - a2_scale, - w2_wscale, - None, - None, - w2_bias, - routing_data, - scatter_indx=scatter_idx, - gammas=g2_gammas, - swizzle_mx_scale=None, - unpadded_N=unpadded_n_w2, - unpadded_K=unpadded_k_w2, - ) - return out - - -def _aiter_mxfp4_fused_experts( - output: torch.Tensor, - hidden_states: torch.Tensor, - w1, - w2, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - quant_config: FusedMoEQuantConfig, - activation: MoEActivation, - apply_router_weight_on_input: bool, - moe_config: FusedMoEConfig | None, -) -> torch.Tensor | None: - """aiter-native MXFP4 MoE for the modular expert path (ROCm fast path). - - Builds aiter routing from ``topk_ids``/``topk_weights``, runs the two-GEMM - aiter compute, restores the per-token router-weight sum that routing - reconstruction normalized away, and writes ``output``. - - Returns ``output`` on success, or ``None`` (caller falls back to ``matmul_ogs``) - when aiter is unavailable, the weights are not MXFP4, or the requested - activation/quant combination is out of the ported scope. - """ - if not rocm_aiter_ops.is_enabled(): - return None - act_mode = _aiter_act_quant_mode(quant_config) - if act_mode is None: - return None - if quant_config.w1_precision is None or quant_config.w2_precision is None: - return None - # SiLU served via a8w4 for FP8 and BF16 (dynamic FP8 quant); only FP4 (a4w4) - # SiLU is out of scope (reference a4w4 SiLU path incomplete). This keeps w4a16 - # SiLU on the aiter-native fast path instead of the vendored-routing fallback. - if activation != MoEActivation.SWIGLUOAI and act_mode == _AITER_ACT_FP4: - return None - aiter_ops = _import_aiter_moe_ops() - if aiter_ops is None: - return None - - num_experts = w1.shape[0] - routing_data, gather_idx, scatter_idx, tw = _aiter_routing_from_topk( - aiter_ops, topk_ids, topk_weights, num_experts, hidden_states.device - ) - swiglu_alpha, swiglu_limit = _aiter_swiglu_params(quant_config) - out = _aiter_moe_two_gemm( - aiter_ops, - hidden_states, - w1, - w2, - routing_data, - gather_idx, - scatter_idx, - topk_ids.size(1), - quant_config, - act_mode, - activation, - swiglu_alpha, - swiglu_limit, - apply_router_weight_on_input, - moe_config, - ) - # Restore the per-token weight sum that routing reconstruction normalized out. - out = out.to(output.dtype) * tw.sum(dim=1, keepdim=True) - output.copy_(out.reshape(output.shape)) - return output - - @triton.jit def _masked_topk_sum_kernel( inp_ptr, # (M, topk, K) contiguous @@ -1563,49 +1024,9 @@ def apply( self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG if expert_map is not None: - topk_ids = expert_map[topk_ids] - - # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible - # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. - if ( - not apply_router_weight_on_input - and rocm_aiter_ops.is_enabled() - and _aiter_mxfp4_fused_experts( - output, - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - self.quant_config, - activation, - apply_router_weight_on_input, - getattr(self, "moe_config", None), - ) - is not None - ): - return - - # ROCm aiter-native MXFP4 fast path (bypasses the gfx1250-incompatible - # matmul_ogs unswizzle). Returns None -> fall through to matmul_ogs. - if ( - not apply_router_weight_on_input - and rocm_aiter_ops.is_enabled() - and _aiter_mxfp4_fused_experts( - output, - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - self.quant_config, - activation, - apply_router_weight_on_input, - getattr(self, "moe_config", None), - ) - is not None - ): - return + # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch): + # make_routing_data treats -1 as the skip sentinel. + topk_ids = remap_topk_to_local(topk_ids, expert_map) local_num_experts = w1.shape[0] if global_num_experts == -1: @@ -1689,7 +1110,7 @@ def activation( alpha = ( quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None - else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default + else 1.702 ) limit = ( quant_config.gemm1_clamp_limit @@ -1722,44 +1143,6 @@ def activation( else: super().activation(activation, output, input) - def _try_apply_aiter_w4a8( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - quant_config: FusedMoEQuantConfig, - activation: MoEActivation, - apply_router_weight_on_input: bool, - ) -> torch.Tensor | None: - """ROCm aiter-native MXFP4 MoE fast path; ``None`` -> caller uses matmul_ogs. - - The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls - ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and - fails to compile on gfx1250. aiter's ``.moe.*`` gluon kernels run on gfx1250, - so we route the MoE through them. - - Delegates to :func:`_aiter_mxfp4_fused_experts` (ported from ROCm/ATOM - ``fused_moe_triton.py``): a8w4 SiLU via dynamic MXFP8 + ``fused_clamp_act_mul``, - and the fused-SwiGLU modes. Routing is built aiter-native from - ``topk_ids``/``topk_weights`` and the router-weight sum is restored afterward; - see that function for the routing/scaling details. - """ - return _aiter_mxfp4_fused_experts( - output, - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - quant_config, - activation, - apply_router_weight_on_input, - self.moe_config, - ) - def apply( self, output: torch.Tensor, @@ -1820,45 +1203,6 @@ def apply( if global_num_experts == -1: global_num_experts = E - # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from - # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs - # unswizzle. Output-side router weighting only (the post-scale below assumes - # it); LoRA still uses the matmul_ogs path further down. - if ( - ( - quant_config.use_mxfp4_w4a8 - or quant_config.use_mxfp4_w4a16 - or quant_config.use_mxfp4_w4a4 - ) - and not apply_router_weight_on_input - and self._lora_context is None - and rocm_aiter_ops.is_enabled() - ): - if ( - self._try_apply_aiter_w4a8( - output, - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - quant_config, - activation, - apply_router_weight_on_input, - ) - is not None - ): - return - - # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids - # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for - # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). - # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not - # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA - # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below - # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA - # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. - # Note that the output tensor might be in workspace13 intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) @@ -1867,108 +1211,18 @@ def apply( gammas = routing_data.gate_scal if routing_data else None - hidden_bf16_for_lora = hidden_states - - _moe_gemm_a8w4_fn = None - _downcast_static_fp8_fn = None - # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's - # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also - # serve the w4a16 case by dynamically quantizing the bf16 activations to - # FP8 here and passing the scale. This routes the MoE through aiter's - # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. - _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 - if _mxfp4_weights and rocm_aiter_ops.is_enabled(): - try: - from aiter.ops.triton.moe.moe_op_gemm_a8w4 import ( - moe_gemm_a8w4 as _moe_gemm_a8w4_fn, - ) - from aiter.ops.triton.moe.quant_moe import ( - downcast_to_static_fp8 as _downcast_static_fp8_fn, - ) - except ImportError: - pass - use_aiter_unfused_a8w4 = ( - _mxfp4_weights - and _moe_gemm_a8w4_fn is not None - and _downcast_static_fp8_fn is not None + matmul_ogs( + hidden_states, + w1, + quant_config.w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=quant_config.w1_precision, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=None, + y=intermediate_cache1, ) - if use_aiter_unfused_a8w4: - assert quant_config.w1_precision is not None - assert quant_config.w2_precision is not None - unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( - _mxfp4_w4a8_unpadded_dims(self.moe_config) - ) - swiglu_alpha = ( - quant_config.gemm1_alpha - if quant_config.gemm1_alpha is not None - else 1.0 # SILU (DeepSeek-V4); was 1.702 gpt-oss default - ) - swiglu_limit = ( - quant_config.gemm1_clamp_limit - if quant_config.gemm1_clamp_limit is not None - else 7.0 - ) - gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, - gemm_num=1, - activations=hidden_bf16_for_lora, - ) - # GEMM1 passes W2's static scale through to the kernel; resolve it before - # GEMM2 activations exist (fallback uses token hidden states for amax). - gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, - gemm_num=2, - activations=hidden_bf16_for_lora, - ) - hidden_fp8 = _downcast_static_fp8_fn( - hidden_bf16_for_lora, - gemm1_lhs_scale, - ) - gemm1_out = _moe_gemm_a8w4_fn( - hidden_fp8, - w1.storage.data, - None, - quant_config.w1_precision.weight_scale.storage.data, - gemm1_lhs_scale, - gemm2_lhs_scale_proxy, - 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.bfloat16, - apply_swiglu=False, - alpha=swiglu_alpha, - limit=swiglu_limit, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) - else: - hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( - _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( - hidden_bf16_for_lora, - quant_config, - ) - ) - w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( - quant_config.w1_precision, - gemm1_lhs_scale_fp8, - ) - - matmul_ogs( - hidden_for_gemm1, - w1, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=w1_precision_for_gemm, - gammas=gammas if apply_router_weight_on_input else None, - fused_activation=None, - y=intermediate_cache1, - ) - # w13 LoRA: gather the activation input from expert-sorted # intermediate_cache1, then add the LoRA delta in-place on that copy # before passing it to activation — exactly mirroring the old @@ -1989,7 +1243,7 @@ def apply( ) = self.apply_w13_lora( lora_context, y=act_input, - x=hidden_bf16_for_lora, + x=hidden_states, topk_ids=global_topk_ids, topk_weights=topk_weights, expert_map=expert_map, @@ -2010,56 +1264,16 @@ def apply( # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. routing_data.n_expts_act = 1 - gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] - if use_aiter_unfused_a8w4: - gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, - gemm_num=2, - activations=gemm2_input_bf16, - ) - gemm2_fp8 = _downcast_static_fp8_fn( - gemm2_input_bf16, - gemm2_lhs_scale, - ) - gemm2_out = _moe_gemm_a8w4_fn( - gemm2_fp8, - w2.storage.data, - None, - quant_config.w2_precision.weight_scale.storage.data, - gemm2_lhs_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, - ) - intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) - else: - gemm2_input, gemm2_lhs_scale_fp8 = ( - _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( - gemm2_input_bf16, - quant_config, - gemm_num=2, - ) - ) - w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( - quant_config.w2_precision, - gemm2_lhs_scale_fp8, - ) - - matmul_ogs( - gemm2_input, - w2, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - precision_config=w2_precision_for_gemm, - gammas=None if apply_router_weight_on_input else gammas, - y=intermediate_cache3, - ) + matmul_ogs( + intermediate_cache2[gather_indx.src_indx], + w2, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_indx, + precision_config=quant_config.w2_precision, + gammas=None if apply_router_weight_on_input else gammas, + y=intermediate_cache3, + ) # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 5dc152be7f9a..ba78573fa3de 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -219,11 +219,14 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( + AiterW4A16ExpertsMonolithic, + ) from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( AiterExperts, ) - return [AiterExperts] + return [AiterExperts, AiterW4A16ExpertsMonolithic] elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( From 34b9c0fc05d680078ba04681ac4c207e6d5a7b06 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Thu, 9 Jul 2026 16:31:12 -0400 Subject: [PATCH 66/68] Remove .bak Signed-off-by: Jaden Mathias --- .../experts/gpt_oss_triton_kernels_moe.py.bak | 1548 ----------------- 1 file changed, 1548 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak deleted file mode 100644 index 189a5ca8fd9e..000000000000 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py.bak +++ /dev/null @@ -1,1548 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from dataclasses import replace - -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import _custom_ops as ops -from vllm._aiter_ops import rocm_aiter_ops -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.config import ( - FUSED_MOE_UNQUANTIZED_CONFIG, - FusedMoEConfig, - FusedMoEParallelConfig, - FusedMoEQuantConfig, - RoutingMethodType, -) -from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( - LoRAExpertsMixin, -) -from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceNoOP, -) -from vllm.model_executor.layers.fused_moe.utils import _resize_cache -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - QuantKey, - kMxfp4Static, -) -from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_triton_kernels - -from ..utils import swiglu_limit_func - -logger = init_logger(__name__) - - -def _triton_kernel_moe_supports_current_device() -> bool: - # Shared device gate for the OAI Triton MoE expert classes. - # Platform-aware to avoid ROCm capability aliasing — cap (9, 0) - # matches both gfx90a (verified) and gfx906 (unverified), so we - # dispatch on gfx-string helpers instead of the cap tuple on ROCm. - p = current_platform - if p.is_cuda(): - cap = p.get_device_capability() - # Keep the original `(9, 0) <= cap < (11, 0)` window on - # CUDA (covers Hopper SM90 and Blackwell SM100, excludes - # SM120) — this PR is ROCm-scoped and the broader CUDA - # range was not validated. - return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0) - if p.is_rocm(): - from vllm.platforms.rocm import on_gfx1x, on_gfx9 - - # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx); - # on_gfx9() already excludes gfx906/gfx908. - # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4); - # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2). - return on_gfx9() or on_gfx1x() - return False - - -def _patch_make_bitmatrix_metadata() -> None: - """Monkey-patch make_bitmatrix_metadata to support non-power-of-2 top_k. - - triton's tl.arange requires a power-of-2 range. The original kernel - computes BLOCK_SIZE = BLOCK_PER_TOK * TOKS_PER_ROW (= 32 * top_k). For - DeepSeek-V4 with top_k=6 this gives 192, which is not a power of 2 and - causes a compile error at the first forward pass. - - Fix: define a drop-in replacement kernel that accepts an extra constexpr - BLOCK_SIZE_PADDED (next power of 2 >= BLOCK_SIZE) and uses it for the - tl.arange call while keeping the actual BLOCK_SIZE as the stride between - thread-blocks so that all flat indices into NonzeroIndx stay correct. - Elements beyond BLOCK_SIZE are masked out (col_indx = 0xffff) and ignored. - - This function is called once at module load time and patches the function - inside the triton_kernels tensor module so that SparseMatrix.__post_init__ - picks up the fixed version transparently. - """ - import torch - import triton - import triton.language as tl - - try: - if current_platform.is_rocm(): - from triton_kernels.tensor_details import bitmatrix as _bm - from triton_kernels.tensor_details.bitmatrix import ( - BitmatrixMetadata, - _keyed_add, - cdiv, - ) - from triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 - sum_bitmatrix_rows, - ) - else: - from vllm.third_party.triton_kernels.tensor_details import ( - bitmatrix as _bm, - ) - from vllm.third_party.triton_kernels.tensor_details.bitmatrix import ( - BitmatrixMetadata, - _keyed_add, - cdiv, - ) - from vllm.third_party.triton_kernels.tensor_details.bitmatrix_details.sum_bitmatrix_rows import ( # noqa: E501 - sum_bitmatrix_rows, - ) - except ImportError: - return - - @triton.jit - def _stage2_pow2( - ColSortedIndx, - RowSortedIndx, - NonzeroIndx, - n_tokens, - ColPartialSum, - stride_pm, - stride_pn, - ColOffs, - TOKS_PER_ROW: tl.constexpr, - BLOCK_PER_TOK: tl.constexpr, - BLOCK_SIZE_PADDED: tl.constexpr, - ): - # Actual number of elements per block (may not be a power of 2). - BLOCK_SIZE: tl.constexpr = BLOCK_PER_TOK * TOKS_PER_ROW - tl.static_assert(BLOCK_SIZE_PADDED <= 32768) - if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr(): - n_tokens = tl.load(n_tokens) - nonzero_indx_size = n_tokens * TOKS_PER_ROW - pid_m = tl.program_id(0) - # Use BLOCK_SIZE_PADDED (a power of 2) for tl.arange, but stride by - # the actual BLOCK_SIZE so flat positions in NonzeroIndx are correct. - # Elements with offs_local >= BLOCK_SIZE have offs_global beyond the - # valid range, get col_indx = 0xffff, and are filtered by the mask - # below without producing any output. - offs_local = tl.arange(0, BLOCK_SIZE_PADDED) - offs_global = pid_m * BLOCK_SIZE + offs_local - mask = offs_global < nonzero_indx_size - col_indx = tl.load(NonzeroIndx + offs_global, mask=mask, other=-1).to(tl.uint32) - kv_pairs = ((col_indx << 16) | offs_local).to(tl.uint32) - kv_pairs = tl.sort(kv_pairs, 0) - col_indx = kv_pairs >> 16 - offs_global = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF) - mask = col_indx != 0xFFFF - x = kv_pairs & 0xFFFF0000 | 0x00000001 - cols_and_inclusive_run_lengths = tl.associative_scan(x, 0, _keyed_add) - exclusive_run_lengths = (cols_and_inclusive_run_lengths - 1) & 0xFFFF - row_sorted_indx = tl.load( - ColPartialSum + pid_m * stride_pm + col_indx * stride_pn, mask=mask - ) - row_sorted_indx += tl.load(ColOffs + col_indx, mask=mask) - row_sorted_indx += exclusive_run_lengths - tl.store(RowSortedIndx + offs_global, row_sorted_indx, mask=mask) - tl.store(ColSortedIndx + row_sorted_indx, offs_global, mask=mask) - - def _make_bitmatrix_metadata_pow2_safe(nonzero_indx, bitmatrix): - assert nonzero_indx.ndim == 2 - PARTIAL_BLOCK_M = 32 - col_sum, col_partial_sum = sum_bitmatrix_rows( - bitmatrix, partials_block_size=PARTIAL_BLOCK_M - ) - device = bitmatrix.device - n_indx = nonzero_indx.numel() - n_cols = bitmatrix.shape[1] - col_offs = torch.empty(n_cols, dtype=torch.int32, device=device) - combined_indx = torch.empty(n_indx * 2, dtype=torch.int32, device=device) - col_sorted_indx = combined_indx[:n_indx] - row_sorted_indx = combined_indx[n_indx:] - MEMSET_BLOCK = 1024 - memset_grid = (cdiv(n_indx * 2, MEMSET_BLOCK) + n_cols + 1,) - _bm._bitmatrix_metadata_compute_stage1[memset_grid]( - combined_indx, - n_indx * 2, - -1, - MEMSET_BLOCK, - col_sum, - col_offs, - col_sum.shape[0], - col_partial_sum, - col_partial_sum.shape[0], - col_partial_sum.stride(0), - col_partial_sum.stride(1), - BLOCK_M=512, - BLOCK_N=512, - ) - toks_per_row = nonzero_indx.shape[-1] - block_size = PARTIAL_BLOCK_M * toks_per_row - # Next power of 2 >= block_size (required by tl.arange). - block_size_padded = 1 << (max(block_size, 1) - 1).bit_length() - compute_grid = (cdiv(bitmatrix.shape_max[0], PARTIAL_BLOCK_M),) - _stage2_pow2[compute_grid]( - col_sorted_indx, - row_sorted_indx, - nonzero_indx, - bitmatrix.shape[0], - col_partial_sum, - col_partial_sum.stride(0), - col_partial_sum.stride(1), - col_offs, - TOKS_PER_ROW=toks_per_row, - BLOCK_PER_TOK=PARTIAL_BLOCK_M, - BLOCK_SIZE_PADDED=block_size_padded, - ) - return BitmatrixMetadata( - col_sum=col_sum, - col_sorted_indx=col_sorted_indx, - row_sorted_indx=row_sorted_indx, - ) - - # The most reliable patch point: SparseMatrix.__post_init__ looks up - # make_bitmatrix_metadata via its own __globals__ dict (the tensor.py - # module dict). Patching through __globals__ works regardless of how - # sys.modules maps "triton_kernels.tensor" vs - # "vllm.third_party.triton_kernels.tensor". - from triton_kernels.tensor import SparseMatrix as _SparseMatrix - - _SparseMatrix.__post_init__.__globals__["make_bitmatrix_metadata"] = ( - _make_bitmatrix_metadata_pow2_safe - ) - # Also patch the bitmatrix module itself in case it is imported directly. - _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe - - -# Two API generations of triton_kernels are supported: -# - v3.5.1 (the version bundled with vLLM): exposes `routing()` and -# `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix` -# constructor takes a `scratchpad` argument. -# - v3.6.0+: removes the `routing` module in favor of a `SparseMatrix` -# based path, and adds a `dtype=BIT` kwarg to `Bitmatrix`. Used only -# when the user has triton_kernels installed system-wide at v3.6.0+. -# -# `use_legacy_triton_kernels` selects between them at import time based on -# whether `SparseMatrix` is importable. -use_legacy_triton_kernels = False - -if has_triton_kernels(): - try: - import triton_kernels.swiglu - from triton_kernels.matmul_ogs import ( - FnSpecs, - FusedActivation, - GatherIndx, - RoutingData, - ScatterIndx, - matmul_ogs, - ) - from triton_kernels.tensor import ( - BIT, - Bitmatrix, - ) - - try: - from triton_kernels.tensor import ( - SparseMatrix, - make_ragged_tensor_metadata, - ) - except ImportError: - # TODO(mgoin): drop the v3.5.1 pin and remove this fallback once - # the gpt-oss perf regression in v3.6.0+ is resolved upstream. - # Tracking: https://github.com/triton-lang/triton/issues/9969 - use_legacy_triton_kernels = True - if not use_legacy_triton_kernels: - _patch_make_bitmatrix_metadata() - except (AttributeError, ImportError) as e: - logger.error( - "Failed to import Triton kernels. Please make sure your triton " - "version is compatible. Error: %s", - e, - ) - - -@triton.jit -def pack_bitmatrix( - bitmatrix, - topk_ids, - n_rows, # n_rows in bitmatrix / topk_ids - bm_cols: tl.constexpr, # n int32_t bitpacks in bitmatrix - n_expts_act, # num_topk - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, -): - """ - Packs topk_ids into a bitmatrix. - code reference: - https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264 - """ - pid_m = tl.program_id(0) - offsets_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offsets_k = tl.arange(0, BLOCK_SIZE_K) - offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :] - mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :] - indices = tl.load(topk_ids + offsets, mask=mask, other=-1) - valid = indices >= 0 - div = indices // 32 - rem = indices % 32 - one = tl.cast(1, tl.uint32) - - # Iterate through all the relevant bitmatrix columns. - for i in range(bm_cols): - # When BLOCK_SIZE_K=32, offs is just the column index. - offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) - # All topks that need to go into this column has the correct bit set. - # Other bits are 0. x is a 2D tensor. - # Guard with `valid` to prevent negative indices from producing - # spurious bits (on HIP, -1 // 32 == 0 and 1 << (-1 % 32) sets - # bit 31). - x = tl.where( - valid[:, :, None] & (div[:, :, None] == offs[None, None, :]), - (one << rem)[:, :, None], - 0, - ) - # Reduce x to get a single int32_t bitpack. - y = tl.reduce_or(x, axis=1) - bitmatrix_ptrs = bitmatrix + offsets_m[:, None] * bm_cols + offs[None, :] - tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) - - -def triton_kernel_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, -) -> torch.Tensor: - sm_first = not renormalize - - # When no expert map is provided (no EP), call the fused `routing()` - # kernel directly. It combines softmax, topk, bitmatrix packing, and - # routing-metadata construction in a single launch, instead of the - # three separate kernels used by the generic path below. - # Only available in the legacy (v3.5.1) API; the v3.6.0+ path inlines - # equivalent logic via SparseMatrix in `make_routing_data`. - if use_legacy_triton_kernels and expert_map is None: - from triton_kernels.routing import routing as fused_routing - - routing_data, gather_idx, scatter_idx = fused_routing( - gating_output, topk, sm_first=sm_first - ) - effective_expert_map = None - effective_global_num_experts = global_num_experts - else: - from triton_kernels.topk import topk as topk_fn - - logits = gating_output - if sm_first: - logits = torch.softmax(logits, dim=-1) - topk_result = topk_fn(logits, topk, apply_softmax=not sm_first) - # topk may return a tuple (vals, indx, bitmatrix) or a - # SparseMatrix depending on the triton_kernels version. - if isinstance(topk_result, tuple): - topk_weights, topk_ids_raw, _ = topk_result - else: - topk_weights = topk_result.vals - topk_ids_raw = topk_result.indx - - if expert_map is not None: - # topk_ids_raw contains global expert IDs - remap to local. - topk_ids = expert_map[topk_ids_raw.to(torch.long)] - local_num_experts = w1.shape[0] - routing_data, gather_idx, scatter_idx = make_routing_data( - topk_ids, topk_weights, local_num_experts - ) - # expert_map already applied; pass None downstream. - effective_expert_map = None - effective_global_num_experts = local_num_experts - else: - topk_ids = topk_ids_raw.to(torch.long) - routing_data, gather_idx, scatter_idx = make_routing_data( - topk_ids, topk_weights, gating_output.shape[-1] - ) - effective_expert_map = expert_map - effective_global_num_experts = global_num_experts - - output = torch.empty_like(hidden_states) - effective_quant_config = ( - quant_config if quant_config is not None else FUSED_MOE_UNQUANTIZED_CONFIG - ) - - return triton_kernel_fused_experts( - output, - hidden_states, - w1, - w2, - routing_data, - gather_idx, - scatter_idx, - topk=topk, - activation=activation, - quant_config=effective_quant_config, - apply_router_weight_on_input=apply_router_weight_on_input, - global_num_experts=effective_global_num_experts, - expert_map=effective_expert_map, - ) - - -# This is a triton implementation of the fused_experts function -def triton_kernel_fused_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 - topk: int, - activation: MoEActivation = MoEActivation.SWIGLUOAI, - 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, - intermediate_cache: torch.Tensor | None = None, - a1q_scale: torch.Tensor | None = None, -) -> torch.Tensor: - """Triton implementation of fused expert computation using OAI kernels.""" - assert activation == MoEActivation.SWIGLUOAI, ( - "Only SWIGLUOAI activation is supported" - ) - 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, only check non-mxfp4 - assert hidden_states.ndim == 2 - assert hidden_states.shape[-1] == w1.shape[-2] - assert w2.shape[-1] == w1.shape[1] - - batch_dim = 1 - M, K = hidden_states.shape[-2:] - E, _, N = w1.shape - - if global_num_experts == -1: - global_num_experts = E - - if intermediate_cache is None: - intermediate_cache = torch.empty( - (batch_dim, M * topk, N // 2), - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - - # Add batch_dim to output buffer because matmul_ogs expects 3D output - intermediate_cache = _resize_cache( - intermediate_cache, (batch_dim, M * topk, N // 2) - ) - output_tensor = _resize_cache(output_tensor, (batch_dim, M, K)) - - act = ( - FusedActivation( - FnSpecs( - "swiglu", - triton_kernels.swiglu.swiglu_fn, - ("alpha", "limit"), - reduction_n=2, - ), - (swiglu_alpha, swiglu_limit), - ) - if not use_legacy_triton_kernels - else FusedActivation( - FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), - (swiglu_alpha, swiglu_limit), - 2, - ) - ) - gammas = routing_data.gate_scal if routing_data else None - - matmul_ogs( - hidden_states, - w1, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=quant_config.w1_precision, - gammas=gammas if apply_router_weight_on_input else None, - fused_activation=act, - y=intermediate_cache, - ) - - matmul_ogs( - intermediate_cache.view(M * topk, N // 2), - w2, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - precision_config=quant_config.w2_precision, - gammas=None if apply_router_weight_on_input else gammas, - y=output_tensor, - ) - output_tensor = output_tensor.view(M, K) - return output_tensor - - -def make_routing_data( - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - num_local_experts: int, -) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: - topk_ids = topk_ids.to(torch.int16) - topk_weights = topk_weights.to(torch.bfloat16) - - n_rows, num_topk = topk_ids.size() - - BLOCK_SIZE_M = 512 - BLOCK_SIZE_K = 32 - - bm_cols = triton.cdiv(num_local_experts, BLOCK_SIZE_K) # n_bitpacks - bitmatrix = torch.zeros( - (n_rows, bm_cols), dtype=torch.uint32, device=topk_ids.device - ) - - grid = (triton.cdiv(n_rows, BLOCK_SIZE_M),) - pack_bitmatrix[grid]( - bitmatrix, - topk_ids, - n_rows, - bm_cols, - num_topk, - BLOCK_SIZE_M=BLOCK_SIZE_M, - BLOCK_SIZE_K=BLOCK_SIZE_K, - ) - - bitmatrix_shape = [n_rows, bm_cols * 32] - bitmatrix_shape_max = [n_rows, None] - bitmatrix = ( - Bitmatrix( - bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max - ) - if not use_legacy_triton_kernels - else Bitmatrix( - bitmatrix, - shape=bitmatrix_shape, - shape_max=bitmatrix_shape_max, - scratchpad=None, - ) - ) - - # matmul_ogs expects invalid topk_weights to be -1s - topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights) - - if use_legacy_triton_kernels: - from triton_kernels.routing import routing_from_bitmatrix - - return routing_from_bitmatrix( - bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk - ) - - sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix) - dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx - combine_indx = sparse_logits.mask_metadata.col_sorted_indx - ragged_batch_metadata = make_ragged_tensor_metadata( - sparse_logits.mask_metadata.col_sum, - dispatch_indx.shape[0], - ) - gate_scal = sparse_logits.vals.flatten()[combine_indx] - routing_data = RoutingData( - gate_scal, - ragged_batch_metadata.block_sizes, - num_local_experts, - num_topk, - ragged_batch_metadata, - ) - gather_indx = GatherIndx(combine_indx, dispatch_indx) - scatter_indx = ScatterIndx(dispatch_indx, combine_indx) - return routing_data, gather_indx, scatter_indx - - -def _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( - hidden_states: torch.Tensor, - quant_config: FusedMoEQuantConfig, - *, - gemm_num: int = 1, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """BF16 activations -> FP8 before unfused ``matmul_ogs`` for MXFP4 w4a16 (ROCm). - - Uses ``aiter.ops.triton.quant_moe.downcast_to_static_fp8`` with a scalar scale: - for GEMM1: ``w1_precision.flex_ctx.lhs_data.scale`` or ``a1_scale`` when - single-element; for GEMM2: ``w2_precision.flex_ctx.lhs_data.scale`` or - ``a2_scale``; otherwise ``max(|x|) / 448`` for e4m3 range (matches AITER static - FP8 path). - - Returns: - (activations, lhs_scale_or_none): ``lhs_scale`` is the scalar tensor used for - FP8 downcast when this path runs; otherwise ``None`` (activations unchanged). - """ - if gemm_num not in (1, 2): - raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") - if not quant_config.use_mxfp4_w4a16: - return hidden_states, None - try: - from aiter.ops.triton.quant_moe import downcast_to_static_fp8 - except ImportError: - return hidden_states, None - if not rocm_aiter_ops.is_enabled(): - return hidden_states, None - - qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision - a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale - - lhs_scale = None - if qp is not None and qp.flex_ctx.lhs_data.scale is not None: - s = qp.flex_ctx.lhs_data.scale - if s.numel() == 1: - lhs_scale = s - if lhs_scale is None and a_scale is not None: - s = a_scale - if s.numel() == 1: - lhs_scale = s - if lhs_scale is None: - amax = hidden_states.abs().max().clamp(min=1e-12) - lhs_scale = (amax / 448.0).to(dtype=torch.float32) - - return downcast_to_static_fp8(hidden_states, lhs_scale), lhs_scale - - -def _mxfp4_w4a8_unpadded_dims( - moe_cfg: FusedMoEConfig, -) -> tuple[int | None, int | None, int | None, int | None]: - """Logical unpadded GEMM sizes for padded MXFP4 checkpoints (e.g. GFX950 swizzle).""" - if ( - moe_cfg.intermediate_size_per_partition_unpadded is None - or moe_cfg.hidden_dim_unpadded is None - ): - return None, None, None, None - unpadded_n_w1 = moe_cfg.intermediate_size_per_partition_unpadded * 2 - unpadded_k_w1 = moe_cfg.hidden_dim_unpadded - unpadded_n_w2 = moe_cfg.hidden_dim_unpadded - unpadded_k_w2 = moe_cfg.intermediate_size_per_partition_unpadded - return unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 - - -def _mxfp4_w4a8_resolve_lhs_scale( - quant_config: FusedMoEQuantConfig, - *, - gemm_num: int, - activations: torch.Tensor, -) -> torch.Tensor: - """Scalar FP8 LHS scale for ``downcast_to_static_fp8`` / ``moe_gemm_a8w4``. - - Prefer calibrated ``flex_ctx.lhs_data.scale`` or ``a{1,2}_scale``; otherwise - ``max(|activations|) / 448`` (e4m3), matching - ``_maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused``. - """ - if gemm_num not in (1, 2): - raise ValueError(f"gemm_num must be 1 or 2, got {gemm_num}") - qp = quant_config.w1_precision if gemm_num == 1 else quant_config.w2_precision - a_scale = quant_config.a1_scale if gemm_num == 1 else quant_config.a2_scale - lhs_scale = None - if qp is not None and qp.flex_ctx is not None: - ld = qp.flex_ctx.lhs_data - if ld is not None and ld.scale is not None and ld.scale.numel() == 1: - lhs_scale = ld.scale - if lhs_scale is None and a_scale is not None: - s = a_scale - if s.numel() == 1: - lhs_scale = s - if lhs_scale is None: - amax = activations.abs().max().clamp(min=1e-12) - lhs_scale = (amax / 448.0).to(dtype=torch.float32) - return lhs_scale - - -def _precision_config_with_mxfp4_unfused_fp8_lhs_scale( - precision_config, - lhs_scale: torch.Tensor | None, -): - """Ensure ``matmul_ogs`` sees the same LHS FP8 scale as ``downcast_to_static_fp8``.""" - if precision_config is None or lhs_scale is None: - return precision_config - from triton_kernels.numerics import InFlexData - - new_flex = replace( - precision_config.flex_ctx, - lhs_data=InFlexData(scale=lhs_scale), - ) - return replace(precision_config, flex_ctx=new_flex) - - -class BaseOAITritonExperts(mk.FusedMoEExpertsModular): - @property - def expects_unquantized_inputs(self) -> bool: - return True - - @staticmethod - def _supports_current_device() -> bool: - return _triton_kernel_moe_supports_current_device() and has_triton_kernels() - - @staticmethod - def _supports_no_act_and_mul() -> bool: - return False - - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - SUPPORTED_W_A = [ - (kMxfp4Static, None), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - raise NotImplementedError - - @staticmethod - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return True - - def moe_problem_size( - self, - a1: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_ids: torch.Tensor, - ) -> tuple[int, int, int, int, int]: - """ - Extract the MoE problem size from the given tensor arguments: - - a: The hidden states, input to the MoE layer. - - w1: The first set of expert weights. - - w2: The second set of expert weights. - - topk_ids: The topk ids. - Note: extracting the problem shape from the weight and activation - tensors is not obvious. It needs to be done this way specifically - due to subtle issues with particular kernels, e.g. the int4 kernels - divide the trailing dimension by two, so it's not "correct" to - extract N or K from the trailing dimension of w1 or w2. Similarly, - some kernels transpose the weights, so this needs to be kept in mind. - Note: This implementation covers most cases. However, if experts - require a specialized implementation, like MarlinExperts, they are free - to override this function. - """ - assert len(w1.shape) == 3 and len(w2.shape) == 3 - E, _, N = w1.shape - K = a1.size(-1) - - assert a1.dim() == 2 - assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}" - M = a1.size(0) - - assert topk_ids.dim() == 2 - topk = topk_ids.size(1) - - return E, M, N, K, topk - - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - # Weight application and reduction happens in the fused_experts kernel. - return TopKWeightAndReduceNoOP() - - def _make_routing_data( - self, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - num_local_experts: int, - ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]: - return make_routing_data(topk_ids, topk_weights, num_local_experts) - - -class OAITritonExperts(BaseOAITritonExperts): - """OAI Triton-based fused MoE expert implementation.""" - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SWIGLUOAI - - @staticmethod - def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - def workspace_shapes( - self, - M: int, - N: int, - K: int, - topk: int, - global_num_experts: int, - local_num_experts: int, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: MoEActivation, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # workspace are allocated inside the kernel - activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (0, 0) - workspace2 = (M * topk, activation_out_dim) - output = (M, K) - return (workspace1, workspace2, output) - - def apply( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, - ): - if self.quant_config is None: - self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG - - if expert_map is not None: - topk_ids = expert_map[topk_ids] - - local_num_experts = w1.shape[0] - if global_num_experts == -1: - global_num_experts = local_num_experts - - routing_data, gather_indx, scatter_indx = self._make_routing_data( - topk_ids, topk_weights, local_num_experts - ) - - topk = topk_ids.size(1) - triton_kernel_fused_experts( - output, - hidden_states, - w1, - w2, - routing_data, - gather_indx, - scatter_indx, - topk=topk, - activation=activation, - quant_config=self.quant_config, - apply_router_weight_on_input=False, - global_num_experts=local_num_experts, - expert_map=None, # applied already - intermediate_cache=workspace2, - a1q_scale=a1q_scale, - ) - - -class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): - """ - A Triton based MoE expert class that operates on expert standard - format and explicitly keeps the activation and reduction (moe_sum) steps - unfused from the matmul_ogs kernel. This exposes injection points - for activation and moe_sum. - - One use case for it is to inject LoRA modules on the activation and moe_sum. - """ - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - return activation in [ - MoEActivation.SILU, - MoEActivation.GELU, - MoEActivation.SWIGLUOAI, - MoEActivation.SWIGLUSTEP, - MoEActivation.SWIGLUOAI_UNINTERLEAVE, - ] - - @staticmethod - def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - def workspace_shapes( - self, - M: int, - N: int, - K: int, - topk: int, - global_num_experts: int, - local_num_experts: int, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: MoEActivation, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # workspace are allocated inside the kernel - activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (M * topk, activation_out_dim) - workspace2 = (M * topk, max(N, K)) - output = (M, K) - return (workspace1, workspace2, output) - - def moe_sum(self, input: torch.Tensor, output: torch.Tensor): - ops.moe_sum(input, output) - - def activation( - self, - activation: MoEActivation, - output: torch.Tensor, - input: torch.Tensor, - **kwargs, - ) -> None: - quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG - if activation == MoEActivation.SWIGLUOAI: - alpha = ( - quant_config.gemm1_alpha - if quant_config.gemm1_alpha is not None - else 1.702 - ) - limit = ( - quant_config.gemm1_clamp_limit - if quant_config.gemm1_clamp_limit is not None - else 7.0 - ) - torch.ops._C.swigluoai_and_mul(output, input, alpha, limit) - elif ( - activation == MoEActivation.SILU - and quant_config.gemm1_clamp_limit is not None - ): - swiglu_limit_func( - output, - input, - quant_config.gemm1_clamp_limit, - ) - elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: - assert quant_config.gemm1_clamp_limit is not None - alpha = ( - quant_config.gemm1_alpha - if quant_config.gemm1_alpha is not None - else 1.0 - ) - beta = ( - quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0 - ) - torch.ops._C.silu_and_mul_with_clamp( - output, input, quant_config.gemm1_clamp_limit, alpha, beta - ) - else: - super().activation(activation, output, input) - - def _try_apply_aiter_w4a8( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - quant_config: FusedMoEQuantConfig, - apply_router_weight_on_input: bool, - ) -> torch.Tensor | None: - """ROCm/gfx1250 MXFP4 MoE via aiter's W4A8 ``moe_gemm_a8w4``. - - The vendored ``triton_kernels`` ``matmul_ogs`` MXFP4 path calls - ``unswizzle_mx_scale_cdna4``, which has no gfx1250 scale-layout variant and - fails to compile on gfx1250. aiter's gluon ``moe_gemm_a8w4`` (FP8 act x MXFP4 - weight) does run on gfx1250, so we route the MoE through it. - - Routing is built aiter-native directly from ``topk_ids``/``topk_weights`` so we - never touch the diverged vendored ``triton_kernels`` routing structs (whose - ``RoutingData``/``GatherIndx``/``ExptData`` field names and layout differ from - aiter's). This mirrors ``triton_kernel_fused_mxfp4_w4a8_experts`` - (aiter_mxfp4_w4a8_moe.py) but resolves the FP8 LHS scales dynamically because - DeepSeek-V4 is ``activation_scheme: dynamic`` (no calibrated static scale). - - aiter's ``routing`` always softmaxes the gate weights, so we feed - ``log(topk_weights)``: softmax-over-selected then recovers - ``topk_weights / sum(topk_weights)``. DeepSeek-V4 weights are - norm_topk_prob-normalized then scaled by ``routed_scaling_factor`` (so they sum - to that factor, not 1); we restore the exact weighting with a per-token output - rescale by the weight sum (== 1 for plain renormalize routing, a no-op there). - - Returns ``output`` on success, or ``None`` if aiter is unavailable (caller then - falls back to ``matmul_ogs``). - """ - try: - from aiter.ops.triton.moe.moe_routing.routing import ( - routing as aiter_routing, - ) - from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 - from aiter.ops.triton.quant_moe import downcast_to_static_fp8 - except ImportError: - return None - - assert quant_config.w1_precision is not None - assert quant_config.w2_precision is not None - - M = hidden_states.shape[0] - num_experts = w1.shape[0] - topk = topk_ids.size(1) - - # Reconstruct dense gating logits from the sparse topk selection. log(weight) - # padded with a large-negative sentinel makes aiter's softmax reproduce both - # this exact expert selection and weight = topk_weights / per-token-sum. - tw = topk_weights.to(torch.float32) - logits = torch.full( - (M, num_experts), -1e30, device=hidden_states.device, dtype=torch.float32 - ) - logits.scatter_( - 1, - topk_ids.long().clamp(min=0, max=num_experts - 1), - torch.log(tw.clamp(min=1e-20)), - ) - logger.info( - "[aiter-w4a8] ENTER M=%d K=%d num_experts=%d topk=%d " - "w1=%s w2=%s w1_wscale=%s w2_wscale=%s", - M, - hidden_states.shape[1], - num_experts, - topk, - tuple(w1.storage.data.shape), - tuple(w2.storage.data.shape), - tuple(quant_config.w1_precision.weight_scale.storage.data.shape), - tuple(quant_config.w2_precision.weight_scale.storage.data.shape), - ) - routing_data, gather_idx, scatter_idx = aiter_routing( - logits, topk, sm_first=False - ) - gammas = routing_data.gate_scal - logger.info( - "[aiter-w4a8] routing OK block_m=%s gate_scal=%s gather=%s scatter=%s", - getattr(routing_data, "block_m", None), - tuple(gammas.shape), - tuple(gather_idx.shape), - tuple(scatter_idx.shape), - ) - - unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( - _mxfp4_w4a8_unpadded_dims(self.moe_config) - ) - swiglu_alpha = ( - quant_config.gemm1_alpha - if quant_config.gemm1_alpha is not None - else 1.702 - ) - swiglu_limit = ( - quant_config.gemm1_clamp_limit - if quant_config.gemm1_clamp_limit is not None - else 7.0 - ) - - # FP8 LHS scales (static if calibrated, else dynamic amax/448). GEMM1 emits FP8 - # directly (apply_swiglu=True, out_dtype fp8) quantized with GEMM2's LHS scale, - # so resolve GEMM2's scale up front (hidden states as the amax proxy) and reuse - # it as GEMM2's input scale so quantize/dequantize are consistent. - gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, gemm_num=1, activations=hidden_states - ) - gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, gemm_num=2, activations=hidden_states - ) - # aiter's in-kernel gather is numerically broken on gfx1250 (validated on the - # FFM sim: do_gather=True -> maxrel ~2.4), so gather rows into expert-sorted - # order in torch and pass gather_indx=None. Per aiter's moe_gemm_torch, sorted - # row i reads source token gather_idx[i] // n_expts_act, so this reproduces the - # in-kernel gather exactly (validated: manual gather -> maxrel ~5e-3). - gather_src = gather_idx.to(torch.long) // topk - hidden_sorted = hidden_states[gather_src] - hidden_fp8 = downcast_to_static_fp8(hidden_sorted, gemm1_lhs_scale) - logger.info( - "[aiter-w4a8] scales g1_lhs=%s g2_lhs=%s hidden_fp8=%s; calling gemm1 " - "(unpadded N/K w1=%s/%s w2=%s/%s, alpha=%.4f limit=%.2f)", - float(gemm1_lhs_scale), - float(gemm2_lhs_scale), - tuple(hidden_fp8.shape), - unpadded_n_w1, - unpadded_k_w1, - unpadded_n_w2, - unpadded_k_w2, - swiglu_alpha, - swiglu_limit, - ) - - intermediate_cache1 = moe_gemm_a8w4( - hidden_fp8, - w1.storage.data, - None, - quant_config.w1_precision.weight_scale.storage.data, - gemm1_lhs_scale, - gemm2_lhs_scale, - quant_config.w1_bias, - routing_data, - gather_indx=None, - gammas=gammas if apply_router_weight_on_input else None, - swizzle_mx_scale=None, - out_dtype=torch.float8_e4m3fn, - apply_swiglu=True, - alpha=swiglu_alpha, - limit=swiglu_limit, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - logger.info( - "[aiter-w4a8] gemm1 OK ic1=%s %s; calling gemm2", - tuple(intermediate_cache1.shape), - intermediate_cache1.dtype, - ) - intermediate_cache3 = moe_gemm_a8w4( - intermediate_cache1, - w2.storage.data, - None, - quant_config.w2_precision.weight_scale.storage.data, - gemm2_lhs_scale, - None, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_idx, - gammas=None if apply_router_weight_on_input else gammas, - swizzle_mx_scale=None, - unpadded_N=unpadded_n_w2, - unpadded_K=unpadded_k_w2, - ) - - # Restore the per-token weight sum (e.g. routed_scaling_factor) that the - # log+softmax normalization divided out. Output weighting (gammas on GEMM2) is - # linear in the gate, so a single post-scale is exact. - logger.info( - "[aiter-w4a8] gemm2 OK ic3=%s %s; rescaling + writing output=%s", - tuple(intermediate_cache3.shape), - intermediate_cache3.dtype, - tuple(output.shape), - ) - out = intermediate_cache3.to(output.dtype) * tw.sum(dim=1, keepdim=True) - output.copy_(out.reshape(output.shape)) - logger.info("[aiter-w4a8] DONE") - return output - - def apply( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, - ): - # Use local variable to help mypy narrow the type after None check - quant_config = self.quant_config - if quant_config is None: - quant_config = FUSED_MOE_UNQUANTIZED_CONFIG - - global_topk_ids = topk_ids - if expert_map is not None: - topk_ids = expert_map[topk_ids] - - local_num_experts = w1.shape[0] - if global_num_experts == -1: - global_num_experts = local_num_experts - - routing_data, gather_indx, scatter_indx = self._make_routing_data( - topk_ids, topk_weights, local_num_experts - ) - - topk = topk_ids.size(1) - - # 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, only check non-mxfp4 - assert hidden_states.ndim == 2 - assert hidden_states.shape[-1] == w1.shape[-2] - assert w2.shape[-1] == w1.shape[1] - - batch_dim = 1 - M, K = hidden_states.shape - E, _, N = w1.shape - - if global_num_experts == -1: - global_num_experts = E - - # gfx1250 fast path: aiter-native W4A8 moe_gemm_a8w4 (built straight from - # topk_ids/topk_weights), bypassing the gfx1250-incompatible matmul_ogs - # unswizzle. Output-side router weighting only (the post-scale below assumes - # it); LoRA still uses the matmul_ogs path further down. - if ( - (quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16) - and not apply_router_weight_on_input - and self._lora_context is None - and rocm_aiter_ops.is_enabled() - ): - if ( - self._try_apply_aiter_w4a8( - output, - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - quant_config, - apply_router_weight_on_input, - ) - is not None - ): - return - - # FP8 activations + AITER Triton moe_gemm_a8w4 (MXFP4 w4a8), then moe_sum — avoids - # matmul_ogs on ROCm. Without AITER the unfused path below uses matmul_ogs (and for - # MXFP4 w4a16, optional FP8 activations via _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused). - # NOTE: the fused `triton_kernel_fused_mxfp4_w4a8_experts` helper is not - # present in this tree (it lives in the dsv4_455 WIP), so the no-LoRA - # fused branch was removed. The unfused AITER `moe_gemm_a8w4` path below - # (`use_aiter_unfused_a8w4`) handles MXFP4 w4a8 for both LoRA and non-LoRA - # and avoids the gfx1250-incompatible `matmul_ogs` unswizzle. - - # Note that the output tensor might be in workspace13 - intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N)) - intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K)) - activation_out_dim = self.adjust_N_for_activation(N, activation) - intermediate_cache2 = _resize_cache(workspace13, (M * topk, activation_out_dim)) - - gammas = routing_data.gate_scal if routing_data else None - - hidden_bf16_for_lora = hidden_states - - _moe_gemm_a8w4_fn = None - _downcast_static_fp8_fn = None - # The weights are MXFP4 in both the w4a8 and w4a16 cases; aiter's - # moe_gemm_a8w4 is a W4A8 (FP8 act x MXFP4 weight) kernel, so we can also - # serve the w4a16 case by dynamically quantizing the bf16 activations to - # FP8 here and passing the scale. This routes the MoE through aiter's - # gfx1250 kernel instead of the gfx1250-incompatible matmul_ogs. - _mxfp4_weights = quant_config.use_mxfp4_w4a8 or quant_config.use_mxfp4_w4a16 - if _mxfp4_weights and rocm_aiter_ops.is_enabled(): - try: - from aiter.ops.triton.moe_op_gemm_a8w4 import ( - moe_gemm_a8w4 as _moe_gemm_a8w4_fn, - ) - from aiter.ops.triton.quant_moe import ( - downcast_to_static_fp8 as _downcast_static_fp8_fn, - ) - except ImportError: - pass - use_aiter_unfused_a8w4 = ( - _mxfp4_weights - and _moe_gemm_a8w4_fn is not None - and _downcast_static_fp8_fn is not None - ) - - if use_aiter_unfused_a8w4: - assert quant_config.w1_precision is not None - assert quant_config.w2_precision is not None - unpadded_n_w1, unpadded_k_w1, unpadded_n_w2, unpadded_k_w2 = ( - _mxfp4_w4a8_unpadded_dims(self.moe_config) - ) - swiglu_alpha = ( - quant_config.gemm1_alpha - if quant_config.gemm1_alpha is not None - else 1.702 - ) - swiglu_limit = ( - quant_config.gemm1_clamp_limit - if quant_config.gemm1_clamp_limit is not None - else 7.0 - ) - gemm1_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, - gemm_num=1, - activations=hidden_bf16_for_lora, - ) - # GEMM1 passes W2's static scale through to the kernel; resolve it before - # GEMM2 activations exist (fallback uses token hidden states for amax). - gemm2_lhs_scale_proxy = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, - gemm_num=2, - activations=hidden_bf16_for_lora, - ) - hidden_fp8 = _downcast_static_fp8_fn( - hidden_bf16_for_lora, - gemm1_lhs_scale, - ) - gemm1_out = _moe_gemm_a8w4_fn( - hidden_fp8, - w1.storage.data, - None, - quant_config.w1_precision.weight_scale.storage.data, - gemm1_lhs_scale, - gemm2_lhs_scale_proxy, - 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.bfloat16, - apply_swiglu=False, - alpha=swiglu_alpha, - limit=swiglu_limit, - unpadded_N=unpadded_n_w1, - unpadded_K=unpadded_k_w1, - ) - intermediate_cache1.copy_(gemm1_out.reshape(intermediate_cache1.shape)) - else: - hidden_for_gemm1, gemm1_lhs_scale_fp8 = ( - _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( - hidden_bf16_for_lora, - quant_config, - ) - ) - w1_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( - quant_config.w1_precision, - gemm1_lhs_scale_fp8, - ) - - matmul_ogs( - hidden_for_gemm1, - w1, - quant_config.w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=w1_precision_for_gemm, - gammas=gammas if apply_router_weight_on_input else None, - fused_activation=None, - y=intermediate_cache1, - ) - - # w13 LoRA: gather the activation input from expert-sorted - # intermediate_cache1, then add the LoRA delta in-place on that copy - # before passing it to activation — exactly mirroring the old - # decorator approach which modified the gathered tensor in-place. - act_input = intermediate_cache1.view(-1, N)[gather_indx.dst_indx] - - sorted_token_ids_lora = None - expert_ids_lora = None - num_tokens_post_padded_lora = None - token_lora_mapping = None - lora_context = self._lora_context - if lora_context is not None: - ( - sorted_token_ids_lora, - expert_ids_lora, - num_tokens_post_padded_lora, - token_lora_mapping, - ) = self.apply_w13_lora( - lora_context, - y=act_input, - x=hidden_bf16_for_lora, - topk_ids=global_topk_ids, - topk_weights=topk_weights, - expert_map=expert_map, - w1=w1, - w2=w2, - num_tokens=M, - top_k_num=topk, - ) - - self.activation( - activation, - intermediate_cache2, - act_input, - ) - - # matmul_ogs grouped reduction fuses sum across multiple experts: - # y[dst_indx // n_expts_act, :] += x - # Set n_expts_act to 1 to unfuse the sum so we can do it manually via moe_sum. - routing_data.n_expts_act = 1 - - gemm2_input_bf16 = intermediate_cache2[gather_indx.src_indx] - if use_aiter_unfused_a8w4: - gemm2_lhs_scale = _mxfp4_w4a8_resolve_lhs_scale( - quant_config, - gemm_num=2, - activations=gemm2_input_bf16, - ) - gemm2_fp8 = _downcast_static_fp8_fn( - gemm2_input_bf16, - gemm2_lhs_scale, - ) - gemm2_out = _moe_gemm_a8w4_fn( - gemm2_fp8, - w2.storage.data, - None, - quant_config.w2_precision.weight_scale.storage.data, - gemm2_lhs_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, - ) - intermediate_cache3.copy_(gemm2_out.reshape(intermediate_cache3.shape)) - else: - gemm2_input, gemm2_lhs_scale_fp8 = ( - _maybe_downcast_bf16_hidden_to_fp8_mxfp4_unfused( - gemm2_input_bf16, - quant_config, - gemm_num=2, - ) - ) - w2_precision_for_gemm = _precision_config_with_mxfp4_unfused_fp8_lhs_scale( - quant_config.w2_precision, - gemm2_lhs_scale_fp8, - ) - - matmul_ogs( - gemm2_input, - w2, - quant_config.w2_bias, - routing_data, - scatter_indx=scatter_indx, - precision_config=w2_precision_for_gemm, - gammas=None if apply_router_weight_on_input else gammas, - y=intermediate_cache3, - ) - - # w2 LoRA: after matmul_ogs with scatter_indx, intermediate_cache3 is - # in token-topk order, matching the (M, topk, K) layout add_lora_w2 expects. - if lora_context is not None: - self.apply_w2_lora( - lora_context, - y=intermediate_cache3.view(-1, topk, K), - x=intermediate_cache2, - topk_weights=topk_weights, - sorted_token_ids_lora=sorted_token_ids_lora, - expert_ids_lora=expert_ids_lora, - num_tokens_post_padded_lora=num_tokens_post_padded_lora, - token_lora_mapping=token_lora_mapping, - num_tokens=M, - w1=w1, - w2=w2, - top_k_num=topk, - ) - - self.moe_sum(intermediate_cache3.view(-1, topk, K), output) - - -class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): - """Monolithic Triton MXFP4 expert. Wraps triton_kernel_moe_forward().""" - - 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: - return _triton_kernel_moe_supports_current_device() and has_triton_kernels() - - @staticmethod - def _supports_no_act_and_mul() -> bool: - return False - - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - SUPPORTED_W_A = [ - (kMxfp4Static, None), - ] - return (weight_key, activation_key) in SUPPORTED_W_A - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - 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 - - @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: - return triton_kernel_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, - ) From 316c685d50939f6a04d1d0b6f0c3e044f15928e8 Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Fri, 10 Jul 2026 13:44:17 -0400 Subject: [PATCH 67/68] Use triton weight prep for 1250 Signed-off-by: Jaden Mathias --- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 151 +++++++++++++++++- .../layers/fused_moe/oracle/mxfp4.py | 7 +- .../layers/quantization/mxfp4.py | 18 ++- 3 files changed, 166 insertions(+), 10 deletions(-) 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 index dffcb94a1192..35cfa8be756b 100644 --- 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 @@ -329,9 +329,95 @@ def apply( def _aiter_raw(t): + if t is None or isinstance(t, torch.Tensor): + return t return t.storage.data if hasattr(t, "storage") else t +def _aiter_w4a16_silu_via_a8w4( + hidden_states: torch.Tensor, + w1_data, + w2_data, + w1_wscale, + w2_wscale, + w1_bias, + w2_bias, + routing_data, + gather_idx, + scatter_idx, + gammas, + apply_router_weight_on_input: bool, + swiglu_limit: float, + unpadded_N_w1, + unpadded_K_w1, + unpadded_N_w2, + unpadded_K_w2, +) -> torch.Tensor: + """ + MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation. + """ + from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant import dynamic_mxfp8_quant + + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + should_use_cdna4_mx_scale_swizzle, + ) + + swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + quant_dtype = torch.float8_e4m3fn + + g1_gammas = gammas if apply_router_weight_on_input else None + g2_gammas = None if apply_router_weight_on_input else gammas + + hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype) + raw_gate_up = moe_gemm_a8w4( + hidden_q, + w1_data, + a1_scale, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gather_idx, + gammas=g1_gammas, + swizzle_mx_scale=swz, + out_dtype=torch.bfloat16, + apply_swiglu=False, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + if unpadded_N_w1 is not None: + raw_gate_up = raw_gate_up[:, :unpadded_N_w1] + + interm_fp8, a2_scale = fused_clamp_act_mul( + raw_gate_up, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=quant_dtype, + scale_dtype_fmt="ue8m0", + quant_block_size=32, + ) + + out = moe_gemm_a8w4( + interm_fp8, + w2_data, + a2_scale, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=swz, + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + return out + + def aiter_triton_kernel_w4a16_moe_forward( hidden_states: torch.Tensor, w1, @@ -348,13 +434,18 @@ def aiter_triton_kernel_w4a16_moe_forward( unpadded_K_w1=None, unpadded_N_w2=None, unpadded_K_w2=None, + num_expert_group: int | None = None, + topk_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + score_mode: str | None = None, ): assert quant_config is not None and rocm_aiter_ops.is_enabled() from vllm.platforms.rocm import on_gfx1250 try: from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 - from aiter.ops.triton.moe.moe_routing.routing import routing as _routing_mod + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod except ImportError: from aiter.ops.triton.moe_routing import routing as _routing_mod @@ -364,9 +455,25 @@ def aiter_triton_kernel_w4a16_moe_forward( _routing_mod.is_tdm_avail = lambda: False aiter_routing = _routing_mod.routing - routing_data, gather_idx, scatter_idx = aiter_routing( - gating_output, topk, sm_first=not renormalize - ) + if score_mode is not None: + use_grouped_topk = num_expert_group is not None and num_expert_group > 1 + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, + topk, + score_mode=score_mode, + bias=e_score_correction_bias, + renorm=renormalize, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) + else: + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) if on_gfx1250(): gather_src = gather_idx.to(torch.long) // topk @@ -394,6 +501,27 @@ def aiter_triton_kernel_w4a16_moe_forward( else 7.0 ) + if activation == MoEActivation.SILU: + return _aiter_w4a16_silu_via_a8w4( + hidden_states, + w1_data, + w2_data, + w1_wscale, + w2_wscale, + quant_config.w1_bias, + quant_config.w2_bias, + routing_data, + gather_idx, + scatter_idx, + gammas, + apply_router_weight_on_input, + swiglu_limit, + unpadded_N_w1, + unpadded_K_w1, + unpadded_N_w2, + unpadded_K_w2, + ) + intermediate = moe_gemm_a16w4( hidden_states, w1_data, @@ -445,6 +573,7 @@ def __init__( self.renormalize = moe_config.routing_method in ( RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, ) @staticmethod @@ -472,7 +601,7 @@ def _supports_quant_scheme( @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SWIGLUOAI + return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) @staticmethod def _supports_parallel_config( @@ -493,6 +622,7 @@ def _supports_routing_method( return routing_method in [ RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, ] @staticmethod @@ -524,6 +654,11 @@ def apply( ) -> torch.Tensor: assert self.moe_config.intermediate_size_per_partition_unpadded is not None assert self.moe_config.hidden_dim_unpadded is not None + score_mode = ( + "sqrtsoftplus" + if self.moe_config.routing_method == RoutingMethodType.DeepseekV4 + else None + ) return aiter_triton_kernel_w4a16_moe_forward( hidden_states=hidden_states, w1=w1, @@ -531,6 +666,7 @@ def apply( gating_output=router_logits, topk=self.topk, renormalize=self.renormalize, + activation=activation, global_num_experts=global_num_experts, expert_map=expert_map, quant_config=self.quant_config, @@ -539,4 +675,9 @@ def apply( 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, + num_expert_group=num_expert_group, + topk_group=topk_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + score_mode=score_mode, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index ba78573fa3de..ef930a77758c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1278,6 +1278,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends. """ + from vllm.platforms.rocm import on_gfx1250 if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( @@ -1450,7 +1451,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not on_gfx1250(): # Initially introduced for DeepSeekV4 if w13_bias is not None: @@ -1507,7 +1508,9 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend in TRITON_BACKENDS: + elif mxfp4_backend in TRITON_BACKENDS or ( + mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + ): from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig if mxfp4_backend == Mxfp4MoeBackend.TRITON: diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 5ef5fd40d5eb..75cae02f15c9 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -690,7 +690,12 @@ def _setup_kernel( # 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: + from vllm.platforms.rocm import on_gfx1250 + uses_triton_weight_format = ( + self.mxfp4_backend in TRITON_BACKENDS or + (self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()) + ) + if not uses_triton_weight_format: replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w13_weight_scale", w13_scale) @@ -702,7 +707,10 @@ def _setup_kernel( self.w2_precision_config = w2_scale # AITER backend requires weights to be marked as shuffled. - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + if ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 + and not uses_triton_weight_format + ): layer.w13_weight.is_shuffled = True layer.w2_weight.is_shuffled = True @@ -745,7 +753,11 @@ def get_fused_moe_quant_config( w2_bias = getattr(layer, "w2_bias", None) swiglu_limit = getattr(layer, "swiglu_limit", None) - if self.mxfp4_backend in TRITON_BACKENDS: + from vllm.platforms.rocm import on_gfx1250 + + if self.mxfp4_backend in TRITON_BACKENDS or ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + ): # TRITON backends free w13/w2_weight_scale after swizzling; the # swizzled scales live inside the precision configs instead. assert self.w13_precision_config is not None From 9a428d09ee91febdabb06391c98bc0a4ca04442f Mon Sep 17 00:00:00 2001 From: Jaden Mathias Date: Mon, 13 Jul 2026 15:03:37 -0400 Subject: [PATCH 68/68] Quick fix 1 --- .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py | 27 +++++-------------- 1 file changed, 6 insertions(+), 21 deletions(-) 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 index 35cfa8be756b..cccddcc9e79e 100644 --- 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 @@ -501,26 +501,11 @@ def aiter_triton_kernel_w4a16_moe_forward( else 7.0 ) - if activation == MoEActivation.SILU: - return _aiter_w4a16_silu_via_a8w4( - hidden_states, - w1_data, - w2_data, - w1_wscale, - w2_wscale, - quant_config.w1_bias, - quant_config.w2_bias, - routing_data, - gather_idx, - scatter_idx, - gammas, - apply_router_weight_on_input, - swiglu_limit, - unpadded_N_w1, - unpadded_K_w1, - unpadded_N_w2, - unpadded_K_w2, - ) + # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu. + # _aiter_w4a16_silu_via_a8w4 routes through moe_gemm_a8w4 which on gfx1250 + # transposes the weight tensor for the gluon kernel layout, but our weights + # are prepared for moe_gemm_a16w4 — the double-transpose produces garbage. + swiglu_add_residual = activation != MoEActivation.SILU intermediate = moe_gemm_a16w4( hidden_states, @@ -537,7 +522,7 @@ def aiter_triton_kernel_w4a16_moe_forward( apply_swiglu=True, alpha=swiglu_alpha, limit=swiglu_limit, - swiglu_add_residual=True, + swiglu_add_residual=swiglu_add_residual, unpadded_N=unpadded_N_w1, unpadded_K=unpadded_K_w1, )