[Bugfix][Model] Remove SciPy dependency from Inkling scale planning (#49485)
Signed-off-by: Michael Gschwind <mgschwind@nvidia.com> Co-authored-by: Michael Gschwind <mgschwind@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Mike G committed
Jul 22, 2026 at 21:34 UTC
4080263bb2c5d10deac17aaeb88e0823bc35bca9
2 files changed
+32
-2
tests/models/inkling/test_contract_validation.py
+17
@@ -6,6 +6,7 @@ import pytest
6
7
from vllm.config.compilation import CompilationConfig, CUDAGraphMode
8
from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser
9
+from vllm.models.inkling.common.towers import plan_out_scales
10
from vllm.models.inkling.configs import (
11
InklingAudioConfig,
12
InklingModelConfig,
@@ -17,6 +18,22 @@ from vllm.models.inkling.nvidia.sconv_swa_attn import (
18
from vllm.v1.attention.backend import AttentionCGSupport
19
20
21
+def test_vision_scale_plan_matches_released_config():
22
+ assert plan_out_scales(2, 40, 4) == [
23
+ (1, 1, 1, 3),
24
+ (1, 5, 5, 128),
25
+ (1, 10, 10, 320),
26
+ (1, 40, 40, 4800),
27
+ (2, 40, 40, 9600),
28
+ ]
29
+
30
+
31
+def test_vision_scale_plan_breaks_assignment_ties_in_order():
32
+ reductions = [np.prod(scale[:-1]) for scale in plan_out_scales(2, 52, 4)]
33
+
34
+ assert reductions == sorted(set(reductions))
35
+
36
+
37
@pytest.mark.parametrize(
38
("config_cls", "kwargs", "missing"),
39
[
vllm/models/inkling/common/towers.py
+15
-2
@@ -9,6 +9,7 @@ Both use vLLM's standard ``RMSNorm`` (CPU-friendly, with a native fallback).
9
10
from __future__ import annotations
11
12
+from itertools import combinations
13
from typing import cast
14
15
import numpy as np
@@ -45,6 +46,20 @@ def _prime_factors(n: int) -> list[int]:
46
return factors
47
48
49
+def linear_sum_assignment(
50
+ cost_matrix: np.ndarray,
51
+) -> tuple[np.ndarray, np.ndarray]:
52
+ """Implement SciPy's assignment for Inkling's ordered L1 cost matrix."""
53
+ rows = np.arange(cost_matrix.shape[0])
54
+ cols = np.array(
55
+ min(
56
+ combinations(range(cost_matrix.shape[1]), len(rows)),
57
+ key=lambda candidate: cost_matrix[rows, candidate].sum(),
58
+ )
59
+ )
60
+ return rows, cols
61
+
62
+
63
def plan_out_scales(
64
temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3
65
) -> list[tuple[int, int, int, int]]:
@@ -97,8 +112,6 @@ def plan_out_scales(
112
if n_layers >= len(scales):
113
idxs = np.argmin(cost_matrix, axis=1)
114
else:
100
- from scipy.optimize import linear_sum_assignment
101
-
115
idxs = linear_sum_assignment(cost_matrix)[1]
116
117
assert len(idxs) >= 2