chore: vLLM, kernels, and flashinfer upgrades

README.md

@@ -21,14 +21,14 @@ Read detailed usage below, or jump to the documentation with [Getting Started](h
 uv pip install "nemo-safe-synthesizer[cu129,engine]" \
   --index https://flashinfer.ai/whl/cu129 \
   --index https://download.pytorch.org/whl/cu129 \
-  --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 \
+  --index https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129 \
   --index-strategy unsafe-best-match
 
 # With pip:
 pip install "nemo-safe-synthesizer[cu129,engine]" \
   --extra-index-url https://download.pytorch.org/whl/cu129 \
   --extra-index-url https://flashinfer.ai/whl/cu129 \
-  --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129
+  --extra-index-url https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129
 ```
 
 Or install from source:

docs/dev-notes/posts/introducing-nemo-safe-synthesizer.md

@@ -87,7 +87,7 @@ Each score maps to concrete remediation guidance in the documentation. The [Prod
 pip install "nemo-safe-synthesizer[cu129,engine]" \
   --extra-index-url https://download.pytorch.org/whl/cu129 \
   --extra-index-url https://flashinfer.ai/whl/cu129 \
-  --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129
+  --extra-index-url https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129
 ```
 
 The quickest way to run your first pipeline is the CLI:

docs/tutorials/differential-privacy.ipynb

@@ -39,10 +39,10 @@
     "# SPDX-License-Identifier: Apache-2.0\n",
     "\n",
     "if command -v uv > /dev/null 2>&1; then\n",
-    "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 --index-strategy unsafe-best-match\n",
+    "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129 --index-strategy unsafe-best-match\n",
     "    uv pip install kagglehub\n",
     "else\n",
-    "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129\n",
+    "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129\n",
     "    pip install kagglehub\n",
     "fi\n"
    ]

docs/tutorials/safe-synthesizer-101.ipynb

@@ -45,10 +45,10 @@
         "# SPDX-License-Identifier: Apache-2.0\n",
         "\n",
         "if command -v uv > /dev/null 2>&1; then\n",
-        "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 --index-strategy unsafe-best-match\n",
+        "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129 --index-strategy unsafe-best-match\n",
         "    uv pip install datasets\n",
         "else\n",
-        "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129\n",
+        "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129\n",
         "    pip install datasets\n",
         "fi\n"
       ]
@@ -217,7 +217,7 @@
   ],
   "metadata": {
     "kernelspec": {
-      "display_name": "Python 3",
+      "display_name": ".venv",
       "language": "python",
       "name": "python3"
     },
@@ -231,7 +231,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.13.9"
+      "version": "3.13.12"
     }
   },
   "nbformat": 4,

docs/tutorials/time-series-financial-transactions.ipynb

@@ -57,9 +57,9 @@
         "# SPDX-License-Identifier: Apache-2.0\n",
         "\n",
         "if command -v uv > /dev/null 2>&1; then\n",
-        "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 --index-strategy unsafe-best-match\n",
+        "    uv pip install \"nemo-safe-synthesizer[engine,cu129]\" --index https://flashinfer.ai/whl/cu129 --index https://download.pytorch.org/whl/cu129 --index https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129 --index-strategy unsafe-best-match\n",
         "else\n",
-        "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129\n",
+        "    pip install \"nemo-safe-synthesizer[engine,cu129]\" --extra-index-url https://flashinfer.ai/whl/cu129 --extra-index-url https://download.pytorch.org/whl/cu129 --extra-index-url https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129\n",
         "fi"
       ]
     },

docs/user-guide/getting-started.md

@@ -37,7 +37,7 @@ indexes outside PyPI. You must pass the extra index URLs shown below.
         pip install "nemo-safe-synthesizer[cu129,engine]" \
           --extra-index-url https://download.pytorch.org/whl/cu129 \
           --extra-index-url https://flashinfer.ai/whl/cu129 \
