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 14, 05:30 UTC
[TPU] Deprecate `xm.mark_step` in favor of ``torch_xla.sync` (#25254)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Ye (Charlotte) Qi <yeq@meta.com>
Nicolò Lucchesi committed
Sep 22, 2025 at 12:12 UTC
4cf71cc88a19f97bf6f5db9802395edd646375ef
5 files changed
+31
-29
tests/tpu/test_moe_pallas.py
+2
-1
@@ -6,6 +6,7 @@ Run `pytest tests/kernels/moe/test_moe_pallas.py`.
6
"""
7
import pytest
8
import torch
9
+import torch_xla
10
11
# yapf conflicts with isort for this block
12
# yapf: disable
@@ -77,7 +78,7 @@ def test_pallas_moe(
78
expert_map=e_map,
79
renormalize=False,
80
)
80
- xm.mark_step()
81
+ torch_xla.sync(wait=False)
82
83
# Compare outputs
84
torch.testing.assert_close(
tests/v1/tpu/test_topk_topp_sampler.py
+6
-5
@@ -4,6 +4,7 @@ import math
4
5
import pytest
6
import torch
7
+import torch_xla
8
9
from vllm.platforms import current_platform
10
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
@@ -63,7 +64,7 @@ def test_topp_result_sums_past_p():
64
probs.masked_fill_(logits_masked.isinf(), 0)
65
masked_prob_sum = probs.sum(dim=-1)
66
66
- xm.mark_step()
67
+ torch_xla.sync()
68
69
# Perform assertion on CPU.
70
assert torch.all(torch.ge(masked_prob_sum.cpu() + TOLERANCE, p.cpu()))
@@ -82,7 +83,7 @@ def test_topp_basic():
83
k=torch.tensor([3, 3]),
84
p=torch.tensor([0.79, 0.79]))
85
85
- xm.mark_step()
86
+ torch_xla.sync()
87
88
# Expect the smallest elements to be dropped.
89
expected_result = logits.clone().cpu()
@@ -104,7 +105,7 @@ def test_topp_select_all():
105
k=torch.tensor([3, 3]),
106
p=torch.tensor([1.0, 1.0]))
107
107
- xm.mark_step()
108
+ torch_xla.sync()
109
110
assert torch.allclose(logits.cpu(), result.cpu())
111
@@ -122,7 +123,7 @@ def test_topp_with_ties():
123
k=torch.tensor([4]),
124
p=torch.tensor([0.2]))
125
125
- xm.mark_step()
126
+ torch_xla.sync()
127
128
# All tie values are included in the top-p set. Tie breaking is left
129
# to be done during final sampling (all tie tokens have equal
@@ -146,7 +147,7 @@ def test_both_topk_topp():
147
k=torch.tensor([1, 3]),
148
p=torch.tensor([0.79, 0.79]))
149
149
- xm.mark_step()
150
+ torch_xla.sync()
151
152
# Since for the first batch k=1, expect only the largest element gets
153
# selected.
vllm/lora/punica_wrapper/punica_tpu.py
+2
-2
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional, Union
6
7
import torch
8
import torch.nn.functional as F
9
-import torch_xla.core.xla_model as xm
9
+import torch_xla
10
11
from vllm.lora.ops.xla_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
12
from vllm.lora.punica_wrapper.utils import convert_mapping
@@ -323,7 +323,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):
323
extra_vocab_size: int,
324
):
325
# Make sure we don't accidentally collect outside operations
326
- xm.mark_step()
326
+ torch_xla.sync()
327
328
# Pad the prompt mapping to avoid running into recompiles on the TPU
329
# TODO: Should this happen inside mapping internally? If so how can we
vllm/model_executor/model_loader/default_loader.py
+4
-5
@@ -211,16 +211,15 @@ class DefaultModelLoader(BaseModelLoader):
211
from vllm.platforms.tpu import USE_TPU_COMMONS
212
213
if not USE_TPU_COMMONS:
214
- # In PyTorch XLA, we should call `xm.mark_step`
214
+ # In PyTorch XLA, we should call `torch_xla.sync`
215
# frequently so that not too many ops are accumulated
216
- # in the XLA program. import torch_xla.core.xla_model
217
- # as xm
218
- import torch_xla.core.xla_model as xm
216
+ # in the XLA program.
217
+ import torch_xla
218
219
def _xla_weights_iterator(iterator: Generator):
220
for weights in iterator:
221
yield weights
223
- xm.mark_step()
222
+ torch_xla.sync(wait=False)
223
224
weights_iterator = _xla_weights_iterator(weights_iterator)
225
vllm/v1/worker/tpu_model_runner.py
+17
-16
@@ -10,6 +10,7 @@ import numpy as np
10
import torch
11
import torch.nn as nn
12
# TPU XLA related
13
+import torch_xla
14
import torch_xla.core.xla_model as xm
15
import torch_xla.distributed.spmd as xs
16
import torch_xla.runtime as xr
@@ -846,10 +847,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
847
# 2. A list or tuple (length: num_items) of tensors, each of shape
848
# (feature_size, hidden_size) in case the feature size is dynamic
849
# depending on the input multimodal items.
849
- xm.mark_step()
850
+ torch_xla.sync(wait=False)
851
curr_group_outputs = self.model.get_multimodal_embeddings(
852
**mm_kwargs_group)
852
- xm.mark_step()
853
+ torch_xla.sync(wait=False)
854
855
sanity_check_mm_encoder_outputs(
856
curr_group_outputs,
@@ -952,7 +953,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
953
mm_embeds = self._gather_mm_embeddings(scheduler_output)
954
else:
955
mm_embeds = []
955
- xm.mark_step()
956
+ torch_xla.sync(wait=False)
957
# Prepare inputs, the requests might be split into multiple
958
# executions, combine the result of each execution.
959
start_index = 0
@@ -969,7 +970,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
970
end_index = self._prepare_inputs(scheduler_output, start_index)
971
input_ids, inputs_embeds = self._get_model_inputs(
972
self.input_ids, mm_embeds)
972
- xm.mark_step()
973
+ torch_xla.sync(wait=False)
974
# Run the decoder
975
with set_forward_context(
976
attn_metadata,
@@ -1183,7 +1184,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1184
1185
# Sync all pending XLA execution during model initialization and weight
1186
# loading.
1186
- xm.mark_step()
1187
+ torch_xla.sync(wait=False)
1188
xm.wait_device_ops()
1189
if not hasattr(self, "model"):
1190
self.model = model
@@ -1267,10 +1268,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1268
1269
def _set_active_loras(self, prompt_lora_mapping, token_lora_mapping,
1270
lora_requests) -> None:
1270
- xm.mark_step() # Captures input updates
1271
+ torch_xla.sync(wait=False) # Captures input updates
1272
super()._set_active_loras(prompt_lora_mapping, token_lora_mapping,
1273
lora_requests)
1273
- xm.mark_step() # Captures metadata updates
1274
+ torch_xla.sync(wait=False) # Captures metadata updates
1275
1276
def _precompile_mm_encoder(self) -> None:
1277
if not self.supports_mm_inputs:
@@ -1297,10 +1298,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1298
num_items,
1299
)
1300
# Run multimodal encoder.
1300
- xm.mark_step()
1301
+ torch_xla.sync(wait=False)
1302
mm_embeds = self.model.get_multimodal_embeddings(
1303
**batched_dummy_mm_inputs)
1303
- xm.mark_step()
1304
+ torch_xla.sync(wait=False)
1305
num_patches = mm_embeds[0].shape[0]
1306
items_size = num_patches * num_items
1307
@@ -1325,7 +1326,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1326
a, b = self._get_model_inputs(placeholders_ids,
1327
[mm_embeds])
1328
assert a is None
1328
- xm.mark_step()
1329
+ torch_xla.sync(wait=False)
1330
1331
# Pre-compile `get_input_embeddings` when mm_embeddings are not
1332
# present. Chunk is only made of text, no mm_placeholders.
@@ -1336,7 +1337,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1337
placeholders_ids = placeholders_ids.to(self.device)
1338
a, b = self._get_model_inputs(placeholders_ids, [])
1339
assert a is None
1339
- xm.mark_step()
1340
+ torch_xla.sync(wait=False)
1341
1342
xm.wait_device_ops()
1343
end = time.perf_counter()
@@ -1532,11 +1533,11 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1533
# Isolate encoder graph from post-processing to minimize
1534
# impact of recompilation until it's fixed.
1535
start = time.perf_counter()
1535
- xm.mark_step()
1536
+ torch_xla.sync(wait=False)
1537
dummy_encoder_outputs = \
1538
self.model.get_multimodal_embeddings(
1539
**batched_dummy_mm_inputs)
1539
- xm.mark_step()
1540
+ torch_xla.sync(wait=False)
1541
xm.wait_device_ops()
1542
end = time.perf_counter()
1543
logger.info(
@@ -1559,7 +1560,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1560
self._dummy_run(num_tokens, self.num_reqs_most_model_len,
1561
self.num_blocks_per_most_len_req)
1562
1562
- xm.mark_step()
1563
+ torch_xla.sync(wait=False)
1564
xm.wait_device_ops()
1565
self.encoder_cache.clear()
1566
gc.collect()
@@ -1927,11 +1928,11 @@ def replace_set_lora(model):
1928
# to a tensor doesn't seem to work anymore. This might be fixed with a
1929
# later release of torch_xla.
1930
self._original_set_lora(index, lora_a, lora_b, embeddings_tensor, bias)
1930
- xm.mark_step()
1931
+ torch_xla.sync(wait=False)
1932
1933
def _tpu_reset_lora(self, index: int):
1934
self._original_reset_lora(index)
1934
- xm.mark_step()
1935
+ torch_xla.sync(wait=False)
1936
1937
for _, module in model.named_modules():
1938
if isinstance(module, BaseLayerWithLoRA):