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 2, 18:15 UTC
Fix early CUDA init (#43791)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
May 27, 2026 at 17:30 UTC
41688e2dc7f52b4f0c22ebe5470e340bbc7e0d6f
5 files changed
+12
-21
vllm/models/deepseek_v4/common/ops/cache_utils.py
+1
-1
@@ -366,7 +366,7 @@ 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 vllm.models.deepseek_v4.nvidia.ops import (
369
+ from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
370
dequantize_and_gather_k_cache_cutedsl,
371
)
372
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+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 vllm.models.deepseek_v4.nvidia.ops 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 vllm.models.deepseek_v4.nvidia.ops 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/compressor.py
+3
-1
@@ -338,7 +338,9 @@ class DeepseekCompressor(nn.Module):
338
if current_platform.is_cuda():
339
# NVIDIA GPUs.
340
if self.head_dim == 512:
341
- from .nvidia.ops import compress_norm_rope_store_cutedsl
341
+ from .nvidia.ops.sparse_attn_compress_cutedsl import (
342
+ compress_norm_rope_store_cutedsl,
343
+ )
344
345
# Main compressor path.
346
# Use a cutedsl kernel for better performance.
vllm/models/deepseek_v4/nvidia/model.py
+1
-1
@@ -59,7 +59,7 @@ from vllm.models.deepseek_v4.attention import (
59
DeepseekV4MLAModules,
60
DeepseekV4MultiHeadLatentAttentionWrapper,
61
)
62
-from vllm.models.deepseek_v4.nvidia.ops import prepare_megamoe_inputs
62
+from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
63
from vllm.platforms import current_platform
64
from vllm.sequence import IntermediateTensors
65
from vllm.utils.torch_utils import direct_register_custom_op
vllm/models/deepseek_v4/nvidia/ops/__init__.py
+5
-16
@@ -5,20 +5,9 @@
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
-"""
9
-
10
-from .dequant_gather_k_cutedsl import dequantize_and_gather_k_cache_cutedsl
11
-from .fused_indexer_q_cutedsl import (
12
- fused_indexer_q_rope_quant_fp8_cutedsl,
13
- fused_indexer_q_rope_quant_mxfp4_cutedsl,
14
-)
15
-from .prepare_megamoe import prepare_megamoe_inputs
16
-from .sparse_attn_compress_cutedsl import compress_norm_rope_store_cutedsl
8
18
-__all__ = [
19
- "compress_norm_rope_store_cutedsl",
20
- "dequantize_and_gather_k_cache_cutedsl",
21
- "fused_indexer_q_rope_quant_fp8_cutedsl",
22
- "fused_indexer_q_rope_quant_mxfp4_cutedsl",
23
- "prepare_megamoe_inputs",
24
-]
9
+This ``__init__`` deliberately imports nothing: re-exporting the cutedsl
10
+modules here would eagerly ``import cutlass`` (initializing the CUDA driver) for
11
+anyone who imports ``vllm.models.deepseek_v4``, breaking forked subprocesses.
12
+Import the leaf modules directly under a ``has_cutedsl()``/``is_cuda()`` gate.
13
+"""