83 lines
3.5 KiB
Python
83 lines
3.5 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
"""Regression tests for the inductor factorable-divisibility patch.
|
||
|
|
|
||
|
|
torch 2.13 proves symbolic divisibility with torch's own unevaluated Mod,
|
||
|
|
so ``15360*s31 + 15360*s87`` is no longer provably divisible by
|
||
|
|
``s31 + s87`` and inductor raises CantSplit on FLUX's fp8 graph (nightly
|
||
|
|
builds 2953/2954). The patch adds an exact-polynomial-division proof that
|
||
|
|
is sound for every symbol assignment; on torch versions whose original
|
||
|
|
check already proves the case (<= 2.11, via sympy.Mod), the patch
|
||
|
|
self-extinguishes.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import sympy
|
||
|
|
from torch._inductor.sizevars import SizeVarAllocator
|
||
|
|
|
||
|
|
import vllm_omni.patch as patch_module
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||
|
|
|
||
|
|
_S1, _S2 = sympy.symbols("s31 s87", positive=True, integer=True)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
("numerator", "denominator", "expected"),
|
||
|
|
[
|
||
|
|
(15360 * _S1 + 15360 * _S2, _S1 + _S2, True), # the CantSplit expression
|
||
|
|
(7 * _S1 + 7 * _S2, _S1 + _S2, True),
|
||
|
|
(_S1 * _S2, _S2, True),
|
||
|
|
(_S1 * _S2 + _S2, _S2, True), # s2*(s1 + 1)
|
||
|
|
(3 * _S1, 2 * _S1, False),
|
||
|
|
(_S1 + 1, _S1, False),
|
||
|
|
(15360 * _S1 + 15359 * _S2, _S1 + _S2, False),
|
||
|
|
(sympy.Integer(12), sympy.Integer(4), True),
|
||
|
|
(sympy.Integer(12), sympy.Integer(5), False),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_provably_factorable_multiple(numerator, denominator, expected):
|
||
|
|
assert patch_module._provably_factorable_multiple(numerator, denominator) is expected
|
||
|
|
|
||
|
|
|
||
|
|
def test_non_polynomial_atoms_never_reach_cancel():
|
||
|
|
from torch.utils._sympy.functions import FloorDiv, ModularIndexing
|
||
|
|
|
||
|
|
assert patch_module._provably_factorable_multiple(FloorDiv(_S1, 2) * 4, _S1) is False
|
||
|
|
assert patch_module._provably_factorable_multiple(ModularIndexing(_S1, 1, 8) * _S2, _S2) is False
|
||
|
|
|
||
|
|
|
||
|
|
def test_divisibility_holds_after_patch_regardless_of_torch_version():
|
||
|
|
"""The end contract on any torch: the CantSplit expression proves divisible.
|
||
|
|
|
||
|
|
On torch <= 2.11 the original method proves it (patch self-extinguished);
|
||
|
|
on torch >= 2.13 the installed wrapper proves it via exact division.
|
||
|
|
"""
|
||
|
|
sv = SizeVarAllocator()
|
||
|
|
method = SizeVarAllocator.statically_known_multiple_of
|
||
|
|
patched = getattr(method, "_vllm_omni_factorable_divisibility", False)
|
||
|
|
assert method(sv, 15360 * _S1 + 15360 * _S2, _S1 + _S2) is True
|
||
|
|
# Negative control must stay non-divisible either way.
|
||
|
|
assert method(sv, 3 * _S1, 2 * _S1) is False
|
||
|
|
if patched:
|
||
|
|
# Wrapper installed: the original alone must have failed the probe.
|
||
|
|
assert patch_module._provably_factorable_multiple(15360 * _S1 + 15360 * _S2, _S1 + _S2) is True
|
||
|
|
|
||
|
|
|
||
|
|
def test_patch_installs_when_original_cannot_prove(monkeypatch):
|
||
|
|
"""Simulate torch 2.13's unevaluated-Mod prover: the wrapper must install."""
|
||
|
|
|
||
|
|
def weak(self, numerator, denominator):
|
||
|
|
return False
|
||
|
|
|
||
|
|
monkeypatch.setattr(SizeVarAllocator, "statically_known_multiple_of", weak)
|
||
|
|
patch_module._patch_inductor_factorable_divisibility()
|
||
|
|
method = SizeVarAllocator.statically_known_multiple_of
|
||
|
|
assert getattr(method, "_vllm_omni_factorable_divisibility", False) is True
|
||
|
|
|
||
|
|
sv = SizeVarAllocator()
|
||
|
|
assert method(sv, 15360 * _S1 + 15360 * _S2, _S1 + _S2) is True
|
||
|
|
assert method(sv, 3 * _S1, 2 * _S1) is False
|
||
|
|
# Idempotency: a second install call must not wrap the wrapper.
|
||
|
|
patch_module._patch_inductor_factorable_divisibility()
|
||
|
|
assert SizeVarAllocator.statically_known_multiple_of is method
|