[Model][NVIDIA] Route DSA models to the CUDA non-compiled path (#52861)
This commit is contained in:
@@ -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"] = []
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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
@@ -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"
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user