yangzhuxinyzx
c9d825d5cf
Bind linear implementations during preparation, partition native tuning and engine scratch ownership, and retain legacy imports and operator defaults. Extend shared configuration and execution explanations. Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm.config import set_current_vllm_config
|
|
from vllm.config.kernel import KernelConfig
|
|
from vllm.model_executor.warmup import sm70_native_cache as cache
|
|
|
|
pytestmark = pytest.mark.cpu_test
|
|
|
|
|
|
def _config(tune=None):
|
|
kernel = KernelConfig()
|
|
kernel.sm70_fp8.native.fp8_dense_tune_max_m = tune
|
|
kernel.sm70_fp8.native.resolve("fp8")
|
|
return SimpleNamespace(kernel_config=kernel)
|
|
|
|
|
|
@pytest.fixture
|
|
def native_tables(monkeypatch):
|
|
imported = []
|
|
exported: dict[tuple, bytes] = {}
|
|
|
|
class Native:
|
|
def __init__(self, values=()):
|
|
self.values = values
|
|
|
|
def sm70_gemm_export_cache(self, device, path):
|
|
payload = exported[self.values]
|
|
Path(path).write_bytes(payload)
|
|
return 1
|
|
|
|
def sm70_gemm_import_cache(self, device, path):
|
|
imported.append((self.values, Path(path).read_bytes()))
|
|
return 1
|
|
|
|
monkeypatch.setattr(cache, "NativeBindings", Native)
|
|
return exported, imported
|
|
|
|
|
|
def test_partition_roundtrip_preserves_legacy_file_and_only_imports_matching_policy(
|
|
tmp_path, native_tables
|
|
):
|
|
exported, imported = native_tables
|
|
first = _config(8)
|
|
second = _config(16)
|
|
one = first.kernel_config.sm70_fp8.native.values
|
|
two = second.kernel_config.sm70_fp8.native.values
|
|
path = str(tmp_path / "lut.bin")
|
|
exported.update({(): b"old-runtime-table", one: b"policy-eight"})
|
|
with set_current_vllm_config(first):
|
|
assert cache.export_cache(None, path) == 2
|
|
assert Path(path).read_bytes() == b"old-runtime-table"
|
|
assert cache.packed_cache_bytes(Path(path)).startswith(cache._MAGIC)
|
|
with set_current_vllm_config(second):
|
|
assert cache.import_cache(None, path) == 1
|
|
assert imported == [((), b"old-runtime-table")]
|
|
assert all(policy != two for policy, payload in imported)
|
|
imported.clear()
|
|
with set_current_vllm_config(first):
|
|
assert cache.import_cache(None, path) == 2
|
|
assert imported == [((), b"old-runtime-table"), (one, b"policy-eight")]
|
|
|
|
|
|
def test_untagged_legacy_seed_does_not_cross_typed_override(tmp_path, native_tables):
|
|
exported, imported = native_tables
|
|
path = str(tmp_path / "legacy.bin")
|
|
Path(path).write_bytes(b"historic-table")
|
|
with set_current_vllm_config(_config(8)):
|
|
assert cache.import_cache(None, path) == 1
|
|
assert imported == [((), b"historic-table")]
|
|
imported.clear()
|
|
legacy = _config()
|
|
with set_current_vllm_config(legacy):
|
|
assert cache.import_cache(None, path) == 2
|
|
assert imported[-1] == (
|
|
legacy.kernel_config.sm70_fp8.native.values,
|
|
b"historic-table",
|
|
)
|
|
|
|
|
|
def test_diagnostic_change_reuses_same_partition(tmp_path, native_tables, monkeypatch):
|
|
exported, imported = native_tables
|
|
first = _config(8)
|
|
monkeypatch.setenv("TM_GEMM_TRACE", "1")
|
|
second = _config(8)
|
|
one, two = (cfg.kernel_config.sm70_fp8.native.values for cfg in (first, second))
|
|
assert one != two
|
|
exported.update({(): b"old", one: b"same-calculation"})
|
|
path = str(tmp_path / "lut.bin")
|
|
with set_current_vllm_config(first):
|
|
cache.export_cache(None, path)
|
|
# Exercise the distributed packed-byte import, without a sidecar filename.
|
|
packed = str(tmp_path / "received.bin")
|
|
Path(packed).write_bytes(cache.packed_cache_bytes(Path(path)))
|
|
with set_current_vllm_config(second):
|
|
cache.import_cache(None, packed)
|
|
assert imported[-1] == (two, b"same-calculation")
|