4e4d017b6f · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/core/block/naive_block.py

+1 -1

@@ -207,7 +207,7 @@ class NaiveBlockAllocator(BlockAllocator):

207

208 Args:

209 absolute_id (int): The absolute block id for the block

210 - in whole allocator.

210 + in whole allocator.

211

212 Returns:

213 int: The zero-offset block id on certain device.

vllm/core/block/prefix_caching_block.py

+1 -1

@@ -61,7 +61,7 @@ class PrefixCachingBlockAllocator(BlockAllocator):

61 Args:

62 num_blocks (int): The total number of blocks to manage.

63 block_size (int): The size of each block in tokens.

64 - block_ids(Optional[Iterable[int]], optional): An optional iterable of

64 + block_ids (Optional[Iterable[int]], optional): An optional iterable of

65 block IDs. If not provided, block IDs will be assigned sequentially

66 from 0 to num_blocks - 1.

67 """

vllm/core/scheduler.py

+1 -1

@@ -657,7 +657,7 @@ class Scheduler:

657 `budget.num_batched_tokens` has not enough capacity to schedule

658 all tokens.

659 partial_prefill_metadata: information about the partial prefills

660 - that are currently running

660 + that are currently running

661

662 Returns:

663 SchedulerRunningOutputs.

vllm/v1/attention/backends/cpu_attn.py

+2 -1

@@ -491,7 +491,8 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):

491 query: shape = [num_tokens, num_heads * head_size]

492 key: shape = [num_tokens, num_kv_heads * head_size]

493 value: shape = [num_tokens, num_kv_heads * head_size]

494 - kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]

494 + kv_cache: shape =

495 + [2, num_blocks, block_size * num_kv_heads * head_size]

496 NOTE: kv_cache will be an empty tensor with shape [0]

497 for profiling run.

498 attn_metadata: Metadata for attention.

vllm/v1/attention/backends/flash_attn.py

+2 -1

@@ -438,7 +438,8 @@ class FlashAttentionImpl(AttentionImpl):

438 query: shape = [num_tokens, num_heads, head_size]

439 key: shape = [num_tokens, num_kv_heads, head_size]

440 value: shape = [num_tokens, num_kv_heads, head_size]

441 - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

441 + kv_cache: shape =

442 + [2, num_blocks, block_size, num_kv_heads, head_size]

443 attn_metadata: Metadata for attention.

444 Returns:

445 shape = [num_tokens, num_heads * head_size]

vllm/v1/attention/backends/flashinfer.py

+3 -5

@@ -637,11 +637,9 @@ class FlashInferImpl(AttentionImpl):

637 query: shape = [num_tokens, num_heads, head_size]

638 key: shape = [num_tokens, num_kv_heads, head_size]

639 value: shape = [num_tokens, num_kv_heads, head_size]

640 - kv_cache: shape -

641 - # NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]

642 - # HND: [num_blocks, 2, num_kv_heads, block_size, head_size]

643 -

644 -

640 + kv_cache: KV cache tensor with different possible shapes:

641 + - NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]

642 + - HND: [num_blocks, 2, num_kv_heads, block_size, head_size]

643 attn_metadata: Metadata for attention.

644 Returns:

645 shape = [num_tokens, num_heads * head_size]

vllm/v1/attention/backends/flex_attention.py

+2 -1

@@ -689,7 +689,8 @@ class FlexAttentionImpl(AttentionImpl):

689 query: shape = [num_tokens, num_heads, head_size]

690 key: shape = [num_tokens, num_kv_heads, head_size]

691 value: shape = [num_tokens, num_kv_heads, head_size]

692 - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

692 + kv_cache: shape =

693 + [2, num_blocks, block_size, num_kv_heads, head_size]

694 attn_metadata: Metadata for attention.

695 Returns:

696 shape = [num_tokens, num_heads * head_size]

vllm/v1/attention/backends/pallas.py

+3 -2

@@ -235,7 +235,8 @@ class PallasAttentionBackendImpl(AttentionImpl):

235 query: shape = [num_tokens, num_heads * head_size]

236 key: shape = [num_tokens, num_kv_heads * head_size]

237 value: shape = [num_tokens, num_kv_heads * head_size]

238 - kv_cache = [num_blocks, block_size, num_kv_heads * 2, head_size]

238 + kv_cache: shape =

239 + [num_blocks, block_size, num_kv_heads * 2, head_size]

240 attn_metadata: Metadata for attention.

241 Returns:

242 shape = [num_tokens, num_heads * head_size]

@@ -329,7 +330,7 @@ def write_to_kv_cache(

330 Args:

331 key: shape = [num_tokens, num_kv_heads, head_size]

332 value: shape = [num_tokens, num_kv_heads, head_size]

332 - kv_cache = [num_blocks, block_size, num_kv_heads * 2, head_size]

333 + kv_cache: shape = [num_blocks, block_size, num_kv_heads * 2, head_size]

334 num_slices_per_kv_cache_update_block: int

335 """

