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 13, 08:15 UTC
[platform] Move get_cu_count to utils (#27005)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
wangxiyuan committed
Nov 13, 2025 at 08:48 UTC
2dacd5739409847e91299e7747a142e200fdff6c
6 files changed
+28
-18
tests/kernels/quantization/test_rocm_skinny_gemms.py
+19
-5
@@ -8,6 +8,7 @@ import torch
8
import vllm._custom_ops as ops
9
from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
10
from vllm.platforms import current_platform
11
+from vllm.utils.platform_utils import get_cu_count
12
13
DTYPES = [torch.bfloat16, torch.float16]
14
# Specific (N, K, M) combinations for targeted testing
@@ -85,7 +86,7 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
86
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
87
def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
88
torch.manual_seed(seed)
88
- cu_count = current_platform.get_cu_count()
89
+ cu_count = get_cu_count()
90
91
A = torch.rand(n, k, dtype=dtype, device="cuda") - 0.5
92
B = torch.rand(m, k, dtype=dtype, device="cuda") - 0.5
@@ -102,7 +103,7 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
103
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
104
def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
105
torch.manual_seed(seed)
105
- cu_count = current_platform.get_cu_count()
106
+ cu_count = get_cu_count()
107
108
xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
109
A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
@@ -121,7 +122,7 @@ def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
122
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
123
def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
124
torch.manual_seed(seed)
124
- cu_count = current_platform.get_cu_count()
125
+ cu_count = get_cu_count()
126
127
xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
128
A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier
@@ -153,7 +154,14 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
154
ref_out = torch._scaled_mm(
155
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b
156
)
156
- out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, current_platform.get_cu_count())
157
+ out = ops.wvSplitKQ(
158
+ B,
159
+ A,
160
+ dtype,
161
+ scale_a,
162
+ scale_b,
163
+ get_cu_count(),
164
+ )
165
166
assert torch.allclose(out, ref_out, rtol=0.01)
167
@@ -180,7 +188,13 @@ def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
188
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
189
)
190
out = ops.wvSplitKQ(
183
- B, A, dtype, scale_a, scale_b, current_platform.get_cu_count(), BIAS
191
+ B,
192
+ A,
193
+ dtype,
194
+ scale_a,
195
+ scale_b,
196
+ get_cu_count(),
197
+ BIAS,
198
)
199
200
assert torch.allclose(out, ref_out, rtol=0.01)
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+2
-1
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
13
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
14
from vllm.platforms import current_platform
15
from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm, has_flashinfer
16
+from vllm.utils.platform_utils import get_cu_count
17
from vllm.utils.torch_utils import direct_register_custom_op
18
19
# Input scaling factors are no longer optional in _scaled_mm starting
@@ -200,7 +201,7 @@ def rocm_per_tensor_w8a8_scaled_mm_impl(
201
out_dtype,
202
scale_a,
203
scale_b,
203
- current_platform.get_cu_count(),
204
+ get_cu_count(),
205
bias,
206
)
207
else:
vllm/model_executor/layers/utils.py
+2
-1
@@ -11,6 +11,7 @@ from vllm import envs
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.logger import init_logger
13
from vllm.platforms import CpuArchEnum, current_platform
14
+from vllm.utils.platform_utils import get_cu_count
15
from vllm.utils.torch_utils import direct_register_custom_op
16
17
logger = init_logger(__name__)
@@ -151,7 +152,7 @@ def rocm_unquantized_gemm_impl(
152
153
x_view = x.reshape(-1, x.size(-1))
154
if m > 8 and 0 < n <= 4:
154
- cu_count = current_platform.get_cu_count()
155
+ cu_count = get_cu_count()
156
out = ops.wvSplitK(weight, x_view, cu_count, bias)
157
return out.reshape(*x.shape[:-1], weight.shape[0])
158
elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
vllm/platforms/interface.py
-7
@@ -545,13 +545,6 @@ class Platform:
545
cls._global_graph_pool = self.graph_pool_handle()
546
return cls._global_graph_pool
547
548
- @classmethod
549
- def get_cu_count(cls, device_id: int = 0) -> int:
550
- """
551
- Returns the total number of compute units (CU) on single GPU.
552
- """
553
- raise NotImplementedError
554
-
548
@classmethod
549
def get_static_graph_wrapper_cls(cls) -> str:
550
"""
vllm/platforms/rocm.py
-4
@@ -423,10 +423,6 @@ class RocmPlatform(Platform):
423
def opaque_attention_op(cls) -> bool:
424
return True
425
426
- @classmethod
427
- def get_cu_count(cls, device_id: int = 0) -> int:
428
- return torch.cuda.get_device_properties(device_id).multi_processor_count
429
-
426
@classmethod
427
def is_navi(cls) -> bool:
428
return "gfx1" in torch.cuda.get_device_properties(0).gcnArchName
vllm/utils/platform_utils.py
+5
@@ -24,6 +24,11 @@ def xpu_is_initialized() -> bool:
24
return torch.xpu.is_initialized()
25
26
27
+def get_cu_count(cls, device_id: int = 0) -> int:
28
+ """Returns the total number of compute units (CU) on single GPU."""
29
+ return torch.cuda.get_device_properties(device_id).multi_processor_count
30
+
31
+
32
def cuda_get_device_properties(
33
device, names: Sequence[str], init_cuda=False
34
) -> tuple[Any, ...]: