@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 730d0ac8b9

Fork 0

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 14, 05:30 UTC

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

Signed-off-by: Zerohertz <ohg3417@gmail.com> Signed-off-by: Hyogeun Oh (오효근) <ohg3417@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Hyogeun Oh (오효근) committed Aug 27, 2025 at 03:19 UTC 730d0ac8b9678d64294ddc1e3431a27a50b5e42f

14 files changed +66 -58

vllm/attention/backends/differential_flash_attn.py

+9 -5

@@ -805,14 +805,18 @@ class DifferentialFlashAttentionImpl(AttentionImpl):

805 """Forward pass with FlashAttention.

806

807 Args:

808 - query: shape = [num_tokens, num_heads, head_size]

809 - key: shape = [num_tokens, num_kv_heads, head_size]

810 - value: shape = [num_tokens, num_kv_heads, head_size]

811 - output: shape = [num_tokens, num_heads, head_size]

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

808 + layer: Attention layer instance.

809 + q: Query tensor with shape = [num_tokens, num_heads, head_size]

810 + k: Key tensor with shape = [num_tokens, num_kv_heads, head_size]

811 + v: Value tensor with shape = [num_tokens, num_kv_heads, head_size]

812 + kv_cache: KV cache tensor with shape

813 + [2, num_blocks, block_size, num_kv_heads, head_size].

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

815 for profiling run.

816 attn_metadata: Metadata for attention.

817 + output: Output tensor with shape [num_tokens, num_heads, head_size]

818 + output_scale: Optional output scale tensor.

819 + output_block_scale: Optional output block scale tensor.

820 NOTE: It in-place updates the output tensor.

821 NOTE: FP8 quantization, flash-attn expect the size of

822 {q,k,v}_descale to be (num_sequences, num_kv_heads).

vllm/attention/backends/flash_attn.py

+3 -2

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

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

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

607 output: shape = [num_tokens, num_heads, head_size]

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

608 + kv_cache: KV cache tensor with shape

609 + [2, num_blocks, block_size, num_kv_heads, head_size].

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

611 for profiling run.

612 attn_metadata: Metadata for attention.

@@ -850,7 +851,7 @@ class FlashAttentionImpl(AttentionImpl):

851

852

853 def _get_query_key_seq_metadata(

853 - attn_metadata,

854 + attn_metadata: FlashAttentionMetadata,

855 is_prompt: bool,

856 attn_type: str,

857 ) -> tuple:

vllm/attention/backends/rocm_flash_attn.py

+6 -5

@@ -584,17 +584,18 @@ class ROCmFlashAttentionImpl(AttentionImpl):

584 use prefill sequence attributes

585

586 Args:

587 + layer: Attention layer instance.

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

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

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

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

591 + kv_cache: KV cache tensor with shape

592 + [2, num_blocks, block_size * num_kv_heads * head_size].

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

594 for profiling run.

595 attn_metadata: Metadata for attention.

594 - attn_type: Select attention type, between encoder attention,

595 - decoder self-attention, or encoder/decoder cross-

596 - attention. Defaults to decoder self-attention,

597 - which is the vLLM default generally

596 + output: Optional output tensor.

597 + output_scale: Optional output scale tensor.

598 + output_block_scale: Optional output block scale tensor.

599 Returns:

600 shape = [num_tokens, num_heads * head_size]

601 """

vllm/attention/backends/utils.py

+1 -1

