feat(vllm): support vLLM 0.29 (torch 2.13) + GPU /pooling fixes

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,

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(

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:

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"