diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py index 42313e7d677..4bdb75dcc49 100644 --- a/tests/models/inkling/test_contract_validation.py +++ b/tests/models/inkling/test_contract_validation.py @@ -6,6 +6,7 @@ import pytest from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser +from vllm.models.inkling.common.towers import plan_out_scales from vllm.models.inkling.configs import ( InklingAudioConfig, InklingModelConfig, @@ -17,6 +18,22 @@ from vllm.models.inkling.nvidia.sconv_swa_attn import ( from vllm.v1.attention.backend import AttentionCGSupport +def test_vision_scale_plan_matches_released_config(): + assert plan_out_scales(2, 40, 4) == [ + (1, 1, 1, 3), + (1, 5, 5, 128), + (1, 10, 10, 320), + (1, 40, 40, 4800), + (2, 40, 40, 9600), + ] + + +def test_vision_scale_plan_breaks_assignment_ties_in_order(): + reductions = [np.prod(scale[:-1]) for scale in plan_out_scales(2, 52, 4)] + + assert reductions == sorted(set(reductions)) + + @pytest.mark.parametrize( ("config_cls", "kwargs", "missing"), [ diff --git a/vllm/models/inkling/common/towers.py b/vllm/models/inkling/common/towers.py index 5ac3dbc20e6..1c8739d0e8b 100644 --- a/vllm/models/inkling/common/towers.py +++ b/vllm/models/inkling/common/towers.py @@ -9,6 +9,7 @@ Both use vLLM's standard ``RMSNorm`` (CPU-friendly, with a native fallback). from __future__ import annotations +from itertools import combinations from typing import cast import numpy as np @@ -45,6 +46,20 @@ def _prime_factors(n: int) -> list[int]: return factors +def linear_sum_assignment( + cost_matrix: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Implement SciPy's assignment for Inkling's ordered L1 cost matrix.""" + rows = np.arange(cost_matrix.shape[0]) + cols = np.array( + min( + combinations(range(cost_matrix.shape[1]), len(rows)), + key=lambda candidate: cost_matrix[rows, candidate].sum(), + ) + ) + return rows, cols + + def plan_out_scales( temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3 ) -> list[tuple[int, int, int, int]]: @@ -97,8 +112,6 @@ def plan_out_scales( if n_layers >= len(scales): idxs = np.argmin(cost_matrix, axis=1) else: - from scipy.optimize import linear_sum_assignment - idxs = linear_sum_assignment(cost_matrix)[1] assert len(idxs) >= 2