A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 08:00 UTC
[Model Refactoring] Move deepseek_v4_ops to models/deepseek_v4 [3/N] (#43073)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 19, 2026 at 00:52 UTC
b14be81c1f63b70668d26d65a377b6383fbca936
20 files changed
+34
-29
.github/CODEOWNERS
-1
@@ -154,7 +154,6 @@ mkdocs.yaml @hmellor
154
155
# DeepseekV4-specific files
156
/vllm/models/deepseek_v4 @zyongye
157
-/vllm/v1/attention/ops/deepseek_v4_ops @zyongye
157
/vllm/model_executor/layers/sparse_attn_indexer.py @zyongye
158
159
# Mistral-specific files
tests/kernels/core/test_fused_q_kv_rmsnorm.py
+1
-1
@@ -13,8 +13,8 @@ from __future__ import annotations
13
import pytest
14
import torch
15
16
+from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm
17
from vllm.platforms import current_platform
17
-from vllm.v1.attention.ops.deepseek_v4_ops import fused_q_kv_rmsnorm
18
19
pytestmark = pytest.mark.skipif(
20
not current_platform.is_cuda_alike(),
tests/kernels/test_compressor_kv_cache.py
+2
-2
@@ -17,11 +17,11 @@ import pytest
17
import torch
18
19
from vllm import _custom_ops as ops
20
-from vllm.v1.attention.ops.deepseek_v4_ops import (
20
+from vllm.models.deepseek_v4.common.ops import (
21
dequantize_and_gather_k_cache,
22
quantize_and_insert_k_cache,
23
)
24
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
24
+from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
25
_fused_kv_compress_norm_rope_insert_indexer_attn,
26
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
27
)
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+1
-1
@@ -19,7 +19,7 @@ The kernel is imported via
19
import pytest
20
import torch
21
22
-from vllm.v1.attention.ops.deepseek_v4_ops import (
22
+from vllm.models.deepseek_v4.common.ops import (
23
dequantize_and_gather_k_cache,
24
quantize_and_insert_k_cache,
25
)
tests/kernels/test_fused_indexer_q_rope_quant.py
+2
-4
@@ -23,10 +23,8 @@ from vllm import _custom_ops as ops
23
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
24
per_token_group_quant_fp8,
25
)
26
+from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant
27
from vllm.utils.import_utils import has_cutedsl
27
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import (
28
- fused_indexer_q_rope_quant,
29
-)
28
29
HEAD_DIM = 128
30
ROPE_DIM = 64
@@ -156,7 +154,7 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
154
# by patching the dispatcher's has_cutedsl() binding to return False.
155
cutedsl_patch = (
156
mock.patch(
159
- "vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q.has_cutedsl",
157
+ "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl",
158
return_value=False,
159
)
160
if not use_cutedsl
tests/kernels/test_fused_inv_rope_fp8_quant.py
+1
-1
@@ -21,7 +21,7 @@ Usage:
21
import pytest
22
import torch
23
24
-from vllm.v1.attention.ops.deepseek_v4_ops import fused_inv_rope_fp8_quant
24
+from vllm.models.deepseek_v4.common.ops import fused_inv_rope_fp8_quant
25
26
# -- Default dimensions matching DeepSeek V3/V4 --------------------------
27
HEAD_DIM = 512
vllm/models/deepseek_v4/attention.py
+3
-3
@@ -19,9 +19,7 @@ from vllm.model_executor.layers.linear import (
19
ReplicatedLinear,
20
)
21
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
22
-from vllm.utils.deep_gemm import fp8_einsum
23
-from vllm.utils.torch_utils import direct_register_custom_op
24
-from vllm.v1.attention.ops.deepseek_v4_ops import (
22
+from vllm.models.deepseek_v4.common.ops import (
23
combine_topk_swa_indices,
24
compute_global_topk_indices_and_lens,
25
dequantize_and_gather_k_cache,
@@ -29,6 +27,8 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
27
fused_inv_rope_fp8_quant,
28
fused_q_kv_rmsnorm,
29
)
30
+from vllm.utils.deep_gemm import fp8_einsum
31
+from vllm.utils.torch_utils import direct_register_custom_op
32
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
33
34
if TYPE_CHECKING:
vllm/models/deepseek_v4/common/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v4/common/ops/__init__.py
renamed
vllm/models/deepseek_v4/common/ops/cache_utils.py
renamed
+3
-1
@@ -366,7 +366,9 @@ def dequantize_and_gather_k_cache(
366
) -> None:
367
if has_cutedsl():
368
# lazily import, otherwise some tests fail due to CUDA driver init failure.
369
- from .dequant_gather_k_cutedsl import dequantize_and_gather_k_cache_cutedsl
369
+ from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
370
+ dequantize_and_gather_k_cache_cutedsl,
371
+ )
372
373
dequantize_and_gather_k_cache_cutedsl(
374
out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
renamed
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
renamed
+2
-2
@@ -346,7 +346,7 @@ def fused_indexer_q_rope_quant(
346
)
347
if has_cutedsl():
348
# lazily import, otherwise some tests fail due to CUDA driver init failure.
349
- from .fused_indexer_q_cutedsl import (
349
+ from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
350
fused_indexer_q_rope_quant_mxfp4_cutedsl,
351
)
352
@@ -400,7 +400,7 @@ def fused_indexer_q_rope_quant(
400
index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
401
if has_cutedsl():
402
# lazily import, otherwise some tests fail due to CUDA driver init failure.
403
- from .fused_indexer_q_cutedsl import (
403
+ from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
404
fused_indexer_q_rope_quant_fp8_cutedsl,
405
)
406
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
renamed
vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py
renamed
vllm/models/deepseek_v4/compressor.py
+6
-10
@@ -11,9 +11,13 @@ from vllm.config import VllmConfig, get_current_vllm_config
11
from vllm.forward_context import get_forward_context
12
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13
from vllm.model_executor.layers.layernorm import RMSNorm
14
-from vllm.model_executor.layers.linear import (
15
- MergedColumnParallelLinear,
14
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear
15
+from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
16
+ _fused_kv_compress_norm_rope_insert_indexer_attn,
17
+ _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
18
+ _fused_kv_compress_norm_rope_insert_sparse_attn,
19
)
20
+from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
21
from vllm.platforms import current_platform
22
from vllm.triton_utils import tl, triton
23
from vllm.v1.attention.backend import (
@@ -23,14 +27,6 @@ from vllm.v1.attention.backend import (
27
CommonAttentionMetadata,
28
MultipleOf,
29
)
26
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
27
- _fused_kv_compress_norm_rope_insert_indexer_attn,
28
- _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
29
- _fused_kv_compress_norm_rope_insert_sparse_attn,
30
-)
31
-from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import (
32
- MXFP4_BLOCK_SIZE,
33
-)
30
from vllm.v1.kv_cache_interface import (
31
KVCacheSpec,
32
MLAAttentionSpec,
vllm/models/deepseek_v4/nvidia/ops/__init__.py
new
+8
@@ -0,0 +1,8 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NVIDIA-only (cutedsl/cutlass) kernels for DeepSeek V4.
4
+
5
+These modules import ``cutlass``/``cutedsl`` at module top level, so they must
6
+not be imported on non-CUDA platforms. Callers should gate on
7
+``vllm.utils.import_utils.has_cutedsl()`` before importing from here.
8
+"""
vllm/models/deepseek_v4/nvidia/ops/cutedsl_utils.py
renamed
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
renamed
+1
-1
@@ -11,7 +11,7 @@ from cutlass import BFloat16, Int32, Uint8, Uint32
11
from cutlass.cute.nvgpu import cpasync
12
from quack.compile_utils import make_fake_tensor
13
14
-from vllm.v1.attention.ops.deepseek_v4_ops.cutedsl_utils import (
14
+from vllm.models.deepseek_v4.nvidia.ops.cutedsl_utils import (
15
_bf16x2_mul,
16
_fp8x4_to_bf16x4,
17
)
vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
renamed
+1
-1
@@ -9,7 +9,7 @@ from cuda.bindings.driver import CUstream
9
from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr
10
from quack.compile_utils import make_fake_tensor
11
12
-from vllm.v1.attention.ops.deepseek_v4_ops.cutedsl_utils import (
12
+from vllm.models.deepseek_v4.nvidia.ops.cutedsl_utils import (
13
_bf16x2_abs,
14
_bf16x2_max,
15
_bf16x2_to_fp32,
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse_dsv4.py
+1
-1
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, cast
7
import torch
8
9
from vllm.forward_context import get_forward_context
10
+from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache
11
from vllm.triton_utils import tl, triton
12
from vllm.v1.attention.backend import (
13
AttentionLayer,
@@ -22,7 +23,6 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
23
DeepseekSparseSWAMetadata,
24
DeepseekSparseSWAMetadataBuilder,
25
)
25
-from vllm.v1.attention.ops.deepseek_v4_ops import dequantize_and_gather_k_cache
26
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
27
build_ragged_indices_from_dense,
28
rocm_sparse_attn_decode,