diff --git a/tests/config/test_virtual_tp.py b/tests/config/test_virtual_tp.py index a1d67ddb1..c1251eee4 100644 --- a/tests/config/test_virtual_tp.py +++ b/tests/config/test_virtual_tp.py @@ -314,11 +314,12 @@ def test_b12x_virtual_tp_padding_glm_dsa_tp6(): } -def test_b12x_virtual_tp_padding_glm_dsa_draft_tp6(): +@pytest.mark.parametrize("method", ["mtp", "dspark"]) +def test_b12x_virtual_tp_padding_glm_dsa_draft_tp6(method: str): target_model_config = FakeGlmDsaModelConfig() draft_model_config = FakeGlmDsaModelConfig() spec_config = SimpleNamespace( - method="mtp", + method=method, target_model_config=target_model_config, draft_model_config=draft_model_config, draft_parallel_config=ParallelConfig( diff --git a/tests/v1/spec_decode/test_dspark.py b/tests/v1/spec_decode/test_dspark.py index 1078f9d7f..fb60c8708 100644 --- a/tests/v1/spec_decode/test_dspark.py +++ b/tests/v1/spec_decode/test_dspark.py @@ -14,7 +14,6 @@ from vllm.models.deepseek_v4.nvidia.dspark import ( _read_dspark_num_layers, ) from vllm.models.deepseek_v4.nvidia.ops.dspark_sparse_attn_tilelang import ( - _get_tilelang_block_size, _get_tilelang_padded_heads, ) from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_target_layer_ids @@ -138,13 +137,14 @@ def test_build_dspark_topk_idxs_uses_rolling_target_window_then_draft_block(): @pytest.mark.parametrize( ("heads", "expected"), [ - (1, 64), - (16, 64), + (1, 16), + (16, 16), + (24, 32), (32, 32), (64, 64), ], ) -def test_dspark_tilelang_sparse_attn_pads_tp4_head_count(heads, expected): +def test_dspark_tilelang_sparse_attn_aligns_local_head_count(heads, expected): assert _get_tilelang_padded_heads(heads) == expected @@ -153,19 +153,6 @@ def test_dspark_tilelang_sparse_attn_rejects_oversized_head_count(): _get_tilelang_padded_heads(128) -@pytest.mark.parametrize( - ("heads", "expected"), - [ - (32, 32), - (64, 16), - ], -) -def test_dspark_tilelang_sparse_attn_uses_small_tile_for_padded_heads( - heads, expected -): - assert _get_tilelang_block_size(heads) == expected - - def test_dspark_b12x_output_projection_uses_attention_helper(): layer = object.__new__(DeepSeekV4DSparkLayer) layer.prefix = "model.layers.40" diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 58b60c97f..96ccc01d0 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1220,7 +1220,7 @@ class SpeculativeConfig: def _maybe_apply_virtual_tp_to_draft(self) -> None: if ( - self.method != "mtp" + self.method not in {"mtp", "dspark"} or self.draft_model_config is None or self.draft_parallel_config is None or self.draft_model_config is self.target_model_config diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 75a057b1f..0a91720d7 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -99,6 +99,16 @@ def _resolve_dsv4_kv_cache_dtype( return kv_cache_dtype, torch.bfloat16 +def resolve_layer_compress_ratio(config, layer_id: int) -> tuple[int, bool]: + """Resolve the operational ratio and whether a layer uses plain RoPE.""" + if layer_id < config.num_hidden_layers: + return max(1, config.compress_ratios[layer_id]), False + if layer_id < len(config.compress_ratios): + raw_compress_ratio = config.compress_ratios[layer_id] + return max(1, raw_compress_ratio), raw_compress_ratio == 0 + return 1, False + + class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): """DeepseekV4 MLA attention layer. @@ -165,6 +175,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): prefix: str, topk_indices_buffer: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, + config_layer_id: int | None = None, ) -> None: super().__init__() self.vllm_config = vllm_config @@ -174,7 +185,11 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self._use_b12x_wo = bool(envs.VLLM_USE_B12X_WO_PROJECTION) self._b12x_wo_projection_weights: Any | None = None tp_size = get_tensor_model_parallel_world_size() - layer_id = extract_layer_index(prefix) + layer_id = ( + config_layer_id + if config_layer_id is not None + else extract_layer_index(prefix) + ) self.prefix = prefix # Alias for compatibility with compressor self.hidden_size = config.hidden_size @@ -189,15 +204,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.n_groups = config.o_groups self.n_local_groups = self.n_groups // tp_size self.window_size = config.sliding_window - # NOTE(zyongye) Compress ratio can't be 0 - # we do this for because MTP layer is not included - # in the compress ratio list - if layer_id < config.num_hidden_layers: - self.compress_ratio = max(1, config.compress_ratios[layer_id]) - elif layer_id < len(config.compress_ratios): - self.compress_ratio = config.compress_ratios[layer_id] - else: - self.compress_ratio = 1 + self.compress_ratio, use_unscaled_rope = resolve_layer_compress_ratio( + config, layer_id + ) self.eps = config.rms_norm_eps self.scale = self.head_dim**-0.5 @@ -269,6 +278,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): rope_head_dim=self.rope_head_dim, max_position_embeddings=config.max_position_embeddings, compress_ratio=self.compress_ratio, + use_unscaled_rope=use_unscaled_rope, ) self.indexer_rotary_emb = self.rotary_emb self.topk_indices_buffer = topk_indices_buffer diff --git a/vllm/models/deepseek_v4/common/rope.py b/vllm/models/deepseek_v4/common/rope.py index 894e1ef31..9aa84d8e8 100644 --- a/vllm/models/deepseek_v4/common/rope.py +++ b/vllm/models/deepseek_v4/common/rope.py @@ -15,12 +15,13 @@ def build_deepseek_v4_rope( rope_head_dim: int, max_position_embeddings: int, compress_ratio: int, + use_unscaled_rope: bool = False, ) -> RotaryEmbedding: rope_parameters = dict(config.rope_parameters) rope_parameters["rope_theta"] = ( config.compress_rope_theta if compress_ratio > 1 else config.rope_theta ) - if compress_ratio == 0: + if use_unscaled_rope: rope_parameters["rope_type"] = "default" if rope_parameters["rope_type"] != "default": rope_parameters["rope_type"] = ( diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 9d23ef67c..875a0e805 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -574,14 +574,15 @@ class DeepSeekV4DSparkLayer(DeepseekV4DecoderLayer): topk_indices_buffer: torch.Tensor, aux_stream_list: list[torch.cuda.Stream] | None, ) -> None: + config = vllm_config.model_config.hf_config super().__init__( vllm_config, prefix=prefix, topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, + config_layer_id=config.num_hidden_layers + stage_id, ) self.prefix = prefix - config = vllm_config.model_config.hf_config self._dspark_use_reference_fp8_linears = bool( int(os.getenv("VLLM_DSPARK_REFERENCE_FP8_LINEARS", "0")) ) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index eef48ba76..58c27ea03 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -848,6 +848,7 @@ class DeepseekV4DecoderLayer(nn.Module): prefix, topk_indices_buffer: torch.Tensor | None = None, aux_stream_list: list[torch.cuda.Stream] | None = None, + config_layer_id: int | None = None, ): super().__init__() @@ -872,6 +873,7 @@ class DeepseekV4DecoderLayer(nn.Module): prefix=f"{prefix}.attn", topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, + config_layer_id=config_layer_id, ) self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") diff --git a/vllm/models/deepseek_v4/nvidia/ops/dspark_sparse_attn_tilelang.py b/vllm/models/deepseek_v4/nvidia/ops/dspark_sparse_attn_tilelang.py index 4fdcc5431..0aea01f94 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dspark_sparse_attn_tilelang.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dspark_sparse_attn_tilelang.py @@ -14,6 +14,21 @@ from typing import Any import torch +_TILELANG_HEAD_ALIGNMENT = 16 +_TILELANG_MAX_SUPPORTED_HEADS = 64 + + +def _get_tilelang_padded_heads(heads: int) -> int: + if heads <= 0 or heads > _TILELANG_MAX_SUPPORTED_HEADS: + raise ValueError( + "DSpark TileLang sparse attention requires between 1 and " + f"{_TILELANG_MAX_SUPPORTED_HEADS} local heads, got {heads}" + ) + return ( + (heads + _TILELANG_HEAD_ALIGNMENT - 1) // _TILELANG_HEAD_ALIGNMENT + ) * _TILELANG_HEAD_ALIGNMENT + + @cache def _build_dspark_sparse_attn_kernel(num_heads: int, head_dim: int, scale: float): import tilelang @@ -156,9 +171,8 @@ def dspark_sparse_attn( topk_idxs = topk_idxs.to(torch.int32) batch, draft_tokens, heads, head_dim = q.shape - padded_heads = heads - if heads < 16: - padded_heads = 16 + padded_heads = _get_tilelang_padded_heads(heads) + if padded_heads != heads: q = torch.cat( [q, q.new_zeros(batch, draft_tokens, padded_heads - heads, head_dim)], dim=2, @@ -170,7 +184,7 @@ def dspark_sparse_attn( out = torch.empty_like(q) kernel = _build_dspark_sparse_attn_kernel(padded_heads, head_dim, softmax_scale) kernel(q.contiguous(), kv.contiguous(), out, attn_sink.contiguous(), topk_idxs) - if heads < 16: + if padded_heads != heads: return out.narrow(2, 0, heads).contiguous() return out diff --git a/tests/models/test_deepseek_v4_dspark_rope.py b/tests/models/test_deepseek_v4_dspark_rope.py new file mode 100644 index 000000000..6ac0ed518 --- /dev/null +++ b/tests/models/test_deepseek_v4_dspark_rope.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for DSpark draft-layer RoPE selection.""" + +import inspect +import types + +import pytest + +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.models.deepseek_v4.attention import ( + DeepseekV4Attention, + resolve_layer_compress_ratio, +) +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v4.nvidia import dspark as dspark_mod +from vllm.models.deepseek_v4.nvidia import model as model_mod + + +def _rope_config(compress_ratios: list[int]) -> types.SimpleNamespace: + return types.SimpleNamespace( + num_hidden_layers=43, + compress_ratios=compress_ratios, + rope_theta=10000.0, + compress_rope_theta=1000000.0, + max_position_embeddings=4096, + rope_parameters={ + "rope_type": "yarn", + "factor": 8.0, + "original_max_position_embeddings": 512, + "beta_fast": 32, + "beta_slow": 1, + }, + ) + + +@pytest.mark.parametrize("stage_id", [0, 1, 2]) +def test_dspark_stage_uses_its_checkpoint_config_entry(monkeypatch, stage_id): + captured: dict[str, object] = {} + + class ParentInitReached(Exception): + pass + + def capture_parent_init(self, *args, **kwargs): + captured.update(kwargs) + raise ParentInitReached + + monkeypatch.setattr( + dspark_mod.DeepseekV4DecoderLayer, "__init__", capture_parent_init + ) + config = types.SimpleNamespace(num_hidden_layers=43) + vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace(hf_config=config) + ) + + with pytest.raises(ParentInitReached): + dspark_mod.DeepSeekV4DSparkLayer( + vllm_config, + stage_id=stage_id, + num_dspark_layers=3, + prefix=f"mtp.{stage_id}", + topk_indices_buffer=None, + aux_stream_list=None, + ) + + assert captured["config_layer_id"] == config.num_hidden_layers + stage_id + + +def test_decoder_forwards_config_layer_id_to_attention(monkeypatch): + captured: dict[str, object] = {} + + class AttentionInitReached(Exception): + pass + + class CaptureAttention: + def __init__(self, *args, **kwargs): + captured.update(kwargs) + raise AttentionInitReached + + monkeypatch.setattr(model_mod, "_use_b12x_mhc", lambda: False) + monkeypatch.setattr( + model_mod, "_select_dsv4_attn_cls", lambda vllm_config: CaptureAttention + ) + config = types.SimpleNamespace(hidden_size=256, rms_norm_eps=1e-6) + vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace(hf_config=config) + ) + + with pytest.raises(AttentionInitReached): + model_mod.DeepseekV4DecoderLayer( + vllm_config, + prefix="mtp.2", + config_layer_id=45, + ) + + assert captured["config_layer_id"] == 45 + + +def test_dspark_zero_entries_select_unscaled_rope(default_vllm_config): + config = _rope_config([0] * 46) + + for layer_id in (43, 44, 45): + ratio, use_unscaled_rope = resolve_layer_compress_ratio(config, layer_id) + assert (ratio, use_unscaled_rope) == (1, True) + + rope = build_deepseek_v4_rope( + config, + head_dim=64, + rope_head_dim=64, + max_position_embeddings=config.max_position_embeddings, + compress_ratio=1, + use_unscaled_rope=True, + ) + assert type(rope) is RotaryEmbedding + + +def test_attention_init_wires_config_layer_override_and_resolver(): + signature = inspect.signature(DeepseekV4Attention.__init__) + assert "config_layer_id" in signature.parameters + assert ( + "resolve_layer_compress_ratio" in DeepseekV4Attention.__init__.__code__.co_names + )