fe85a92e86 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/attention/utils.py

+1

@@ -107,6 +107,7 @@ def create_common_attn_metadata(

107 query_start_loc=query_start_loc,

108 query_start_loc_cpu=query_start_loc_cpu,

109 seq_lens=seq_lens,

110 + seq_lens_cpu_upper_bound=seq_lens_cpu,

111 _seq_lens_cpu=seq_lens_cpu,

112 _num_computed_tokens_cpu=num_computed_tokens_cpu,

113 num_reqs=batch_spec.batch_size,

tests/v1/spec_decode/test_tree_attention.py

+3 -1

@@ -241,11 +241,13 @@ def forward_attention(

241 )

242 kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

243 builder = builder_cls(kv_cache_spec, [], vllm_config, q.device)

244 + seq_lens_cpu = seq_lens.cpu()

245 common_attn_metadata = CommonAttentionMetadata(

246 query_start_loc=query_start_loc,

247 query_start_loc_cpu=query_start_loc.cpu(),

248 seq_lens=seq_lens,

248 - _seq_lens_cpu=seq_lens.cpu(),

249 + seq_lens_cpu_upper_bound=seq_lens_cpu,

250 + _seq_lens_cpu=seq_lens_cpu,

251 _num_computed_tokens_cpu=context_lens.cpu(),

252 num_reqs=batch_size,

253 num_actual_tokens=num_actual_tokens,

vllm/model_executor/layers/attention/cross_attention.py

+12 -4

@@ -90,15 +90,23 @@ def create_cross_attention_backend(

90 assert new_metadata.encoder_seq_lens_cpu is not None

91 max_encoder_len = int(new_metadata.encoder_seq_lens_cpu.max())

92 new_metadata.max_seq_len = max_encoder_len

93 - # Any computed tokens indicated decode step>1 (no chunked prefill)

94 - num_cache_decodes = (

95 - (common_attn_metadata.num_computed_tokens_cpu > 0).sum().item()

93 + # Any computed tokens indicates decode step>1 (no chunked prefill).

94 + # The upper bound is exact for this `> 0` test - prefill rows have

95 + # num_computed == 0 and decode rows have num_computed > 0.

96 + query_lens_cpu = (

97 + common_attn_metadata.query_start_loc_cpu[1:]

98 + - common_attn_metadata.query_start_loc_cpu[:-1]

99 )

100 + assert common_attn_metadata.seq_lens_cpu_upper_bound is not None

101 + num_computed_tokens_cpu = (

102 + common_attn_metadata.seq_lens_cpu_upper_bound - query_lens_cpu

103 + )

104 + num_cache_decodes = (num_computed_tokens_cpu > 0).sum().item()

105 if num_cache_decodes > 0:

106 # CrossAttn KV cache has already been populated on first decoder step,

107 # skip slot_mapping calculation for requests that do not need

108 # reshape_and_cache.

101 - num_tokens = common_attn_metadata.num_computed_tokens_cpu.numpy()

109 + num_tokens = num_computed_tokens_cpu.numpy()

110 new_metadata.encoder_seq_lens_cpu = np.where(

111 num_tokens > 0, 0, new_metadata.encoder_seq_lens_cpu

112 )

vllm/model_executor/layers/attention/mla_attention.py

+10 -5

@@ -1822,13 +1822,18 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1822

1823 prefill_metadata = None

1824 if num_prefills > 0:

1825 - num_computed_tokens_cpu = (

1826 - common_attn_metadata.compute_num_computed_tokens().cpu()

1827 - )

1828 -

1825 reqs_start = num_decodes # prefill_start

1826

1831 - context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]

1827 + # Upper bound is exact for prefill rows (no D2H sync).

1828 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

1829 + assert seq_lens_cpu is not None

1830 + prefill_query_lens_cpu = (

1831 + query_start_loc_cpu[reqs_start + 1 : num_reqs + 1]

1832 + - query_start_loc_cpu[reqs_start:num_reqs]

1833 + )

1834 + context_lens_cpu = (

1835 + seq_lens_cpu[reqs_start:num_reqs] - prefill_query_lens_cpu

1836 + )

1837 max_context_len_cpu = context_lens_cpu.max().item()

1838 num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()

1839 prefill_query_start_loc = (

vllm/v1/attention/backend.py

+6

@@ -397,6 +397,12 @@ class CommonAttentionMetadata:

397 (num_computed_tokens < num_prompt_tokens). Used by some backends to

398 distinguish actual decodes from short extends."""

399

400 + seq_lens_cpu_upper_bound: torch.Tensor | None = None

401 + """(batch_size,) CPU upper bound on seq_lens. Precise for prefill rows

402 + and for all rows outside async spec decode; optimistic for async-spec

403 + decode rows (assumes every draft was accepted). Not safe for kernels

404 + that need exact per-row context lengths on decode rows."""

405 +

406 # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)

407 _seq_lens_cpu: torch.Tensor | None = None

408 _num_computed_tokens_cpu: torch.Tensor | None = None

vllm/v1/attention/backends/flex_attention.py

+5 -4

@@ -782,10 +782,11 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat

782 def build_for_cudagraph_capture(

783 self, common_attn_metadata: CommonAttentionMetadata

784 ) -> FlexAttentionMetadata:

785 - # Use actual max_seq_len instead of max_model_len to avoid

786 - # torch.compile recompilation during CUDA graph capture.

787 - common_attn_metadata.max_seq_len = (

788 - common_attn_metadata.seq_lens_cpu.max().item()

785 + # Use actual max_seq_len (not max_model_len) to avoid torch.compile

786 + # recompilation during CUDA graph capture.

787 + assert common_attn_metadata.seq_lens_cpu_upper_bound is not None

788 + common_attn_metadata.max_seq_len = int(

789 + common_attn_metadata.seq_lens_cpu_upper_bound.max().item()

790 )

791 return self.build(

792 common_prefix_len=0, common_attn_metadata=common_attn_metadata

vllm/v1/attention/backends/mla/flashmla_sparse.py

+4 -1

@@ -364,7 +364,10 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

364 # For pure decode batches, prefill_request_id will be None

365 # For mixed batches, it will have -1 for decode and request_id for prefill

366 if num_prefills > 0:

367 - seq_lens_cpu = common_attn_metadata.seq_lens.cpu()

367 + # Upper bound is exact for prefill rows (the `[num_decodes:]`

368 + # slice below), so no D2H sync is needed.

369 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

370 + assert seq_lens_cpu is not None

371 seq_lens = common_attn_metadata.seq_lens

372 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

373

vllm/v1/attention/backends/mla/indexer.py

+6 -2

@@ -554,8 +554,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

554 query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]

555 )

556 max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024

557 + # Upper bound is exact for prefill rows (the `[num_decodes:]`

558 + # slice below).

559 + assert common_attn_metadata.seq_lens_cpu_upper_bound is not None

560 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

561 chunk_specs = split_indexer_prefill_chunks(

558 - common_attn_metadata.seq_lens_cpu[num_decodes:],

562 + seq_lens_cpu[num_decodes:],

563 prefill_query_lens_cpu,

564 self.max_prefill_buffer_size,

565 max_logits_bytes,

@@ -566,7 +570,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

570 req_slice,

571 query_slice,

572 query_start_loc_cpu,

569 - common_attn_metadata.seq_lens_cpu,

573 + seq_lens_cpu,

574 common_attn_metadata.block_table_tensor,

575 skip_kv_gather=query_slice.start > 0,

576 )

vllm/v1/attention/backends/utils.py

+7 -1

@@ -356,6 +356,7 @@ def make_local_attention_virtual_batches(

356 block_table_tensor=block_table_local,

357 slot_mapping=common_attn_metadata.slot_mapping,

358 causal=True,

359 + seq_lens_cpu_upper_bound=seq_lens_cpu,

360 _seq_lens_cpu=seq_lens_cpu,

361 _num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),

362 ), make_block_table

@@ -414,6 +415,7 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(

415 block_table_tensor=common_attn_metadata.block_table_tensor,

416 slot_mapping=common_attn_metadata.slot_mapping,

417 causal=True,

418 + seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,

419 _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,

420 _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

421 )

@@ -445,7 +447,11 @@ def split_decodes_prefills_and_extends(

447 num_reqs = common_attn_metadata.num_reqs

448 num_tokens = common_attn_metadata.num_actual_tokens

449 query_start_loc = common_attn_metadata.query_start_loc_cpu

448 - seq_lens = common_attn_metadata.seq_lens_cpu

450 + # Upper bound is exact for prefill rows; decode rows still satisfy

451 + # seq_len > query_len under the optimistic bound, so `seq_lens ==

452 + # query_lens` identifies prefills correctly either way.

453 + assert common_attn_metadata.seq_lens_cpu_upper_bound is not None

454 + seq_lens = common_attn_metadata.seq_lens_cpu_upper_bound

455

456 if max_query_len <= decode_threshold:

457 return num_reqs, 0, 0, num_tokens, 0, 0

vllm/v1/spec_decode/dflash.py

+7

@@ -151,6 +151,12 @@ class DFlashProposer(SpecDecodeBaseProposer):

151 if has_num_rejected:

152 effective_seq_lens = effective_seq_lens - num_rejected_tokens_gpu

153

154 + # Skip num_rejected_tokens (GPU-only); overestimating is fine here.

155 + new_seq_lens_cpu_upper_bound = (

156 + cad.seq_lens_cpu_upper_bound + num_query_per_req

157 + if cad.seq_lens_cpu_upper_bound is not None

158 + else None

159 + )

160 new_cad = CommonAttentionMetadata(

161 query_start_loc=new_query_start_loc,

162 seq_lens=effective_seq_lens + num_query_per_req,

@@ -160,6 +166,7 @@ class DFlashProposer(SpecDecodeBaseProposer):

166 ),

167 _seq_lens_cpu=None,

168 _num_computed_tokens_cpu=None,

169 + seq_lens_cpu_upper_bound=new_seq_lens_cpu_upper_bound,

170 num_reqs=cad.num_reqs,

171 num_actual_tokens=num_query_total,

172 max_query_len=num_query_per_req,

vllm/v1/spec_decode/llm_base_proposer.py

+9 -1

@@ -593,6 +593,8 @@ class SpecDecodeBaseProposer:

593 common_attn_metadata._seq_lens_cpu += 1

594 if common_attn_metadata._num_computed_tokens_cpu is not None:

595 common_attn_metadata._num_computed_tokens_cpu += 1

596 + if common_attn_metadata.seq_lens_cpu_upper_bound is not None:

597 + common_attn_metadata.seq_lens_cpu_upper_bound += 1

598

599 # Rebuild attention metadata

600 _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(

@@ -959,6 +961,7 @@ class SpecDecodeBaseProposer:

961 query_start_loc_cpu=query_start_loc_cpu,

962 _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,

963 _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

964 + seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,

965 num_reqs=common_attn_metadata.num_reqs,

966 num_actual_tokens=total_num_tokens,

967 max_query_len=new_query_len_per_req.max().item(),

@@ -1183,7 +1186,11 @@ class SpecDecodeBaseProposer:

1186

1187 device = common_attn_metadata.query_start_loc.device

1188 query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

1186 - new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu - num_rejected_tokens

1189 + # upper_bound - rejected = actual post-rejection seq_lens (no D2H sync).

1190 + assert common_attn_metadata.seq_lens_cpu_upper_bound is not None

1191 + new_seq_lens_cpu = (

1192 + common_attn_metadata.seq_lens_cpu_upper_bound - num_rejected_tokens

1193 + )

1194

1195 # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]

1196 new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

@@ -1237,6 +1244,7 @@ class SpecDecodeBaseProposer:

1244 query_start_loc_cpu=new_query_start_loc_cpu,

1245 _seq_lens_cpu=new_seq_lens_cpu,

1246 _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

1247 + seq_lens_cpu_upper_bound=new_seq_lens_cpu,

1248 num_reqs=common_attn_metadata.num_reqs,

1249 num_actual_tokens=total_num_tokens,

1250 max_query_len=new_query_len_per_req.max().item(),

vllm/v1/worker/gpu/attn_utils.py

+4

@@ -227,12 +227,15 @@ def build_attn_metadata(

227 block_tables: Sequence[torch.Tensor],

228 slot_mappings: torch.Tensor,

229 kv_cache_config: KVCacheConfig,

230 + seq_lens_cpu_upper_bound: torch.Tensor | None = None,

231 dcp_local_seq_lens: torch.Tensor | None = None,

232 encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] | None = None,

233 ) -> dict[str, Any]:

234 seq_lens = seq_lens[:num_reqs]

235 if dcp_local_seq_lens is not None:

236 dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]

237 + if seq_lens_cpu_upper_bound is not None:

238 + seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs]

239

240 attn_metadata: dict[str, Any] = {}

241 num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)

@@ -244,6 +247,7 @@ def build_attn_metadata(

247 query_start_loc=query_start_loc_gpu,

248 query_start_loc_cpu=query_start_loc_cpu,

249 seq_lens=seq_lens,

250 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

251 max_seq_len=max_seq_len,

252 num_reqs=num_reqs,

253 num_actual_tokens=num_tokens,

vllm/v1/worker/gpu/input_batch.py

+5

@@ -60,6 +60,8 @@ class InputBatch:

60 query_start_loc_np: np.ndarray

61 # [num_reqs]

62 seq_lens: torch.Tensor

63 + # [num_reqs] CPU upper bound on seq_lens (see CommonAttentionMetadata).

64 + seq_lens_cpu_upper_bound: torch.Tensor

65 # [num_reqs]

66 dcp_local_seq_lens: torch.Tensor | None

67

@@ -121,6 +123,8 @@ class InputBatch:

123 logits_indices = query_start_loc[1:] - 1

124 cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)

125 cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)

126 + # Dummy: seq_len == query_len (fresh-prefill shape).

127 + seq_lens_cpu_upper_bound = torch.from_numpy(num_scheduled_tokens.copy())

128 return cls(

129 req_ids=req_ids,

130 num_reqs=num_reqs,

@@ -136,6 +140,7 @@ class InputBatch:

140 query_start_loc=query_start_loc,

141 query_start_loc_np=query_start_loc_np,

142 seq_lens=seq_lens,

143 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

144 dcp_local_seq_lens=None,

145 input_ids=input_ids,

146 positions=positions,

vllm/v1/worker/gpu/model_runner.py

+18

@@ -799,6 +799,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

799 total_num_logits,

800 )

801

802 + # CPU upper bound on seq_lens; padded entries left at zero.

803 + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)

804 + np.add(

805 + self.req_states.num_computed_tokens_np[idx_mapping_np],

806 + num_scheduled_tokens,

807 + out=seq_lens_cpu_upper_bound_np[:num_reqs],

808 + )

809 + seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)

810 +

811 return InputBatch(

812 req_ids=req_ids,

813 num_reqs=num_reqs,

@@ -814,6 +823,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

823 query_start_loc=query_start_loc,

824 query_start_loc_np=query_start_loc_np,

825 seq_lens=seq_lens,

826 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

827 dcp_local_seq_lens=dcp_local_seq_lens,

828 input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],

829 positions=self.input_buffers.positions[:num_tokens_after_padding],

@@ -927,6 +937,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

937 np.minimum(

938 computed_prefill, self.req_states.prefill_len.np, out=computed_prefill

939 )

940 + # Advance the CPU mirror optimistically (assume all scheduled accepted).

941 + self.req_states.num_computed_tokens_np[idx_mapping_np] += (

942 + input_batch.num_scheduled_tokens

943 + )

944

945 @torch.inference_mode()

946 def execute_model(

@@ -1297,6 +1311,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1311 np.minimum(

1312 computed_prefill, self.req_states.prefill_len.np, out=computed_prefill

1313 )

1314 + # Advance the CPU mirror optimistically (assume all scheduled accepted).

1315 + self.req_states.num_computed_tokens_np[idx_mapping_np] += (

1316 + input_batch.num_scheduled_tokens

1317 + )

1318

1319 ########### EPLB methods start ###########

1320 @property

vllm/v1/worker/gpu/model_states/default.py

+8 -1

@@ -173,6 +173,12 @@ class DefaultModelState(ModelState):

173 num_tokens = input_batch.num_tokens

174 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)

175 max_query_len = input_batch.num_scheduled_tokens.max().item()

176 + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound

177 + if for_capture:

178 + # Capture with worst-case max_seq_len so the graph is valid at any replay.

179 + max_seq_len = self.max_model_len

180 + else:

181 + max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item())

182 attn_metadata = build_attn_metadata(

183 attn_groups=attn_groups,

184 num_reqs=num_reqs,

@@ -181,10 +187,11 @@ class DefaultModelState(ModelState):

187 query_start_loc_cpu=query_start_loc_cpu,

188 max_query_len=max_query_len,

189 seq_lens=input_batch.seq_lens,

184 - max_seq_len=self.max_model_len,

190 + max_seq_len=max_seq_len,

191 block_tables=block_tables,

192 slot_mappings=slot_mappings,

193 kv_cache_config=kv_cache_config,

194 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

195 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,

196 )

197 return attn_metadata

vllm/v1/worker/gpu/model_states/whisper.py

+7 -1

@@ -117,6 +117,11 @@ class WhisperModelState(ModelState):

117

118 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)

119 max_query_len = input_batch.num_scheduled_tokens.max().item()

120 + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound

121 + if for_capture:

122 + max_seq_len = self.max_model_len

123 + else:

124 + max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item())

125 attn_metadata = build_attn_metadata(

126 attn_groups=attn_groups,

127 num_reqs=num_reqs,

@@ -125,10 +130,11 @@ class WhisperModelState(ModelState):

130 query_start_loc_cpu=query_start_loc_cpu,

131 max_query_len=max_query_len,

132 seq_lens=input_batch.seq_lens,

128 - max_seq_len=self.max_model_len,

133 + max_seq_len=max_seq_len,

134 block_tables=block_tables,

135 slot_mappings=slot_mappings,

136 kv_cache_config=kv_cache_config,

137 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

138 dcp_local_seq_lens=input_batch.dcp_local_seq_lens,

139 encoder_seq_lens=encoder_seq_lens,

140 )

vllm/v1/worker/gpu/states.py

+3

@@ -57,6 +57,8 @@ class RequestState:

57 self.num_computed_tokens = StagedWriteTensor(

58 self.max_num_reqs, dtype=torch.int32, device=device

59 )

60 + # Optimistic CPU mirror of num_computed_tokens (upper bound on GPU value).

61 + self.num_computed_tokens_np = np.zeros(self.max_num_reqs, dtype=np.int32)

62

63 # Last sampled tokens.

64 self.last_sampled_tokens = torch.zeros(

@@ -100,6 +102,7 @@ class RequestState:

102 self.total_len.stage_write_elem(req_idx, prefill_len)

103 self.all_token_ids.stage_write(req_idx, 0, all_token_ids)

104 self.num_computed_prefill_tokens[req_idx] = num_computed_tokens

105 + self.num_computed_tokens_np[req_idx] = num_computed_tokens

106 self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)

107

108 if num_computed_tokens > 0 and num_computed_tokens <= prefill_len:

vllm/v1/worker/gpu_model_runner.py

+2

@@ -2155,6 +2155,7 @@ class GPUModelRunner(

2155 :num_reqs_padded

2156 ]

2157 seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs_padded]

2158 + seq_lens_cpu_upper_bound = seq_lens_cpu

2159

2160 # is_prefilling: True if request is still in prefill phase.

2161 # Used by mamba backends to distinguish actual decodes from

@@ -2172,6 +2173,7 @@ class GPUModelRunner(

2173 seq_lens=self.seq_lens[:num_reqs_padded],

2174 _seq_lens_cpu=seq_lens_cpu,

2175 _num_computed_tokens_cpu=num_computed_tokens_cpu,

2176 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

2177 num_reqs=num_reqs_padded,

2178 num_actual_tokens=num_tokens_padded,

2179 max_query_len=max_query_len,

vllm/v1/worker/ubatch_utils.py

+25 -5

@@ -177,7 +177,22 @@ def _make_metadata_with_slice(

177 query_start_loc[1:] -= tokens_skipped

178 query_start_loc_cpu[1:] -= tokens_skipped

179 seq_lens = attn_metadata.seq_lens[request_slice]

180 - seq_lens_cpu = attn_metadata.seq_lens_cpu[request_slice]

180 + # Read raw fields to avoid triggering the deprecated D2H-syncing properties.

181 + seq_lens_cpu = (

182 + attn_metadata._seq_lens_cpu[request_slice]

183 + if attn_metadata._seq_lens_cpu is not None

184 + else None

185 + )

186 + seq_lens_cpu_upper_bound = (

187 + attn_metadata.seq_lens_cpu_upper_bound[request_slice]

188 + if attn_metadata.seq_lens_cpu_upper_bound is not None

189 + else None

190 + )

191 + num_computed_tokens_cpu = (

192 + attn_metadata._num_computed_tokens_cpu[request_slice]

193 + if attn_metadata._num_computed_tokens_cpu is not None

194 + else None

195 + )

196

197 if splits_last_request:

198 # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate

@@ -190,12 +205,16 @@ def _make_metadata_with_slice(

205 # Make sure we don't modify the seq_lens tensors

206 # (not cudagraph compatible)

207 seq_lens = seq_lens.clone()

193 - seq_lens_cpu = seq_lens_cpu.clone()

208 seq_lens[-1] -= tokens_skipped

195 - seq_lens_cpu[-1] -= tokens_skipped

209 + if seq_lens_cpu is not None:

210 + seq_lens_cpu = seq_lens_cpu.clone()

211 + seq_lens_cpu[-1] -= tokens_skipped

212 + if seq_lens_cpu_upper_bound is not None:

213 + seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound.clone()

214 + seq_lens_cpu_upper_bound[-1] -= tokens_skipped

215

197 - max_seq_len = int(seq_lens_cpu.max())

198 - num_computed_tokens_cpu = attn_metadata.num_computed_tokens_cpu[request_slice]

216 + assert seq_lens_cpu_upper_bound is not None

217 + max_seq_len = int(seq_lens_cpu_upper_bound.max())

218

219 num_requests = request_slice.stop - request_slice.start

220 num_actual_tokens = token_slice.stop - token_slice.start

@@ -221,6 +240,7 @@ def _make_metadata_with_slice(

240 max_seq_len=max_seq_len,

241 block_table_tensor=block_table_tensor,

242 slot_mapping=slot_mapping,

243 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

244 _seq_lens_cpu=seq_lens_cpu,

245 _num_computed_tokens_cpu=num_computed_tokens_cpu,

246 )

[Core] Avoid seq_lens_cpu GPU->CPU sync (#40654)