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, 17:15 UTC
Apply torchfix (#15532)
Signed-off-by: cyy <cyyever@outlook.com>
cyyever committed
Mar 26, 2025 at 20:09 UTC
1aa162e030d62fcf476388ac77c141cb5b52957b
5 files changed
+15
-11
vllm/attention/backends/rocm_flash_attn.py
+2
-3
@@ -884,9 +884,8 @@ def _sdpa_attention(
884
885
for i, seq_len in enumerate(seq_lens):
886
end = start + seq_len
887
- with torch.backends.cuda.sdp_kernel(enable_math=True,
888
- enable_flash=False,
889
- enable_mem_efficient=False):
887
+ with torch.nn.attention.sdpa_kernel(
888
+ torch.nn.attention.SDPBackend.MATH):
889
sub_out = torch.nn.functional.scaled_dot_product_attention(
890
query[:, start:end, :],
891
key[:, start:end, :],
vllm/lora/models.py
+3
-1
@@ -272,7 +272,9 @@ class LoRAModel(AdapterModel):
272
f" target modules in {expected_lora_modules}"
273
f" but received {unexpected_modules}."
274
f" Please verify that the loaded LoRA module is correct")
275
- tensors = torch.load(lora_bin_file_path, map_location=device)
275
+ tensors = torch.load(lora_bin_file_path,
276
+ map_location=device,
277
+ weights_only=True)
278
else:
279
raise ValueError(f"{lora_dir} doesn't contain tensors")
280
vllm/model_executor/models/nemotron.py
+3
-3
@@ -63,8 +63,8 @@ def _cast_if_autocast_enabled(*args):
63
if not torch.is_autocast_enabled():
64
return args
65
else:
66
- return torch.cuda.amp.autocast_mode._cast(
67
- args, torch.get_autocast_gpu_dtype())
66
+ return torch.amp.autocast_mode._cast(
67
+ args, device_type="cuda", dtype=torch.get_autocast_gpu_dtype())
68
69
70
class NemotronLayerNorm1P(nn.LayerNorm):
@@ -89,7 +89,7 @@ class NemotronLayerNorm1P(nn.LayerNorm):
89
residual = x
90
args = _cast_if_autocast_enabled(x, self.normalized_shape,
91
self.weight + 1, self.bias, self.eps)
92
- with torch.cuda.amp.autocast(enabled=False):
92
+ with torch.amp.autocast("cuda", enabled=False):
93
x = torch.nn.functional.layer_norm(*args)
94
return x if residual is None else (x, residual)
95
vllm/model_executor/models/phi4mm_utils.py
+6
-3
@@ -1766,9 +1766,12 @@ class MultiHeadedAttention(nn.Module):
1766
if mask.dtype != q.dtype:
1767
attn_mask = attn_mask.to(q.dtype)
1768
1769
- with torch.backends.cuda.sdp_kernel(enable_flash=True,
1770
- enable_math=True,
1771
- enable_mem_efficient=True):
1769
+ with torch.nn.attention.sdpa_kernel([
1770
+ torch.nn.attention.SDPBackend.FLASH_ATTENTION,
1771
+ torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
1772
+ torch.nn.attention.SDPBackend.MATH,
1773
+ torch.nn.attention.SDPBackend.CUDNN_ATTENTION,
1774
+ ]):
1775
x = torch.nn.functional.scaled_dot_product_attention(
1776
q,
1777
k,
vllm/multimodal/image.py
+1
-1
@@ -149,7 +149,7 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
149
return self.load_bytes(base64.b64decode(data))
150
151
def load_file(self, filepath: Path) -> torch.Tensor:
152
- return torch.load(filepath)
152
+ return torch.load(filepath, weights_only=True)
153
154
def encode_base64(self, media: torch.Tensor) -> str:
155
return base64.b64encode(media.numpy()).decode('utf-8')