8ad54a991b · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/attention/test_cutlass_mla_decode.py

+2 -2

@@ -9,6 +9,7 @@ import torch

9 import vllm._custom_ops as ops

10 from vllm.platforms import current_platform

11 from vllm.triton_utils import triton

12 +from vllm.utils.platform_utils import num_compute_units

13

14

15 def cal_diff(

@@ -124,8 +125,7 @@ def test_cutlass_mla_decode(

125 q_pe = q_pe_padded

126

127 kv_cache_flat = blocked_k.squeeze(2)

127 - device_properties = torch.cuda.get_device_properties(torch.device("cuda:0"))

128 - sm_count = device_properties.multi_processor_count

128 + sm_count = num_compute_units(device.index)

129 workspace_size = ops.sm100_cutlass_mla_get_workspace_size(

130 max_seqlen * block_size, b, sm_count, num_kv_splits=1

131 )

tests/kernels/quantization/test_allspark_gemm.py

+2 -1

@@ -13,6 +13,7 @@ from vllm.model_executor.layers.quantization.utils.allspark_utils import (

13 from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights

14 from vllm.platforms import current_platform

15 from vllm.scalar_type import scalar_types

16 +from vllm.utils.platform_utils import num_compute_units

17

18

19 def is_gptq_allspark_supported(min_capability: int, max_capability: int) -> bool:

@@ -78,7 +79,7 @@ def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype):

79 if has_zp:

80 zp = zp.to(dtype)

81 properties = torch.cuda.get_device_properties(qw.device.index)

81 - sm_count = properties.multi_processor_count

82 + sm_count = num_compute_units(qw.device.index)

83 sm_version = properties.major * 10 + properties.minor

84

85 n_32align = (n + 32 - 1) // 32 * 32

tests/kernels/quantization/test_rocm_skinny_gemms.py

+6 -6

@@ -9,7 +9,7 @@ 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.platforms.rocm import on_gfx950

12 -from vllm.utils.platform_utils import get_cu_count

12 +from vllm.utils.platform_utils import num_compute_units

13

14 DTYPES = [torch.bfloat16, torch.float16]

15 BIAS_MODES = [0, 1, 2]

@@ -121,7 +121,7 @@ def pad_fp8(weight):

121 @pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")

122 def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):

123 torch.manual_seed(seed)

124 - cu_count = get_cu_count()

124 + cu_count = num_compute_units()

125

126 # Next ^2 of n

127 N_p2 = 1 << (n - 1).bit_length()

@@ -186,7 +186,7 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):

186 @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")

187 def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):

188 torch.manual_seed(seed)

189 - cu_count = get_cu_count()

189 + cu_count = num_compute_units()

190

191 A = torch.rand(n, k, dtype=dtype, device="cuda") - 0.5

192 B = torch.rand(m, k, dtype=dtype, device="cuda") - 0.5

@@ -203,7 +203,7 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):

203 @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")

204 def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):

205 torch.manual_seed(seed)

206 - cu_count = get_cu_count()

206 + cu_count = num_compute_units()

207

208 xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas

209 A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier

@@ -222,7 +222,7 @@ def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):

222 @pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")

223 def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):

224 torch.manual_seed(seed)

225 - cu_count = get_cu_count()

225 + cu_count = num_compute_units()

226

227 xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas

228 A = (torch.rand(n, k, dtype=dtype, device="cuda") - 0.5) * xavier

@@ -267,7 +267,7 @@ def test_rocm_wvsplitk_fp8_kernel(

267 ref_out = torch._scaled_mm(

268 A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS

269 )

270 - out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, get_cu_count(), BIAS)

270 + out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, num_compute_units(), BIAS)

271

272 if xnorm:

273 torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)

vllm/model_executor/kernels/linear/mixed_precision/allspark.py

+2 -1

@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.allspark_utils import (

11 check_allspark_supported_dtype_shape,

12 )

13 from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_

14 +from vllm.utils.platform_utils import num_compute_units

15

16 from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig

17

