diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 483a9aceac..288e7621ca 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -205,12 +205,12 @@ steps: device: h200_35gb source_file_dependencies: - vllm/models/qwen4_exp/ - - vllm/transformers_utils/configs/qwen4_exp.py - vllm/config/compilation.py - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ commands: - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_prepare.py models/qwen4_exp/test_qsa_reference.py + - pytest -v -s models/qwen4_exp/test_ple.py -k test_fused_conv_correctness mirror: amd: label: ":amd: (MI355 DPX) Qwen4 Exp" @@ -231,7 +231,6 @@ steps: device: cpu-small source_file_dependencies: - vllm/models/qwen4_exp/ - - vllm/transformers_utils/configs/qwen4_exp.py - vllm/config/compilation.py - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 9ab75cfc6e..8d2bc5dbaf 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -2,6 +2,23 @@ group: Models - Distributed depends_on: - image-build steps: +- label: ":nvidia: (L4) Qwen4 Exp SP" + key: qwen4-exp-sp-tests-4-gpus + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + device: l4 + num_devices: 4 + source_file_dependencies: + - vllm/config/parallel.py + - vllm/models/common/ops/sequence_parallel.py + - vllm/models/qwen4_exp/ + - vllm/model_executor/layers/vocab_parallel_embedding.py + - vllm/model_executor/models/qwen3_next.py + - tests/models/qwen4_exp/test_ple.py + - tests/models/qwen4_exp/test_tpsp.py + commands: + - pytest -v -s models/qwen4_exp/test_tpsp.py + - label: ":nvidia: (L4) Distributed Models Shard %N" key: distributed-model-tests-2-gpus timeout_in_minutes: 35 diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py index b0038c084e..b50b5e5688 100644 --- a/tests/kernels/core/test_vocab_parallel_embedding.py +++ b/tests/kernels/core/test_vocab_parallel_embedding.py @@ -9,6 +9,8 @@ each rank, and summing the per-rank outputs (what the all-reduce does) reconstructs a full-table lookup. """ +from unittest.mock import Mock + import pytest import torch @@ -124,3 +126,43 @@ def test_rank_sum_is_full_lookup(tp_size): total += _fused(ids, _shard(full_table, si, rows, dtype), si) torch.testing.assert_close(total, full_table[ids.long()], atol=0.0, rtol=0.0) + + +@requires_cuda +@torch.inference_mode() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("use_fused", [False, True]) +def test_embedding_optionally_reduces_output(dtype, use_fused, default_vllm_config): + """SP's deferred reduction preserves masked contributions and their dtype.""" + with torch.device("cuda"): + values = torch.arange(128 * 16).reshape(128, 16) + table = ((values % 29 - 14) / 16).to(dtype) + ids = torch.tensor([0, 63, 64, 127, 12]) + expected = table[ids].float() + for rank in range(2): + local = expected.clone() + local[ids // 64 != rank] = 0 + remote = (expected - local).to(dtype) + group = Mock(world_size=2, rank_in_group=rank) + # Model the other vocabulary owner's contribution, including FP8 bytes. + group.all_reduce.side_effect = lambda output, remote=remote: ( + output + remote.view(output.dtype) + ) + for reduce_results in (True, False): + embedding = VocabParallelEmbedding( + 128, + 16, + params_dtype=dtype, + parallel_group=group, + reduce_results=reduce_results, + ) + embedding.use_fused_embedding = use_fused + embedding.weight_loader(embedding.get_parameter("weight"), table) + actual = embedding(ids) + assert actual.dtype == dtype + torch.testing.assert_close( + actual.float(), + expected if reduce_results else local, + atol=0, + rtol=0, + ) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 6fde77eda2..818f515438 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -9,6 +9,7 @@ import pytest import torch from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig +from vllm.config import CacheConfig, ParallelConfig from vllm.config.speculative import SpeculativeConfig from vllm.model_executor.models.config import ( Qwen3_5ForConditionalGenerationConfig, @@ -48,23 +49,54 @@ def _text_config(**kwargs) -> Qwen4ExpTextConfig: return Qwen4ExpTextConfig(**values) -def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: +def test_moe_sp_rejects_pipeline_parallel() -> None: + """Automatic MoE SP must validate HC topology even without --enable-hc-sp.""" + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=_text_config(num_experts=4, ple_layer_ids=[]), + ), + cache_config=CacheConfig(), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + pipeline_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=True, + all2all_backend="allgather_reducescatter", + ), + ) + with ( + patch("vllm.platforms.current_platform.is_cuda", return_value=True), + pytest.raises(ValueError, match="MoE SP requires PP=1"), + ): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(config) + + +@pytest.mark.parametrize("use_sp", [False, True]) +def test_qwen4_exp_mtp_returns_sample_and_multi_streams(use_sp: bool) -> None: + """MTP returns complete, contiguous outputs after gathering padded SP shards.""" from vllm.models.qwen4_exp.nvidia.mtp import ( Qwen4ExpMultiTokenPredictor, ) model = object.__new__(Qwen4ExpMultiTokenPredictor) + model.use_hc_sequence_parallel = use_sp + model.use_sequence_parallel = False torch.nn.Module.__init__(model) model.hc_count = 2 model.hidden_size = 4 model.num_mtp_layers = 1 - model.layers = [ - lambda **kwargs: ( - kwargs["hidden_states"], - kwargs["hidden_states"], - torch.zeros(kwargs["hidden_states"].shape[0], 2), - ), - ] + model.pre_fc_norm_embedding = torch.nn.Identity() + model.pre_fc_norm_hidden = torch.nn.Identity() + model.fc_embedding = torch.nn.Identity() + model.fc_hidden = torch.nn.Identity() + + def decoder(**kwargs): + """Check that the embedding residual shares the hidden state's token slice.""" + hidden = kwargs["hidden_states"] + assert kwargs["prev_block_output"].shape == (hidden.shape[0], model.hidden_size) + return hidden, hidden, torch.zeros(hidden.shape[0], model.hc_count) + + model.layers = [decoder] model.hyper_connection_mixer = SimpleNamespace( combine_and_mix=lambda hidden_states, block_output, injection: ( hidden_states, @@ -72,24 +104,32 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: None, ), ) - multi_hidden = torch.arange(16, dtype=torch.float32).reshape(2, 8) - pp_group = SimpleNamespace(is_first_rank=False, is_last_rank=True) + multi_hidden = torch.arange(24, dtype=torch.float32).reshape(3, 8) + expected_sample = multi_hidden.unflatten(-1, (2, 4)).mean(-2) + pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) - with patch( - "vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", - return_value=pp_group, + def gather(local): + """Supply rank 1's last token and padding for a simulated two-rank gather.""" + assert local.shape[0] == 2 + remote = torch.cat([expected_sample[2:], multi_hidden[2:]], dim=-1) + return torch.cat([local, remote, torch.full_like(remote, float("nan"))]) + + with ( + patch("vllm.models.qwen4_exp.nvidia.mtp.get_pp_group", return_value=pp_group), + patch("vllm.models.qwen4_exp.nvidia.mtp.sp_shard", side_effect=lambda x: x[:2]), + patch("vllm.models.qwen4_exp.nvidia.mtp.sp_all_gather", side_effect=gather), ): sample_hidden, returned_multi_hidden = model.forward( input_ids=None, - positions=torch.arange(2), - intermediate_tensors={"hidden_states": multi_hidden}, + positions=torch.arange(3), + hidden_states=multi_hidden, + inputs_embeds=torch.zeros(3, 4), ) - torch.testing.assert_close( - sample_hidden, - multi_hidden.unflatten(-1, (2, 4)).mean(dim=-2), - ) - assert returned_multi_hidden is multi_hidden + torch.testing.assert_close(sample_hidden, expected_sample) + torch.testing.assert_close(returned_multi_hidden, multi_hidden) + assert sample_hidden.is_contiguous() + assert returned_multi_hidden.is_contiguous() @spawn_new_process_for_each_test @@ -150,9 +190,7 @@ def test_qwen4_exp_rejects_pipeline_parallel_only_with_ple(ple_layer_ids) -> Non hf_text_config=_text_config(ple_layer_ids=ple_layer_ids), multimodal_config=None, ), - parallel_config=SimpleNamespace( - pipeline_parallel_size=2, enable_dbo=False, ubatch_size=1 - ), + parallel_config=ParallelConfig(pipeline_parallel_size=2), speculative_config=None, ) with patch.object( diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index d7ee461d8f..a4a428eccc 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -611,46 +611,65 @@ def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None: ) -def test_pinned_embedding_forward_finalizes_prefetched_output( +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("reduce_results", [False, True]) +def test_pinned_embedding_finalize_prefetch_optionally_reduces_output( monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + reduce_results: bool, ) -> None: + """Deferring reduction must preserve full token rows and wait for the lookup.""" + total_tokens = 3 embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding) nn.Module.__init__(embedding) - embedding._prefetch_buffer = torch.empty(4, 2, 3, dtype=torch.float8_e4m3fn) + values = torch.arange(1, 13).reshape(4, 3) + embedding._prefetch_buffer = torch.stack( + (values, torch.zeros_like(values)), dim=1 + ).to(dtype) embedding._output_dim = 6 - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16) - expected = torch.arange(12).reshape(2, 6).to(torch.float8_e4m3fn) - - def finalize_prefetched( - prefetch_output: torch.Tensor, - output: torch.Tensor, - ) -> None: - assert prefetch_output is embedding._prefetch_buffer - output.copy_(expected) - + embedding.tp_size = 2 + embedding.reduce_results = reduce_results + embedding.etp_data_parallel_size = 1 + embedding._prefetch_stream = object() + waited_streams: list[object] = [] + reduced_dtypes = [] monkeypatch.setattr( - embedding, - "_finalize_prefetch", - finalize_prefetched, + torch.cuda, + "current_stream", + lambda: SimpleNamespace(wait_stream=waited_streams.append), ) - output = embedding(hidden_states) + def all_reduce(tensor: torch.Tensor) -> torch.Tensor: + assert waited_streams == [embedding._prefetch_stream] + reduced_dtypes.append(tensor.dtype) + # The remote rank owns the other head's vocabulary rows. + return tensor + tensor.flip(1) - assert output.dtype == embedding._prefetch_buffer.dtype - assert torch.equal(output, expected) + embedding.parallel_group = SimpleNamespace(all_reduce=all_reduce) + output = embedding.finalize_prefetch(total_tokens=total_tokens) + + expected = ( + values.unsqueeze(1).expand(-1, 2, -1).to(dtype) + if reduce_results + else embedding._prefetch_buffer + ) + expected = expected[:total_tokens].flatten(-2) + comm_dtype = torch.int8 if dtype == torch.float8_e4m3fn else dtype + assert reduced_dtypes == ([comm_dtype] if reduce_results else []) + assert waited_streams == [embedding._prefetch_stream] + assert output.dtype == dtype + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) -def test_pinned_embedding_forward_requires_prior_start_prefetch() -> None: - """Forward before any start_prefetch must fail loudly, not read garbage.""" +def test_pinned_embedding_finalize_prefetch_requires_prior_start_prefetch() -> None: + """Finalizing before start_prefetch must fail instead of reading garbage.""" embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding) nn.Module.__init__(embedding) embedding._prefetch_buffer = None embedding._prefetch_stream = None embedding._output_dim = 6 - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="prior start_prefetch"): - embedding(hidden_states) + embedding.finalize_prefetch(total_tokens=2) def test_pinned_embedding_finalize_requires_prior_start_prefetch() -> None: @@ -781,15 +800,10 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( +def test_pinned_embedding_start_prefetch_allocates_lazily( monkeypatch: pytest.MonkeyPatch, ) -> None: - """start_prefetch lazily allocates the stream/buffer, then forward consumes it. - - Exercises the full pinned-prefetch contract end to end: first call - allocates the side stream and buffer with the constructor geometry, - and forward finalizes the lookup into the returned output. - """ + """Allocate prefetch storage on first use and finalize the lookup results.""" _mock_etp_group(monkeypatch) monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr( @@ -825,8 +839,7 @@ def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( ) ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") - hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16, device="cuda:0") - embedding.start_prefetch(hidden_states, ngram_ids) + embedding.start_prefetch(ngram_ids) buffer = embedding._prefetch_buffer assert buffer is not None @@ -835,7 +848,7 @@ def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward( assert buffer.device.type == "cuda" assert embedding._prefetch_stream is not None - output = embedding(hidden_states) + output = embedding.finalize_prefetch(total_tokens=ngram_ids.shape[0]) expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) torch.testing.assert_close(output, expected, rtol=0, atol=0) @@ -1149,7 +1162,7 @@ def test_ngram_prefetch_ids_outlive_start_prefetch() -> None: prefetched: list[torch.Tensor] = [] module.ngram_embedding = SimpleNamespace( supports_prefetch=True, - start_prefetch=lambda _, ids: prefetched.append(weak_ref_tensor(ids)), + start_prefetch=lambda ids: prefetched.append(weak_ref_tensor(ids)), ) expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params) @@ -1157,7 +1170,7 @@ def test_ngram_prefetch_ids_outlive_start_prefetch() -> None: # like a later CUDA graph segment would. pool = torch.cuda.MemPool() with torch.cuda.use_mem_pool(pool): - module.start_prefetch(None, input_ids, query_start_loc, ngram_context) + module.start_prefetch(input_ids, query_start_loc, ngram_context) torch.full_like(expected, -1) assert torch.equal(prefetched[0], expected) @@ -1859,12 +1872,18 @@ def _make_conv_metadata( ], ) @pytest.mark.parametrize( - "layer_cls", [Qwen4ExpPLELayer, Qwen4ExpPLELayerAMD], ids=["nvidia", "amd"] + ("layer_cls", "fuse_residual"), + [ + pytest.param(Qwen4ExpPLELayer, True, id="nvidia-fused"), + pytest.param(Qwen4ExpPLELayer, False, id="nvidia-unfused"), + pytest.param(Qwen4ExpPLELayerAMD, True, id="amd-fused"), + ], ) def test_fused_conv_correctness( case: _ConvBatchCase, state_layout: str, layer_cls: type[nn.Module], + fuse_residual: bool, ) -> None: device = torch.device("cuda") metadata, num_real_tokens = _make_conv_metadata(case, device) @@ -1906,12 +1925,15 @@ def test_fused_conv_correctness( ) null_state = conv_state[NULL_BLOCK_ID].clone() residual_kernel = residual.clone() - residual_reference = residual.clone() + # The unfused kernel must overwrite the buffer without reading its contents. + residual_reference = ( + residual.clone() if fuse_residual else torch.zeros_like(residual) + ) module._short_conv_dilated_dispatch( inputs=inputs, residual=residual_kernel, - outer_residual=outer_residual, + outer_residual=outer_residual if fuse_residual else None, metadata=metadata, conv_state=conv_state, conv_weights=weights, @@ -1925,7 +1947,8 @@ def test_fused_conv_correctness( conv_state_len=module.conv_state_len, dilation=module.short_conv_dilation, ) - residual_reference = outer_residual + residual_reference + if fuse_residual: + residual_reference = outer_residual + residual_reference assert torch.equal( residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens] @@ -1935,7 +1958,7 @@ def test_fused_conv_correctness( if case.graph_padding: assert torch.equal( residual_kernel[num_real_tokens:], - (outer_residual + residual)[num_real_tokens:], + residual_reference[num_real_tokens:], ) diff --git a/tests/models/qwen4_exp/test_tpsp.py b/tests/models/qwen4_exp/test_tpsp.py new file mode 100644 index 0000000000..5182ab09c3 --- /dev/null +++ b/tests/models/qwen4_exp/test_tpsp.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multi-GPU equivalence tests for Qwen4Exp HC and MoE sequence parallelism.""" + +from dataclasses import MISSING, fields +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.multiprocessing as mp +from torch import nn +from transformers import Qwen4ExpTextConfig + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("Qwen4Exp SP requires CUDA", allow_module_level=True) + +from tests.utils import multi_gpu_test +from vllm.config import ( + EngramConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + tensor_model_parallel_all_reduce, +) +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.linear import RowParallelLinear +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) +from vllm.models.qwen4_exp.nvidia import model as qwen4_model +from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer +from vllm.utils.network_utils import get_open_port +from vllm.v1.attention.backends.short_conv_attn import PleShortConvAttentionMetadata +from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager + +from .test_ple import _ConvBatchCase, _make_conv_metadata + + +class _TPAttentionProjection(nn.Module): + """Keep real TP communication without setting up an attention KV cache.""" + + def __init__(self, config, *, reduce_results: bool, **kwargs) -> None: + super().__init__() + self.proj = RowParallelLinear( + config.hidden_size, + config.hidden_size, + bias=False, + input_is_parallel=False, + reduce_results=reduce_results, + ) + + def forward(self, hidden_states: torch.Tensor): + return self.proj(hidden_states)[0] + + +def _make_config(rank: int, *, hc_sp: bool = False, moe_sp: bool = False) -> VllmConfig: + """Configure HC SP explicitly or MoE SP through TP=2, DP=2 and EP.""" + config = VllmConfig( + parallel_config=ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=2 if moe_sp else 1, + data_parallel_rank=rank // 2, + enable_expert_parallel=moe_sp, + enable_hc_sp=hc_sp, + all2all_backend="allgather_reducescatter", + ), + ) + config.model_config = SimpleNamespace( + architecture="Qwen4ExpForCausalLM", + is_moe=True, + sleep_mode_offload_cudagraph=False, + hf_text_config=Qwen4ExpTextConfig( + vocab_size=128, + eos_token_id=0, + hidden_size=128, + hc_count=2, + hc_lowrank=16, + num_hidden_layers=1, + layer_types=["linear_attention"], + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=256, + shared_expert_intermediate_size=256, + ple_layer_ids=[] if moe_sp else [1], + ple_embed_dim=128, + heads_per_ngram=2, + ngram_vocab_size_base=32, + make_ngram_vocab_size_divisible_by=64, + ), + dtype=torch.bfloat16, + ) + config.engram_config = EngramConfig(cpu_offload=False) + return config + + +def _make_decoder(config: VllmConfig, rank: int) -> Qwen4ExpDecoderLayer: + """Use the production constructor and identical initial weights for both modes.""" + with ( + set_current_vllm_config(config), + patch.object(qwen4_model, "QwenGatedDeltaNetAttention", _TPAttentionProjection), + ): + layer = Qwen4ExpDecoderLayer( + config, config.model_config.hf_text_config.layer_types[0], "model.layers.0" + ) + torch.manual_seed(12) + for param in layer.parameters(): + param.normal_(std=0.08) + # Distinct rank contributions expose missing or duplicate reductions. + layer.linear_attn.proj.weight.add_(rank * 0.01) + experts = layer.mlp.experts.routed_experts + if config.parallel_config.enable_expert_parallel: + experts.w13_weight.add_(rank * 0.01) + experts.quant_method.process_weights_after_loading(experts) + return layer + + +def _check_decoder_ple(rank: int) -> None: + """Compare HC SP with full tokens across a PLE prefill and cached decode.""" + configs = [_make_config(rank, hc_sp=enabled) for enabled in (False, True)] + layers = [_make_decoder(config, rank) for config in configs] + states = [] + for layer in layers: + ple = layer.ple + assert ple is not None + state = torch.zeros(2, ple.hc_hidden_size, ple.conv_state_len) + if not is_conv_state_dim_first(): + state = state.transpose(-1, -2).contiguous() + ple.kv_cache = (state,) + states.append(state) + + # Three prefill tokens straddle the TP=2 boundary and require SP padding. + for case in ( + _ConvBatchCase(prefill_query_lens=(3,)), + _ConvBatchCase(num_decodes=1), + ): + layout, num_tokens = _make_conv_metadata(case, torch.device("cuda", rank)) + metadata_args = { + f.name: None + for f in fields(PleShortConvAttentionMetadata) + if f.default is MISSING and f.default_factory is MISSING + } + metadata = PleShortConvAttentionMetadata(**(metadata_args | vars(layout))) + ids = torch.randint(0, 128, (num_tokens,)) + positions = torch.arange(num_tokens).expand(3, -1) + qsl = torch.tensor([0, num_tokens], dtype=torch.int32) + context = torch.zeros(1, layers[0].config.ngram_size - 1, dtype=torch.int64) + hidden = torch.randn(num_tokens, 256) + block = torch.randn(num_tokens, 128) + injection = torch.randn(num_tokens, 2) + outputs = [] + for config, layer in zip(configs, layers): + assert layer.ple is not None + inputs = (hidden, block, injection) + if config.parallel_config.enable_hc_sp: + inputs = tuple(sp_shard(x) for x in inputs) + with ( + set_current_vllm_config(config), + set_forward_context({layer.ple.prefix: metadata}, config), + ): + layer.ple.start_prefetch(ids, qsl, context) + output = layer( + *inputs, + positions, + input_ids=ids, + query_start_loc=qsl, + ngram_context=context, + ) + if config.parallel_config.enable_hc_sp: + output = tuple(sp_all_gather(x)[:num_tokens] for x in output) + # A later MoE call may reuse workspace-backed output storage. + outputs.append(tuple(x.clone() for x in output)) + torch.testing.assert_close(outputs[1], outputs[0], atol=0.002, rtol=0.02) + torch.testing.assert_close(states[1], states[0], atol=0.002, rtol=0.02) + + +def _check_moe_sp(rank: int, config: VllmConfig) -> None: + """Compare native SP with full-token GR and the MoE wrapper's public interface.""" + layer = _make_decoder(config, rank) + dp_rank = config.parallel_config.data_parallel_rank + counts = torch.tensor([3, 5], device="cpu") + num_tokens = int(counts[dp_rank]) + torch.manual_seed(30 + dp_rank) + hidden = torch.randn(num_tokens, 256) + positions = torch.arange(num_tokens) + # Include runner padding in addition to the padding introduced by SP. + padding = sp_padding_mask(positions == num_tokens - 1, hidden) + with ( + set_current_vllm_config(config), + set_forward_context( + None, + config, + num_tokens=num_tokens, + num_tokens_across_dp=counts, + is_padding=padding, + ), + ): + # The reference keeps full GR rows. MoE handles its own chunk and gather. + full_hidden, attn_input, injection = layer.attn_hyper_connection.mix(hidden) + attn_output = tensor_model_parallel_all_reduce(layer.linear_attn(attn_input)) + full_hidden, moe_input, injection = layer.mlp_hyper_connection.combine_and_mix( + full_hidden, attn_output, injection + ) + expected = tuple( + x.clone() for x in (full_hidden, layer.mlp(moe_input), injection) + ) + output = layer( + sp_shard(hidden), + None, + None, + positions, + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + actual = tuple(sp_all_gather(x)[:num_tokens] for x in output) + torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.02) + + +def _run_sp(rank: int, port: int, moe_sp: bool) -> None: + """Initialize rank-local runtime state and run one equivalence check.""" + torch.accelerator.set_device_index(rank) + torch.set_num_threads(1) + config = _make_config(rank, moe_sp=moe_sp) + init_distributed_environment( + world_size=4 if moe_sp else 2, + rank=rank, + local_rank=rank, + distributed_init_method=f"tcp://127.0.0.1:{port}", + ) + try: + # Spawned ranks bypass the GPU worker's workspace initialization. + init_workspace_manager(torch.device(f"cuda:{rank}")) + with set_current_vllm_config(config): + initialize_model_parallel(tensor_model_parallel_size=2) + with torch.device(f"cuda:{rank}"), torch.inference_mode(): + torch.set_default_dtype(torch.bfloat16) + if moe_sp: + _check_moe_sp(rank, config) + else: + _check_decoder_ple(rank) + finally: + reset_workspace_manager() + cleanup_dist_env_and_memory() + + +@multi_gpu_test(num_gpus=2) +def test_hc_sp_decoder_and_ple() -> None: + """HC SP preserves decoder output and PLE cache across a token-shard boundary.""" + mp.spawn(_run_sp, args=(get_open_port(), False), nprocs=2) + + +@multi_gpu_test(num_gpus=4) +def test_moe_sp_equivalence(monkeypatch) -> None: + """Native MoE SP preserves outputs with unequal DP batches and padding.""" + monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1") + mp.spawn(_run_sp, args=(get_open_port(), True), nprocs=4) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index d269aeec13..cffc728043 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -127,6 +127,22 @@ def test_registry_is_pp(model_arch, is_pp): assert model_info.supports_pp is is_pp +@pytest.mark.parametrize( + "model_arch,supported", + [ + ("Qwen4ExpForCausalLM", not current_platform.is_rocm()), + ("Qwen4ExpForConditionalGeneration", not current_platform.is_rocm()), + ("Qwen4ExpMTP", not current_platform.is_rocm()), + ("LlamaForCausalLM", False), + ], +) +def test_registry_supports_hc_sp(model_arch, supported): + """HC SP is opt-in for each model implementation, including MTP.""" + model_info = ModelRegistry._try_inspect_model_cls(model_arch) + assert model_info is not None + assert model_info.supports_hc_sp is supported + + @create_new_process_for_each_test() @pytest.mark.parametrize( "model_arch", diff --git a/tests/test_config.py b/tests/test_config.py index 734253edc9..a7e36ada62 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4109,6 +4109,47 @@ def test_revision_resolved_for_model(mock_resolve): mock_resolve.assert_any_call(model, None, config.hf_token) +@pytest.mark.parametrize( + "enabled,supported,is_cuda,tp_size,pp_size,error", + [ + (True, True, True, 2, 1, None), + (True, False, True, 2, 1, "not supported for this model"), + (True, True, False, 2, 1, "requires CUDA"), + (True, True, True, 1, 1, None), + (True, False, False, 1, 2, None), + (True, True, True, 2, 2, "requires PP=1"), + (False, False, False, 2, 1, None), + ], +) +def test_hc_sp_model_support(enabled, supported, is_cuda, tp_size, pp_size, error): + """Validate active HC SP configurations and ignore the flag with TP=1.""" + from vllm.model_executor.models.registry import ModelRegistry + + model = SimpleNamespace( + architectures=["HCModel"], + registry=ModelRegistry, + model_arch_config=SimpleNamespace(total_num_attention_heads=8), + multimodal_config=None, + ) + parallel = ParallelConfig( + enable_hc_sp=enabled, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + # Avoid backend auto-detection in this config-only test. + distributed_executor_backend="mp", + ) + with ( + patch("vllm.config.model.current_platform.is_cuda", return_value=is_cuda), + patch.object(ModelRegistry, "is_hc_sp_supported_model", return_value=supported), + patch.object(ModelRegistry, "is_pp_supported_model", return_value=True), + ): + if error: + with pytest.raises(ValueError, match=error): + ModelConfig.verify_with_parallel_config(model, parallel) + else: + ModelConfig.verify_with_parallel_config(model, parallel) + + @pytest.mark.parametrize( ("layer_types", "expected_attention"), [ diff --git a/vllm/config/model.py b/vllm/config/model.py index d4bafad5d6..545c306e1e 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1537,6 +1537,14 @@ class ModelConfig: self, parallel_config: ParallelConfig, ) -> None: + if parallel_config.enable_hc_sp and parallel_config.tensor_parallel_size > 1: + if not current_platform.is_cuda(): + raise ValueError("--enable-hc-sp currently requires CUDA") + if not self.registry.is_hc_sp_supported_model(self.architectures, self): + raise ValueError("--enable-hc-sp is not supported for this model") + if parallel_config.pipeline_parallel_size != 1: + raise ValueError("--enable-hc-sp requires PP=1") + total_num_attention_heads = self.model_arch_config.total_num_attention_heads tensor_parallel_size = parallel_config.tensor_parallel_size if total_num_attention_heads % tensor_parallel_size != 0: diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6960661e0c..e103c9ffd7 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -181,6 +181,11 @@ class ParallelConfig: --data-parallel-start-rank.""" is_moe_model: bool | None = None """Whether the deployed model is MoE (if known).""" + enable_hc_sp: bool = False + """Enable sequence parallelism for HyperConnection (HC) computation. + Currently supported by Qwen4Exp on CUDA with TP>1 and PP=1 (ignored when TP=1). + Does not change MoE parallelism. MoE SP may enable HC SP independently of this flag. + """ enable_expert_parallel: bool = False """Use expert parallelism instead of tensor parallelism for MoE layers.""" enable_batch_sharded_sampling: bool | None = None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 31777e266f..a892900816 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -538,6 +538,7 @@ class EngineArgs: data_parallel_external_lb: bool = False data_parallel_multi_port_external_lb: bool = False data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend + enable_hc_sp: bool = ParallelConfig.enable_hc_sp enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel enable_batch_sharded_sampling: bool | None = ( ParallelConfig.enable_batch_sharded_sampling @@ -1258,6 +1259,7 @@ class EngineArgs: "server per local data parallel rank and exposes aggregated health on " "a supervisor port.", ) + parallel_group.add_argument("--enable-hc-sp", **parallel_kwargs["enable_hc_sp"]) parallel_group.add_argument( "--enable-expert-parallel", "-ep", @@ -2532,6 +2534,7 @@ class EngineArgs: data_parallel_backend=self.data_parallel_backend, data_parallel_hybrid_lb=self.data_parallel_hybrid_lb, is_moe_model=model_config.is_moe, + enable_hc_sp=self.enable_hc_sp, enable_expert_parallel=self.enable_expert_parallel, enable_batch_sharded_sampling=self.enable_batch_sharded_sampling, enable_ep_weight_filter=self.enable_ep_weight_filter, diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index be20c41243..92e65c839d 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.base_config import ( method_has_implemented_embedding, resolve_quant_method, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.layers.utils import dispatch_unquantized_gemm from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs @@ -240,6 +241,7 @@ class VocabParallelEmbedding(PluggableLayer): quant_config: quant config for the layer prefix: full name of the layer in the state dict disable_tp: If true, tensor parallelism will be disabled for this layer. + reduce_results: Whether to all-reduce contributions from vocabulary shards. quant_method: Preselected quantization method for model-specific layers. parallel_group: Process group used to shard and reduce the embedding. @@ -258,6 +260,7 @@ class VocabParallelEmbedding(PluggableLayer): prefix: str = "", *, disable_tp: bool = False, + reduce_results: bool = True, quant_method: QuantizeMethodBase | None = None, parallel_group: GroupCoordinator | None = None, ): @@ -266,6 +269,7 @@ class VocabParallelEmbedding(PluggableLayer): # Keep the input dimensions. self.disable_tp = disable_tp self.parallel_group = parallel_group + self.reduce_results = reduce_results if disable_tp: tp_rank, self.tp_size = 0, 1 elif parallel_group is not None: @@ -516,6 +520,7 @@ class VocabParallelEmbedding(PluggableLayer): param[loaded_weight.shape[0] :].data.fill_(0) def forward(self, input_): + """Look up IDs, optionally reducing contributions from vocabulary shards.""" if self.tp_size == 1: return self.quant_method.embedding(self, input_.long()) @@ -543,24 +548,20 @@ class VocabParallelEmbedding(PluggableLayer): self.shard_indices.added_vocab_end_index, ) output_parallel = self.quant_method.embedding(self, masked_input.long()) - if output_parallel.dtype in ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ): - # Each vocab token has one owner, so FP8 bytes can use int8 SUM. - comm_output = output_parallel.view(torch.int8) - comm_output.masked_fill_(input_mask.unsqueeze(-1), 0) - output = ( - self.parallel_group.all_reduce(comm_output) - if self.parallel_group is not None - else tensor_model_parallel_all_reduce(comm_output) - ) - return output.view(output_parallel.dtype) + + output_dtype = output_parallel.dtype + if is_fp8(output_dtype): + # Each vocab token has one owner, so FP8 bytes can use int8 SUM. + output_parallel = output_parallel.view(torch.int8) + if not self.use_fused_embedding: output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) - # Reduce across all the model parallel GPUs. - if self.parallel_group is not None: - return self.parallel_group.all_reduce(output_parallel) - return tensor_model_parallel_all_reduce(output_parallel) + if self.reduce_results: + output_parallel = ( + self.parallel_group.all_reduce(output_parallel) + if self.parallel_group is not None + else tensor_model_parallel_all_reduce(output_parallel) + ) + return output_parallel.view(output_dtype) def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings}" diff --git a/vllm/model_executor/models/__init__.py b/vllm/model_executor/models/__init__.py index 478b00d0af..78fa39ee41 100644 --- a/vllm/model_executor/models/__init__.py +++ b/vllm/model_executor/models/__init__.py @@ -3,6 +3,7 @@ from .interfaces import ( HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsMultiModal, @@ -10,6 +11,7 @@ from .interfaces import ( SupportsPP, SupportsTranscription, has_inner_state, + supports_hc_sp, supports_lora, supports_mrope, supports_multimodal, @@ -33,6 +35,8 @@ __all__ = [ "is_text_generation_model", "HasInnerState", "has_inner_state", + "SupportsHCSP", + "supports_hc_sp", "SupportsLoRA", "supports_lora", "SupportsMultiModal", diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 7f6a3aef40..d64e51bfa8 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -912,11 +912,21 @@ class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConf @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: + from vllm.platforms import current_platform + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) text_config = vllm_config.model_config.hf_text_config if text_config.hc_count <= 1: raise ValueError("Qwen4Exp requires hc_count > 1") parallel_config = vllm_config.parallel_config + # MoE SP also enables HC SP without the explicit HC flag. + if ( + current_platform.is_cuda() + and parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size != 1 + ): + raise ValueError("Qwen4Exp MoE SP requires PP=1") + uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( getattr(text_config, "indexer_n_heads", None) is not None ) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 272b78da50..4dc67ff34d 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -858,6 +858,27 @@ def _supports_pp_inspect(model: type[object] | object) -> bool: return supports_kw(model_forward, "intermediate_tensors") +@runtime_checkable +class SupportsHCSP(Protocol): + """Support sequence-parallel HC without changing MoE parallelism.""" + + supports_hc_sp: ClassVar[Literal[True]] = True + + +@overload +def supports_hc_sp(model: type[object]) -> TypeIs[type[SupportsHCSP]]: ... + + +@overload +def supports_hc_sp(model: object) -> TypeIs[SupportsHCSP]: ... + + +def supports_hc_sp( + model: type[object] | object, +) -> TypeIs[type[SupportsHCSP]] | TypeIs[SupportsHCSP]: + return getattr(model, "supports_hc_sp", False) + + @runtime_checkable class HasInnerState(Protocol): """The interface required for all models that has inner state.""" diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 04d20c5741..ee5c05d0a4 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -129,7 +129,12 @@ def _should_replicate_misaligned_shared_expert( class Qwen3NextSparseMoeBlock(nn.Module): - def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ): super().__init__() config = vllm_config.model_config.hf_text_config @@ -231,6 +236,7 @@ class Qwen3NextSparseMoeBlock(nn.Module): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, + reduce_results=reduce_results, fuse_shared_experts=self.is_fused_shared_expert_enabled, shared_expert_prefix=f"{prefix}.shared_expert", shared_expert_gate=self.shared_expert_gate diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index fa29244413..427648045f 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -49,6 +49,7 @@ from .interfaces import ( is_attention_free, is_hybrid, requires_raw_input_tokens, + supports_hc_sp, supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, @@ -873,6 +874,7 @@ class _ModelInfo: requires_raw_input_tokens: bool supports_multimodal_encoder_tp_data: bool supports_pp: bool + supports_hc_sp: bool has_inner_state: bool is_attention_free: bool is_hybrid: bool @@ -903,6 +905,7 @@ class _ModelInfo: model ), supports_pp=supports_pp(model), + supports_hc_sp=supports_hc_sp(model), has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), @@ -1463,6 +1466,15 @@ class _ModelRegistry: model_info, _ = self.inspect_model_cls(architectures, model_config) return model_info.supports_pp + def is_hc_sp_supported_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + """Check support for HC SP without changing MoE parallelism.""" + model_info, _ = self.inspect_model_cls(architectures, model_config) + return model_info.supports_hc_sp + def model_has_inner_state( self, architectures: str | list[str], diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 6dad6b7b5a..85a382c5ec 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -70,6 +70,7 @@ class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): num_ngram_heads: int = 1, max_total_tokens: int = 0, data_parallel_rank: int = 0, + reduce_results: bool = True, ) -> None: del num_ngram_heads, max_total_tokens super().__init__( @@ -80,6 +81,7 @@ class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): prefix=prefix, quant_method=embedding_method, parallel_group=get_etp_group(), + reduce_results=reduce_results, ) self.embedding_method = embedding_method self.data_parallel_rank = data_parallel_rank @@ -153,14 +155,15 @@ class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): return embeddings.narrow(0, slot_offset, local_num_tokens) @abstractmethod - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Start an asynchronous lookup when supported.""" raise NotImplementedError + @abstractmethod + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Return prefetched embeddings for the full input token set.""" + raise NotImplementedError + class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): """Quantization interface shared by resident and pinned PLE tables.""" @@ -351,16 +354,16 @@ class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): """Allocate the complete PLE weight on the active device.""" return torch.empty(num_embeddings, embedding_dim, dtype=dtype) - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Resident embedding prefetch is a no-op.""" return None + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Reject prefetch finalization for resident embeddings.""" + raise RuntimeError("Device PLE embedding does not support prefetch") + def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: - """Gather ETP inputs, look up embeddings, and select local rows.""" + """Gather ETP IDs, optionally reduce lookup results, and select local rows.""" slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) embeddings = super().forward(gathered_ids) @@ -422,6 +425,7 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): num_ngram_heads: int = 1, max_total_tokens: int = 0, data_parallel_rank: int = 0, + reduce_results: bool = True, ) -> None: if not is_uva_available(): raise RuntimeError("Engram CPU offload requires UVA support") @@ -435,6 +439,7 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): num_ngram_heads=num_ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=reduce_results, ) self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) self._row_bytes = self.embedding_dim * self.weight.element_size() @@ -521,11 +526,7 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): return self.parallel_group.all_reduce(embeddings) @eager_break_during_capture - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: + def start_prefetch(self, ngram_ids: torch.Tensor) -> None: """Gather ETP IDs and launch their UVA lookup on the side stream.""" buffer = self._prefetch_buffer if buffer is None: @@ -574,14 +575,18 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): prefetch_output: torch.Tensor, output: torch.Tensor, ) -> None: - """Join the side stream, reduce ETP shards, and select local rows.""" + """Join the side stream, optionally reduce ETP shards, and select rows.""" prefetch_stream = self._prefetch_stream if prefetch_stream is None: raise RuntimeError("pinned PLE finalize requires a prior start_prefetch") torch.cuda.current_stream().wait_stream(prefetch_stream) slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0]) - active_output = prefetch_output[: slot_size * self.etp_data_parallel_size] - embeddings = self._reduce_etp_embeddings(active_output) + embeddings = prefetch_output[: slot_size * self.etp_data_parallel_size] + if self.reduce_results: + embeddings = self._reduce_etp_embeddings(embeddings) + else: + # The outer TP reduce-scatter requires ETP=TP. + assert self.etp_data_parallel_size == 1 embeddings = self._select_embeddings( embeddings, output.shape[0], @@ -589,11 +594,18 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): ) output.copy_(embeddings.flatten(-2)) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Finish the pinned lookup into graph-owned output storage.""" + def finalize_prefetch(self, total_tokens: int) -> torch.Tensor: + """Finalize the full-token pinned lookup with optional ETP reduction.""" buffer = self._prefetch_buffer if buffer is None: raise RuntimeError("pinned PLE lookup requires a prior start_prefetch") - output = buffer.new_empty((hidden_states.shape[0], self._output_dim)) + # total_tokens counts this DP rank's tokens before SP sharding. + output = buffer.new_empty((total_tokens, self._output_dim)) self._finalize_prefetch(buffer, output) return output + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + raise RuntimeError( + "Pinned-host PLE embedding requires " + "start_prefetch() and finalize_prefetch()." + ) diff --git a/vllm/models/qwen4_exp/common/ops/ple.py b/vllm/models/qwen4_exp/common/ops/ple.py index bc4c99957d..41ce082c9e 100644 --- a/vllm/models/qwen4_exp/common/ops/ple.py +++ b/vllm/models/qwen4_exp/common/ops/ple.py @@ -344,6 +344,7 @@ def _ple_conv_kernel( MODE: tl.constexpr, HAS_INIT: tl.constexpr, NULL_STATE_ID: tl.constexpr, + FUSE_RESIDUAL: tl.constexpr, launch_pdl: tl.constexpr, ): t = tl.program_id(0) @@ -438,21 +439,23 @@ def _ple_conv_kernel( conv = acc.to(residual_ptr.dtype.element_ty).to(tl.float32) y = conv * tl.sigmoid(conv) conv_output = tl.where(out_ok, y, 0.0).to(residual_ptr.dtype.element_ty) - residual = tl.load( - residual_ptr + output_t * C + c_offs, - mask=c_mask, - other=0.0, - ) - # Preserve the original eager operation boundaries: short convolution is - # first accumulated into the BF16/FP16 PLE output, then the outer residual - # is added to that rounded value. - ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) - outer_residual = tl.load( - outer_residual_ptr + output_t * C + c_offs, - mask=c_mask, - other=0.0, - ) - ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) + ple_output = conv_output + if FUSE_RESIDUAL: + residual = tl.load( + residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + # Preserve the original eager operation boundaries: short convolution is + # first accumulated into the BF16/FP16 PLE output, then the outer residual + # is added to that rounded value. + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) + outer_residual = tl.load( + outer_residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) if launch_pdl: tl.extra.cuda.gdc_launch_dependents() tl.store( @@ -575,7 +578,7 @@ def ple_conv( conv_state: torch.Tensor, conv_weights: torch.Tensor, state_indices: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, *, mode: Literal["decode", "spec", "prefill"], dilation: int, @@ -585,7 +588,9 @@ def ple_conv( spec_query_len: int = 1, token_indices: torch.Tensor | None = None, ) -> None: - """Add short convolution and the outer residual; update state.""" + """Update state and write activated convolution into residual. + A non-None outer_residual enables both fused residual additions. + """ BLOCK_C = 512 kernel_spec_query_len = spec_query_len if mode == "spec" else 1 T, C = inputs.shape @@ -657,6 +662,7 @@ def ple_conv( SPEC_QUERY_LEN=kernel_spec_query_len, MODE=mode, HAS_INIT=has_initial_states_arg, + FUSE_RESIDUAL=outer_residual is not None, NULL_STATE_ID=NULL_BLOCK_ID, launch_pdl=launch_pdl, num_warps=num_warps, diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 2de883cc47..2fb056bde4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -9,8 +9,10 @@ import torch from torch import nn from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig +from vllm import envs from vllm.config import VllmConfig from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -36,6 +38,7 @@ from vllm.model_executor.models.interfaces import ( IsHybrid, MixtureOfExperts, MultiModalEmbeddings, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -66,6 +69,12 @@ from vllm.model_executor.models.utils import ( maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -84,6 +93,14 @@ from .qsa import Qwen4ExpQSAAttention _QSA_LAYER_TYPES = ("qwen_sparse_attention", "indexed_attention") +def is_hc_sequence_parallel_enabled(vllm_config: VllmConfig) -> bool: + """Enable GR/PLE SP explicitly or automatically with sequence-parallel MoE.""" + parallel_config = vllm_config.parallel_config + return parallel_config.tensor_parallel_size > 1 and ( + parallel_config.enable_hc_sp or parallel_config.use_sequence_parallel_moe + ) + + def without_modelopt_fp4( quant_config: QuantizationConfig | None, ) -> QuantizationConfig | None: @@ -161,13 +178,19 @@ _EXTRA_WEIGHTS_MAPPER = WeightsMapper( class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): """Qwen3Next MoE with Qwen4Exp HC validation.""" - def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: - parallel_config = vllm_config.parallel_config - if parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) - super().__init__(vllm_config=vllm_config, prefix=prefix) + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__( + vllm_config=vllm_config, prefix=prefix, reduce_results=reduce_results + ) + if self.replicate_shared_expert and not self.is_sequence_parallel: + # Reduce routed outputs before adding replicated shared outputs. + # SP then slices the result to avoid summing shared outputs again via RS. + self.experts.moe_config.skip_final_all_reduce = False config = vllm_config.model_config.hf_text_config self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) @@ -185,11 +208,12 @@ class Qwen4ExpDecoderLayer(nn.Module): self.config = config self.layer_type = layer_type + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) self.layer_idx = extract_layer_index(prefix) - if vllm_config.parallel_config.use_sequence_parallel_moe: - raise NotImplementedError( - "Qwen4Exp HC does not support sequence-parallel MoE" - ) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.ple: Qwen4ExpPLELayer | None = None ple_layer_ids = config.ple_layer_ids if (self.layer_idx + 1) in ple_layer_ids: @@ -204,6 +228,7 @@ class Qwen4ExpDecoderLayer(nn.Module): layer_idx=self.layer_idx, ple_dense_layer_id=ple_dense_layer_id, prefix=f"{prefix}.ple", + use_sequence_parallel=self.use_hc_sequence_parallel, ) if layer_type == "linear_attention": @@ -212,6 +237,7 @@ class Qwen4ExpDecoderLayer(nn.Module): vllm_config=vllm_config, prefix=f"{prefix}.linear_attn", gqa_interleaved_layout=False, + reduce_results=not self.use_hc_sequence_parallel, ) elif layer_type in _QSA_LAYER_TYPES: use_qsa = getattr(config, "indexer_n_heads", None) is not None @@ -222,6 +248,7 @@ class Qwen4ExpDecoderLayer(nn.Module): cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: self.self_attn = Qwen4ExpQSAAttention( @@ -230,12 +257,15 @@ class Qwen4ExpDecoderLayer(nn.Module): layer_id=self.layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn", + reduce_results=not self.use_hc_sequence_parallel, ) else: raise ValueError(f"Invalid layer_type {layer_type}") self.mlp = Qwen4ExpSparseMoeBlock( - vllm_config=vllm_config, prefix=f"{prefix}.mlp" + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + reduce_results=not self.use_hc_sequence_parallel, ) hc_config = HyperConnectionConfig( @@ -295,6 +325,9 @@ class Qwen4ExpDecoderLayer(nn.Module): else: hidden_states, block_input, injection = attn_hc.mix(hidden_states) + if self.use_hc_sequence_parallel: + # Positions retain the full token axis for both text and MRoPE inputs. + block_input = sp_all_gather(block_input)[: positions.shape[-1]] if self.layer_type == "linear_attention": attn_out = self.linear_attn(hidden_states=block_input) elif self.layer_type in _QSA_LAYER_TYPES: @@ -306,10 +339,25 @@ class Qwen4ExpDecoderLayer(nn.Module): raise ValueError("Invalid layer_type") mlp_hc = self.mlp_hyper_connection + if self.use_hc_sequence_parallel: + attn_out = sp_reduce_scatter(attn_out) hidden_states, block_input, injection = mlp_hc.combine_and_mix( hidden_states, attn_out, injection ) - mlp_out = self.mlp(block_input) + if self.use_sequence_parallel: + # MoE consumes local tokens and returns complete local outputs. + mlp_out = self.mlp(block_input, already_sequence_parallel=True) + else: + if self.use_hc_sequence_parallel: + block_input = sp_all_gather(block_input)[: positions.shape[-1]] + mlp_out = self.mlp(block_input) + if self.use_hc_sequence_parallel: + assert isinstance(self.mlp, Qwen4ExpSparseMoeBlock) + reduce_output = self.mlp.experts.moe_config.skip_final_all_reduce + # Only sum partial outputs. Reduced outputs already include all ranks. + mlp_out = ( + sp_reduce_scatter(mlp_out) if reduce_output else sp_shard(mlp_out) + ) return hidden_states, mlp_out, injection @@ -373,6 +421,11 @@ class Qwen4ExpModel(nn.Module): self.num_redundant_experts = ( vllm_config.parallel_config.eplb_config.num_redundant_experts ) + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.vocab_size = config.vocab_size self._qsa_layer_ids = frozenset( layer_idx @@ -447,7 +500,6 @@ class Qwen4ExpModel(nn.Module): @staticmethod def _start_layer_ple_prefetch( layer: nn.Module, - hidden_states: torch.Tensor, input_ids: torch.Tensor | None, query_start_loc: torch.Tensor | None, ngram_context: torch.Tensor | None, @@ -458,12 +510,7 @@ class Qwen4ExpModel(nn.Module): return if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - ple.start_prefetch( - hidden_states, - input_ids, - query_start_loc, - ngram_context, - ) + ple.start_prefetch(input_ids, query_start_loc, ngram_context) def forward( self, @@ -482,19 +529,31 @@ class Qwen4ExpModel(nn.Module): if input_ids is None: raise ValueError("input_ids or inputs_embeds is required") hidden_states = self.embed_input_ids(input_ids) + if self.use_hc_sequence_parallel: + if ( + self.use_sequence_parallel + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + # Expand only the local token rows when SP is enabled. hidden_states = hidden_states.repeat(1, self.config.hc_count) else: if intermediate_tensors is None: raise ValueError("pipeline stage requires intermediate tensors") hidden_states = intermediate_tensors["hidden_states"] + full_num_tokens = positions.shape[-1] block_output = None injection = None last_layer = None if self.start_layer < self.end_layer: self._start_layer_ple_prefetch( self.layers[self.start_layer], - hidden_states, input_ids, query_start_loc, ngram_context, @@ -506,7 +565,6 @@ class Qwen4ExpModel(nn.Module): if layer_idx + 1 < self.end_layer: self._start_layer_ple_prefetch( self.layers[layer_idx + 1], - hidden_states, input_ids, query_start_loc, ngram_context, @@ -526,6 +584,8 @@ class Qwen4ExpModel(nn.Module): deepstack_embed = deepstack_input_embeds[ f"deepstack_input_embeds_{layer_idx}" ] + if self.use_hc_sequence_parallel: + deepstack_embed = sp_shard(deepstack_embed) deepstack_embed = ( deepstack_embed.unsqueeze(-2) .expand( @@ -560,6 +620,24 @@ class Qwen4ExpModel(nn.Module): multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix( hidden_states, block_output, injection ) + if self.use_hc_sequence_parallel: + if self._mtp_hidden_buffer is not None: + # Gather LM-head and MTP states together when both are needed. + hidden_size = self.config.hidden_size + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[ + :full_num_tokens + ] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, self.config.hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + else: + sample_hidden_states = sp_all_gather(sample_hidden_states)[ + :full_num_tokens + ] if self._mtp_hidden_buffer is not None: # Capture the pre-final-mixer multi-stream hidden state # [T, hc_count*H] for the MTP drafter (zero extra compute: @@ -615,6 +693,7 @@ class Qwen4ExpModel(nn.Module): class Qwen4ExpForCausalLM( nn.Module, HasInnerState, + SupportsHCSP, SupportsLoRA, SupportsMRoPE, SupportsPP, @@ -865,6 +944,7 @@ class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo): class Qwen4ExpForConditionalGeneration( Qwen3_5ForConditionalGeneration, HasInnerState, + SupportsHCSP, Qwen4ExpMixtureOfExperts, ): """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 61290617ae..32e95faa0d 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -20,8 +20,10 @@ import torch from torch import nn from transformers import Qwen4ExpTextConfig +from vllm import envs from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) @@ -33,7 +35,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import SupportsHCSP, SupportsPP from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -44,6 +46,11 @@ from vllm.model_executor.models.utils import ( maybe_fuse_shared_experts, maybe_prefix, ) +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) from vllm.sequence import IntermediateTensors from .hyperconnection import GatedResidual, HyperConnectionConfig @@ -54,6 +61,7 @@ from .model import ( Qwen4ExpDecoderLayer, Qwen4ExpMixtureOfExperts, Qwen4ExpSparseMoeBlock, + is_hc_sequence_parallel_enabled, ) @@ -171,6 +179,11 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): config: Qwen4ExpTextConfig = model_config.hf_text_config self.config = config + # MoE SP also enables HC SP, while HC SP can be enabled independently. + self.use_hc_sequence_parallel = is_hc_sequence_parallel_enabled(vllm_config) + self.use_sequence_parallel = ( + vllm_config.parallel_config.use_sequence_parallel_moe + ) self.vocab_size = config.vocab_size self.mtp_start_layer_idx = config.num_hidden_layers @@ -310,6 +323,22 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): hidden_states = intermediate_tensors["hidden_states"] current_step_idx = spec_step_idx % self.num_mtp_layers + full_num_tokens = positions.shape[-1] + if self.use_hc_sequence_parallel: + # SP requires PP=1, so the first-rank branch sets prev_block_output. + assert prev_block_output is not None + if ( + self.use_sequence_parallel + and envs.VLLM_MOE_SKIP_PADDING + and is_forward_context_available() + ): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + # The FC output is already complete and needs only a slice. + prev_block_output = sp_shard(prev_block_output) layer = self.layers[current_step_idx] hidden_states, block_output, injection = layer( hidden_states=hidden_states, @@ -337,6 +366,17 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): hidden_states, block_output, injection ) ) + if self.use_hc_sequence_parallel: + # Gather both outputs together, then restore their contiguous layouts. + packed_hidden_states = torch.cat( + [sample_hidden_states, multi_hidden], dim=-1 + ) + packed_hidden_states = sp_all_gather(packed_hidden_states)[:full_num_tokens] + sample_hidden_states, multi_hidden = packed_hidden_states.split( + [hidden_size, hc_count * hidden_size], dim=-1 + ) + sample_hidden_states = sample_hidden_states.contiguous() + multi_hidden = multi_hidden.contiguous() return sample_hidden_states, multi_hidden def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -357,7 +397,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): return loader.load_weights(weights, mapper=mapper) -class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP(nn.Module, SupportsHCSP, SupportsPP, Qwen4ExpMixtureOfExperts): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index cf03cbb40d..794eea87a6 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -9,11 +9,14 @@ from torch import nn from transformers import Qwen4ExpTextConfig from vllm.config import get_current_vllm_config +from vllm.distributed import get_etp_group, get_tp_group from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter, sp_shard from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, @@ -146,8 +149,10 @@ class Qwen4ExpNGramEmbedding(nn.Module): prefix: str, quant_config: QuantizationConfig | None = None, params_dtype: torch.dtype | None = None, + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel self.embedding_dim = embedding_dim self.ngram_size = int(config.ngram_size) self.heads_per_ngram = int(config.heads_per_ngram) @@ -211,6 +216,11 @@ class Qwen4ExpNGramEmbedding(nn.Module): if engram_config is not None and engram_config.cpu_offload else Qwen4ExpPLEDeviceEmbedding ) + # A TP reduce-scatter covers all vocabulary owners only when ETP=TP. + self.use_reduce_scatter = ( + self.use_sequence_parallel + and get_etp_group().world_size == get_tp_group().world_size + ) self.ngram_embedding = embedding_cls( padded_vocab_size, self.head_dim, @@ -221,6 +231,7 @@ class Qwen4ExpNGramEmbedding(nn.Module): num_ngram_heads=self.ngram_heads, max_total_tokens=max_total_tokens, data_parallel_rank=data_parallel_rank, + reduce_results=not self.use_reduce_scatter, ) if self.ngram_embedding.supports_prefetch: # The side-stream lookup outlives eager-break args, whose @@ -348,6 +359,18 @@ class Qwen4ExpNGramEmbedding(nn.Module): id_blocks.append(ids[request_indices, adjusted_columns]) return torch.cat(id_blocks, dim=-1) + def _sp_shard_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: + """Shard embedding token rows, optionally reducing vocabulary contributions.""" + # Each embedding element has one vocabulary owner, including FP8 bytes. + output_dtype = embeddings.dtype + if is_fp8(embeddings): + embeddings = embeddings.view(torch.int8) + if self.use_reduce_scatter: + embeddings = sp_reduce_scatter(embeddings) + else: + embeddings = sp_shard(embeddings) + return embeddings.view(output_dtype) + def forward( self, hidden_states: torch.Tensor, @@ -355,15 +378,22 @@ class Qwen4ExpNGramEmbedding(nn.Module): query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> torch.Tensor: + """Return full embeddings or a token shard after combining vocabulary owners.""" embedding = self.ngram_embedding if embedding.supports_prefetch: - return embedding(hidden_states) - ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) - return self.ngram_embedding(ngram_ids).flatten(-2) + # Prefetch contains full tokens even when hidden states are sharded. + embeddings = embedding.finalize_prefetch(total_tokens=input_ids.shape[0]) + else: + ngram_ids = self.compute_ngram_ids( + input_ids, query_start_loc, ngram_context + ) + embeddings = embedding(ngram_ids).flatten(-2) + if not self.use_sequence_parallel: + return embeddings + return self._sp_shard_embeddings(embeddings) def start_prefetch( self, - hidden_states: torch.Tensor, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, @@ -378,7 +408,7 @@ class Qwen4ExpNGramEmbedding(nn.Module): ngram_context, output=self._prefetch_ids[: input_ids.numel()], ) - embedding.start_prefetch(hidden_states, ngram_ids) + embedding.start_prefetch(ngram_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load hash buffers and checkpoint-split embedding rows.""" @@ -444,8 +474,3 @@ class Qwen4ExpNGramEmbedding(nn.Module): if regular_weights: loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) return loaded - - -__all__ = [ - "Qwen4ExpNGramEmbedding", -] diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index bdb16b3ade..a595dba18d 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -18,6 +18,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateShapeCalculator, is_conv_state_dim_first, ) +from vllm.models.common.ops.sequence_parallel import sp_all_gather, sp_shard from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionBackend, @@ -69,8 +70,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): layer_idx: int = 0, ple_dense_layer_id: int | None = None, prefix: str = "", + use_sequence_parallel: bool = False, ) -> None: super().__init__() + self.use_sequence_parallel = use_sequence_parallel model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config @@ -100,8 +103,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): prefix=f"{prefix}.ple_embedding", quant_config=quant_config, params_dtype=model_config.dtype, + use_sequence_parallel=self.use_sequence_parallel, ) - # The PLE cache is TP-replicated, so this merged projection is too. + # Projection weights stay replicated even when their token rows are sharded. self.kv_proj = MergedColumnParallelLinear( int(config.ple_embed_dim), [self.hc_hidden_size, self.hidden_size], @@ -151,18 +155,12 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): def start_prefetch( self, - hidden_states: torch.Tensor, input_ids: torch.Tensor, query_start_loc: torch.Tensor, ngram_context: torch.Tensor, ) -> None: """Start the pinned PLE lookup while the preceding decoder layer runs.""" - self.ple_embedding.start_prefetch( - hidden_states, - input_ids, - query_start_loc, - ngram_context, - ) + self.ple_embedding.start_prefetch(input_ids, query_start_loc, ngram_context) @property def mamba_type(self) -> MambaAttentionBackendEnum: @@ -192,7 +190,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -207,7 +205,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] - outer_residual = outer_residual[: metadata.num_actual_tokens] + if outer_residual is not None: + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -259,11 +258,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) - outer_residual_d, outer_residual_p = torch.split( - outer_residual, - [num_decode_tokens, num_prefill_tokens], - dim=0, - ) + outer_residual_d: torch.Tensor | None = None + outer_residual_p: torch.Tensor | None = None + if outer_residual is not None: + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: @@ -336,14 +338,16 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self, inputs: torch.Tensor, residual: torch.Tensor, - outer_residual: torch.Tensor, + outer_residual: torch.Tensor | None, ) -> None: + """Update convolution state, optionally fusing both residual additions.""" forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata # Profiling omits all metadata or this Mamba entry. Short convolution # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -354,7 +358,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: - residual.add_(outer_residual) + if outer_residual is not None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -396,7 +401,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): ngram_context: torch.Tensor, ) -> torch.Tensor: input_ids = input_ids.reshape(-1) - if input_ids.shape[0] != hidden_states.shape[0]: + if ( + not self.use_sequence_parallel + and input_ids.shape[0] != hidden_states.shape[0] + ): raise ValueError( "PLE expects input_ids and hidden_states to have the same " f"token length, got {input_ids.shape[0]} and " @@ -408,6 +416,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): query_start_loc, ngram_context, ) + assert embeddings.shape[0] == hidden_states.shape[0] embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) kv, _ = self.kv_proj(embeddings) key, value = kv.split(self.kv_proj.output_sizes, dim=-1) @@ -420,9 +429,16 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self.norm_conv.weight, self.norm_key.eps, ) - # State routing depends on runtime request metadata and remains outside - # the graph; short convolution accumulates into gated_output. - self._short_conv(conv_input, gated_output, hidden_states) + if self.use_sequence_parallel: + # Only convolution needs full tokens. Add residuals on local shards. + conv_input = sp_all_gather(conv_input)[: input_ids.shape[0]] + # Metadata may omit padding tokens or all tokens during profiling. + conv_output = torch.zeros_like(conv_input) + self._short_conv(conv_input, conv_output, outer_residual=None) + gated_output.add_(sp_shard(conv_output)) + gated_output.add_(hidden_states) + else: + self._short_conv(conv_input, gated_output, hidden_states) return gated_output