4d36693687 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/moe/test_deepep_deepgemm_moe.py

+7 -13

@@ -6,7 +6,6 @@ fp8 block-quantized case.

6 """

7

8 import dataclasses

9 -import importlib

9 from typing import Optional

10

11 import pytest

@@ -21,18 +20,11 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

20 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

21 per_token_group_quant_fp8)

22 from vllm.platforms import current_platform

23 +from vllm.utils import has_deep_ep, has_deep_gemm

24

25 from .utils import ProcessGroupInfo, parallel_launch

26

27 -has_deep_ep = importlib.util.find_spec("deep_ep") is not None

28 -

29 -try:

30 - import deep_gemm

31 - has_deep_gemm = True

32 -except ImportError:

33 - has_deep_gemm = False

34 -

35 -if has_deep_ep:

27 +if has_deep_ep():

28 from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501

29 DeepEPHTPrepareAndFinalize)

30 from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501

@@ -40,19 +32,21 @@ if has_deep_ep:

32

33 from .deepep_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a

34

43 -if has_deep_gemm:

35 +if has_deep_gemm():

36 + import deep_gemm

37 +

38 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (

39 BatchedDeepGemmExperts)

40 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

41 DeepGemmExperts)

42

43 requires_deep_ep = pytest.mark.skipif(

50 - not has_deep_ep,

44 + not has_deep_ep(),

45 reason="Requires deep_ep kernels",

46 )

47

48 requires_deep_gemm = pytest.mark.skipif(

55 - not has_deep_gemm,

49 + not has_deep_gemm(),

50 reason="Requires deep_gemm kernels",

51 )

52

tests/kernels/moe/test_deepep_moe.py

+3 -5

@@ -4,7 +4,6 @@ Test deepep dispatch-combine logic

4 """

5

6 import dataclasses

7 -import importlib

7 from typing import Optional, Union

8

9 import pytest

@@ -22,12 +21,11 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

21 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

22 per_token_group_quant_fp8)

23 from vllm.platforms import current_platform

24 +from vllm.utils import has_deep_ep

25

26 from .utils import ProcessGroupInfo, parallel_launch

27

28 -has_deep_ep = importlib.util.find_spec("deep_ep") is not None

29 -

30 -if has_deep_ep:

28 +if has_deep_ep():

29 from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501

30 DeepEPHTPrepareAndFinalize)

31 from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501

@@ -36,7 +34,7 @@ if has_deep_ep:

34 from .deepep_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a

35

36 requires_deep_ep = pytest.mark.skipif(

39 - not has_deep_ep,

37 + not has_deep_ep(),

38 reason="Requires deep_ep kernels",

39 )

40

vllm/distributed/device_communicators/all2all.py

+5 -5

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import importlib.util

3 from typing import TYPE_CHECKING, Any

4

5 import torch

@@ -8,6 +7,7 @@ import torch.distributed as dist

7

8 from vllm.forward_context import get_forward_context

9 from vllm.logger import init_logger

10 +from vllm.utils import has_deep_ep, has_pplx

11

12 from .base_device_communicator import All2AllManagerBase, Cache

13

@@ -80,8 +80,8 @@ class PPLXAll2AllManager(All2AllManagerBase):

80 """

81

82 def __init__(self, cpu_group):

83 - has_pplx = importlib.util.find_spec("pplx_kernels") is not None

84 - assert has_pplx, "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install pplx_kernels." # noqa

83 + assert has_pplx(

84 + ), "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install pplx_kernels." # noqa

85 super().__init__(cpu_group)

86

87 if self.internode:

@@ -133,8 +133,8 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):

133 """

134

135 def __init__(self, cpu_group):

136 - has_deepep = importlib.util.find_spec("deep_ep") is not None

137 - assert has_deepep, "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install DeepEP kernels." # noqa

136 + assert has_deep_ep(

137 + ), "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install DeepEP kernels." # noqa

138 super().__init__(cpu_group)

139 self.handle_cache = Cache()

140

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

-3

@@ -1,5 +1,4 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -import importlib.util

2 from typing import Optional

3

4 import torch