@@ -45,7 +46,7 @@ class AllSparkLinearKernel(MPLinearKernel):

46

47 # prepare the parameters required for the kernel

48 properties = torch.cuda.get_device_properties(device.index)

48 - sm_count = properties.multi_processor_count

49 + sm_count = num_compute_units(device.index)

50 sm_version = properties.major * 10 + properties.minor

51 gemm_args = {}

52 gemm_args["sm_count"] = sm_count

vllm/model_executor/kernels/linear/scaled_mm/rocm.py

+2 -2

@@ -7,7 +7,7 @@ import torch

7 import vllm.envs as envs

8 from vllm import _custom_ops as ops

9 from vllm.platforms import current_platform

10 -from vllm.utils.platform_utils import get_cu_count

10 +from vllm.utils.platform_utils import num_compute_units

11 from vllm.utils.torch_utils import direct_register_custom_op

12

13 from .ScaledMMLinearKernel import (

@@ -36,7 +36,7 @@ def rocm_per_tensor_float_w8a8_scaled_mm_impl(

36 out_dtype,

37 As,

38 Bs,

39 - get_cu_count(),

39 + num_compute_units(),

40 bias,

41 )

42 # Fallback

vllm/model_executor/layers/batch_invariant.py

+2 -1

@@ -9,6 +9,7 @@ import torch

9 from vllm.logger import init_logger

10 from vllm.platforms import current_platform

11 from vllm.triton_utils import tl, triton

12 +from vllm.utils.platform_utils import num_compute_units

13 from vllm.utils.torch_utils import is_torch_equal_or_newer

14 from vllm.v1.attention.backends.registry import AttentionBackendEnum

15

@@ -147,7 +148,7 @@ def matmul_persistent(

148 assert bias is None or bias.dim() == 1, (

149 "Currently assuming bias is 1D, let Horace know if you run into this"

150 )

150 - NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count

151 + NUM_SMS = num_compute_units(a.device.index)

152 M, K = a.shape

153 K, N = b.shape

154 dtype = a.dtype

vllm/model_executor/layers/fla/ops/layernorm_guard.py

+2 -10

@@ -13,8 +13,6 @@

13 # This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling.

14 # The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.

15

16 -from functools import lru_cache

17 -

16 import torch

17 import torch.nn as nn

18 import torch.nn.functional as F

@@ -22,6 +20,7 @@ from einops import rearrange

20

21 from vllm.triton_utils import tl, triton

22 from vllm.utils.math_utils import cdiv, next_power_of_2

23 +from vllm.utils.platform_utils import num_compute_units

24

25 from .utils import input_guard

26

@@ -162,15 +161,8 @@ def layer_norm_fwd_kernel(

161 tl.store(Y_base, y, mask=mask)

162

163

165 -@lru_cache

166 -def _get_sm_count(device: torch.device) -> int:

167 - """Get and cache the SM count for a given device."""

168 - props = torch.cuda.get_device_properties(device)

169 - return props.multi_processor_count

170 -

171 -

164 def calc_rows_per_block(M: int, device: torch.device) -> int:

173 - sm_count = _get_sm_count(device)

165 + sm_count = num_compute_units(device.index)

166 rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count))

167 rows_per_block = min(rows_per_block, 4)

168 return rows_per_block

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

+2 -1

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (

16 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

17 from vllm.platforms import current_platform

18 from vllm.scalar_type import ScalarType, scalar_types

19 +from vllm.utils.platform_utils import num_compute_units

20

21 from .quant_utils import pack_cols, unpack_cols

22

@@ -271,7 +272,7 @@ def marlin_make_workspace_new(

272 ) -> torch.Tensor:

273 # In the new marlin kernel, we use the num of threadblocks as workspace

274 # size. The num of threadblocks is sms_count * max_blocks_per_sm.

274 - sms = torch.cuda.get_device_properties(device).multi_processor_count

275 + sms = num_compute_units(device.index)

276 return torch.zeros(

277 sms * max_blocks_per_sm, dtype=torch.int, device=device, requires_grad=False

278 )

vllm/model_executor/layers/utils.py

+3 -3

@@ -11,7 +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

14 +from vllm.utils.platform_utils import num_compute_units

15 from vllm.utils.torch_utils import direct_register_custom_op

16

17 logger = init_logger(__name__)

@@ -149,7 +149,7 @@ def rocm_unquantized_gemm_impl(

149 m = weight.shape[0]

150 k = weight.shape[1]

151

152 - cu_count = get_cu_count()

152 + cu_count = num_compute_units()

153 if use_aiter_triton_gemm(n, m, k, x.dtype):

154 from aiter.ops.triton.gemm_a16w16 import gemm_a16w16

155

@@ -199,7 +199,7 @@ def rocm_unquantized_gemm_impl(

199

200 x_view = x.reshape(-1, x.size(-1))

201 if m > 8 and 0 < n <= 4:

202 - cu_count = get_cu_count()

202 + cu_count = num_compute_units()

203 out = ops.wvSplitK(weight, x_view, cu_count, bias)

204 return out.reshape(*x.shape[:-1], weight.shape[0])

205 elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:

vllm/model_executor/warmup/deep_gemm_warmup.py

+2 -1

@@ -26,6 +26,7 @@ from vllm.utils.deep_gemm import (

26 m_grouped_fp8_gemm_nt_contiguous,

27 )

28 from vllm.utils.math_utils import cdiv

29 +from vllm.utils.platform_utils import num_compute_units

30

31

32 def _generate_optimal_warmup_m_values(

@@ -44,7 +45,7 @@ def _generate_optimal_warmup_m_values(

45 # DeepGEMM's possible block sizes

46 block_ms = [64, 128, 256]

47 block_ns = list(range(16, min(257, n + 1), 16))

47 - num_sms = torch.cuda.get_device_properties(device).multi_processor_count

48 + num_sms = num_compute_units(device.index)

49

50 m_values = set()

51

vllm/platforms/cuda.py

+4

@@ -538,6 +538,10 @@ class CudaPlatformBase(Platform):

538 def support_static_graph_mode(cls) -> bool:

539 return True

540

541 + @classmethod

542 + def num_compute_units(cls, device_id=0):

543 + return torch.cuda.get_device_properties(device_id).multi_processor_count

544 +

545

546 # NVML utils

547 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,

vllm/platforms/interface.py

+10

@@ -692,6 +692,16 @@ class Platform:

692 """

693 return {}

694

695 + @classmethod

696 + def num_compute_units(cls, device_id: int = 0) -> int:

697 + """

698 + Get the number of compute units for the current platform.

699 + (NVIDIA SM / AMD CU / Intel EU)

700 + """

701 + raise NotImplementedError(

702 + "num_compute_units is not implemented for the current platform."

703 + )

704 +

705

706 class UnspecifiedPlatform(Platform):

707 _enum = PlatformEnum.UNSPECIFIED

vllm/platforms/rocm.py

+4

@@ -682,3 +682,7 @@ class RocmPlatform(Platform):

682 @classmethod

683 def support_static_graph_mode(cls) -> bool:

684 return True

685 +

686 + @classmethod

687 + def num_compute_units(cls, device_id=0):

688 + return torch.cuda.get_device_properties(device_id).multi_processor_count

vllm/platforms/xpu.py

+4

@@ -277,3 +277,7 @@ class XPUPlatform(Platform):

277 """Copy blocks from XPU to host (CPU)."""

278 _src_cache = src_cache[:, src_block_indices]

279 dst_cache[:, dst_block_indices] = _src_cache.cpu()

280 +

281 + @classmethod

282 + def num_compute_units(cls, device_id: int = 0) -> int:

283 + return torch.xpu.get_device_properties(device_id).max_compute_units

vllm/utils/platform_utils.py

+8 -5

@@ -24,11 +24,6 @@ def xpu_is_initialized() -> bool:

24 return torch.xpu.is_initialized()

25

26

27 -def get_cu_count(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 -

27 def cuda_get_device_properties(

28 device, names: Sequence[str], init_cuda=False

29 ) -> tuple[Any, ...]:

@@ -57,3 +52,11 @@ def is_uva_available() -> bool:

52 # UVA requires pinned memory.

53 # TODO: Add more requirements for UVA if needed.

54 return is_pin_memory_available()

55 +

56 +

57 +@cache

58 +def num_compute_units(device_id: int = 0) -> int:

59 + """Get the number of compute units of the current device."""

60 + from vllm.platforms import current_platform

61 +

62 + return current_platform.num_compute_units(device_id)

vllm/v1/attention/backends/mla/cutlass_mla.py

+2 -2

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

16 MLACommonMetadataBuilder,

17 )

18 from vllm.platforms.interface import DeviceCapability

19 +from vllm.utils.platform_utils import num_compute_units

20 from vllm.v1.attention.backend import (

21 AttentionCGSupport,

22 AttentionLayer,

@@ -74,8 +75,7 @@ class SM100Workspace:

75

76 # Pre-compute sm_count to avoid recomputing it. Use device 0 as a proxy

77 # (assumes all devices are similar)

77 - properties = torch.cuda.get_device_properties(torch.device("cuda:0"))

78 - self._sm_count = properties.multi_processor_count

78 + self._sm_count = num_compute_units(0)

79

80 def get_buf(self):

81 return self._workspace_buf

vllm/v1/attention/backends/mla/flashmla.py

+2 -2

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.batch_invariant import (

21 vllm_is_batch_invariant,

22 )

23 from vllm.platforms.interface import DeviceCapability

24 +from vllm.utils.platform_utils import num_compute_units

25 from vllm.v1.attention.backend import (

26 AttentionCGSupport,

27 AttentionLayer,

@@ -130,8 +131,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

131 self.cg_buf_num_splits = None

132 self.is_fp8_kvcache = vllm_config.cache_config.cache_dtype.startswith("fp8")

133

133 - device_properties = torch.cuda.get_device_properties(self.device)

134 - num_sms = device_properties.multi_processor_count

134 + num_sms = num_compute_units(self.device.index)

135

136 if self.compilation_config.cudagraph_mode.has_full_cudagraphs():

137 self.cg_buf_tile_scheduler_metadata = torch.zeros(

vllm/v1/attention/backends/mla/flashmla_sparse.py

+2 -2

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

15 )

16 from vllm.platforms import current_platform

17 from vllm.platforms.interface import DeviceCapability

18 +from vllm.utils.platform_utils import num_compute_units

19 from vllm.v1.attention.backend import (

20 AttentionBackend,

21 AttentionCGSupport,

@@ -237,8 +238,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

238 # DeepGEMM indexer constraint (fp8_paged_mqa_logits only supports next_n <= 2)

239 self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

240

240 - props = torch.cuda.get_device_properties(device)

241 - sm_count = props.multi_processor_count

241 + sm_count = num_compute_units(device.index)

242

243 self.num_heads = self.model_config.get_num_attention_heads(parallel_config)

244 self.mla_dims = get_mla_dims(self.model_config)

vllm/v1/attention/backends/mla/indexer.py

+2 -2

@@ -9,6 +9,7 @@ from vllm.config import VllmConfig

9 from vllm.logger import init_logger

10 from vllm.platforms import current_platform

11 from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata, has_deep_gemm

12 +from vllm.utils.platform_utils import num_compute_units

13 from vllm.v1.attention.backend import (

14 AttentionBackend,

15 AttentionCGSupport,

@@ -219,8 +220,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

220 )

221 self.reorder_batch_threshold += self.num_speculative_tokens

222

222 - props = torch.cuda.get_device_properties(self.device)

223 - sm_count = props.multi_processor_count

223 + sm_count = num_compute_units(self.device.index)

224 self.num_sms = sm_count

225

226 self.decode_lens_buffer = torch.empty(

vllm/v1/attention/backends/rocm_aiter_fa.py

+2 -2

@@ -13,7 +13,7 @@ from vllm.logger import init_logger

13 from vllm.model_executor.layers.attention import Attention

14 from vllm.platforms import current_platform

15 from vllm.utils.math_utils import cdiv

16 -from vllm.utils.platform_utils import get_cu_count

16 +from vllm.utils.platform_utils import num_compute_units

17 from vllm.v1.attention.backend import (

18 AttentionBackend,

19 AttentionCGSupport,

@@ -38,7 +38,7 @@ if current_platform.is_rocm():

38 return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))

39

40 def num_programs(total_tokens):

41 - return min(total_tokens, get_cu_count())

41 + return min(total_tokens, num_compute_units())

42

43 @triton.jit

44 def cp_mha_gather_cache_kernel(

vllm/v1/sample/ops/topk_topp_triton.py

+2 -1

@@ -13,6 +13,7 @@ import torch

13

14 from vllm.triton_utils import tl, triton

15 from vllm.utils.math_utils import next_power_of_2

16 +from vllm.utils.platform_utils import num_compute_units

17

18 _TRITON_TABLE_CACHE: dict[tuple[torch.device], tuple[torch.Tensor, torch.Tensor]] = {}

19 _TRITON_BUFFER_CACHE: dict[tuple[torch.device, torch.dtype, int], torch.Tensor] = {}

@@ -988,7 +989,7 @@ def apply_top_k_top_p_triton(

989 else:

990 p_ptr = logits # Dummy pointer (won't be read)

991

991 - num_sm = torch.cuda.get_device_properties(logits.device).multi_processor_count

992 + num_sm = num_compute_units(logits.device.index)

993 NUM_PROGRAMS = min(num_sm, batch_size)

994

995 # Cache per-Triton Program buffer on each device.

vllm/v1/worker/gpu_model_runner.py

+3 -3

@@ -98,7 +98,7 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds

98 from vllm.utils.math_utils import cdiv, round_up

99 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib

100 from vllm.utils.nvtx_pytorch_hooks import PytHooks

101 -from vllm.utils.platform_utils import is_pin_memory_available

101 +from vllm.utils.platform_utils import is_pin_memory_available, num_compute_units

102 from vllm.utils.torch_utils import (

103 get_dtype_size,

104 kv_cache_dtype_str_to_dtype,

@@ -909,8 +909,8 @@ class GPUModelRunner(

909 # Note: used for model runner override.

910 def _init_device_properties(self) -> None:

911 """Initialize attributes from torch.cuda.get_device_properties"""

912 - self.device_properties = torch.cuda.get_device_properties(self.device)

913 - self.num_sms = self.device_properties.multi_processor_count

912 +

913 + self.num_sms = num_compute_units(self.device.index)

914

915 # Note: used for model runner override.

916 def _sync_device(self) -> None:

vllm/v1/worker/gpu_ubatch_wrapper.py

+2 -2

@@ -23,6 +23,7 @@ from vllm.logger import init_logger

23 from vllm.platforms import current_platform

24 from vllm.sequence import IntermediateTensors

25 from vllm.utils.import_utils import has_deep_gemm

26 +from vllm.utils.platform_utils import num_compute_units

27 from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts

28

29 logger = init_logger(__name__)

@@ -72,8 +73,7 @@ class SMControlContextManager:

73 "SM control is currently only supported on CUDA"

74 )

75

75 - props = torch.cuda.get_device_properties(torch.cuda.current_device())

76 - total_sms = props.multi_processor_count

76 + total_sms = num_compute_units(torch.cuda.current_device().index)

77

78 assert comm_sms < total_sms

79 self.total_sms = total_sms

vllm/v1/worker/xpu_model_runner.py

-3

@@ -28,9 +28,6 @@ class XPUModelRunner(GPUModelRunner):

28 # FIXME: To be verified.

29 self.cascade_attn_enabled = False

30

31 - def _init_device_properties(self) -> None:

32 - self.num_sms = None

33 -

31 def _sync_device(self) -> None:

32 torch.xpu.synchronize()

33

[Platform] Add current_platform.num_compute_units interface (#35042)