@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 23eca9cf68

Fork 0

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 15, 22:15 UTC

[model][refactor] remove cuda hard code in models and layers (#13658)

Mengqing Cao committed Feb 24, 2025 at 22:10 UTC 23eca9cf688f2ff142a9997e9b3a300be53a13e5

7 files changed +29 -14

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+2 -1

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

8 from vllm.model_executor.layers.fused_moe.fused_moe import (

9 fused_topk, moe_align_block_size, try_get_optimal_moe_config)

10 +from vllm.platforms import current_platform

11 from vllm.scalar_type import scalar_types

12 from vllm.utils import direct_register_custom_op

@@ -238,7 +239,7 @@ def fused_marlin_moe(

239 max_workspace_size = (max(2 * N, K) // 64) * 16

240 workspace = torch.zeros(max_workspace_size,

241 dtype=torch.int,

241 - device="cuda",

242 + device=current_platform.device_type,

243 requires_grad=False)

244

245 if has_no_zp:

vllm/model_executor/layers/rotary_embedding.py

+9 -4

@@ -30,6 +30,7 @@ import torch.nn as nn

30 from transformers import PretrainedConfig

32 from vllm.model_executor.custom_op import CustomOp

33 +from vllm.platforms import current_platform

36 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:

@@ -650,9 +651,13 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

651 is_neox_style, dtype)

652

653 def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:

653 - pos_freqs = self.base**(torch.arange(

654 - 0, self.rotary_dim, 2, dtype=torch.float, device="cuda") /

655 - self.rotary_dim)

654 + pos_freqs = self.base**(

655 + torch.arange(0,

656 + self.rotary_dim,

657 + 2,

658 + dtype=torch.float,

659 + device=current_platform.device_type) /

660 + self.rotary_dim)

661 inv_freq_extrapolation = 1.0 / pos_freqs

662 inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)

663

@@ -670,7 +675,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

675 def _compute_cos_sin_cache(self) -> torch.Tensor:

676 inv_freq = self._compute_inv_freq(self.scaling_factor)

677 t = torch.arange(self.max_position_embeddings * self.scaling_factor,

673 - device="cuda",

678 + device=current_platform.device_type,

679 dtype=torch.float32)

680 freqs = torch.einsum("i,j -> ij", t, inv_freq)

681 cos = (freqs.cos() * self.mscale)

vllm/model_executor/layers/spec_decode_base_sampler.py

+3 -1

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

7 import torch.jit

8 import torch.nn as nn

10 +from vllm.platforms import current_platform

11 +

13 class SpecDecodeBaseSampler(nn.Module):

14 """Base class for samplers used for Speculative Decoding verification

@@ -35,7 +37,7 @@ class SpecDecodeBaseSampler(nn.Module):

37 def init_gpu_tensors(self, device: Union[int, str]) -> None:

38 assert self.num_accepted_tokens is None

39 if isinstance(device, int):

38 - device = f"cuda:{device}"

40 + device = f"{current_platform.device_type}:{device}"

41 elif not isinstance(device, str):

42 raise ValueError(f"Device must be int or str, get {type(device)}")

43 self.num_accepted_tokens = torch.tensor(0,

vllm/model_executor/model_loader/loader.py

+2 -1

@@ -914,7 +914,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):

914 if param_name + "." in k:

915 quant_state[k] = temp_state_dict[k]

916

917 - return QuantState.from_dict(quant_state, device="cuda")

917 + return QuantState.from_dict(quant_state,

918 + device=current_platform.device_type)

919

920 # Second iterate over all prequant and normal weights

921 # pre quantized weights would have a quant_state

vllm/model_executor/models/arctic.py

+3 -2

@@ -30,6 +30,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

30 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

31 from vllm.model_executor.sampling_metadata import SamplingMetadata

32 from vllm.model_executor.utils import set_weight_attrs

33 +from vllm.platforms import current_platform

34 from vllm.sequence import IntermediateTensors

35 from vllm.transformers_utils.configs.arctic import ArcticConfig

@@ -138,13 +139,13 @@ class ArcticMoE(nn.Module):

139 torch.empty(self.num_experts,

140 2 * self.intermediate_size,

141 self.hidden_size,

141 - device="cuda",

142 + device=current_platform.device_type,

143 dtype=self.params_dtype))

144 self.w2s = nn.Parameter(

145 torch.empty(self.num_experts,

146 self.hidden_size,

147 self.intermediate_size,

147 - device="cuda",

148 + device=current_platform.device_type,

149 dtype=self.params_dtype))

150 set_weight_attrs(self.ws, {

151 "weight_loader": self.weight_loader,

vllm/model_executor/models/minicpm.py

+3 -2

@@ -51,6 +51,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

51 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

52 from vllm.model_executor.sampling_metadata import SamplingMetadata

53 from vllm.model_executor.utils import set_weight_attrs

54 +from vllm.platforms import current_platform

55 from vllm.sequence import IntermediateTensors

57 from .interfaces import SupportsLoRA, SupportsPP

@@ -98,13 +99,13 @@ class MiniCPMMoE(nn.Module):

99 torch.empty(self.num_total_experts,

100 2 * self.intermediate_size,

101 self.hidden_size,

101 - device="cuda",

102 + device=current_platform.device_type,

103 dtype=self.params_dtype))

104 self.w2s = nn.Parameter(

105 torch.empty(self.num_total_experts,

106 self.hidden_size,

107 self.intermediate_size,

107 - device="cuda",

108 + device=current_platform.device_type,

109 dtype=self.params_dtype))

110

111 set_weight_attrs(self.ws, {

vllm/model_executor/models/minicpmv.py

+7 -3

@@ -59,6 +59,7 @@ from vllm.multimodal.parse import (DictEmbeddingItems, ImageItem, ImageSize,

59 from vllm.multimodal.processing import (BaseMultiModalProcessor,

60 BaseProcessingInfo, PromptReplacement)

61 from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

62 +from vllm.platforms import current_platform

63 from vllm.sequence import IntermediateTensors

65 from .idefics2_vision_model import Idefics2VisionTransformer

@@ -1184,7 +1185,8 @@ class MiniCPMV2_0(MiniCPMVBaseModel):

1185 quant_config=quant_config,

1186 prefix=prefix)

1187

1187 - return resampler.to(device="cuda", dtype=torch.get_default_dtype())

1188 + return resampler.to(device=current_platform.device_type,

1189 + dtype=torch.get_default_dtype())

1190

1191 def get_vision_embedding(

1192 self,

@@ -1266,7 +1268,8 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):

1268 quant_config=quant_config,

1269 prefix=prefix)

1270

1269 - return resampler.to(device="cuda", dtype=torch.get_default_dtype())

1271 + return resampler.to(device=current_platform.device_type,

1272 + dtype=torch.get_default_dtype())

1273

1274 def get_vision_embedding(

1275 self,

@@ -1360,7 +1363,8 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):

1363 quant_config=quant_config,

1364 prefix=prefix)

1365

1363 - return resampler.to(device="cuda", dtype=torch.get_default_dtype())

1366 + return resampler.to(device=current_platform.device_type,

1367 + dtype=torch.get_default_dtype())

1368

1369 def get_vision_embedding(

1370 self,