Optimize PLE MTP metadata transfers (#55054)
bhsueh_NV committed
Sep 3, 2026 at 23:02 UTC
edc0fb7e03b43fc17265c117f92021813711c6d8
1 file changed
+12
-9
vllm/v1/attention/backends/short_conv_attn.py
+12
-9
@@ -324,8 +324,8 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
324
decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0]
325
prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0]
326
non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu))
327
- spec_req_idx = spec_req_idx_cpu.to(query_start_loc.device)
328
- non_spec_req_idx = non_spec_req_idx_cpu.to(query_start_loc.device)
327
+ spec_req_idx = async_tensor_h2d(spec_req_idx_cpu, device=query_start_loc.device)
328
+ non_spec_req_idx: torch.Tensor | None = None
329
330
if num_decodes == 0 and num_prefills == 0:
331
# Pure speculative-decode batch: all real tokens are spec tokens.
@@ -348,6 +348,12 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
348
# stable sort, so tokens of each request stay contiguous and in
349
# request order. This yields spec tokens first, then the non-spec
350
# [decode, prefill] tokens.
351
+ non_spec_req_idx = async_tensor_h2d(
352
+ non_spec_req_idx_cpu, device=query_start_loc.device
353
+ )
354
+ decode_req_idx = async_tensor_h2d(
355
+ decode_req_idx_cpu, device=query_start_loc.device
356
+ )
357
req_group = torch.full(
358
(m.num_reqs,),
359
2,
@@ -355,7 +361,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
361
device=query_start_loc.device,
362
)
363
req_group[spec_req_idx] = 0
358
- req_group[decode_req_idx_cpu.to(query_start_loc.device)] = 1
364
+ req_group[decode_req_idx] = 1
365
token_group = torch.repeat_interleave(req_group, query_lens)
366
token_perm = torch.argsort(token_group, stable=True)
367
spec_token_indx = token_perm[:num_spec_decode_tokens]
@@ -391,9 +397,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
397
assert num_accepted_tokens is not None
398
# Accepted-token counts must follow the same request order as the
399
# speculative state indices.
394
- num_accepted_tokens = num_accepted_tokens[
395
- spec_req_idx_cpu.to(num_accepted_tokens.device)
396
- ]
400
+ num_accepted_tokens = num_accepted_tokens[spec_req_idx]
401
402
# Compute the conv-state slots for the non-spec decode/prefill split,
403
# plus the initial-state masks and Triton causal_conv1d metadata.
@@ -410,9 +414,8 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
414
state_indices_tensor_d = None
415
if num_decodes > 0 or num_prefills > 0:
416
num_computed_tokens = m.compute_num_computed_tokens()
413
- if non_spec_req_idx_cpu is not None:
414
- non_spec_req_idx = non_spec_req_idx_cpu.to(num_computed_tokens.device)
415
- num_computed_tokens = num_computed_tokens[non_spec_req_idx]
417
+ assert non_spec_req_idx is not None
418
+ num_computed_tokens = num_computed_tokens[non_spec_req_idx]
419
420
state_indices_tensor_d = state_indices_tensor[:num_decodes]
421
state_indices_tensor_p = state_indices_tensor[