diff --git a/gridfm_graphkit/vllm/model.py b/gridfm_graphkit/vllm/model.py index 1b8d62e1..8840c8cd 100644 --- a/gridfm_graphkit/vllm/model.py +++ b/gridfm_graphkit/vllm/model.py @@ -282,7 +282,14 @@ def apply( ) with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes(self.info.model_id) + # vLLM 0.29 added a required hash-algorithm argument to + # ``get_mm_hashes`` (was single-arg on the 0.26 line). Source the + # algorithm from the multimodal config, mirroring vLLM's own + # Terratorch wrapper. + mm_hashes = inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) mm_placeholders = {_MODALITY: [PlaceholderRange(offset=0, length=0)]} @@ -355,16 +362,23 @@ def forward( **kwargs: object, ) -> torch.Tensor: # vLLM's multimodal collation prepends an "items" dimension to every - # field: one graph per prompt token. Real requests carry a single graph - # (leading dim 1); vLLM's warmup ``_dummy_run`` replicates the dummy - # graph across ``max_num_reqs`` items. vLLM then treats our output as - # ``[num_tokens, hidden]`` — it slices ``hidden_states[:num_tokens]`` in - # ``_pool`` and indexes ``hidden_states[logit_indices]`` (up to - # ``num_tokens - 1``) during warmup. So we must return one packed row per - # item, keeping the leading dim equal to the token/item count and all - # graph data in the trailing dimension (mirroring vLLM's Terratorch - # wrapper). Collapsing to a single graph would truncate real output and - # blow the warmup index out of bounds. + # field: one graph per collated item. We produce one packed row per item + # and vLLM treats our output as ``[num_tokens, hidden]`` — it slices + # ``hidden_states[:num_tokens]`` when pooling and gathers + # ``hidden_states[logit_indices]`` with indices up to ``num_tokens - 1``. + # + # In real serving each request is exactly one sentinel prompt token + # carrying one graph, so ``num_tokens == n_items`` and the two axes + # coincide. During warmup they do NOT: ``_dummy_run`` derives + # ``num_reqs = min(num_tokens, max_num_seqs)`` and collates only + # ``num_reqs`` graph items, but packs ``num_tokens`` tokens across those + # requests (several tokens each when ``max_num_seqs < num_tokens``). The + # profiling and flashinfer-autotune warmup runs both call + # ``_dummy_run(max_num_batched_tokens)``, so ``num_tokens`` can far + # exceed ``n_items`` and the downstream ``hidden_states[logit_indices]`` + # gather would run out of bounds. Pad the (discarded) warmup output up to + # the scheduled token count to keep that gather in bounds; on real + # requests ``num_tokens == n_items`` and the pad is a no-op. fields = { name: kwargs[name] for name in graph_codec.GRAPH_FIELDS if name in kwargs } @@ -385,7 +399,17 @@ def forward( ), ) - return torch.stack(packed_rows, dim=0) + packed = torch.stack(packed_rows, dim=0) + + # Align the leading (token) axis with the scheduled token count so + # vLLM's ``hidden_states[logit_indices]`` gather stays in bounds during + # warmup dummy runs; identity on real requests (num_tokens == n_items). + num_tokens = positions.shape[0] + if packed.shape[0] < num_tokens: + pad = packed.new_zeros((num_tokens - packed.shape[0], packed.shape[1])) + packed = torch.cat([packed, pad], dim=0) + + return packed def load_weights( self, diff --git a/gridfm_graphkit/vllm/plugins/pf_reconstruction/io_processor.py b/gridfm_graphkit/vllm/plugins/pf_reconstruction/io_processor.py index d2c05443..b5f47b69 100644 --- a/gridfm_graphkit/vllm/plugins/pf_reconstruction/io_processor.py +++ b/gridfm_graphkit/vllm/plugins/pf_reconstruction/io_processor.py @@ -11,7 +11,6 @@ from __future__ import annotations -import asyncio import logging from collections.abc import Sequence from typing import Any, Optional @@ -104,14 +103,12 @@ def pre_process( request_id: str | None = None, **kwargs, ) -> PromptType | Sequence[PromptType]: - return asyncio.run(self.pre_process_async(prompt, request_id, **kwargs)) - - async def pre_process_async( - self, - prompt: IOProcessorInput, - request_id: str | None = None, - **kwargs, - ) -> PromptType | Sequence[PromptType]: + # vLLM's online pooling path (>=0.29) calls this synchronous method from + # inside the running server event loop (``get_request_factory_online`` → + # ``pre_process``), so it must do the work directly and must NOT spin up + # its own loop. Graph construction/normalization is pure CPU work (no + # awaits), so the synchronous implementation lives here and the async + # variant simply delegates. request: GridFMRequest = prompt case = request.case @@ -136,6 +133,15 @@ async def pre_process_async( return {"prompt_token_ids": [1], "multi_modal_data": multi_modal_data} + async def pre_process_async( + self, + prompt: IOProcessorInput, + request_id: str | None = None, + **kwargs, + ) -> PromptType | Sequence[PromptType]: + # No blocking I/O in pre_process; delegate to the synchronous path. + return self.pre_process(prompt, request_id, **kwargs) + # --- post-processing -------------------------------------------------- def post_process( diff --git a/gridfm_graphkit/vllm/utils.py b/gridfm_graphkit/vllm/utils.py index e857101f..29a968cb 100644 --- a/gridfm_graphkit/vllm/utils.py +++ b/gridfm_graphkit/vllm/utils.py @@ -6,9 +6,9 @@ # The interface this integration targets. vLLM's IO-processor and pooling-model # APIs shift between minor releases (renderer constructor arg, multimodal input -# nesting, pooling model mixins), so the plugin is written against this line and -# refuses to load silently against an untested one. -SUPPORTED_VLLM = ">=0.26,<0.27" +# nesting, get_mm_hashes arity, pooling model mixins), so the plugin is written +# against this line and refuses to load silently against an untested one. +SUPPORTED_VLLM = ">=0.29,<0.30" def check_vllm_version(target_version: str, comparison: str) -> bool: diff --git a/pyproject.toml b/pyproject.toml index 63200162..a898f7bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ namespaces = false [project] name = "gridfm-graphkit" description = "Grid Foundation Model" -version = "0.9.0.2" +version = "0.9.1" readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10,<3.13" @@ -48,7 +48,7 @@ dependencies = [ "pandas>=2.3.0", "plotly>=6.1.2", "pyyaml>=6.0.2", - "torch>=2.10,<2.13", # add upper bound + "torch>=2.10,<2.14", # upper bound; <2.14 admits torch 2.13, required by vllm>=0.29 (see [vllm] extra) "torch-geometric>=2.6.1", "torch-scatter>=2.1.2", "torchaudio>=2.10", @@ -76,10 +76,10 @@ test = [ # Optional dependency for serving a trained GridFM model through vLLM's # /pooling endpoint. Aligned with the vLLM pin used downstream (algorithm-nexus -# pins vllm==0.26.0). The plugin interfaces this integration targets are stable -# across the 0.26 line; see gridfm_graphkit.vllm.utils.SUPPORTED_VLLM. +# candidate pins vllm==0.29.0). The plugin interfaces this integration targets +# shift between minor releases; see gridfm_graphkit.vllm.utils.SUPPORTED_VLLM. vllm = [ - "vllm>=0.26,<0.27", + "vllm>=0.29,<0.30", "safetensors>=0.4.0", # vLLM parses the served config.json through HuggingFace AutoConfig; a config # that carries no `model_type` (see gridfm_graphkit.vllm.export.build_hf_config) @@ -88,16 +88,32 @@ vllm = [ ] [tool.uv] -# torch-scatter is not on PyPI — it requires a PyG wheel index matching your torch version. -# The find-links URL below must be updated whenever the torch version in [project.dependencies] changes. +# torch-scatter is not on PyPI as a universal wheel — prebuilt wheels live on a +# PyG wheel index matching your torch version. The find-links below must be +# updated whenever the torch version in [project.dependencies] changes. # CPU: https://data.pyg.org/whl/torch-X.Y.Z+cpu.html # CUDA: https://data.pyg.org/whl/torch-X.Y.Z+cuNNN.html # After updating, regenerate the lock file with: uv lock +# NOTE: torch 2.13 (pulled by the vllm>=0.29 extra) has no prebuilt torch-scatter +# wheels on the PyG index yet, so torch-scatter is built from its PyPI sdist. That +# sdist imports torch at build time without declaring it, so we inject torch into +# its build environment below (extra-build-dependencies). Kept the torch 2.12 +# links too for the non-vllm torch range. find-links = [ + "https://data.pyg.org/whl/torch-2.13.0+cpu.html", + "https://data.pyg.org/whl/torch-2.13.0+cu130.html", "https://data.pyg.org/whl/torch-2.12.0+cpu.html", "https://data.pyg.org/whl/torch-2.12.0+cu126.html", ] +# torch-scatter's sdist imports torch at build time without declaring it as a +# build dependency, so `uv lock`/`uv sync` fail to build it (e.g. on torch 2.13, +# which has no prebuilt wheel). Inject torch into its build environment. (uv only +# reads this from the top-level project, so downstream consumers still need their +# own copy; the real fix belongs in torch-scatter's build-system.requires.) +[tool.uv.extra-build-dependencies] +torch-scatter = ["torch"] + [project.scripts] gridfm_graphkit = "gridfm_graphkit.__main__:main"