-          --extra-index-url https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129
+          --extra-index-url https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129
         ```
 
     === "uv"
@@ -46,7 +46,7 @@ indexes outside PyPI. You must pass the extra index URLs shown below.
         uv pip install "nemo-safe-synthesizer[cu129,engine]" \
           --index https://flashinfer.ai/whl/cu129 \
           --index https://download.pytorch.org/whl/cu129 \
-          --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 \
+          --index https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129 \
           --index-strategy unsafe-best-match
         ```
 

pyproject.toml

@@ -113,8 +113,8 @@ engine = [
 cpu = [
   "accelerate>=1.1.0",
   "bitsandbytes>=0.46.1",
-  "flashinfer-python==0.6.8.post1; sys_platform=='linux'",
-  "flashinfer-cubin==0.6.8.post1; sys_platform=='linux'",
+  "flashinfer-python==0.6.13; sys_platform=='linux'",
+  "flashinfer-cubin==0.6.13; sys_platform=='linux'",
   "gliner",
   "kernels>=0.12.1,<0.15",
   "peft>=0.18.0",
@@ -127,19 +127,19 @@ cpu = [
   "torchvision==0.26.0; sys_platform == 'darwin'",
   "torchvision==0.26.0+cpu; sys_platform == 'linux'",
   "torchao==0.17.0",
-  "transformers>=5.6,<6",
+  "transformers==5.12.0",
   "triton>=2.0.0; sys_platform=='linux'",
   "trl>=0.23.0",
-  "vllm==0.20.0; sys_platform=='linux'",
+  "vllm==0.24.0; sys_platform=='linux'",
   "xgrammar>=0.2.0; sys_platform=='linux'",
 ]
 
 cu129 = [
   "accelerate>=1.1.0",
   "bitsandbytes>=0.46.1",
-  "flashinfer-python==0.6.8.post1; sys_platform == 'linux'",
-  "flashinfer-cubin==0.6.8.post1; sys_platform == 'linux'",
-  "flashinfer-jit-cache==0.6.8.post1+cu129; sys_platform == 'linux'",
+  "flashinfer-python==0.6.13; sys_platform == 'linux'",
+  "flashinfer-cubin==0.6.13; sys_platform == 'linux'",
+  "flashinfer-jit-cache==0.6.13+cu129; sys_platform == 'linux'",
   "gliner",
   "kernels>=0.12.1,<0.15",
   "nvidia-cublas-cu12; sys_platform == 'linux'",
@@ -152,10 +152,10 @@ cu129 = [
   "torchaudio==2.11.0+cu129; sys_platform == 'linux'",
   "torchvision==0.26.0+cu129; sys_platform == 'linux'",
   "torchao==0.17.0+cu129; sys_platform == 'linux' and platform_machine == 'x86_64'",
-  "transformers>=5.6,<6",
+  "transformers>=5.12,<5.12.1",
   "triton>=2.0.0; sys_platform == 'linux'",
   "trl>=0.23.0",
-  "vllm==0.20.0+cu129; sys_platform == 'linux'",
+  "vllm==0.24.0+cu129; sys_platform == 'linux'",
   "xgrammar>=0.2.0; sys_platform == 'linux'",
 ]
 
@@ -180,8 +180,8 @@ conflicts = [
 
 
 override-dependencies = [
-    "flashinfer-python==0.6.8.post1; sys_platform != 'darwin'",
-    "flashinfer-cubin==0.6.8.post1; sys_platform != 'darwin'",
+    "flashinfer-python==0.6.13; sys_platform != 'darwin'",
+    "flashinfer-cubin==0.6.13; sys_platform != 'darwin'",
 ]
 
 environments = [
@@ -225,7 +225,7 @@ constraint-dependencies = [
   # here we specifiy the dep's source from which UV should attempt to resolve
   triton = { index = "pytorch-cu129", extra = "cu129" }
   vllm = [
-    { index = "vllm-v0-20-0-cu129", marker = "sys_platform == 'linux'", extra = "cu129" },
+    { index = "vllm-v0-24-0-cu129", marker = "sys_platform == 'linux'", extra = "cu129" },
   ]
 # for torch - we have a "cpu" dep set and "cu129" dep set, specified in the optional dependencies above
   torch = [
@@ -319,8 +319,8 @@ url = "https://flashinfer.ai/whl/cu129"
 explicit = true
 
 [[tool.uv.index]]
-name = "vllm-v0-20-0-cu129"
-url = "https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129"
+name = "vllm-v0-24-0-cu129"
+url = "https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129"
 explicit = true
 
 [[tool.uv.index]]

script/slurm/slurm_nss_matrix.sh

@@ -119,7 +119,7 @@ if [[ -n "${NSS_VERSION:-}" ]]; then
     uv pip install "nemo-safe-synthesizer[cu129,engine]==${NSS_VERSION}" \
         --index https://flashinfer.ai/whl/cu129 \
         --index https://download.pytorch.org/whl/cu129 \
-        --index https://wheels.vllm.ai/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/cu129 \
+        --index https://wheels.vllm.ai/ee0da84ab9e04ac7610e28580af62c365e898389/cu129 \
         --index-strategy unsafe-best-match
     NSS_RUN_CMD="${PYPI_VENV}/bin/safe-synthesizer"
     echo "[NSS SLURM] Using PyPI install: nemo-safe-synthesizer==${NSS_VERSION} on Python ${NSS_PYTHON_VERSION}"

src/nemo_safe_synthesizer/generation/vllm_backend.py

@@ -76,6 +76,28 @@
 os.environ["VLLM_V1_USE_OUTLINES_CACHE"] = "0"
 
 
+def _build_rope_hf_overrides(model_metadata: ModelMetadata) -> dict[str, Any] | None:
+    """Return vLLM ``hf_overrides`` needed for NSS RoPE context extension."""
+    rope_scaling = model_metadata.rope_scaling
+    if rope_scaling is None or rope_scaling.factor <= 1.0:
+        return None
+
+    rope_type = rope_scaling.rope_type
+    if rope_type == "default":
+        rope_type = "linear"
+
+    return {
+        "max_position_embeddings": model_metadata.max_seq_length,
+        "max_model_len": model_metadata.max_seq_length,
+        "rope_parameters": {
+            "rope_type": rope_type,
+            "factor": float(rope_scaling.factor),
+            "original_max_position_embeddings": model_metadata.base_max_seq_length,
+            "rope_theta": float(rope_scaling.theta),
+        },
+    }
+
+
 def _secure_outlines_cache_dir() -> None:
     """Pin ``OUTLINES_CACHE_DIR`` to a per-user path and tighten permissions.
 
@@ -293,16 +315,21 @@ def initialize(self, **kwargs) -> None:
             backend=self.config.generation.structured_generation.backend,
         )
         model_ref = ModelRef.parse(self.config.training.pretrained_model)
+        llm_kwargs: dict[str, Any] = {}
+        if hf_overrides := _build_rope_hf_overrides(self.model_metadata):
+            llm_kwargs["hf_overrides"] = hf_overrides
 
         with heartbeat("Model loading", logger_name=__name__, model=self.config.training.pretrained_model):
             self.llm = vLLM(
                 model=model_ref.target(),
                 gpu_memory_utilization=max_vram,
+                max_model_len=self.model_metadata.max_seq_length,
                 enable_lora=True,
                 max_lora_rank=self.config.training.lora_r,
                 structured_outputs_config=structured_outputs_config,
                 attention_config=attention_config,
                 trust_remote_code=model_ref.trust_remote_code,
+                **llm_kwargs,
             )
 
         # Cache the engine's *effective* runtime config once at init. Read by

tests/generation/test_vllm_backend.py

@@ -20,7 +20,7 @@
 from nemo_safe_synthesizer.errors import ParameterError
 from nemo_safe_synthesizer.generation.processors import TabularDataProcessor
 from nemo_safe_synthesizer.generation.vllm_observability import GenerationObservability
-from nemo_safe_synthesizer.llm.metadata import ModelMetadata
+from nemo_safe_synthesizer.llm.metadata import ModelMetadata, RopeScaling
 
 
 @pytest.fixture
@@ -43,6 +43,8 @@ def mock_model_metadata(fixture_session_cache_dir):
     # helper's fallback when ``max_tokens_per_example`` is unset).
     # Individual tests override per-prompt return values where needed.
     metadata.max_seq_length = 2048
+    metadata.base_max_seq_length = 2048
+    metadata.rope_scaling = None
     metadata.max_tokens_per_example = None
     metadata.generation_max_tokens_for.return_value = 2048
     return metadata
@@ -393,8 +395,52 @@ def test_initialize_passes_cached_snapshot_target_and_trust_to_vllm(
 
         assert backend.llm is mock_llm
         assert mock_vllm.call_args.kwargs["model"] == str(snapshot)
+        assert mock_vllm.call_args.kwargs["max_model_len"] == mock_model_metadata.max_seq_length
+        assert "hf_overrides" not in mock_vllm.call_args.kwargs
         assert mock_vllm.call_args.kwargs["trust_remote_code"] is True
 
+    def test_initialize_passes_rope_hf_overrides_for_extended_context(
+        self,
+        base_params,
+        mock_model_metadata,
+        mock_schema,
+        mock_workdir,
+        fixture_cached_nvidia_snapshot,
+    ):
+        """VLLM 0.24 requires config overrides when max_model_len exceeds the base context."""
+        cache_root, snapshot = fixture_cached_nvidia_snapshot
+        base_params.training.pretrained_model = "nvidia/Nemotron-Mini-4B-Instruct"
+        mock_model_metadata.base_max_seq_length = 2048
+        mock_model_metadata.max_seq_length = 4096
+        mock_model_metadata.rope_scaling = RopeScaling(rope_type="linear", factor=2.0, theta=10000.0)
+        backend = create_backend(base_params, mock_model_metadata, mock_schema, mock_workdir)
+        mock_llm = MagicMock()
+        mock_llm.get_tokenizer.return_value = MagicMock()
+
+        with (
+            patch(
+                "nemo_safe_synthesizer.generation.vllm_backend.ModelRef._default_hf_cache_root",
+                return_value=cache_root,
+            ),
+            patch("nemo_safe_synthesizer.generation.vllm_backend.vLLM", return_value=mock_llm) as mock_vllm,
+            patch("nemo_safe_synthesizer.generation.vllm_backend.get_max_vram", return_value={0: 0.8}),
+            patch("nemo_safe_synthesizer.generation.vllm_backend.create_processor", return_value=MagicMock()),
+        ):
+            backend.initialize()
+
+        assert mock_vllm.call_args.kwargs["model"] == str(snapshot)
+        assert mock_vllm.call_args.kwargs["max_model_len"] == 4096
+        assert mock_vllm.call_args.kwargs["hf_overrides"] == {
+            "max_model_len": 4096,
+            "max_position_embeddings": 4096,
+            "rope_parameters": {
+                "rope_type": "linear",
+                "factor": 2.0,
+                "original_max_position_embeddings": 2048,
+                "rope_theta": 10000.0,
+            },
+        }
+
     def test_initialize_caches_engine_runtime_config(
         self,
         base_params,

tests/generation/test_vllm_shutdown.py

@@ -29,6 +29,7 @@ def backend(_mock_vllm_cleanup, fixture_session_cache_dir):
     mock_metadata.instruction = "Generate"
     mock_metadata.prompt_config = MagicMock()
     mock_metadata.prompt_config.template = "{instruction} {schema}"
+    mock_metadata.rope_scaling = None
 
     mock_config = MagicMock()
     # Pin branching fields so create_processor() selects TabularDataProcessor deterministically.