336 _, page_size, num_combined_kv_heads, head_size = kv_cache.shape

vllm/v1/attention/backends/rocm_aiter_fa.py

+2 -1

@@ -429,7 +429,8 @@ class AiterFlashAttentionImpl(AttentionImpl):

429 query: shape = [num_tokens, num_heads, head_size]

430 key: shape = [num_tokens, num_kv_heads, head_size]

431 value: shape = [num_tokens, num_kv_heads, head_size]

432 - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

432 + kv_cache: shape =

433 + [2, num_blocks, block_size, num_kv_heads, head_size]

434 attn_metadata: Metadata for attention.

435 Returns:

436 shape = [num_tokens, num_heads * head_size]

vllm/v1/attention/backends/tree_attn.py

+2 -1

@@ -362,7 +362,8 @@ class TreeAttentionImpl(AttentionImpl):

362 query: shape = [num_tokens, num_heads, head_size]

363 key: shape = [num_tokens, num_kv_heads, head_size]

364 value: shape = [num_tokens, num_kv_heads, head_size]

365 - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

365 + kv_cache: shape =

366 + [2, num_blocks, block_size, num_kv_heads, head_size]

367 attn_metadata: Metadata for attention.

368 Returns:

369 shape = [num_tokens, num_heads * head_size]

vllm/v1/attention/backends/triton_attn.py

+2 -1

@@ -285,7 +285,8 @@ class TritonAttentionImpl(AttentionImpl):

285 query: shape = [num_tokens, num_heads, head_size]

286 key: shape = [num_tokens, num_kv_heads, head_size]

287 value: shape = [num_tokens, num_kv_heads, head_size]

288 - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

288 + kv_cache: shape =

289 + [2, num_blocks, block_size, num_kv_heads, head_size]

290 attn_metadata: Metadata for attention.

291 Returns:

292 shape = [num_tokens, num_heads * head_size]

vllm/v1/attention/backends/xformers.py

+2 -1

@@ -330,7 +330,8 @@ class XFormersAttentionImpl(AttentionImpl):

330 query: shape = [num_tokens, num_heads, head_size]

331 key: shape = [num_tokens, num_kv_heads, head_size]

332 value: shape = [num_tokens, num_kv_heads, head_size]

333 - kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]

333 + kv_cache: shape =

334 + [2, num_blocks, block_size, num_kv_heads, head_size]

335 attn_metadata: Metadata for attention.

336 Returns:

337 shape = [num_tokens, num_heads * head_size]

vllm/v1/core/encoder_cache_manager.py

+4 -4