@@ -11,8 +10,6 @@ from vllm.triton_utils import tl, triton

10

11 logger = init_logger(__name__)

12

14 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None

15 -

13

14 @triton.jit

15 def _silu_mul_fp8_quant_deep_gemm(

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+5 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import functools

4 -import importlib.util

4 from typing import Optional

5

6 import torch

@@ -12,14 +11,13 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (

11 _moe_permute)

12 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

13 MoEPrepareAndFinalizeNoEP)

15 -from vllm.model_executor.layers.fused_moe.utils import (

16 - _resize_cache, per_token_group_quant_fp8)

17 -from vllm.utils import round_up

14 +from vllm.model_executor.layers.fused_moe.utils import _resize_cache

15 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (

16 + per_token_group_quant_fp8)

17 +from vllm.utils import has_deep_gemm, round_up

18

19 logger = init_logger(__name__)

20

21 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None

22 -

21

22 @functools.cache

23 def deep_gemm_block_shape() -> list[int]:

@@ -41,7 +39,7 @@ def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,

39 gemm kernel. All of M, N, K and the quantization block_shape must be

40 aligned by `dg.get_m_alignment_for_contiguous_layout()`.

41 """

44 - if not has_deep_gemm:

42 + if not has_deep_gemm():

43 logger.debug("DeepGemm disabled: deep_gemm not available.")

44 return False

45

vllm/model_executor/layers/fused_moe/layer.py

+3 -7

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -import importlib

4 from abc import abstractmethod

5 from collections.abc import Iterable

6 from dataclasses import dataclass

@@ -32,10 +31,7 @@ from vllm.model_executor.layers.quantization.base_config import (

31 from vllm.model_executor.utils import set_weight_attrs

32 from vllm.platforms import current_platform

33 from vllm.platforms.interface import CpuArchEnum

35 -from vllm.utils import direct_register_custom_op

36 -

37 -has_pplx = importlib.util.find_spec("pplx_kernels") is not None

38 -has_deepep = importlib.util.find_spec("deep_ep") is not None

34 +from vllm.utils import direct_register_custom_op, has_deep_ep, has_pplx

35

36 if current_platform.is_cuda_alike():

37 from .fused_batched_moe import BatchedTritonExperts

@@ -43,9 +39,9 @@ if current_platform.is_cuda_alike():

39 from .modular_kernel import (FusedMoEModularKernel,

40 FusedMoEPermuteExpertsUnpermute,

41 FusedMoEPrepareAndFinalize)

46 - if has_pplx:

42 + if has_pplx():

43 from .pplx_prepare_finalize import PplxPrepareAndFinalize

48 - if has_deepep:

44 + if has_deep_ep():

45 from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize

46 from .deepep_ll_prepare_finalize import (DEEPEP_QUANT_BLOCK_SIZE,

47 DeepEPLLPrepareAndFinalize)

vllm/model_executor/layers/fused_moe/utils.py

+1 -1

@@ -104,4 +104,4 @@ def find_free_port():

104 with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:

105 s.bind(('', 0))

106 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

107 - return s.getsockname()[1]

\ No newline at end of file

107 + return s.getsockname()[1]

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+3 -5

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import enum

5 -import importlib

5 from enum import Enum

6 from typing import Callable, Optional

7

@@ -29,13 +28,12 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

28 from vllm.model_executor.utils import set_weight_attrs

29 from vllm.platforms import current_platform

30 from vllm.scalar_type import scalar_types

32 -

33 -has_pplx = importlib.util.find_spec("pplx_kernels") is not None

31 +from vllm.utils import has_pplx

32

33 if current_platform.is_cuda_alike():

34 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

35 BatchedPrepareAndFinalize)

38 - if has_pplx:

36 + if has_pplx():

37 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (

38 PplxPrepareAndFinalize)

39

@@ -577,7 +575,7 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):

575 use_batched_format=True,

576 )

577

580 - if has_pplx and isinstance(

578 + if has_pplx() and isinstance(

579 prepare_finalize,

580 (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):

581 # no expert_map support in this case

vllm/model_executor/layers/quantization/deepgemm.py

+2 -4

@@ -1,15 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -import importlib.util

2 import logging

3

4 import torch

5

6 from vllm.platforms import current_platform

7 from vllm.triton_utils import triton

9 -from vllm.utils import direct_register_custom_op

8 +from vllm.utils import direct_register_custom_op, has_deep_gemm

9

11 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None

12 -if has_deep_gemm:

10 +if has_deep_gemm():

11 import deep_gemm

12

13 logger = logging.getLogger(__name__)

vllm/model_executor/layers/quantization/fp8.py

+2 -4

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import functools

5 -import importlib.util

5 from typing import Any, Callable, Optional, Union

6

7 import torch

@@ -38,13 +37,12 @@ from vllm.model_executor.parameter import (BlockQuantScaleParameter,

37 from vllm.model_executor.utils import set_weight_attrs

38 from vllm.platforms import current_platform

39 from vllm.scalar_type import scalar_types

40 +from vllm.utils import has_deep_gemm

41

42 ACTIVATION_SCHEMES = ["static", "dynamic"]

43

44 logger = init_logger(__name__)

45

46 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None

47 -

46

47 def _is_col_major(x: torch.Tensor) -> bool:

48 assert x.dim() == 3

@@ -451,7 +449,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

449 # Check for DeepGemm support.

450 self.allow_deep_gemm = False

451 if envs.VLLM_USE_DEEP_GEMM:

454 - if not has_deep_gemm:

452 + if not has_deep_gemm():

453 logger.warning_once("Failed to import DeepGemm kernels.")

454 elif not self.block_quant:

455 logger.warning_once("Model is not block quantized. Not using "

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+2 -4

@@ -3,7 +3,6 @@

3

4 # Adapted from https://github.com/sgl-project/sglang/pull/2575

5 import functools

6 -import importlib.util

6 import json

7 import os

8 from typing import Any, Callable, Optional, Union

@@ -19,10 +18,9 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

18 CUTLASS_BLOCK_FP8_SUPPORTED)

19 from vllm.platforms import current_platform

20 from vllm.triton_utils import tl, triton

22 -from vllm.utils import cdiv, direct_register_custom_op

21 +from vllm.utils import cdiv, direct_register_custom_op, has_deep_gemm

22

23 logger = init_logger(__name__)

25 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None

24

25

26 def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:

@@ -109,7 +107,7 @@ def should_use_deepgemm(output_dtype: torch.dtype, weight: torch.Tensor):

107 """