@@ -561,7 +561,7 @@ def get_num_prefill_decode_query_kv_tokens(

561

562 Raises:

563 AssertionError: If the number of encoder tokens in `attn_metadata`

564 - is `None` when required for the calculations.

564 + is `None` when required for the calculations.

565 """

566 num_prefill_query_tokens = 0

567 num_decode_query_tokens = 0

vllm/attention/backends/xformers.py

+6 -6

@@ -471,17 +471,18 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):

471 max_encoder_seq_len)

472

473 Args:

474 + layer: Attention layer instance.

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

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

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

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

478 + kv_cache: KV cache tensor with shape

479 + [2, num_blocks, block_size * num_kv_heads * head_size].

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

481 for profiling run.

482 attn_metadata: Metadata for attention.

481 - attn_type: Select attention type, between encoder attention,

482 - decoder self-attention, or encoder/decoder cross-

483 - attention. Defaults to decoder self-attention,

484 - which is the vLLM default generally

483 + output: Optional output tensor.

484 + output_scale: Optional output scale tensor.

485 + output_block_scale: Optional output block scale tensor.

486 Returns:

487 shape = [num_tokens, num_heads * head_size]

488 """

@@ -644,7 +645,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):

645 for API spec.

646

647 Args:

647 - output: shape = [num_prefill_tokens, num_heads, head_size]

648 query: shape = [num_prefill_tokens, num_heads, head_size]

649 key: shape = [num_prefill_tokens, num_kv_heads, head_size]

650 value: shape = [num_prefill_tokens, num_kv_heads, head_size]

vllm/core/block_manager.py

+3 -5

@@ -352,7 +352,7 @@ class SelfAttnBlockSpaceManager(BlockSpaceManager):

352 with num_lookahead_slots.

353

354 Args:

355 - sequence_group (SequenceGroup): The sequence group to swap in.

355 + seq_group (SequenceGroup): The sequence group to swap in.

356 num_lookahead_slots (int): Number of lookahead slots used in

357 speculative decoding, default to 0.

358

@@ -405,8 +405,6 @@ class SelfAttnBlockSpaceManager(BlockSpaceManager):

405

406 Args:

407 seq_group (SequenceGroup): The sequence group to swap out.

408 - num_lookahead_slots (int): Number of lookahead slots used in

409 - speculative decoding, default to 0.

408

409 Returns:

410 bool: Whether it's possible to swap out current sequence group.

@@ -420,7 +418,7 @@ class SelfAttnBlockSpaceManager(BlockSpaceManager):

418 swapping out the given sequence_group with num_lookahead_slots.

419

420 Args:

423 - sequence_group (SequenceGroup): The sequence group to swap out.

421 + seq_group (SequenceGroup): The sequence group to swap out.

422

423 Returns:

424 List[Tuple[int, int]]: The mapping of swapping block from

@@ -473,7 +471,7 @@ class SelfAttnBlockSpaceManager(BlockSpaceManager):

471 on to the 'device'.

472

473 Args:

476 - sequence_group (SequenceGroup): The sequence group to swap in/out.

474 + seq_group (SequenceGroup): The sequence group to swap in/out.

475 device (Device): device to swap the 'seq_group' on.

476 status (SequenceStatus): The status of sequence which is needed

477 for action. RUNNING for swap out and SWAPPED for swap in

vllm/engine/async_llm_engine.py

+2 -2

@@ -486,10 +486,10 @@ class AsyncLLMEngine(EngineClient):

486 _engine_class: Type[_AsyncLLMEngine] = _AsyncLLMEngine

487

488 def __init__(self,

489 - *args,

489 + *args: Any,

490 log_requests: bool = True,

491 start_engine_loop: bool = True,

492 - **kwargs) -> None:

492 + **kwargs: Any) -> None:

493 if envs.VLLM_USE_V1:

494 raise ValueError(

495 "Using V0 AsyncLLMEngine, but envs.VLLM_USE_V1=True. "

vllm/engine/llm_engine.py

+4 -4

@@ -644,10 +644,10 @@ class LLMEngine:

644 Details:

645 - Set arrival_time to the current time if it is None.

646 - Set prompt_token_ids to the encoded prompt if it is None.

647 - - Create `n` number of [Sequence][vllm.Sequence] objects.

648 - - Create a [SequenceGroup][vllm.SequenceGroup] object

649 - from the list of [Sequence][vllm.Sequence].

650 - - Add the [SequenceGroup][vllm.SequenceGroup] object to the

647 + - Create `n` number of [Sequence][vllm.sequence.Sequence] objects.

648 + - Create a [SequenceGroup][vllm.sequence.SequenceGroup] object

649 + from the list of [Sequence][vllm.sequence.Sequence].

650 + - Add the [SequenceGroup][vllm.sequence.SequenceGroup] object to the

651 scheduler.

652

653 Example:

vllm/entrypoints/llm.py

+5 -5

@@ -186,7 +186,7 @@ class LLM:

186 CompilationConfig]] = None,

