[Model][NVIDIA] Route DSA models to the CUDA non-compiled path (#52861)

This commit is contained in:
Woosuk Kwon
2026-08-19 00:13:52 -07:00
committed by GitHub
parent ee11730751
commit b09bd69b5b
16 changed files with 260 additions and 146 deletions
+1 -20
View File
@@ -90,11 +90,8 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
backend_name = attn_backend.backend.name.lower()
requires_mla = "deepseek" in model_name.lower()
is_mla = "mla" in backend_name
# DeepSeek V3.2 uses sparse MLA
requires_sparse = "v3.2" in model_name.lower()
is_sparse = "sparse" in backend_name
if requires_mla != is_mla or requires_sparse != is_sparse:
if requires_mla != is_mla:
pytest.skip(
f"Incompatible model '{model_name}' and "
f"attention backend '{attn_backend.backend.name}'"
@@ -128,22 +125,6 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
# models (e.g. Llama-4-Scout-FP8) at 16384 tokens may trigger OOM.
model_kwargs.setdefault("max_num_batched_tokens", 8192)
# Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
# decompose_auto_functionalized when +rotary_embedding is forced into
# the compile graph. Disable qk_norm+rope fusion (which auto-enables
# +rotary_embedding) for this combo to avoid the known torch bug.
# TODO: remove once upstream torch fix lands.
if requires_sparse:
if "pass_config" in compilation_config:
compilation_config["pass_config"].enable_qk_norm_rope_fusion = False
matches_check = [m for m in matches_check if m != "norm_rope_fusion"]
# DSv3.2 sparse indexer uses persistent_topk with k=config.index_topk
# (2048 for the default config). max_model_len must be >= index_topk
# or the topk kernel raises "k out of range" at runtime.
model_kwargs["max_model_len"] = max(
model_kwargs.get("max_model_len", 0), 2048
)
# Always compile the full graph instead of piecewise
if not compilation_config["use_inductor_graph_partition"]:
compilation_config["splitting_ops"] = []
-24
View File
@@ -58,18 +58,6 @@ TRITON_MLA_ATTN = pytest.param(
id="TRITON_MLA",
)
FLASHMLA_SPARSE_ATTN = pytest.param(
AttentionBackendCase(
backend=AttentionBackendEnum.FLASHMLA_SPARSE,
model_kwargs=dict(kv_cache_dtype="fp8_ds_mla"),
),
id="FLASHMLA_SPARSE",
marks=pytest.mark.skipif(
not is_blackwell(),
reason="FlashMLA Sparse requires Blackwell",
),
)
# Models
llama3_8b = ModelFusionInfo(
model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -197,18 +185,6 @@ deepseek_r1_fp4 = ModelFusionInfo(
),
)
deepseek_v32_fp4 = ModelFusionInfo(
model_name="nvidia/DeepSeek-V3.2-NVFP4",
matches=lambda n_layers: Matches(
rms_quant_fusion=0,
# silu+quant on dense layers only; MoE hides the act+quant site
act_quant_fusion=min(3, n_layers),
# MLA attn + NVFP4 output quant fuses on sparse MLA output path
attn_quant_fusion=n_layers,
ar_rms_fusion=n_layers * 2 + 1,
),
)
gpt_oss_20b = ModelFusionInfo(
model_name="openai/gpt-oss-20b",
matches=lambda n_layers: Matches(
+2 -4
View File
@@ -18,7 +18,6 @@ from .common import (
from .models import (
FLASHINFER_ATTN,
FLASHINFER_MLA_ATTN,
FLASHMLA_SPARSE_ATTN,
ROCM_AITER_UNIFIED_ATTN,
ROCM_ATTN,
TRITON_ATTN,
@@ -26,7 +25,6 @@ from .models import (
deepseek_coder_v2_lite_fp8,
deepseek_r1_fp4,
deepseek_v3_fp8,
deepseek_v32_fp4,
llama3_8b_fp4,
llama3_8b_fp8,
llama4_scout_fp4,
@@ -149,11 +147,11 @@ def test_tp1_fp8_fusions(
@pytest.mark.parametrize(
"model_name, matches_fn, model_kwargs, hf_overrides",
[llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4],
[llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4],
)
@pytest.mark.parametrize(
"attn_backend",
[FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN],
[FLASHINFER_ATTN, FLASHINFER_MLA_ATTN],
)
@pytest.mark.parametrize("n_layers", [6])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
+2 -4
View File
@@ -18,14 +18,12 @@ from .common import (
from .models import (
FLASHINFER_ATTN,
FLASHINFER_MLA_ATTN,
FLASHMLA_SPARSE_ATTN,
ROCM_AITER_UNIFIED_ATTN,
ROCM_ATTN,
TRITON_ATTN,
deepseek_coder_v2_lite_fp8,
deepseek_r1_fp4,
deepseek_v3_fp8,
deepseek_v32_fp4,
gpt_oss_20b,
llama3_8b,
llama3_8b_fp4,
@@ -120,11 +118,11 @@ def test_tp2_ar_rms_fp8_fusions(
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
"model_name, matches_fn, model_kwargs, hf_overrides",
[llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4],
[llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4],
)
@pytest.mark.parametrize(
"attn_backend",
[FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN],
[FLASHINFER_ATTN, FLASHINFER_MLA_ATTN],
)
@pytest.mark.parametrize("n_layers", [4])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
-13
View File
@@ -132,19 +132,6 @@ MODEL_SPECS = [
),
id="gpt_oss_120b",
),
# NOTE: DeepSeek-V3.2 requires sparse MLA (index_topk) which needs
# Hopper+ GPUs. This test must run on H100 (see pytorch.yaml).
pytest.param(
ModelStartupSpec(
model="deepseek-ai/DeepSeek-V3.2",
hf_overrides=_SMALL_MOE_OVERRIDES,
cold_artifacts_saved=9,
# https://github.com/vllm-project/vllm/issues/38051
warm_artifacts_saved=0 if is_torch_equal_or_newer("2.12.0") else 9,
warm_artifacts_loaded=9 if is_torch_equal_or_newer("2.12.0") else 0,
),
id="deepseek_v3.2",
),
pytest.param(
ModelStartupSpec(
model="moonshotai/Kimi-K2.5",
@@ -146,8 +146,8 @@ def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str):
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])
@pytest.mark.parametrize("index_interleave", [True, False])
@pytest.mark.parametrize("mla_fp8", [False, True])
def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool):
@pytest.mark.parametrize("mla_dtype", ["auto", "bfloat16", "fp8"])
def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_dtype: str):
torch.manual_seed(0)
dev = "cuda"
max_pos = 8192
@@ -167,13 +167,12 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool)
bs = max_pos # single block covering all tokens
mla_dim = KV_LORA + ROPE_DIM
mla_fp8 = mla_dtype == "fp8"
if mla_fp8:
mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8)
mla_dtype = "fp8"
mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32)
else:
mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16)
mla_dtype = "auto"
mla_k_scale = None
idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 # 132
idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8)
+5
View File
@@ -1679,6 +1679,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
speculative_model="luccafong/deepseek_mtp_draft_random",
trust_remote_code=True,
),
"DeepseekV32MTPModel": _HfExamplesInfo(
"nvidia/DeepSeek-V3.2-NVFP4",
speculative_model="nvidia/DeepSeek-V3.2-NVFP4",
is_available_online=False,
),
"DeepSeekV4MTPModel": _HfExamplesInfo(
"deepseek-ai/DeepSeek-V4-Flash",
speculative_model="deepseek-ai/DeepSeek-V4-Flash",
+140 -4
View File
@@ -113,18 +113,154 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch):
"""ROCm keeps DeepSeek V4 on MRV1, which is still faster there."""
from vllm.config.vllm import default_v2_model_runner_architectures
def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
"""ROCm keeps DeepSeek V3.2 and V4 on their compiled MRV1 paths."""
from vllm.config.vllm import (
default_breakable_cudagraph_architectures,
default_v2_model_runner_architectures,
)
from vllm.platforms import current_platform
monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
# The lookup is lru_cached against a fixed platform.
default_v2_model_runner_architectures.cache_clear()
default_breakable_cudagraph_architectures.cache_clear()
try:
assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures()
v2_architectures = default_v2_model_runner_architectures()
breakable_architectures = default_breakable_cudagraph_architectures()
assert "DeepseekV32ForCausalLM" not in v2_architectures
assert "DeepseekV4ForCausalLM" not in v2_architectures
assert "DeepseekV32ForCausalLM" not in breakable_architectures
assert "DeepseekV32MTPModel" not in breakable_architectures
finally:
default_v2_model_runner_architectures.cache_clear()
default_breakable_cudagraph_architectures.cache_clear()
@pytest.mark.parametrize(
("model", "architecture"),
[
("nvidia/GLM-5.2-NVFP4", "GlmMoeDsaForCausalLM"),
("zai-org/GLM-5.2-FP8", "GlmMoeDsaForCausalLM"),
("nvidia/DeepSeek-V3.2-NVFP4", "DeepseekV32ForCausalLM"),
],
)
@pytest.mark.parametrize("with_mtp", [False, True], ids=["no-mtp", "mtp"])
def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
monkeypatch, model, architecture, with_mtp
):
from vllm.compilation.breakable_cudagraph import (
is_breakable_cudagraph_enabled,
)
from vllm.config.vllm import (
default_breakable_cudagraph_architectures,
default_v2_model_runner_architectures,
)
from vllm.platforms import current_platform
monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False)
monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)
monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True)
monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
default_v2_model_runner_architectures.cache_clear()
default_breakable_cudagraph_architectures.cache_clear()
model_config = SimpleNamespace(
model=model,
architectures=[architecture],
runner_type="generate",
is_moe=True,
is_hybrid=False,
is_attention_free=False,
is_diffusion=False,
)
config = SimpleNamespace(
model_config=model_config,
speculative_config=SimpleNamespace(method="mtp") if with_mtp else None,
parallel_config=SimpleNamespace(prefill_context_parallel_size=1),
compilation_config=CompilationConfig(
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE
),
)
config._dflash_needs_multi_kv_group = lambda: False
config._is_default_v2_model_runner_model = lambda: (
VllmConfig._is_default_v2_model_runner_model(config)
)
config._get_v2_model_runner_unsupported_features = lambda: []
config._uses_breakable_cudagraph_by_default = lambda: (
VllmConfig._uses_breakable_cudagraph_by_default(config)
)
try:
assert VllmConfig.use_v2_model_runner.fget(config)
assert VllmConfig._maybe_enable_breakable_cudagraph(config)
assert is_breakable_cudagraph_enabled()
assert config.compilation_config.mode == CompilationMode.NONE
assert config.compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
finally:
os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None)
default_v2_model_runner_architectures.cache_clear()
default_breakable_cudagraph_architectures.cache_clear()
@pytest.mark.parametrize(
("architecture", "is_rocm", "expected"),
[
("DeepseekV32ForCausalLM", False, True),
("DeepseekV32ForCausalLM", True, False),
("DeepseekV32MTPModel", False, True),
("DeepseekV32MTPModel", True, False),
("GlmMoeDsaForCausalLM", False, True),
("GlmMoeDsaForCausalLM", True, True),
],
)
def test_dsa_breakable_cudagraph_platform_default(
monkeypatch, architecture, is_rocm, expected
):
from vllm.config.vllm import default_breakable_cudagraph_architectures
from vllm.platforms import current_platform
monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False)
monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm)
default_breakable_cudagraph_architectures.cache_clear()
config = SimpleNamespace(
model_config=SimpleNamespace(architectures=[architecture]),
compilation_config=CompilationConfig(),
)
config._uses_breakable_cudagraph_by_default = lambda: (
VllmConfig._uses_breakable_cudagraph_by_default(config)
)
try:
assert VllmConfig._maybe_enable_breakable_cudagraph(config) is expected
if expected:
assert config.compilation_config.mode == CompilationMode.NONE
finally:
os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None)
default_breakable_cudagraph_architectures.cache_clear()
@pytest.mark.parametrize(
("model_type", "expected_architecture"),
[
("deepseek_v32", "DeepseekV32MTPModel"),
("glm_moe_dsa", "DeepseekV32MTPModel"),
("deepseek_v3", "DeepSeekMTPModel"),
],
)
def test_dsa_models_select_matching_mtp(model_type, expected_architecture):
from transformers import PretrainedConfig
hf_config = PretrainedConfig(
architectures=["DeepseekV32ForCausalLM"],
num_nextn_predict_layers=1,
)
hf_config.model_type = model_type
SpeculativeConfig.hf_config_override(hf_config)
assert hf_config.architectures == [expected_architecture]
@pytest.mark.parametrize(
+7 -5
View File
@@ -341,6 +341,7 @@ class SpeculativeConfig:
@staticmethod
def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
initial_architecture = hf_config.architectures[0]
use_v32_mtp = hf_config.model_type in ("deepseek_v32", "glm_moe_dsa")
if hf_config.model_type == "dots3_note":
n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
mtp_layer_types = getattr(hf_config, "mtp_layer_types", None)
@@ -365,7 +366,12 @@ class SpeculativeConfig:
if hf_config.model_type == "deepseek_mtp":
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
hf_config.update(
{"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]}
{
"n_predict": n_predict,
"architectures": [
"DeepseekV32MTPModel" if use_v32_mtp else "DeepSeekMTPModel"
],
}
)
if hf_config.model_type == "deepseek_v4":
hf_config.model_type = "deepseek_mtp"
@@ -759,10 +765,6 @@ class SpeculativeConfig:
if self.method == "mtp":
if self.target_model_config is None:
raise ValueError("target_model_config must be present for mtp")
if self.target_model_config.hf_text_config.model_type == "deepseek_v32":
# FIXME(luccafong): cudagraph with v32 MTP is not supported,
# remove this when the issue is fixed.
self.enforce_eager = True
# use the draft model from the same model:
self.model = self.target_model_config.model
# Align the quantization of draft model for cases such as
+67 -39
View File
@@ -69,7 +69,9 @@ logger = init_logger(__name__)
DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
{
"DeepseekV2ForCausalLM",
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
"GlmMoeDsaForCausalLM",
"GraniteMoeForCausalLM",
"InklingForCausalLM",
"InklingForConditionalGeneration",
@@ -79,6 +81,23 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
}
)
DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
{
"DeepseekV32MTPModel",
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
"DeepSeekV4MTPModel",
"GlmMoeDsaForCausalLM",
"InklingForCausalLM",
"InklingForConditionalGeneration",
"KimiK3ForConditionalGeneration",
"KimiK3MTPModel",
"KimiLinearForCausalLM",
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
}
)
@lru_cache
def default_v2_model_runner_architectures() -> frozenset[str]:
@@ -86,13 +105,28 @@ def default_v2_model_runner_architectures() -> frozenset[str]:
from vllm.platforms import current_platform
if current_platform.is_rocm():
# TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The
# attention layer picks the eager cudagraph region MRV1 needs, so
# this is a perf default only; drop it once MRV2 catches up.
return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"}
# TODO(rocm): These models are either unsupported by MRV2 or slower with
# MRV2 on AMD GPUs.
return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
}
return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES
@lru_cache
def default_breakable_cudagraph_architectures() -> frozenset[str]:
"""Architectures defaulting to breakable CUDA graphs on this platform."""
from vllm.platforms import current_platform
if current_platform.is_rocm():
return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES - {
"DeepseekV32ForCausalLM",
"DeepseekV32MTPModel",
}
return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES
class OptimizationLevel(IntEnum):
"""Optimization level enum."""
@@ -691,6 +725,34 @@ class VllmConfig:
return False
return is_default_v2_architecture or not model_config.is_moe
def _uses_breakable_cudagraph_by_default(self) -> bool:
model_config = self.model_config
if model_config is None:
return False
architectures = set(model_config.architectures)
return bool(architectures & default_breakable_cudagraph_architectures())
def _maybe_enable_breakable_cudagraph(self) -> bool:
if (
"VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
and self._uses_breakable_cudagraph_by_default()
):
os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
logger.info_once(
"Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
"Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
)
from vllm.compilation.breakable_cudagraph import (
is_breakable_cudagraph_enabled,
)
enabled = is_breakable_cudagraph_enabled()
if enabled:
self.compilation_config.mode = CompilationMode.NONE
return enabled
@property
def needs_dp_coordinator(self) -> bool:
"""
@@ -1288,41 +1350,7 @@ class VllmConfig:
)
self.compilation_config.mode = CompilationMode.NONE
# For model classes don't carry @support_torch_compile —
# the breakable cudagraph is the supported PIECEWISE path. Auto-enable
# it unless the user has explicitly opted out via the env var.
if (
self.model_config is not None
and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
and any(
a
in (
"DeepseekV4ForCausalLM",
"DeepSeekV4MTPModel",
"InklingForCausalLM",
"InklingForConditionalGeneration",
"KimiK3ForConditionalGeneration",
"KimiK3MTPModel",
"KimiLinearForCausalLM",
"MiniMaxM3SparseForCausalLM",
"MiniMaxM3SparseForConditionalGeneration",
)
for a in self.model_config.architectures
)
):
os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
logger.info_once(
"Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
"Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
)
from vllm.compilation.breakable_cudagraph import (
is_breakable_cudagraph_enabled,
)
breakable_cudagraph_enabled = is_breakable_cudagraph_enabled()
if breakable_cudagraph_enabled:
self.compilation_config.mode = CompilationMode.NONE
breakable_cudagraph_enabled = self._maybe_enable_breakable_cudagraph()
if not breakable_cudagraph_enabled and (
self.compilation_config.backend == "eager"
+6 -2
View File
@@ -92,7 +92,10 @@ _TEXT_GENERATION_MODELS = {
"DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"),
"DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
"DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
"DeepseekV32ForCausalLM": (
"vllm.models.deepseek_v32",
"DeepseekV32ForCausalLM",
),
"DeepseekV4ForCausalLM": ("vllm.models.deepseek_v4", "DeepseekV4ForCausalLM"),
"Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),
"Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),
@@ -115,7 +118,7 @@ _TEXT_GENERATION_MODELS = {
"Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
"Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
"Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
"GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
"GlmMoeDsaForCausalLM": ("vllm.models.deepseek_v32", "GlmMoeDsaForCausalLM"),
"GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
"GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
"GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
@@ -661,6 +664,7 @@ _SPECULATIVE_DECODING_MODELS = {
"Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
"EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
"DeepseekV32MTPModel": ("vllm.models.deepseek_v32", "DeepseekV32MTP"),
"Dots3NoteMTPModel": ("vllm.models.dots3_note", "Dots3NoteMTP"),
"DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
"BailingMoeV3MTPModel": ("bailing_moe_v3_mtp", "BailingMoeV3MTPModel"),
+17 -8
View File
@@ -1,26 +1,35 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""DeepSeek V3.2 (``deepseek_v32``) model — hardware-isolated entry point.
"""DeepSeek V3.2 (``deepseek_v32``) platform entry point.
DeepSeek V3.2 introduced the DeepSeek Sparse Attention (DSA) architecture:
MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA
attend. The same model code serves any DSA checkpoint, including GLM-5.2
(``glm_moe_dsa``), which reuses this architecture.
The CUDA implementation selects capability-specific kernels internally and
falls back when an optimization is unavailable. Other platforms use the
generic implementation by default.
"""
from vllm.platforms import current_platform
if current_platform.is_rocm():
from .amd.model import DeepseekV32ForCausalLM
from .amd.mtp import DeepseekV32MTP
elif current_platform.is_xpu():
raise NotImplementedError("deepseek_v32 does not yet support XPU.")
else:
# Covers Blackwell (sm100) and all other CUDA devices.
if current_platform.is_cuda():
# GLM-5.2 (glm_moe_dsa) reuses the CUDA DSA module. Individual optimized
# kernels remain gated on the device capabilities they support.
from .nvidia.model import DeepseekV32ForCausalLM
from .nvidia.model import DeepseekV32ForCausalLM as GlmMoeDsaForCausalLM
from .nvidia.mtp import DeepseekV32MTP
else:
# ROCm, XPU, and CPU keep the generic implementation.
from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP
from vllm.model_executor.models.deepseek_v2 import (
DeepseekV3ForCausalLM as DeepseekV32ForCausalLM,
)
from vllm.model_executor.models.deepseek_v2 import GlmMoeDsaForCausalLM
__all__ = [
"DeepseekV32ForCausalLM",
"DeepseekV32MTP",
"GlmMoeDsaForCausalLM",
]
-1
View File
@@ -38,7 +38,6 @@ class DeepseekV32ROCmIndexer(DeepseekV32Indexer):
class DeepseekV32MLAAttention(DeepseekV32Attention):
require_fp8_kv_cache: bool = False
indexer_cls = DeepseekV32ROCmIndexer
def __init__(self, vllm_config, config, prefix, topk_indices_buffer=None):
+3 -16
View File
@@ -170,7 +170,6 @@ class DeepseekV32Attention(MLAAttention):
indexer: "DeepseekV32Indexer | None"
indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
require_fp8_kv_cache: bool = True
supports_dense_mha_prefill = False
def __init__(
@@ -283,21 +282,9 @@ class DeepseekV32Attention(MLAAttention):
self.layer_name if enable_short_prefill_scoring_skip else ""
)
if self.require_fp8_kv_cache:
assert is_quantized_kv_cache(self.kv_cache_dtype), (
"deepseek_v32 (nvidia) requires an fp8 KV cache served by a sparse "
"MLA backend. Launch with --kv-cache-dtype fp8 (FlashInfer sparse) "
"or --kv-cache-dtype fp8_ds_mla (FlashMLA sparse)."
)
self._fp8_query = self.impl.supports_quant_query_input
if not self._fp8_query:
assert self.kv_cache_dtype == "fp8_ds_mla", (
"deepseek_v32 (nvidia) on a bf16-query sparse MLA backend "
"(FlashMLA sparse) requires the fp8_ds_mla KV cache layout. "
"Launch with --kv-cache-dtype fp8_ds_mla."
)
self._fp8_kv_needs_view = self.kv_cache_dtype != "fp8_ds_mla"
fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
self._fp8_query = fp8_attention and self.impl.supports_quant_query_input
self._fp8_kv_needs_view = fp8_attention and self.kv_cache_dtype != "fp8_ds_mla"
self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False)
+2 -1
View File
@@ -6,6 +6,7 @@ import torch
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.torch_utils import is_quantized_kv_cache
# Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
# has_indexer=False path so the indexer args don't allocate every call.
@@ -473,7 +474,7 @@ def fused_norm_rope(
# --- MLA KV cache setup ---
mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla"
mla_cache_fp8 = mla_kv_cache_dtype not in ("auto", "fp8_ds_mla")
mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla
mla_num_tiles = 1
mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device)
mla_ds_rope_view = torch.empty(0, dtype=torch.bfloat16, device=device)
+5 -1
View File
@@ -1018,7 +1018,11 @@ class SpecDecodeBaseProposer:
# feedback into the next draft step.
architectures = self.draft_model_config.hf_config.architectures or []
return bool(
{"DeepSeekMTPModel", "KimiK3MTPModel"}.intersection(architectures)
{
"DeepSeekMTPModel",
"DeepseekV32MTPModel",
"KimiK3MTPModel",
}.intersection(architectures)
)
return self.method not in ("mtp", "draft_model", "dflash")