c64c356990 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -47,6 +47,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys

47 from vllm.model_executor.models.transformers.utils import recursive_replace_linear

48 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix

49 from vllm.multimodal import MULTIMODAL_REGISTRY

50 +from vllm.platforms import current_platform

51 from vllm.v1.outputs import LogprobsTensors

52 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata

53 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu

@@ -516,11 +517,6 @@ def _compiled_sample_step(

517 num_decode = decode_slots.shape[0]

518 device = decode_slots.device

519

519 - # Clear outputs so prefill / non-decode slots report 0 (decode slots are

520 - # overwritten below).

521 - sampled.zero_()

522 - num_sampled.zero_()

523 -

520 # ---- Phase 1: Temperature schedule ----

521 steps_f = step_tensor[decode_slots].float()

522 remaining = (max_denoising_steps - steps_f).clamp(min=1.0)

@@ -1276,9 +1272,12 @@ class DiffusionSampler:

1272 src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1)

1273 logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype)

1274

1279 - # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0.

1275 + # Clear once: the tiled loop below only scatters its own decode slots,

1276 + # so it must not re-clear earlier tiles' writes.

1277 sampled = self._sampled[:num_reqs]

1278 num_sampled = self._num_sampled[:num_reqs]

1279 + sampled.zero_()

1280 + num_sampled.zero_()

1281

1282 all_slots = input_batch.idx_mapping[:num_reqs]

1283

@@ -1286,94 +1285,109 @@ class DiffusionSampler:

1285 # since it mutates is_encoder_phase (commit→False, converge→True).

1286 is_committing = states.is_encoder_phase[decode_slots].clone()

1287

1289 - # --- Single compiled call: temp → sample → probs → post-process ---

1290 - scaled = _compiled_sample_step(

1291 - logits,

1292 - decode_slots,

1293 - decode_idx,

1294 - all_slots,

1295 - valid_canvas_len,

1296 - # State

1297 - states.canvas,

1298 - states.argmax_canvas,

1299 - states.step,

1300 - states.is_encoder_phase,

1301 - states.confident,

1302 - states.self_conditioning_embeds,

1303 - self.embed_weight,

1304 - self.normalizer,

1305 - states.accepted_canvas_history,

1306 - states.accepted_canvas_history_len,

1307 - # Output

1308 - sampled,

1309 - num_sampled,

1310 - self.req_states.draft_tokens,

1311 - # Config

1312 - max_denoising_steps=float(states.max_denoising_steps),

1313 - t_min=self.t_min,

1314 - t_max=self.t_max,

1315 - confidence_threshold=self.confidence_threshold,

1316 - vocab_size=self.vocab_size,

1317 - CL=self.canvas_length,

1318 - ST=states.stability_threshold,

1319 - entropy_bound=self.entropy_bound,

1320 - sc_vocab_start=self.sc_vocab_start,

1321 - sc_vocab_end=self.sc_vocab_end,

1322 - tp_size=self.tp_size,

1323 - tp_group_name=self.tp_group_name,

1324 - )

1325 -

1326 - # --- Logprobs: stash on convergence, return on commit ---

1288 slots_np = input_batch.idx_mapping_np[:num_reqs]

1289 is_decode_np = per_req_nlogits_np > 0

1290 + max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)

1291

1292 + # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in

1293 + # _compiled_sample_step keeps several live [group * CL, vocab] copies, so

1294 + # size each tile to a fraction of free memory to bound the transient at

1295 + # high concurrency. Tiling is bit-identical to a single pass.

1296 + group = max(num_decode, 1)

1297 + if num_decode > 0:

1298 + free, _ = current_platform.mem_get_info()

1299 + # ~10 transient fp32 copies of [group * CL, vocab] inside the step

1300 + # (eager peaks at ~8; pad for allocator overhead and small tensors).

1301 + bytes_per_req = CL * self.vocab_size * 4 * 10

1302 + budget = int(free * 0.5) // max(bytes_per_req, 1)

1303 + group = max(1, min(num_decode, budget))

1304 +

1305 + for start_req in range(0, num_decode, group):

1306 + end_req = min(start_req + group, num_decode)

1307 + tile = slice(start_req, end_req)

1308 + tile_slots = decode_slots[tile]

1309 +

1310 + scaled = _compiled_sample_step(

1311 + logits[start_req * CL : end_req * CL],

1312 + tile_slots,

1313 + decode_idx[tile],

1314 + all_slots,

1315 + valid_canvas_len[tile],

1316 + # State

1317 + states.canvas,

1318 + states.argmax_canvas,

1319 + states.step,

1320 + states.is_encoder_phase,

1321 + states.confident,

1322 + states.self_conditioning_embeds,

1323 + self.embed_weight,

1324 + self.normalizer,

1325 + states.accepted_canvas_history,

1326 + states.accepted_canvas_history_len,

1327 + # Output

1328 + sampled,

1329 + num_sampled,

1330 + self.req_states.draft_tokens,

1331 + # Config

1332 + max_denoising_steps=float(states.max_denoising_steps),

1333 + t_min=self.t_min,

1334 + t_max=self.t_max,

1335 + confidence_threshold=self.confidence_threshold,

1336 + vocab_size=self.vocab_size,

1337 + CL=CL,

1338 + ST=states.stability_threshold,

1339 + entropy_bound=self.entropy_bound,

1340 + sc_vocab_start=self.sc_vocab_start,

1341 + sc_vocab_end=self.sc_vocab_end,

1342 + tp_size=self.tp_size,

1343 + tp_group_name=self.tp_group_name,

1344 + )