187 logits_processors: Optional[list[Union[str,

188 type[LogitsProcessor]]]] = None,

189 - **kwargs,

189 + **kwargs: Any,

190 ) -> None:

191 """LLM constructor."""

192

@@ -697,8 +697,8 @@ class LLM:

697 Generate responses for a chat conversation.

698

699 The chat conversation is converted into a text prompt using the

700 - tokenizer and calls the [generate][] method to generate the

701 - responses.

700 + tokenizer and calls the [generate][vllm.LLM.generate] method to generate

701 + the responses.

702

703 Multi-modal inputs can be passed in the same way you would pass them

704 to the OpenAI API.

@@ -1334,8 +1334,8 @@ class LLM:

1334

1335 def wake_up(self, tags: Optional[list[str]] = None):

1336 """

1337 - Wake up the engine from sleep mode. See the [sleep][] method

1338 - for more details.

1337 + Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]

1338 + method for more details.

1339

1340 Args:

1341 tags: An optional list of tags to reallocate the engine memory

vllm/entrypoints/openai/tool_parsers/minimax_tool_parser.py

+2 -1

@@ -461,7 +461,8 @@ class MinimaxToolParser(ToolParser):

461 i += 1

462 return boundaries

463

464 - def _extract_tool_args(self, tool_content: str, args_match) -> str:

464 + def _extract_tool_args(self, tool_content: str,

465 + args_match: re.Match[str]) -> str:

466 """

467 Extract tool arguments from tool content.

468

vllm/model_executor/layers/lightning_attn.py

+10 -1

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Optional

4 +

5 import torch

6 from einops import rearrange

@@ -453,7 +455,14 @@ class _attention(torch.autograd.Function):

455 lightning_attention_ = _attention.apply

456

457

456 -def lightning_attention(q, k, v, ed, block_size=256, kv_history=None):

458 +def lightning_attention(

459 + q: torch.Tensor,

460 + k: torch.Tensor,

461 + v: torch.Tensor,

462 + ed: torch.Tensor,

463 + block_size: int = 256,

464 + kv_history: Optional[torch.Tensor] = None

465 +) -> tuple[torch.Tensor, torch.Tensor]:

466 """

467 Apply lightning attention algorithm

468 to compute attention efficiently.

vllm/model_executor/layers/linear.py

+3 -2

@@ -233,10 +233,10 @@ class LinearBase(CustomOp):

233 Args:

234 input_size: input dimension of the linear layer.

235 output_size: output dimension of the linear layer.

236 - bias: If true, add bias.

236 skip_bias_add: If true, skip adding bias but instead return it.

237 params_dtype: Data type for the parameters.

238 quant_config: Quantization configure.

239 + prefix: Prefix for parameter names.

240 return_bias: If true, return bias together with outputs in forward pass.

241 """

242

@@ -378,13 +378,14 @@ class MergedReplicatedLinear(ReplicatedLinear):

378

379 Args:

380 input_size: input dimension of the linear layer.

381 - output_size: output dimension of the linear layer.

381 + output_sizes: list of output dimensions of the linear layer.

382 bias: If true, add bias.

383 skip_bias_add: If true, skip adding bias but instead return it.

384 params_dtype: Data type for the parameters.

385 quant_config: Quantization configure.

386 prefix: The name of the layer in the state dict, including all parents

387 (e.g. model.layers.0.qkv_proj)

388 + return_bias: If true, return bias together with outputs in forward pass.

