iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Qwen4Exp] support Sequence Parallelism for Qwen3.8-Flash-Next (#56322)

Browse Source 
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Harry Huang
2026-10-12 20:56:44 +08:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 258b97ff6b
commit c75cb63cd2
23 changed files with 870 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/models_basic.yaml
+1 -2
View File
@@ -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/
.buildkite/test_areas/models_distributed.yaml
+17
View File
@@ -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
tests/kernels/core/test_vocab_parallel_embedding.py
+42
View File
@@ -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,
)
tests/models/qwen4_exp/test_config.py
+61 -23
View File
@@ -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(
tests/models/qwen4_exp/test_ple.py
+63 -40
View File
@@ -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:],
)
tests/models/qwen4_exp/test_tpsp.py
+267
View File
@@ -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)
tests/models/test_registry.py
+16
View File
@@ -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",
tests/test_config.py
+41
View File
@@ -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"),
[
vllm/config/model.py
+8
View File
@@ -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:
vllm/config/parallel.py
+5
View File
@@ -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
vllm/engine/arg_utils.py
+3
View File
@@ -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,
vllm/model_executor/layers/vocab_parallel_embedding.py
+18 -17
View File
@@ -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}"
vllm/model_executor/models/__init__.py
+4
View File
@@ -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",
vllm/model_executor/models/config.py
+10
View File
@@ -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
)
vllm/model_executor/models/interfaces.py
+21
View File
@@ -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."""
vllm/model_executor/models/qwen3_next.py
+7 -1
View File
@@ -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
vllm/model_executor/models/registry.py
+12
View File
@@ -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],
vllm/models/qwen4_exp/common/ngram_embedding.py
+34 -22
View File
@@ -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()."
)
vllm/models/qwen4_exp/common/ops/ple.py
+23 -17
View File
@@ -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,
vllm/models/qwen4_exp/nvidia/model.py
+102 -22
View File
@@ -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."""
vllm/models/qwen4_exp/nvidia/mtp.py
+42 -2
View File
@@ -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"],
vllm/models/qwen4_exp/nvidia/ngram_embedding.py
+35 -10
View File
@@ -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",
]
vllm/models/qwen4_exp/nvidia/ple_layer.py
+38 -22
View File
@@ -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