1345 +

1346 + # Logprobs for denoise steps that just converged (is_encoder_phase

1347 + # flipped False→True), stashed per tile so `scaled` is freed each tile.

1348 + if max_num_logprobs >= 0:

1349 + converged_mask = states.is_encoder_phase[tile_slots]

1350 + just_converged = converged_mask & ~is_committing[tile]

1351 + if just_converged.any():

1352 + flat_logits = scaled.reshape(-1, scaled.shape[-1])

1353 + argmax_tokens = scaled.argmax(dim=-1)

1354 + for local_idx in just_converged.nonzero(as_tuple=True)[0]:

1355 + li = local_idx.item()

1356 + slot = tile_slots[local_idx]

1357 + # Stash only the real canvas positions (== CL unless this

1358 + # canvas was truncated near max_model_len); padded tail

1359 + # positions are never emitted.

1360 + k_i = int(valid_canvas_len_np[start_req + li])

1361 + pos = li * CL

1362 + self._pending_logprobs[slot.item()] = compute_topk_logprobs(

1363 + flat_logits[pos : pos + k_i],

1364 + max_num_logprobs,

1365 + argmax_tokens[local_idx][:k_i],

1366 + )

1367 +

1368 + # Commit steps: is_committing was True at entry. Reassemble previously

1369 + # stashed logprobs and attach to SamplerOutput.

1370 logprobs_tensors = None

1331 - max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np)

1332 - if max_num_logprobs >= 0:

1333 - # Denoise steps that just converged: the compiled step flipped

1334 - # is_encoder_phase from False→True. Detect as slots where

1335 - # is_encoder_phase is now True but is_committing was False.

1336 - converged_mask = states.is_encoder_phase[decode_slots]

1337 - just_converged = converged_mask & ~is_committing

1338 - if just_converged.any():

1339 - flat_logits = scaled.reshape(-1, scaled.shape[-1])

1340 - argmax_tokens = scaled.argmax(dim=-1)

1341 - for local_idx in just_converged.nonzero(as_tuple=True)[0]:

1342 - li = local_idx.item()

1343 - slot = decode_slots[local_idx]

1344 - # Stash only the real canvas positions (== CL unless this

1345 - # canvas was truncated near max_model_len); padded tail

1346 - # positions are never emitted.

1347 - k_i = int(valid_canvas_len_np[li])

1348 - start = li * CL

1349 - self._pending_logprobs[slot.item()] = compute_topk_logprobs(

1350 - flat_logits[start : start + k_i],

1351 - max_num_logprobs,

1352 - argmax_tokens[local_idx][:k_i],

1353 - )

1354 -

1355 - # Commit steps: is_committing was True at entry. Reassemble

1356 - # previously stashed logprobs and attach to SamplerOutput.

1357 - if is_committing.any() and self._pending_logprobs:

1358 - parts_ids, parts_lp, parts_ranks = [], [], []

1359 - cu_gen: list[int] = []

1360 - flat_offset = 0

1361 - for i in range(num_reqs):

1362 - cu_gen.append(flat_offset)

1363 - slot = int(slots_np[i])

1364 - if is_decode_np[i] and slot in self._pending_logprobs:

1365 - lp = self._pending_logprobs.pop(slot)

1366 - parts_ids.append(lp.logprob_token_ids)

1367 - parts_lp.append(lp.logprobs)

1368 - parts_ranks.append(lp.selected_token_ranks)

1369 - flat_offset += lp.logprobs.shape[0]

1370 - if parts_ids:

1371 - logprobs_tensors = LogprobsTensors(

1372 - logprob_token_ids=torch.cat(parts_ids),

1373 - logprobs=torch.cat(parts_lp),

1374 - selected_token_ranks=torch.cat(parts_ranks),

1375 - cu_num_generated_tokens=cu_gen,

1376 - )

1371 + if max_num_logprobs >= 0 and is_committing.any() and self._pending_logprobs:

1372 + parts_ids, parts_lp, parts_ranks = [], [], []

1373 + cu_gen: list[int] = []

1374 + flat_offset = 0

1375 + for i in range(num_reqs):

1376 + cu_gen.append(flat_offset)

1377 + slot = int(slots_np[i])

1378 + if is_decode_np[i] and slot in self._pending_logprobs:

1379 + lp = self._pending_logprobs.pop(slot)

1380 + parts_ids.append(lp.logprob_token_ids)

1381 + parts_lp.append(lp.logprobs)

1382 + parts_ranks.append(lp.selected_token_ranks)

1383 + flat_offset += lp.logprobs.shape[0]

1384 + if parts_ids:

1385 + logprobs_tensors = LogprobsTensors(

1386 + logprob_token_ids=torch.cat(parts_ids),

1387 + logprobs=torch.cat(parts_lp),

1388 + selected_token_ranks=torch.cat(parts_ranks),

1389 + cu_num_generated_tokens=cu_gen,

1390 + )

1391

1392 return self._build_output(

1393 input_batch,

[Perf] Bound DiffusionGemma sampler transient via request-tiled logits (#45672)