@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a80af24356
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 1, 20:00 UTC

Speed up docs build (#44635)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Jun 5, 2026 at 15:51 UTC a80af243560e25f915010011db76ae97c0cef5b5
32 files changed +234 -159
AGENTS.md
+2
@@ -103,6 +103,8 @@ pre-commit run mypy-3.10 --all-files --hook-stage manual
103
104 The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
105
106 +Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
107 +
108 ### Commit messages
109
110 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
docs/features/speculative_decoding/README.md
+1 -1
@@ -169,7 +169,7 @@ speculative decoding, breaking down the guarantees into three key areas:
169 > distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
170 > - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
171 > without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
172 - > provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
172 + > provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../../tests/v1/spec_decode).
173 > verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
174
175 3. **vLLM Logprob Stability**
mkdocs.yaml
+9 -9
@@ -83,22 +83,22 @@ plugins:
83 - "re:vllm\\._.*" # Internal modules
84 - "vllm.third_party"
85 - "vllm.vllm_flash_attn"
86 - - "re:vllm\\.grpc\\..*_pb2.*" # Auto-generated protobuf files
86 + - "vllm.transformers_utils.configs"
87 + - "vllm.transformers_utils.processors"
88 - !ENV [API_AUTONAV_EXCLUDE, "re:^$"] # Match nothing by default
89 - mkdocstrings:
90 handlers:
91 python:
92 options:
92 - show_symbol_type_heading: true
93 - show_symbol_type_toc: true
94 - filters:
95 - - "!.*_pb2_grpc" # Exclude auto-generated gRPC stubs
96 - summary:
97 - modules: true
98 - show_signature_annotations: true
99 - separate_signature: true
93 + filters: []
94 show_overloads: true
95 signature_crossrefs: true
96 + # Recommendations from api-autonav
97 + docstring_section_style: list
98 + parameter_headings: true
99 + show_symbol_type_heading: true
100 + show_symbol_type_toc: true
101 + summary: true
102 inventories:
103 - https://docs.python.org/3/objects.inv
104 - https://typing-extensions.readthedocs.io/en/latest/objects.inv
vllm/_custom_ops.py
+10 -6
@@ -868,9 +868,10 @@ def cutlass_scaled_mm_azp(
868 bias: torch.Tensor | None = None,
869 ) -> torch.Tensor:
870 """
871 - :param azp_adj: In the per-tensor case, this should include the azp.
872 - Always per-channel.
873 - :param azp: Only set in the per-token case. Per-token if set.
871 + Args:
872 + azp_adj: In the per-tensor case, this should include the azp.
873 + Always per-channel.
874 + azp: Only set in the per-token case. Per-token if set.
875 """
876 assert b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0
877 assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
@@ -3886,9 +3887,12 @@ def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
3887 Note that sylvester hadamard transforms are also symmetric, which means that
3888 this function is also applies the (transpose <=> inverse) transform.
3889
3889 - :param x: value to be transformed inplace
3890 - :param inplace: modify value in place
3891 - :return: value after transformation
3890 + Args:
3891 + x: value to be transformed inplace
3892 + inplace: modify value in place
3893 +
3894 + Returns:
3895 + value after transformation
3896 """
3897 return torch.ops._C.hadacore_transform(x, inplace)
3898
vllm/compilation/passes/inductor_pass.py
+9 -6
@@ -82,11 +82,12 @@ class InductorPass(CustomGraphPass): # type: ignore[misc]
82 def hash_source(*srcs: str | Any) -> str:
83 """
84 Utility method to hash the sources of functions or objects.
85 - :param srcs: strings or objects to add to the hash.
86 - Objects and functions have their source inspected.
87 - Results are cached by resolved types to avoid repeated
88 - inspect.getsource() calls.
89 - :return:
85 +
86 + Args:
87 + srcs: strings or objects to add to the hash.
88 + Objects and functions have their source inspected.
89 + Results are cached by resolved types to avoid repeated
90 + inspect.getsource() calls.
91 """
92 # Resolve instances to their class for a hashable cache key.
93 cache_key = tuple(
@@ -99,7 +100,9 @@ class InductorPass(CustomGraphPass): # type: ignore[misc]
100 def hash_dict(dict_: dict[Any, Any]) -> str:
101 """
102 Utility method to hash a dictionary, can alternatively be used for uuid.
102 - :return: A sha256 hash of the json rep of the dictionary.
103 +
104 + Returns:
105 + A sha256 hash of the json rep of the dictionary.
106 """
107 encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
108 return hashlib.sha256(encoded).hexdigest()
vllm/compilation/passes/utility/fix_functionalization.py
+10 -7
@@ -276,9 +276,11 @@ class FixFunctionalizationPass(VllmInductorPass):
276 """
277 Replace mutated getitem users of the auto-functionalized node with the
278 mutated arguments.
279 - :param node: The auto-functionalized node
280 - :param mutated_args: The mutated arguments, indexed by getitem index.
281 - If the value of an arg is a string, `node.kwargs[arg]` is used.
279 +
280 + Args:
281 + node: The auto-functionalized node
282 + mutated_args: The mutated arguments, indexed by getitem index.
283 + If the value of an arg is a string, `node.kwargs[arg]` is used.
284 """
285 for idx, user in self.getitem_users(node).items():
286 # Some functionalized nodes may return both a result at getitem[0]
@@ -317,10 +319,11 @@ class FixFunctionalizationPass(VllmInductorPass):
319 as node.kwargs cannot be used.
320 See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351
321
320 - :param graph: Graph to insert the defunctionalized node into
321 - :param node: The auto-functionalized node to defunctionalize
322 - :param args: If we cannot use kwargs, specify args directly.
323 - If an arg is a string, `node.kwargs[arg]` is used.
322 + Args:
323 + graph: Graph to insert the defunctionalized node into
324 + node: The auto-functionalized node to defunctionalize
325 + args: If we cannot use kwargs, specify args directly.
326 + If an arg is a string, `node.kwargs[arg]` is used.
327 """ # noqa: E501
328 assert is_func(node, auto_functionalized), (
329 f"node must be auto-functionalized, is {node} instead"
vllm/compilation/passes/utility/noop_elimination.py
+7 -3
@@ -108,9 +108,13 @@ class NoOpEliminationPass(VllmInductorPass):
108 def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool:
109 """
110 This function checks if two dimensions are equivalent.
111 - :param dim: The dimension arg to reshape/slice
112 - :param i_dim: The corresponding dimension in the input tensor
113 - :return: Are the dimensions equivalent?
111 +
112 + Args:
113 + dim: The dimension arg to reshape/slice
114 + i_dim: The corresponding dimension in the input tensor
115 +
116 + Returns:
117 + Are the dimensions equivalent?
118
119 There are two cases in which the dimensions are equivalent:
120 1. The dimensions are equal (both integers)
vllm/device_allocator/cumem.py
+10 -7
@@ -180,8 +180,9 @@ class CuMemAllocator:
180 All data in the memory allocation with the specified tag will be
181 offloaded to CPU memory, and others will be discarded.
182
183 - :param offload_tags: The tags of the memory allocation that will be
184 - offloaded. The rest of the memory allocation will be discarded.
183 + Args:
184 + offload_tags: The tags of the memory allocation that will be
185 + offloaded. The rest of the memory allocation will be discarded.
186 """
187 if offload_tags is None:
188 # by default, allocated tensors are offloaded
@@ -230,9 +231,10 @@ class CuMemAllocator:
231 All data that is previously offloaded will be loaded back to GPU
232 memory, and the rest of the data will have empty memory.
233
233 - :param tags: The tags of the memory allocation that will be loaded
234 - back to GPU memory. If None, all memory allocation will be loaded
235 - back to GPU memory.
234 + Args:
235 + tags: The tags of the memory allocation that will be loaded
236 + back to GPU memory. If None, all memory allocation will be loaded
237 + back to GPU memory.
238 """
239 for ptr, data in self.pointer_to_data.items():
240 if tags is None or data.tag in tags:
@@ -255,8 +257,9 @@ class CuMemAllocator:
257 All memory allocation created inside the context will be allocated
258 in the memory pool, and has the specified tag.
259
258 - :param tag: The tag of the memory allocation. If None, the default tag
259 - will be used.
260 + Args:
261 + tag: The tag of the memory allocation. If None, the default tag
262 + will be used.
263 """
264 if tag is None:
265 tag = CuMemAllocator.default_tag
vllm/distributed/kv_events.py
+10 -5
@@ -132,7 +132,8 @@ class KVEventAggregator:
132 """
133 Add events from a worker batch.
134
135 - :param events: List of KVCacheEvent objects.
135 + Args:
136 + events: List of KVCacheEvent objects.
137 """
138 if not isinstance(events, list):
139 raise TypeError("events must be a list of KVCacheEvent.")
@@ -142,7 +143,8 @@ class KVEventAggregator:
143 """
144 Return events that appeared in all workers.
145
145 - :return: List of events present in all workers.
146 + Returns:
147 + List of events present in all workers.
148 """
149 return [
150 event
@@ -154,7 +156,8 @@ class KVEventAggregator:
156 """
157 Return all events for all workers.
158
157 - :return: List of events for all workers.
159 + Returns:
160 + List of events for all workers.
161 """
162 return list(self._event_counter.elements())
163
@@ -168,7 +171,8 @@ class KVEventAggregator:
171 """
172 Increment the number of workers contributing events.
173
171 - :param count: Number to increment the workers by.
174 + Args:
175 + count: Number to increment the workers by.
176 """
177 if count <= 0:
178 raise ValueError("count must be positive.")
@@ -184,7 +188,8 @@ class KVEventAggregator:
188 """
189 Return the number of workers.
190
187 - :return: int number of workers.
191 + Returns:
192 + int number of workers.
193 """
194 return self._num_workers
195
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+5 -6
@@ -439,13 +439,12 @@ def _init_lmcache_engine(
439 `LMCACHE_CONFIG_FILE` to load the configuration file. If that environment
440 variable is not set, this function will return None.
441
442 - :param lmcache_config: The LMCache configuration.
443 - :type lmcache_config: LMCacheEngineConfig
444 - :param vllm_config: The vLLM configuration.
445 - :type vllm_config: VllmConfig
442 + Args:
443 + lmcache_config: The LMCache configuration.
444 + vllm_config: The vLLM configuration.
445
447 - :return: The initialized LMCache engine
448 - :rtype: LMCacheEngine
446 + Returns:
447 + The initialized LMCache engine
448 """
449 if curr_engine := LMCacheEngineBuilder.get(ENGINE_NAME):
450 return curr_engine
vllm/ir/op.py
+17 -11
@@ -113,11 +113,14 @@ def register_op(
113 """
114 Register a new vLLM IR op.
115
116 - :param f: the native implementation of the op
117 - :param name: the name of the op, defaults to the function name
118 - :param activations: list of activation params, defaults to params starting with 'x'
119 - :param allow_inplace: add a maybe_inplace overload that allows inplace impls
120 - :return: the IrOp object if f is provided, otherwise a decorator
116 + Args:
117 + f: the native implementation of the op
118 + name: the name of the op, defaults to the function name
119 + activations: list of activation params, defaults to params starting with 'x'
120 + allow_inplace: add a maybe_inplace overload that allows inplace impls
121 +
122 + Returns:
123 + the IrOp object if f is provided, otherwise a decorator
124
125 Example usage:
126 ```python
@@ -245,14 +248,17 @@ class IrOp:
248 supported: bool = True,
249 supports_args: Callable[..., bool] | None = None,
250 inplace: bool = False,
248 - ):
251 + ) -> Callable[[Callable[..., Any]], "IrOpImpl"]:
252 """
253 Register an implementation for this custom op.
251 - :param provider: The name of the provider, must be unique.
252 - :param supported: Static support check, use this to check platform support.
253 - :param supports_args: Dynamic arg support check, used for types and shapes.
254 - :param inplace: Does this op reuse activation input memory for outputs
255 - :return: A decorator that registers the implementation.
254 + Args:
255 + provider: The name of the provider, must be unique.
256 + supported: Static support check, use this to check platform support.
257 + supports_args: Dynamic arg support check, used for types and shapes.
258 + inplace: Does this op reuse activation input memory for outputs
259 +
260 + Returns:
261 + A decorator that registers the implementation.
262
263 The decorated function must have the same semantics and signature as
264 the native implementation.
vllm/ir/util.py
+3 -3
@@ -12,9 +12,9 @@ from typing import Any
12 def hash_source(*srcs: str | Any) -> str:
13 """
14 Utility method to hash the sources of functions or objects.
15 - :param srcs: strings or objects to add to the hash.
16 - Objects and functions have their source inspected.
17 - :return:
15 + Args:
16 + srcs: strings or objects to add to the hash.
17 + Objects and functions have their source inspected.
18 """
19 hasher = hashlib.sha256()
20 for src in srcs:
vllm/model_executor/layers/fused_moe/layer.py
+8 -6
@@ -599,12 +599,14 @@ class FusedMoE(PluggableLayer):
599 ):
600 """
601 Load grouped weight scales for group quantization or model weights
602 - :param shard_dim: dimension to shard
603 - :param expert_data: parameter for a particular expert
604 - :param shard_id: either w1, w2, or w3
605 - :param loaded_weight: checkpoint weight to load into the param
606 - :param tp_rank: tensor parallel rank
607 - :param load_full_w2: whether or not the w2 loaded should be sharded.
602 +
603 + Args:
604 + shard_dim: dimension to shard
605 + expert_data: parameter for a particular expert
606 + shard_id: either w1, w2, or w3
607 + loaded_weight: checkpoint weight to load into the param
608 + tp_rank: tensor parallel rank
609 + load_full_w2: whether or not the w2 loaded should be sharded.
610 """
611 if shard_id == "w2":
612 # In the case where we have actorder/g_idx, we do not partition the
vllm/model_executor/layers/quantization/base_config.py
+3 -2
@@ -178,8 +178,9 @@ class QuantizationConfig(ABC):
178 Interface for models to update module names referenced in
179 quantization configs in order to reflect the vllm model structure
180
181 - :param hf_to_vllm_mapper: maps from hf model structure (the assumed
182 - structure of the qconfig) to vllm model structure
181 + Args:
182 + hf_to_vllm_mapper: maps from hf model structure (the assumed
183 + structure of the qconfig) to vllm model structure
184 """
185 # TODO (@kylesayrs): add implementations for all subclasses
186 pass
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+13 -5
@@ -267,8 +267,11 @@ class CompressedTensorsConfig(QuantizationConfig):
267 cls, config: dict[str, Any]
268 ) -> tuple[dict[str, SparsityCompressionConfig], list[str]]:
269 """
270 - :param config: The `quantization_config` dictionary from config.json
271 - :return: A tuple with two elements
270 + Args:
271 + config: The `quantization_config` dictionary from config.json
272 +
273 + Returns:
274 + A tuple with two elements
275 1. A dictionary mapping target layer names to their corresponding
276 sparsity_config
277 2. A list of layer names to ignore for sparsity
@@ -296,8 +299,11 @@ class CompressedTensorsConfig(QuantizationConfig):
299 cls, config: dict[str, Any]
300 ) -> QUANTIZATION_SCHEME_MAP_TYPE:
301 """
299 - :param config: The `quantization_config` dictionary from config.json
300 - :return: A dictionary mapping target layer names to their corresponding
302 + Args:
303 + config: The `quantization_config` dictionary from config.json
304 +
305 + Returns:
306 + A dictionary mapping target layer names to their corresponding
307 quantization_args for weights and input activations
308 """
309 target_scheme_map: dict[str, Any] = dict()
@@ -967,7 +973,9 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
973 """
974 Validator for the kv cache scheme. Useful for controlling the
975 kv cache quantization schemes, that are being supported in vLLM
970 - :param kv_cache_scheme: the compressed-tensors kv cache scheme
976 +
977 + Args:
978 + kv_cache_scheme: the compressed-tensors kv cache scheme
979 """
980 if kv_cache_scheme is None:
981 return
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py
+5 -5
@@ -38,11 +38,11 @@ class CompressedTensorsScheme(ABC):
38 Run the forward pass for the particular scheme. This is where
39 scheme-specific dequant/quant steps/kernels should be applied.
40
41 - :param layer: torch.nn.Module with the registered weights and
42 - other parameters relevant to the particular scheme.
43 - :param x: input to the layer
44 - :param bias: bias parameter
45 -
41 + Args:
42 + layer: torch.nn.Module with the registered weights and
43 + other parameters relevant to the particular scheme.
44 + x: input to the layer
45 + bias: bias parameter
46 """
47 raise NotImplementedError()
48
vllm/model_executor/layers/quantization/compressed_tensors/utils.py
+13 -12
@@ -133,12 +133,11 @@ def find_matched_target(
133 *All* component module names must match in order for a match to be
134 successful. A successful match returns the first component target
135
136 - :param layer_name: layer name
137 - :param module: torch.nn.Module
138 - :param targets: list of targets to match the layer against
139 - :param fused_mapping: map from fused layer names to its components
140 - :param fused_strategy: either "all" or "any". If using "all", fused
141 - layers match if "all" of its components match
136 + Args:
137 + layer_name: layer name
138 + module: torch.nn.Module
139 + targets: list of targets to match the layer against
140 + fused_mapping: map from fused layer names to its components
141 """
142
143 if layer_name is None:
@@ -161,9 +160,10 @@ def _find_first_match(
160 exactly or as a regex after 're:'. If check_contains is set to True,
161 additionally checks if the target string is contained within the value.
162
164 - :param value: string to compare the list of targets against
165 - :param targets: list of targets to match the layer against
166 - :param check_contains: whether or not to do a substring match
163 + Args:
164 + value: string to compare the list of targets against
165 + targets: list of targets to match the layer against
166 + check_contains: whether or not to do a substring match
167 """
168
169 for target in targets:
@@ -205,9 +205,10 @@ def _match_fused_layer(
205 Implements an "all" matching strategy where a fused layer matches iff
206 "all" of its components match
207
208 - :param layer_name: layer name
209 - :param target_layers: list of targets to match the layer against
210 - :param fused_mapping: map from fused layer names to its components
208 + Args:
209 + layer_name: layer name
210 + target_layers: list of targets to match the layer against
211 + fused_mapping: map from fused layer names to its components
212
213 Examples:
214 layer_name = "model.layers.0.self_attn.qkv_proj"
vllm/model_executor/layers/quantization/gguf.py
+3 -2
@@ -114,8 +114,9 @@ class GGUFConfig(QuantizationConfig):
114 Interface for models to update module names referenced in
115 quantization configs in order to reflect the vllm model structure
116
117 - :param hf_to_vllm_mapper: maps from hf model structure (the assumed
118 - structure of the qconfig) to vllm model structure
117 + Args:
118 + hf_to_vllm_mapper: maps from hf model structure (the assumed
119 + structure of the qconfig) to vllm model structure
120 """
121 if self.unquantized_modules is not None:
122 self.unquantized_modules = hf_to_vllm_mapper.apply_list(
vllm/model_executor/layers/quantization/input_quant_fp8.py
+11 -10
@@ -46,16 +46,17 @@ class QuantFP8(CustomOp):
46 compile_native: bool = True,
47 ):
48 """
49 - :param static: static or dynamic quantization
50 - :param group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
51 - PER_CHANNEL, or arbitrary block size)
52 - :param num_token_padding: Pad the token dimension of output to this
53 - size
54 - :param tma_aligned_scales: For group quantization, output scales in
55 - TMA-aligned layout
56 - :param column_major_scales: For group quantization, output scales in
57 - column major format
58 - :param compile_native: Manually compile forward_native if compile mode > None
49 + Args:
50 + static: static or dynamic quantization
51 + group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
52 + PER_CHANNEL, or arbitrary block size)
53 + num_token_padding: Pad the token dimension of output to this
54 + size
55 + tma_aligned_scales: For group quantization, output scales in
56 + TMA-aligned layout
57 + column_major_scales: For group quantization, output scales in
58 + column major format
59 + compile_native: Manually compile forward_native if compile mode > None
60 """
61 super().__init__(compile_native=compile_native)
62 self.static = static
vllm/model_executor/layers/quantization/kv_cache.py
+2 -1
@@ -47,7 +47,8 @@ class BaseKVCacheMethod(QuantizeMethodBase):
47 - quantize k/v_cache entries before saving them to the cache
48 - dequantize k/v_cache entries before fetching them from the cache
49
50 - :param quant_config: the appropriate QuantizationConfig
50 + Args:
51 + quant_config: the appropriate QuantizationConfig
52 """
53
54 def __init__(self, quant_config: QuantizationConfig):
vllm/model_executor/layers/quantization/quark/quark.py
+6 -3
@@ -87,8 +87,9 @@ class QuarkConfig(QuantizationConfig):
87 Interface for models to update module names referenced in
88 quantization configs in order to reflect the vllm model structure
89
90 - :param hf_to_vllm_mapper: maps from hf model structure (the assumed
91 - structure of the qconfig) to vllm model structure
90 + Args:
91 + hf_to_vllm_mapper: maps from hf model structure (the assumed
92 + structure of the qconfig) to vllm model structure
93 """
94 quant_config_with_hf_to_vllm_mapper: dict[str, Any] = {}
95
@@ -724,7 +725,9 @@ class QuarkKVCacheMethod(BaseKVCacheMethod):
725 """
726 Validator for the kv cache configuration. Useful for controlling the
727 kv cache quantization schemes, that are being supported in vLLM
727 - :param kv_cache_config: the quark kv cache scheme
728 +
729 + Args:
730 + kv_cache_config: the quark kv cache scheme
731 """
732 if kv_cache_config is None:
733 return
vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py
+5 -5
@@ -38,11 +38,11 @@ class QuarkScheme(ABC):
38 Run the forward pass for the particular scheme. This is where
39 scheme-specific dequant/quant steps/kernels should be applied.
40
41 - :param layer: torch.nn.Module with the registered weights and
42 - other parameters relevant to the particular scheme.
43 - :param x: input to the layer
44 - :param bias: bias parameter
45 -
41 + Args:
42 + layer: torch.nn.Module with the registered weights and
43 + other parameters relevant to the particular scheme.
44 + x: input to the layer
45 + bias: bias parameter
46 """
47 raise NotImplementedError
48
vllm/model_executor/model_loader/reload/layerwise.py
+5 -3
@@ -123,7 +123,8 @@ def initialize_online_processing(layer: torch.nn.Module):
123 Called by either `initialize_layerwise_reload` or an online quantization scheme,
124 prevents double wrapping in the case of online quantization + reloading
125
126 - :param layer: layer whose parameter weight loaders will be wrapped
126 + Args:
127 + layer: layer whose parameter weight loaders will be wrapped
128 """
129 info = get_layerwise_info(layer)
130
@@ -222,8 +223,9 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
223 This function should be applied after `initialize_layerwise_reload` is applied
224 unwrap the layerwise weight loaders.
225
225 - :param model: model to finalize processing for
226 - :param model_config: config needed for applying processing to attention layers
226 + Args:
227 + model: model to finalize processing for
228 + model_config: config needed for applying processing to attention layers
229 """
230 if hasattr(model, "_original_do_torchao_reload"):
231 model._do_torchao_reload = model._original_do_torchao_reload
vllm/model_executor/model_loader/reload/meta.py
+6 -3
@@ -175,9 +175,12 @@ def get_numel_loaded(
175 """
176 Determine how many elements would be loaded by a weight loader call.
177
178 - :param weight loader: used to load weights
179 - :param args: bound arguments to weight loader
180 - :return: number of elements loaded by the weight loader, the return value of the
178 + Args:
179 + weight_loader: used to load weights
180 + args: bound arguments to weight loader
181 +
182 + Returns:
183 + number of elements loaded by the weight loader, the return value of the
184 weight loader
185 """
186 with CopyCounter() as counter:
vllm/model_executor/model_loader/reload/sanitize.py
+11 -6
@@ -20,9 +20,12 @@ def sanitize_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.T
20 tensors will reference layers, and the WeakKeyDictionary will never evict entries,
21 even when the model is deleted.
22
23 - :param tensor: tensor to be sanitized
24 - :param layer: layer whose references should be removed
25 - :return: sanitized tensor
23 + Args:
24 + tensor: tensor to be sanitized
25 + layer: layer whose references should be removed
26 +
27 + Returns:
28 + sanitized tensor
29 """
30 for key, value in tensor.__dict__.items():
31 if isinstance(value, MethodType) and value.__self__ is layer:
@@ -38,10 +41,12 @@ def restore_layer_refs(tensor: torch.Tensor, layer: torch.nn.Module) -> torch.Te
41 Used by `restore_layer_on_meta` to add back layer references, allowing for proper
42 weight loading.
43
41 - :param tensor: tensor to be sanitized
42 - :param layer: layer whose references should be removed
43 - :return: sanitized tensor
44 + Args:
45 + tensor: tensor to be sanitized
46 + layer: layer whose references should be removed
47
48 + Returns:
49 + sanitized tensor
50 """
51 for key, value in tensor.__dict__.items():
52 if isinstance(value, MethodType) and value.__self__ is layer_ref_sentinel:
vllm/model_executor/model_loader/reload/utils.py
+10 -4
@@ -49,8 +49,11 @@ def has_device_tensors(bound_args: BoundArguments) -> bool:
49 """
50 Return True if the loaded weights exist on an accelerator device
51
52 - :param bound_args: args to load weights
53 - :return: True if weights are on accelerator device
52 + Args:
53 + bound_args: args to load weights
54 +
55 + Returns:
56 + True if weights are on accelerator device
57 """
58 return any(
59 isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu")
@@ -62,8 +65,11 @@ def get_info_size(info: LayerReloadingInfo) -> int:
65 """
66 Calculate the number of bytes used by loaded weights for a given layer
67
65 - :param info: layerwise info to get size of
66 - :return: number of bytes used by loaded weights
68 + Args:
69 + info: layerwise info to get size of
70 +
71 + Returns:
72 + number of bytes used by loaded weights
73 """
74 return sum(
75 value.nbytes
vllm/model_executor/models/olmo.py
+2 -1
@@ -277,7 +277,8 @@ class OlmoModel(nn.Module):
277 inputs_embeds: torch.Tensor | None = None,
278 ) -> torch.Tensor | IntermediateTensors:
279 """
280 - :param input_ids: A tensor of shape `(batch_size, seq_len)`.
280 + Args:
281 + input_ids: A tensor of shape `(batch_size, seq_len)`.
282 """
283 if get_pp_group().is_first_rank:
284 if inputs_embeds is not None:
vllm/model_executor/models/olmo2.py
+2 -1
@@ -314,7 +314,8 @@ class Olmo2Model(nn.Module):
314 inputs_embeds: torch.Tensor | None = None,
315 ) -> torch.Tensor | IntermediateTensors:
316 """
317 - :param input_ids: A tensor of shape `(batch_size, seq_len)`.
317 + Args:
318 + input_ids: A tensor of shape `(batch_size, seq_len)`.
319 """
320 if get_pp_group().is_first_rank:
321 if inputs_embeds is not None:
vllm/model_executor/parameter.py
+10 -9
@@ -3,6 +3,7 @@
3
4 from collections.abc import Callable, Hashable
5 from fractions import Fraction
6 +from typing import Any
7 from weakref import WeakValueDictionary
8
9 import torch
@@ -42,10 +43,9 @@ class BasevLLMParameter(Parameter):
43 """
44 Initialize the BasevLLMParameter
45
45 - :param data: torch tensor with the parameter data
46 - :param weight_loader: weight loader callable
47 -
48 - :returns: a torch.nn.parameter
46 + Args:
47 + data: torch tensor with the parameter data
48 + weight_loader: weight loader callable
49 """
50
51 # During weight loading, we often do something like:
@@ -445,15 +445,16 @@ class SharedWeightParameter(BasevLLMParameter):
445 "currently support tensor parallelism"
446 )
447
448 - def add_partition(self, index: int, data_key: Hashable, *args, **kwargs):
448 + def add_partition(self, index: int, data_key: Hashable, *args: Any, **kwargs: Any):
449 """
450 Add a partition to the weight parameter. Partitions whose `data_key`
451 is the same will share tensor data
452
453 - :param index: index of partition to add
454 - :param data_key: hashable key used to key shared tensors
455 - :param *args: arguments for `torch.empty`
456 - :param **kwargs: keyword arguments for `torch.empty`
453 + Args:
454 + index: index of partition to add
455 + data_key: hashable key used to key shared tensors
456 + *args: arguments for `torch.empty`
457 + **kwargs: keyword arguments for `torch.empty`
458 """
459 # load (shared) tensor using `data_key`
460 if data_key not in self.tensors_registry:
vllm/transformers_utils/utils.py
+5 -2
@@ -84,8 +84,11 @@ def maybe_model_redirect(model: str) -> str:
84 """
85 Use model_redirect to redirect the model name to a local folder.
86
87 - :param model: hf model name
88 - :return: maybe redirect to a local folder
87 + Args:
88 + model: hf model name
89 +
90 + Returns:
91 + maybe redirect to a local folder
92 """
93
94 model_redirect_path = envs.VLLM_MODEL_REDIRECT_PATH
vllm/v1/attention/backend.py
+6 -3
@@ -808,14 +808,17 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
808 ) -> torch.Tensor:
809 raise NotImplementedError
810
811 - def fused_output_quant_supported(self, quant_key: "QuantKey"):
811 + def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool:
812 """
813 Does this attention implementation support fused output quantization.
814 This is used by the AttnFusionPass to only fuse output quantization
815 onto implementations that support it.
816
817 - :param quant_key: QuantKey object that describes the quantization op
818 - :return: is fusion supported for this type of quantization
817 + Args:
818 + quant_key: QuantKey object that describes the quantization op
819 +
820 + Returns:
821 + is fusion supported for this type of quantization
822 """
823 return False
824
vllm/v1/worker/gpu_model_runner.py
+15 -12
@@ -1878,9 +1878,8 @@ class GPUModelRunner(
1878 SpecDecodeMetadata | None,
1879 ]:
1880 """
1881 - :return: tuple[
1882 - logits_indices, spec_decode_metadata,
1883 - ]
1881 + Returns:
1882 + tuple[logits_indices, spec_decode_metadata]
1883 """
1884 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1885 assert total_num_scheduled_tokens > 0
@@ -2205,7 +2204,8 @@ class GPUModelRunner(
2204 slot_mappings: dict[int, torch.Tensor] | None = None,
2205 ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:
2206 """
2208 - :return: tuple[attn_metadata, spec_decode_common_attn_metadata]
2207 + Returns:
2208 + tuple[attn_metadata, spec_decode_common_attn_metadata]
2209 """
2210 # Attention metadata is not needed for attention free models
2211 if len(self.kv_cache_config.kv_cache_groups) == 0:
@@ -2503,9 +2503,11 @@ class GPUModelRunner(
2503 num_common_prefix_blocks: list[int],
2504 ) -> list[list[int]] | None:
2505 """
2506 - :return: Optional[cascade_attn_prefix_lens]
2507 - cascade_attn_prefix_lens is 2D: ``[kv_cache_group_id][attn_group_idx]``,
2508 - None if we should not use cascade attention
2506 + Returns:
2507 + Optional[cascade_attn_prefix_lens]
2508 + cascade_attn_prefix_lens is 2D:
2509 + ``[kv_cache_group_id][attn_group_idx]``,
2510 + None if we should not use cascade attention
2511 """
2512
2513 use_cascade_attn = False
@@ -5324,11 +5326,12 @@ class GPUModelRunner(
5326 """
5327 Reload weights from a weights iterator or from disk
5328
5327 - :param weights_iterator: weights to load into model
5328 - :param weights_path: path to load weights from if weights_iterator is not
5329 - provided. Use path of original model if neither is provided.
5330 - :param is_checkpoint_format: set to False if weights have already been processed
5331 - into kernel format (repacking, renaming, etc.)
5329 + Args:
5330 + weights_iterator: weights to load into model
5331 + weights_path: path to load weights from if weights_iterator is not
5332 + provided. Use path of original model if neither is provided.
5333 + is_checkpoint_format: set to False if weights have already been
5334 + processed into kernel format (repacking, renaming, etc.)
5335 """
5336 # TODO(@kylesayrs): generalize to all runners and loaders
5337 # argument validation