389 """

390

391 def __init__(

vllm/outputs.py

+2 -2

@@ -409,7 +409,7 @@ class EmbeddingOutput:

409

410 Args:

411 embedding: The embedding vector, which is a list of floats.

412 - Its length depends on the hidden dimension of the model.

412 + Its length depends on the hidden dimension of the model.

413 """

414 embedding: list[float]

415

@@ -447,7 +447,7 @@ class ClassificationOutput:

447

448 Args:

449 probs: The probability vector, which is a list of floats.

450 - Its length depends on the number of classes.

450 + Its length depends on the number of classes.

451 """

452 probs: list[float]

453

vllm/sequence.py

+10 -17

@@ -147,18 +147,7 @@ class SequenceDataDelta(

147

148 class SequenceData(msgspec.Struct,

149 omit_defaults=True): # type: ignore[call-arg]

150 - """Data associated with a sequence.

151 -

152 - Args:

153 - prompt_token_ids: The token IDs of the prompt.

154 - output_token_ids: The token IDs of the output. Set to an empty list if

155 - None.

156 -

157 - Attributes:

158 - prompt_token_ids: The token IDs of the prompt.

159 - output_token_ids: The token IDs of the output.

160 - cumulative_logprob: The cumulative log probability of the output.

161 - """

150 + """Data associated with a sequence."""

151 # NOTE: we cannot use Union[list, array] because msgspec cannot support

152 # union of 2 list types.

153 _prompt_token_ids: array

@@ -256,10 +245,12 @@ class SequenceData(msgspec.Struct,

245

246 @property

247 def cumulative_logprob(self) -> float:

248 + """The cumulative log probability of the output."""

249 return self._cumulative_logprob

250

251 @property

252 def prompt_token_ids(self) -> tuple[int, ...]:

253 + """The token IDs of the prompt."""

254 return self._prompt_token_ids_tuple

255

256 @prompt_token_ids.setter

@@ -277,6 +268,7 @@ class SequenceData(msgspec.Struct,

268

269 @property

270 def output_token_ids(self) -> tuple[int, ...]:

271 + """The token IDs of the output."""

272 return tuple(self._output_token_ids)

273

274 @output_token_ids.setter

@@ -940,7 +932,7 @@ class SequenceGroupMetadata(

932 omit_defaults=True): # type: ignore[call-arg]

933 """Metadata for a sequence group. Used to create `AttentionMetadata`.

934

943 - Args:

935 + Attributes:

936 request_id: The ID of the request.

937 is_prompt: Whether the request is at prompt stage.

938 seq_data: The sequence data. (Seq id -> sequence data)

@@ -950,14 +942,14 @@ class SequenceGroupMetadata(

942 do_sample: True if sampling is required. Sampling is not required when

943 e.g., prefill is chunked, and the current iteration only computes

944 query tokens for prefill, we don't need sampling.

953 - token_chunk_size: The number of tokens to be processed (per sequence).

954 - None if chunking is not required.

945 + pooling_params: Pooling parameters.

946 lora_request: LoRA request.

947 computed_block_nums: The block numbers that are already computed,

948 used in prefix caching.

949 state: Internal state tied to this sequence group.

950 + token_type_ids: Token type IDs.

951 multi_modal_data: Multi modal data.

960 - mm_processor_kwargs: Multimodal input processor / mapper overrides.

952 + multi_modal_placeholders: Multi modal placeholders.

953 encoder_seq_data: Optional sequence data for encoder prompt

954 (SequenceGroup.encoder_seq). Should be None

955 unless you are working with an encoder/decoder

@@ -1043,12 +1035,13 @@ class SequenceOutput(

1035 array_like=True): # type: ignore[call-arg]

1036 """The model output associated with a sequence.

1037

1046 - Args:

1038 + Attributes:

1039 parent_seq_id: The ID of the parent sequence (for forking in beam

1040 search).

1041 output_token: The output token ID.

1042 logprobs: The logprobs of the output token.

1043 (Token id -> logP(x_i+1 | x_0, ..., x_i))

1044 + output_embed: Optional output embedding tensor.

1045 """

1046 parent_seq_id: int

1047 output_token: int