108

109 return (current_platform.is_cuda()

112 - and current_platform.is_device_capability(90) and has_deep_gemm

110 + and current_platform.is_device_capability(90) and has_deep_gemm()

111 and envs.VLLM_USE_DEEP_GEMM and output_dtype == torch.bfloat16

112 and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)

113

vllm/utils.py

+28

@@ -2929,3 +2929,31 @@ def is_torch_equal_or_newer(target: str) -> bool:

2929 def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool:

2930 torch_version = version.parse(torch_version)

2931 return torch_version >= version.parse(target)

2932 +

2933 +

2934 +@cache

2935 +def _has_module(module_name: str) -> bool:

2936 + """Return True if *module_name* can be found in the current environment.

2937 +

2938 + The result is cached so that subsequent queries for the same module incur

2939 + no additional overhead.

2940 + """

2941 + return importlib.util.find_spec(module_name) is not None

2942 +

2943 +

2944 +def has_pplx() -> bool:

2945 + """Whether the optional `pplx_kernels` package is available."""

2946 +

2947 + return _has_module("pplx_kernels")

2948 +

2949 +

2950 +def has_deep_ep() -> bool:

2951 + """Whether the optional `deep_ep` package is available."""

2952 +

2953 + return _has_module("deep_ep")

2954 +

2955 +

2956 +def has_deep_gemm() -> bool:

2957 + """Whether the optional `deep_gemm` package is available."""

2958 +

2959 + return _has_module("deep_gemm")

\ No newline at end of file

[Refactor] Create a function util and cache the results for `has_deepgemm`, `has_deepep`, `has_pplx` (#20187)