@@ -255,9 +255,9 @@ def compute_encoder_budget(

255

256 Returns:

257 - Compute budget for encoder execution, measured in number of tokens

258 - from the input sequence.

258 + from the input sequence.

259 - Space budget for encoder cache size, measured in number of tokens

260 - from the input sequence.

260 + from the input sequence.

261 """

262 if mm_registry.supports_multimodal_inputs(model_config):

263 max_tokens_by_modality = mm_registry \

@@ -303,9 +303,9 @@ def compute_mm_encoder_budget(

303

304 Returns:

305 - Compute budget for encoder execution, measured in number of tokens

306 - from the input sequence.

306 + from the input sequence.

307 - Space budget for encoder cache size, measured in number of tokens

308 - from the input sequence.

308 + from the input sequence.

309 """

310

311 if not max_tokens_by_modality:

vllm/v1/core/kv_cache_coordinator.py

+2 -1

@@ -119,7 +119,8 @@ class KVCacheCoordinator(ABC):

119

120 Args:

121 request: The request.

122 - num_tokens: The total number of tokens that need to be cached

122 + num_computed_tokens: The total number of tokens

123 + that need to be cached

124 (including tokens that are already cached).

125 """

126 for manager in self.single_type_managers:

vllm/v1/core/kv_cache_manager.py

+6 -5

@@ -54,14 +54,15 @@ class KVCacheBlocks:

54 def get_block_ids(

55 self,

56 allow_none: bool = False,

57 - ):

57 + ) -> Optional[tuple[list[int], ...]]:

58 """

59 Converts the KVCacheBlocks instance to block_ids.

60 -

60 +

61 Returns:

62 - tuple[list[int], ...]: A tuple of lists where

63 - * the outer tuple corresponds to KV cache groups

64 - * each inner list contains the block_ids of the blocks in that group

62 + tuple[list[int], ...]: A tuple of lists where:

63 + - the outer tuple corresponds to KV cache groups

64 + - each inner list contains the block_ids of the blocks in that

65 + group

66 """

67 if allow_none and all(len(group) == 0 for group in self.blocks):

68 return None

vllm/v1/executor/ray_distributed_executor.py

+2 -1

@@ -8,6 +8,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator

8 from vllm.executor.ray_distributed_executor import ( # noqa

9 RayDistributedExecutor as RayDistributedExecutorV0)

10 from vllm.logger import init_logger

11 +from vllm.v1.core.sched.output import SchedulerOutput

12 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

13 from vllm.v1.executor.abstract import Executor

14 from vllm.v1.outputs import ModelRunnerOutput

@@ -64,7 +65,7 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):

65

66 def execute_model(

67 self,

67 - scheduler_output,

68 + scheduler_output: SchedulerOutput,

69 ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:

70 """Execute the model on the Ray workers.

71

vllm/v1/metrics/prometheus.py

+1 -1

@@ -36,7 +36,7 @@ def setup_multiprocess_prometheus():

36 "and vLLM will properly handle cleanup.")

37

38

39 -def get_prometheus_registry():

39 +def get_prometheus_registry() -> CollectorRegistry:

40 """Get the appropriate prometheus registry based on multiprocessing

41 configuration.

42

vllm/v1/sample/logits_processor/interface.py

+2 -2

@@ -91,7 +91,7 @@ class LogitsProcessor(ABC):

91 to each forward pass.

92

93 Args:

94 - batch_update is non-None iff there have been

95 - changes to the batch makeup.

94 + batch_update: Non-None iff there have been changes

95 + to the batch makeup.

96 """

97 raise NotImplementedError

vllm/v1/sample/rejection_sampler.py

+1 -1

@@ -68,7 +68,7 @@ class RejectionSampler(nn.Module):

68 different requests are flattened into a single tensor because

69 this is the shape of the output logits.

70 NOTE: `target_logits` can be updated in place to save memory.

71 - bonus_token_ids_tensor (torch.Tensor):

71 + bonus_token_ids (torch.Tensor):

72 A tensor containing bonus tokens. Shape is [batch_size, 1].

73 Bonus tokens are added to the end of the sequence if all

74 proposed tokens are accepted. We generate the bonus tokens

vllm/v1/sample/tpu/sampler.py

+1 -1

@@ -89,7 +89,7 @@ class Sampler(nn.Module):

89 Gather logprobs for topk and sampled/prompt token.

90

91 Args:

92 - logits: (num tokens) x (vocab) tensor

92 + logprobs: (num tokens) x (vocab) tensor

93 num_logprobs: minimum number of logprobs to

94 retain per token

95 token_ids: prompt tokens (if prompt logprobs)

vllm/v1/structured_output/backend_types.py

+2 -2

@@ -110,7 +110,7 @@ class StructuredOutputBackend(ABC):

110

111 Args:

112 request_type (StructuredOutputOptions): The type of structured

113 - output request.

113 + output request.

114 grammar_spec (str): The grammar specification to compile.

115

116 Returns:

@@ -124,7 +124,7 @@ class StructuredOutputBackend(ABC):

124

125 Args:

126 max_num_seqs (int): The maximum number of sequences for which

127 - to allocate the bitmask.

127 + to allocate the bitmask.

128 """

129

130 @abstractmethod

vllm/v1/worker/gpu_input_batch.py

-3

@@ -525,9 +525,6 @@ class InputBatch:

525 Any consecutive empty indices at the very end of the list are not

526 filled.

527

528 - Args:

529 - empty_req_indices: empty indices which may be filled.

530 -

528 Returns:

529 swaps: list of (from,to) swap tuples for moved requests

530 empty_req_indices: indices not filled by condensation

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -2955,7 +2955,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2955 Args:

2956 kv_cache_config: The KV cache config

2957 kv_cache_raw_tensors: The KV cache buffer of each layer, with

2958 - correct size but uninitialized shape.

2958 + correct size but uninitialized shape.

2959 Returns:

2960 Dict[str, torch.Tensor]: A map between layer names to their

2961 corresponding memory buffer for KV cache.

vllm/v1/worker/tpu_model_runner.py

+5 -5

@@ -552,7 +552,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

552 return kv_cache_spec

553

554 def _get_slot_mapping_metadata(self, num_reqs,

555 - num_scheduled_tokens_per_req):

555 + num_scheduled_tokens_per_req) -> np.ndarray:

556 """

557 Computes metadata for mapping slots to blocks in the key-value (KV)

558 cache for a batch of requests.

@@ -565,15 +565,15 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

565 Args:

566 num_reqs (int): Number of requests in the current batch.

567 num_scheduled_tokens_per_req (int or np.ndarray): Number of tokens

568 - to be scheduled for each request.

568 + to be scheduled for each request.

569

570 Returns:

571 np.ndarray: A 2D array of shape (total_block_len, 3), where each row

572 - contains:

572 + contains:

573 - kv_cache_start_index (int): The starting index in the KV cache

574 - for the corresponding slice.

574 + for the corresponding slice.

575 - new_kv_start_index (int): The starting index in the new KV

576 - cache for the corresponding slice.

576 + cache for the corresponding slice.

577 - slice_len (int): The length of the slice.

578 """

579 slices_start = self.input_batch.num_computed_tokens_cpu[:num_reqs]

vllm/v1/worker/utils.py

+4 -4

@@ -172,10 +172,10 @@ def scatter_mm_placeholders(

172

173 Args:

174 embeds: The multimodal embeddings.

175 - Shape: `(num_embeds, embed_dim)`

175 + Shape: `(num_embeds, embed_dim)`

176 is_embed: A boolean mask indicating which positions in the placeholder

177 - tokens need to be filled with multimodal embeddings.

178 - Shape: `(num_placeholders, num_embeds)`

177 + tokens need to be filled with multimodal embeddings.

178 + Shape: `(num_placeholders, num_embeds)`

179 """

180 if is_embed is None:

181 return embeds

@@ -278,7 +278,7 @@ def bind_kv_cache(

278 Args:

279 kv_caches: The allocated kv_caches with layer names as keys.

280 forward_context: The global forward context containing all Attention

281 - layers with layer names as keys.

281 + layers with layer names as keys.

282 runner_kv_caches: The kv_cache declared by ModelRunner.

283 """

284 # Bind kv_caches to ModelRunner

vllm/v1/worker/worker_base.py

+2 -2

@@ -36,8 +36,8 @@ class WorkerBase(WorkerBaseV0):

36 local_rank: Local device index

37 rank: Global rank in distributed setup

38 distributed_init_method: Distributed initialization method

39 - is_driver_worker: Whether this worker handles driver

40 - responsibilities

39 + is_driver_worker: Whether this worker handles driver

40 + responsibilities

41 """

42 # Configuration storage

43 super().__init__(vllm_config=vllm_config)

[Docs] Fix warnings in `mkdocs build` (continued) (#23743)