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 9, 18:15 UTC
[Feature]: Remove DtoH Copy for lfm2_vl On Default Stream (#32815)
Signed-off-by: Tianshu Yu <tianshuyu.formal@gmail.com>
tianshu-Michael-yu committed
Jan 23, 2026 at 05:20 UTC
13d8746c545576bcb6e0771ee4dc0b6fae694fa1
5 files changed
+260
-158
vllm/model_executor/models/lfm2_siglip2.py
renamed
+89
-92
@@ -40,99 +40,111 @@ class Siglip2VisionEmbeddings(nn.Module):
40
self.position_embedding_size = int(self.num_patches**0.5)
41
self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
42
43
+ def forward(
44
+ self,
45
+ pixel_values_packed: torch.FloatTensor,
46
+ spatial_shapes: torch.LongTensor,
47
+ ) -> torch.Tensor:
48
+ """Embed patchified pixel values in packed (unpadded) form.
49
+
50
+ Args:
51
+ pixel_values_packed: (1, total_tokens, patch_dim) or
52
+ (total_tokens, patch_dim), packed in tile order.
53
+ spatial_shapes: (num_tiles, 2) on CPU (height, width) per tile.
54
+
55
+ Returns:
56
+ (1, total_tokens, embed_dim) packed embeddings.
57
+ """
58
+ assert spatial_shapes.device.type == "cpu", (
59
+ "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "
60
+ "variable-length packing."
61
+ )
62
+
63
+ if pixel_values_packed.dim() == 3:
64
+ assert pixel_values_packed.shape[0] == 1
65
+ pixel_values_flat = pixel_values_packed[0]
66
+ else:
67
+ pixel_values_flat = pixel_values_packed
68
+
69
+ lengths = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(dtype=torch.int64)
70
+ lengths_list = lengths.tolist()
71
+ total_tokens = int(sum(lengths_list))
72
+ if total_tokens != pixel_values_flat.shape[0]:
73
+ raise ValueError(
74
+ "Packed pixel_values token count does not match spatial_shapes: "
75
+ f"{pixel_values_flat.shape[0]} vs {total_tokens}."
76
+ )
77
+
78
+ target_dtype = self.patch_embedding.weight.dtype
79
+ patch_embeds = self.patch_embedding(pixel_values_flat.to(dtype=target_dtype))
80
+
81
+ positional_embeddings = self.position_embedding.weight.reshape(
82
+ self.position_embedding_size, self.position_embedding_size, -1
83
+ )
84
+ packed_pos_embeds = self.resize_positional_embeddings_packed(
85
+ positional_embeddings,
86
+ spatial_shapes,
87
+ lengths_list=lengths_list,
88
+ )
89
+
90
+ embeddings = patch_embeds + packed_pos_embeds
91
+ return embeddings.unsqueeze(0)
92
+
93
@staticmethod
44
- def resize_positional_embeddings(
94
+ def resize_positional_embeddings_packed(
95
positional_embeddings: torch.Tensor,
96
spatial_shapes: torch.LongTensor,
47
- max_length: int,
97
+ lengths_list: list[int],
98
) -> torch.Tensor:
49
- """
50
- Resize positional embeddings to image-specific size and pad to a fixed size.
99
+ """Resize positional embeddings per image and return a packed tensor.
100
101
Args:
53
- positional_embeddings (`torch.Tensor`):
54
- Position embeddings of shape (height, width, embed_dim)
55
- spatial_shapes (`torch.LongTensor`):
56
- Spatial shapes of shape (batch_size, 2) to resize the positional
57
- embeddings to
58
- max_length (`int`):
59
- Maximum length of the positional embeddings to pad resized
60
- positional embeddings to
102
+ positional_embeddings: (height, width, embed_dim) base grid.
103
+ spatial_shapes: (batch_size, 2) on CPU, (height, width) per image.
104
+ lengths_list: flattened token length per image (height * width).
105
106
Returns:
63
- `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
107
+ (total_tokens, embed_dim) packed positional embeddings, concatenated
108
+ in the same order as `lengths_list`.
109
"""
65
- batch_size = spatial_shapes.shape[0]
110
+ assert spatial_shapes.device.type == "cpu"
111
+
112
embed_dim = positional_embeddings.shape[-1]
113
source_dtype = positional_embeddings.dtype
114
69
- resulted_positional_embeddings = torch.empty(
70
- (batch_size, max_length, embed_dim),
115
+ total_tokens = int(sum(lengths_list))
116
+ packed_pos_embeds = torch.empty(
117
+ (total_tokens, embed_dim),
118
device=positional_embeddings.device,
119
dtype=source_dtype,
120
)
121
75
- # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
76
- positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
122
+ # (height, width, embed_dim) -> (1, embed_dim, height, width)
123
+ pos_4d = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
124
125
# Upcast to float32 on CPU because antialias is not supported for
79
- # bfloat16/float16 on CPU
80
- if positional_embeddings.device.type == "cpu":
81
- positional_embeddings = positional_embeddings.to(torch.float32)
82
-
83
- for i in range(batch_size):
84
- # (1, dim, height, width) -> (1, dim, target_height, target_width)
85
- height, width = spatial_shapes[i]
86
- resized_embeddings = F.interpolate(
87
- positional_embeddings,
126
+ # bfloat16/float16 on CPU.
127
+ if pos_4d.device.type == "cpu":
128
+ pos_4d = pos_4d.to(torch.float32)
129
+
130
+ offset = 0
131
+ for i, length in enumerate(lengths_list):
132
+ if length <= 0:
133
+ continue
134
+ height, width = spatial_shapes[i].tolist()
135
+ resized = F.interpolate(
136
+ pos_4d,
137
size=(height, width),
138
mode="bilinear",
139
align_corners=False,
140
antialias=True,
141
)
142
+ resized = resized.reshape(embed_dim, height * width).transpose(0, 1)
143
+ resized = resized.to(source_dtype)
144
+ packed_pos_embeds[offset : offset + length] = resized
145
+ offset += length
146
94
- # (1, dim, target_height, target_width) ->
95
- # (target_height * target_width, dim)
96
- resized_embeddings = resized_embeddings.reshape(
97
- embed_dim, height * width
98
- ).transpose(0, 1)
99
-
100
- # Cast to original dtype
101
- resized_embeddings = resized_embeddings.to(source_dtype)
102
-
103
- resulted_positional_embeddings[i, : height * width] = resized_embeddings
104
- resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
105
-
106
- return resulted_positional_embeddings
107
-
108
- def forward(
109
- self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor
110
- ) -> torch.Tensor:
111
- """
112
- Args:
113
- pixel_values (`torch.FloatTensor`):
114
- Pixel values of shape (batch_size, max_num_patches,
115
- num_channels * patch_size * patch_size)
116
- spatial_shapes (`list[tuple[int, int]]`):
117
- Spatial shapes of shape (batch_size, 2) to resize the positional
118
- embeddings to
119
- """
120
-
121
- # Apply patch embeddings to already patchified pixel values
122
- target_dtype = self.patch_embedding.weight.dtype
123
- patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
124
-
125
- # Get positional resized and padded positional embeddings
126
- positional_embeddings = self.position_embedding.weight.reshape(
127
- self.position_embedding_size, self.position_embedding_size, -1
128
- )
129
- resized_positional_embeddings = self.resize_positional_embeddings(
130
- positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
131
- )
132
-
133
- # Add positional embeddings to patch embeddings
134
- embeddings = patch_embeds + resized_positional_embeddings
135
- return embeddings
147
+ return packed_pos_embeds
148
149
150
class Siglip2Attention(nn.Module):
@@ -402,36 +414,23 @@ class Siglip2VisionTransformer(nn.Module):
414
415
def forward(
416
self,
405
- pixel_values: torch.FloatTensor,
417
+ pixel_values_packed: torch.FloatTensor,
418
spatial_shapes: torch.LongTensor,
407
- packed_mask: torch.Tensor,
419
cu_seqlens: torch.Tensor,
409
- max_seqlen: int | torch.Tensor,
420
+ max_seqlen: torch.Tensor,
421
) -> torch.Tensor:
422
r"""
423
spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
424
Tensor containing the spatial dimensions (height, width)
414
- of the input images.
425
+ of the input images.
426
"""
416
- hidden_states = self.embeddings(pixel_values, spatial_shapes)
417
- flat_mask = packed_mask.view(-1)
418
- packed_indices = flat_mask.nonzero(as_tuple=True)[0]
419
- flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
420
- hidden_states = flat_hidden_states.index_select(0, packed_indices).unsqueeze(0)
427
+ hidden_states = self.embeddings(pixel_values_packed, spatial_shapes)
428
encoder_outputs = self.encoder(
429
inputs_embeds=hidden_states,
430
cu_seqlens=cu_seqlens,
431
max_seqlen=max_seqlen,
432
)
426
- unpacked = encoder_outputs.new_zeros(
427
- packed_mask.numel(), encoder_outputs.shape[-1]
428
- )
429
- unpacked.index_copy_(0, packed_indices, encoder_outputs.squeeze(0))
430
- encoder_outputs = unpacked.view(
431
- packed_mask.shape + (encoder_outputs.shape[-1],)
432
- )
433
- last_hidden_state = self.post_layernorm(encoder_outputs)
434
- return last_hidden_state
433
+ return self.post_layernorm(encoder_outputs)
434
435
436
class Siglip2Model(torch.nn.Module):
@@ -453,16 +452,14 @@ class Siglip2Model(torch.nn.Module):
452
453
def forward(
454
self,
456
- pixel_values: torch.FloatTensor,
455
+ pixel_values_packed: torch.FloatTensor,
456
spatial_shapes: torch.LongTensor,
458
- packed_mask: torch.Tensor,
457
cu_seqlens: torch.Tensor,
460
- max_seqlen: int | torch.Tensor,
458
+ max_seqlen: torch.Tensor,
459
) -> torch.Tensor:
460
return self.vision_model(
463
- pixel_values=pixel_values,
461
+ pixel_values_packed=pixel_values_packed,
462
spatial_shapes=spatial_shapes,
465
- packed_mask=packed_mask,
463
cu_seqlens=cu_seqlens,
464
max_seqlen=max_seqlen,
465
)
vllm/model_executor/models/lfm2_vl.py
+129
-51
@@ -50,7 +50,7 @@ from .interfaces import (
50
SupportsMultiModal,
51
SupportsPP,
52
)
53
-from .siglip2 import Siglip2Model
53
+from .lfm2_siglip2 import Siglip2Model
54
from .utils import (
55
AutoWeightsLoader,
56
WeightsMapper,
@@ -450,29 +450,78 @@ class Lfm2VLMultiModalProjector(nn.Module):
450
bias=config.projector_bias,
451
)
452
453
- def forward(self, image_features: torch.Tensor):
454
- image_features = self.pixel_unshuffle(image_features)
455
- if self.projector_use_layernorm:
456
- image_features = self.layer_norm(image_features)
457
- hidden_states = self.linear_1(image_features)
458
- hidden_states = self.act(hidden_states)
459
- hidden_states = self.linear_2(hidden_states)
460
- return hidden_states
453
+ def forward(
454
+ self,
455
+ vision_features_packed: torch.Tensor,
456
+ spatial_shapes: torch.Tensor,
457
+ ) -> torch.Tensor:
458
+ """Project packed vision features without materializing padded tensors.
459
462
- def pixel_unshuffle(self, hidden_states: torch.Tensor):
463
- batch_size, width, height, channels = hidden_states.size()
464
- hidden_states = hidden_states.reshape(
465
- batch_size, width, height // self.factor, channels * self.factor
466
- )
467
- hidden_states = hidden_states.permute(0, 2, 1, 3)
468
- hidden_states = hidden_states.reshape(
469
- batch_size,
470
- height // self.factor,
471
- width // self.factor,
472
- channels * self.factor**2,
460
+ Args:
461
+ vision_features_packed: (total_tokens, hidden_size) packed in tile order.
462
+ spatial_shapes: (num_tiles, 2) on CPU (height, width) per tile.
463
+
464
+ Returns:
465
+ projected_packed: (total_projected_tokens, text_hidden_size)
466
+ """
467
+ assert spatial_shapes.device.type == "cpu", (
468
+ "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "
469
+ "variable-length packing."
470
)
474
- hidden_states = hidden_states.permute(0, 2, 1, 3)
475
- return hidden_states
471
+ factor = self.factor
472
+ device = vision_features_packed.device
473
+ hidden_size = vision_features_packed.shape[-1]
474
+
475
+ spatial_shapes_list: list[list[int]] = spatial_shapes.tolist()
476
+ lengths_list = [h * w for h, w in spatial_shapes_list]
477
+
478
+ gather_idx_parts: list[torch.Tensor] = []
479
+ offset = 0
480
+
481
+ dh = torch.arange(factor, dtype=torch.int64)
482
+ dw = torch.arange(factor, dtype=torch.int64)
483
+ dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")
484
+ dh_flat = dh_grid.reshape(-1)
485
+ dw_flat = dw_grid.reshape(-1)
486
+
487
+ for (height, width), length in zip(spatial_shapes_list, lengths_list):
488
+ if length <= 0:
489
+ continue
490
+ if height % factor != 0 or width % factor != 0:
491
+ raise ValueError(
492
+ "spatial_shapes must be divisible by downsample_factor: "
493
+ f"got ({height}, {width}) with factor={factor}."
494
+ )
495
+ height_out = height // factor
496
+ width_out = width // factor
497
+
498
+ rows_out = torch.arange(height_out, dtype=torch.int64)
499
+ cols_out = torch.arange(width_out, dtype=torch.int64)
500
+ rr, cc = torch.meshgrid(rows_out, cols_out, indexing="ij")
501
+ rr = rr.reshape(-1)
502
+ cc = cc.reshape(-1)
503
+
504
+ token_idx = (rr[:, None] * factor + dh_flat[None, :]) * width + (
505
+ cc[:, None] * factor + dw_flat[None, :]
506
+ )
507
+ gather_idx_parts.append(token_idx.reshape(-1) + offset)
508
+ offset += length
509
+
510
+ if gather_idx_parts:
511
+ gather_idx = torch.cat(gather_idx_parts).to(device=device)
512
+ gathered = vision_features_packed.index_select(0, gather_idx)
513
+ unshuffled = gathered.reshape(-1, factor * factor * hidden_size)
514
+ else:
515
+ unshuffled = vision_features_packed.new_empty(
516
+ (0, factor * factor * hidden_size)
517
+ )
518
+
519
+ if self.projector_use_layernorm:
520
+ unshuffled = self.layer_norm(unshuffled)
521
+ hidden_states = self.linear_1(unshuffled)
522
+ hidden_states = self.act(hidden_states)
523
+ projected_packed = self.linear_2(hidden_states)
524
+ return projected_packed
525
526
527
@MULTIMODAL_REGISTRY.register_processor(
@@ -598,61 +647,90 @@ class Lfm2VLForConditionalGeneration(
647
pixel_values: torch.FloatTensor,
648
spatial_shapes: torch.Tensor,
649
) -> torch.Tensor:
650
+ assert spatial_shapes.device.type == "cpu", (
651
+ "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "
652
+ "variable-length packing."
653
+ )
654
+
655
pixel_values = pixel_values.to(
656
dtype=self.vision_tower.vision_model.embeddings.patch_embedding.weight.dtype
657
) # fp16 compatibility
658
659
# LFM2-VL's HF processor pads patch sequences with trailing zeros.
606
- # Derive the valid-patch mask from spatial_shapes instead of carrying
607
- # pixel_attention_mask through the vLLM multimodal pipeline.
608
- max_seq_len = pixel_values.shape[1]
660
+ # Pack patch tokens upfront so the vision tower runs entirely unpadded.
661
+ spatial_shapes_list: list[list[int]] = spatial_shapes.tolist()
662
+ lengths_list = [h * w for h, w in spatial_shapes_list]
663
+ total_tokens = int(sum(lengths_list))
664
lengths_cpu = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(
665
dtype=torch.int32
666
)
667
max_seqlen = (
613
- lengths_cpu.max().reshape(1).to(device=pixel_values.device)
668
+ lengths_cpu.max().reshape(1)
669
if lengths_cpu.numel()
615
- else torch.tensor([0], dtype=torch.int32, device=pixel_values.device)
670
+ else torch.tensor([0], dtype=torch.int32)
671
+ )
672
+
673
+ if total_tokens == 0:
674
+ return []
675
+
676
+ packed_pixel_values = pixel_values.new_empty(
677
+ (total_tokens, pixel_values.shape[-1])
678
)
617
- lengths = lengths_cpu.to(device=pixel_values.device)
618
- packed_mask = (
619
- torch.arange(max_seq_len, device=pixel_values.device)[None, :]
620
- < lengths[:, None]
679
+ offset = 0
680
+ for i, length in enumerate(lengths_list):
681
+ if length <= 0:
682
+ continue
683
+ packed_pixel_values[offset : offset + length].copy_(
684
+ pixel_values[i, :length]
685
+ )
686
+ offset += length
687
+ packed_pixel_values = packed_pixel_values.unsqueeze(0)
688
+
689
+ lengths = torch.tensor(
690
+ lengths_list, dtype=torch.int32, device=pixel_values.device
691
)
692
cu_seqlens = torch.zeros(
693
lengths.shape[0] + 1,
694
dtype=torch.int32,
625
- device=lengths.device,
695
+ device=pixel_values.device,
696
)
697
cu_seqlens[1:] = torch.cumsum(lengths, dim=0)
698
699
with set_forward_context(None, self.vllm_config):
700
vision_outputs = self.vision_tower(
631
- pixel_values=pixel_values,
701
+ pixel_values_packed=packed_pixel_values,
702
spatial_shapes=spatial_shapes,
633
- packed_mask=packed_mask,
703
cu_seqlens=cu_seqlens,
704
max_seqlen=max_seqlen,
705
)
637
- image_outputs = getattr(vision_outputs, "last_hidden_state", vision_outputs)
638
-
639
- image_features = []
640
-
641
- # spatial_shapes is on CPU (keep_on_cpu=True), so .tolist() is instant
642
- spatial_shapes_list = spatial_shapes.tolist()
643
- for img_idx, (feature_org_h, feature_org_w) in enumerate(spatial_shapes_list):
644
- feature_len = feature_org_h * feature_org_w
645
- feature = image_outputs[img_idx, :feature_len]
646
-
647
- # reshape to original height and width
648
- feature = feature.reshape(1, feature_org_h, feature_org_w, -1)
706
+ image_outputs_packed = getattr(
707
+ vision_outputs, "last_hidden_state", vision_outputs
708
+ )
709
+ vision_features_packed = image_outputs_packed[0]
710
+
711
+ factor = self.multi_modal_projector.factor
712
+ projected_lengths_list: list[int] = []
713
+ for (height, width), length in zip(spatial_shapes_list, lengths_list):
714
+ if length <= 0:
715
+ projected_lengths_list.append(0)
716
+ continue
717
+ if height % factor != 0 or width % factor != 0:
718
+ raise ValueError(
719
+ "spatial_shapes must be divisible by downsample_factor: "
720
+ f"got ({height}, {width}) with factor={factor}."
721
+ )
722
+ projected_lengths_list.append((height // factor) * (width // factor))
723
650
- # project the image representation
651
- img_embedding = self.multi_modal_projector(feature)
724
+ projected_packed = self.multi_modal_projector(
725
+ vision_features_packed=vision_features_packed,
726
+ spatial_shapes=spatial_shapes,
727
+ )
728
653
- # flatten here to handle variable length in naflex
654
- img_embedding = img_embedding.reshape(-1, img_embedding.size(-1))
655
- image_features.append(img_embedding)
729
+ image_features: list[torch.Tensor] = []
730
+ offset = 0
731
+ for out_len in projected_lengths_list:
732
+ image_features.append(projected_packed[offset : offset + out_len])
733
+ offset += out_len
734
735
return image_features
736
vllm/v1/attention/backends/gdn_attn.py
+24
-4
@@ -155,9 +155,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
155
m = common_attn_metadata
156
157
query_start_loc = m.query_start_loc
158
+ query_start_loc_cpu = m.query_start_loc_cpu
159
context_lens_tensor = m.compute_num_computed_tokens()
160
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
161
162
+ spec_sequence_masks_cpu: torch.Tensor | None = None
163
if (
164
not self.use_spec_decode
165
or num_decode_draft_tokens_cpu is None
@@ -169,12 +171,13 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
171
spec_sequence_masks = None
172
num_spec_decodes = 0
173
else:
172
- spec_sequence_masks = num_decode_draft_tokens_cpu >= 0
173
- num_spec_decodes = spec_sequence_masks.sum().item()
174
+ spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
175
+ num_spec_decodes = spec_sequence_masks_cpu.sum().item()
176
if num_spec_decodes == 0:
177
spec_sequence_masks = None
178
+ spec_sequence_masks_cpu = None
179
else:
177
- spec_sequence_masks = spec_sequence_masks.to(
180
+ spec_sequence_masks = spec_sequence_masks_cpu.to(
181
query_start_loc.device, non_blocking=True
182
)
183
@@ -189,9 +192,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
192
non_spec_state_indices_tensor = m.block_table_tensor[:, 0]
193
spec_query_start_loc = None
194
non_spec_query_start_loc = query_start_loc
195
+ non_spec_query_start_loc_cpu = query_start_loc_cpu
196
num_accepted_tokens = None
197
else:
198
query_lens = query_start_loc[1:] - query_start_loc[:-1]
199
+ assert spec_sequence_masks_cpu is not None
200
+ query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
201
202
non_spec_query_lens = query_lens[~spec_sequence_masks]
203
num_decodes = (non_spec_query_lens == 1).sum().item()
@@ -219,6 +225,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
225
non_spec_state_indices_tensor = None
226
spec_query_start_loc = query_start_loc
227
non_spec_query_start_loc = None
228
+ non_spec_query_start_loc_cpu = None
229
else:
230
spec_token_masks = torch.repeat_interleave(
231
spec_sequence_masks, query_lens
@@ -253,6 +260,15 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
260
dim=0,
261
out=non_spec_query_start_loc[1:],
262
)
263
+ non_spec_query_start_loc_cpu = torch.zeros(
264
+ query_lens_cpu.size(0) - num_spec_decodes + 1,
265
+ dtype=torch.int32,
266
+ )
267
+ torch.cumsum(
268
+ query_lens_cpu[~spec_sequence_masks_cpu],
269
+ dim=0,
270
+ out=non_spec_query_start_loc_cpu[1:],
271
+ )
272
273
assert num_accepted_tokens is not None
274
num_accepted_tokens = num_accepted_tokens[spec_sequence_masks]
@@ -261,8 +277,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
277
has_initial_state = context_lens_tensor > 0
278
if spec_sequence_masks is not None:
279
has_initial_state = has_initial_state[~spec_sequence_masks]
280
+ assert non_spec_query_start_loc_cpu is not None
281
nums_dict, batch_ptr, token_chunk_offset_ptr = (
265
- compute_causal_conv1d_metadata(non_spec_query_start_loc)
282
+ compute_causal_conv1d_metadata(
283
+ non_spec_query_start_loc_cpu,
284
+ device=query_start_loc.device,
285
+ )
286
)
287
else:
288
has_initial_state = None
vllm/v1/attention/backends/mamba_attn.py
+10
-7
@@ -219,21 +219,24 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
219
if num_prefills > 0:
220
if num_computed_tokens is None:
221
num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
222
- num_computed_tokens_cpu = num_computed_tokens.cpu()
222
223
+ query_start_loc_p_cpu = (
224
+ common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]
225
+ - num_decode_tokens
226
+ )
227
query_start_loc_p = (
228
common_attn_metadata.query_start_loc[-num_prefills - 1 :]
229
- num_decode_tokens
230
)
228
- has_initial_states_cpu = (
229
- num_computed_tokens_cpu[num_reqs - num_prefills : num_reqs] > 0
230
- )
231
- has_initial_states_p = has_initial_states_cpu.to(
232
- common_attn_metadata.query_start_loc.device
231
+ has_initial_states_p = (
232
+ num_computed_tokens[num_reqs - num_prefills : num_reqs] > 0
233
)
234
235
nums_dict, batch_ptr, token_chunk_offset_ptr = (
236
- compute_causal_conv1d_metadata(query_start_loc_p)
236
+ compute_causal_conv1d_metadata(
237
+ query_start_loc_p_cpu,
238
+ device=common_attn_metadata.query_start_loc.device,
239
+ )
240
)
241
242
if self.vllm_config.cache_config.enable_prefix_caching:
vllm/v1/attention/backends/utils.py
+8
-4
@@ -732,13 +732,17 @@ def create_fast_prefill_custom_backend(
732
return attn_backend
733
734
735
-def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
736
- # Needed for causal_conv1d
737
- seqlens = query_start_loc_p.diff().to("cpu")
735
+def compute_causal_conv1d_metadata(
736
+ query_start_loc_p_cpu: torch.Tensor,
737
+ *,
738
+ device: torch.device,
739
+):
740
+ # Needed for causal_conv1d. Use the CPU query_start_loc to avoid DtoH sync.
741
+ assert query_start_loc_p_cpu.device.type == "cpu"
742
+ seqlens = query_start_loc_p_cpu.diff()
743
nums_dict = {} # type: ignore
744
batch_ptr = None
745
token_chunk_offset_ptr = None
741
- device = query_start_loc_p.device
746
for BLOCK_M in [8]: # cover all BLOCK_M values
747
nums = -(-seqlens // BLOCK_M)
748
nums_dict[BLOCK_M] = {}