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, 08:30 UTC
[V1][Spec Decode] Small refactors to improve eagle bookkeeping performance (#18424)
Signed-off-by: qizixi <qizixi@meta.com>
qizixi committed
May 23, 2025 at 23:51 UTC
d55e446d1320d0f5f22bc3584f81f18d7924f166
3 files changed
+21
-19
tests/v1/spec_decode/test_eagle.py
+5
-1
@@ -100,8 +100,12 @@ def test_prepare_inputs():
100
dtype=torch.int32,
101
device=device)
102
103
+ # n1 + n2 + n3 - a - b -c
104
+ num_tokens = cu_target_query_lens[-1].item() - num_rejected_tokens.sum(
105
+ ).item()
106
+
107
cu_num_tokens, token_indices = EagleProposer.prepare_inputs(
104
- cu_target_query_lens, num_rejected_tokens)
108
+ cu_target_query_lens, num_rejected_tokens, num_tokens)
109
110
assert torch.equal(cu_num_tokens, expected_cu_num_tokens)
111
assert token_indices.shape[0] == expected_cu_num_tokens[-1].item()
vllm/v1/spec_decode/eagle.py
+3
-7
@@ -271,6 +271,7 @@ class EagleProposer:
271
cu_target_query_lens: torch.Tensor,
272
# [batch_size]
273
num_rejected_tokens: torch.Tensor,
274
+ num_tokens: int,
275
) -> tuple[torch.Tensor, torch.Tensor]:
276
# cu_target_query_lens: [0, a, a + b, a + b + c]
277
# num_rejected_tokens: [n1, n2, n3]
@@ -288,18 +289,13 @@ class EagleProposer:
289
290
# [a - n1, b - n2, c - n3] ->
291
# [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3]
291
- cu_num_tokens = torch.empty_like(cu_target_query_lens)
292
+ cu_num_tokens = torch.zeros_like(cu_target_query_lens)
293
torch.cumsum(num_tokens_per_req, dim=0, out=cu_num_tokens[1:])
293
- cu_num_tokens[0] = 0
294
-
295
- # FIXME(woosuk): Avoid synchronization.
296
- num_tokens = cu_num_tokens[-1].item()
294
token_indices = torch.empty(
295
num_tokens,
296
dtype=torch.int32,
300
- device=cu_num_tokens.device,
297
+ device=cu_target_query_lens.device,
298
)
302
-
299
batch_size = num_rejected_tokens.shape[0]
300
BLOCK_SIZE = 1024
301
prepare_eagle_input_kernel[(batch_size, )](
vllm/v1/worker/gpu_model_runner.py
+13
-11
@@ -34,8 +34,8 @@ from vllm.multimodal.utils import group_mm_inputs_by_modality
34
from vllm.sampling_params import SamplingType
35
from vllm.sequence import IntermediateTensors
36
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
37
- GiB_bytes, LazyLoader, cdiv, check_use_alibi,
38
- is_pin_memory_available)
37
+ GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
38
+ check_use_alibi, is_pin_memory_available)
39
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
40
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
41
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -281,7 +281,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
281
def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> bool:
282
"""
283
Update the order of requests in the batch based on the attention
284
- backend's needs. For example, some attention backends (namely MLA) may
284
+ backend's needs. For example, some attention backends (namely MLA) may
285
want to separate requests based on if the attention computation will be
286
compute-bound or memory-bound.
287
@@ -1360,9 +1360,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1360
scheduler_output.num_scheduled_tokens[req_id])
1361
next_token_id = req_state.get_token_id(seq_len)
1362
next_token_ids.append(next_token_id)
1363
- next_token_ids = torch.tensor(next_token_ids,
1364
- dtype=torch.int32,
1365
- device=self.device)
1363
+ next_token_ids = async_tensor_h2d(next_token_ids,
1364
+ dtype=torch.int32,
1365
+ target_device=self.device,
1366
+ pin_memory=True)
1367
eagle_attn_metadata = attn_metadata[self.drafter.attn_layer_name]
1368
1369
# NOTE: deepseek_mtp uses MLA which does not have `block_table`
@@ -1390,14 +1391,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1391
n + 1 - len(valid_sampled_token_ids[i]) if n > 0 else 0
1392
for i, n in enumerate(num_draft_tokens)
1393
]
1393
- num_rejected_tokens = torch.tensor(
1394
+ num_rejected_tokens_tensor = async_tensor_h2d(
1395
num_rejected_tokens,
1396
dtype=torch.int32,
1396
- device=self.device,
1397
- )
1397
+ target_device=self.device,
1398
+ pin_memory=True)
1399
+ num_tokens = num_scheduled_tokens - sum(num_rejected_tokens)
1400
cu_num_tokens, token_indices = self.drafter.prepare_inputs(
1401
eagle_attn_metadata.query_start_loc,
1400
- num_rejected_tokens,
1402
+ num_rejected_tokens_tensor,
1403
+ num_tokens,
1404
)
1405
target_token_ids = self.input_ids[token_indices]
1406
target_positions = positions[token_indices]
@@ -1408,7 +1411,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1411
target_hidden_states = hidden_states[token_indices]
1412
target_slot_mapping = eagle_attn_metadata.slot_mapping[
1413
token_indices]
1411
-
1414
draft_token_ids = self.drafter.propose(
1415
target_token_ids=target_token_ids,
1416
target_positions=target_positions,