vllm-omni/tests/edge/test_spark_export.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Tests for the Spark-X2.5 decode-step exporter.

The point of the export is the hybrid cache layout, so that is what these
check: sliding layers stay at a fixed 512 entries however long the context
gets, full layers grow, and the window actually rolls.
"""

import pytest
import torch

from vllm_omni.edge.spark_export import (
    FULL,
    SLIDING,
    SparkDecodeStep,
    SparkStepConfig,
    config_from_spark,
    example_inputs,
    input_names,
    kv_cache_bytes,
    output_names,
)

HF_CONFIG = {
    "hidden_size": 64,
    "num_attention_heads": 4,
    "num_key_value_heads": 2,
    "head_dim": 16,
    "intermediate_size": 128,
    "vocab_size": 256,
    "sliding_window": 8,
    "rms_norm_eps": 1e-6,
    "layer_types": [SLIDING, SLIDING, SLIDING, FULL],
    "rope_parameters": {
        "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25},
        "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0},
    },
}


def _cfg(**kw) -> SparkStepConfig:
    return config_from_spark(HF_CONFIG, **kw)


@pytest.mark.core_model
@pytest.mark.cpu
def test_sliding_cache_is_capped_but_full_cache_grows():
    cfg = _cfg(cache_layout="roll")
    assert cfg.cache_len(SLIDING, 4) == 4  # shorter than the window
    assert cfg.cache_len(SLIDING, 1024) == 8  # capped at sliding_window
    assert cfg.cache_len(SLIDING, 1_000_000) == 8
    assert cfg.cache_len(FULL, 1024) == 1024


@pytest.mark.core_model
@pytest.mark.cpu
def test_rotary_dim_differs_by_layer_type():
    cfg = _cfg()
    assert cfg.rotary_dim(SLIDING) == 16  # all head dims
    assert cfg.rotary_dim(FULL) == 4  # first quarter only


@pytest.mark.core_model
@pytest.mark.cpu
def test_kv_bytes_beats_uniform_attention_and_the_gap_grows():
    cfg = _cfg()
    near = kv_cache_bytes(cfg, 16)
    far = kv_cache_bytes(cfg, 4096)
    assert near["hybrid"] < near["uniform"]
    assert far["hybrid"] < far["uniform"]
    # The saving is a growing fraction, not a constant one.
    assert far["saved"] / far["uniform"] > near["saved"] / near["uniform"]


@pytest.mark.core_model
@pytest.mark.cpu
def test_example_inputs_match_the_declared_signature():
    cfg = _cfg()
    args = example_inputs(cfg, context=1024)
    names = input_names(cfg)
    assert len(args) == len(names)
    caches = args[7:]
    assert len(caches) == 2 * cfg.num_layers
    for i, layer_type in enumerate(cfg.layer_types):
        expected = cfg.cache_len(layer_type, 1024)
        assert caches[2 * i].shape == (1, cfg.num_key_value_heads, expected, cfg.head_dim)
    assert len(output_names(cfg)) == 1 + 2 * cfg.num_layers


@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("layout", ["ring", "roll"])
def test_every_layer_returns_what_the_runtime_must_write(layout):
    cfg = _cfg(cache_layout=layout)
    step = SparkDecodeStep(cfg).eval()
    context = 1024
    with torch.no_grad():
        out = step(*example_inputs(cfg, context))
    logits, caches = out[0], out[1:]
    assert logits.shape == (1, cfg.vocab_size)
    for i, layer_type in enumerate(cfg.layer_types):
        k_new = caches[2 * i]
        if layout == "ring":
            # One entry per layer, dropped into the ring slot for this
            # position -- nothing cache-sized crosses the graph boundary.
            assert k_new.shape[2] == 1
        elif layer_type == SLIDING:
            assert k_new.shape[2] == cfg.cache_len(SLIDING, context)
        else:
            assert k_new.shape[2] == 1


@pytest.mark.core_model
@pytest.mark.cpu
def test_ring_and_roll_agree_on_the_same_history():
    """Both layouts must attend over the same window of tokens.

    Roll keeps `sliding_window` entries and drops the oldest after appending;
    ring keeps one fewer and appends in score space. Given a shared history
    they have to produce identical output, or the device numbers measure a
    different model.
    """
    torch.manual_seed(0)
    ring = _cfg(cache_layout="ring", layer_slice=slice(0, 1))
    roll = _cfg(cache_layout="roll", layer_slice=slice(0, 1))
    step_ring = SparkDecodeStep(ring).eval()
    step_roll = SparkDecodeStep(roll).eval()
    step_roll.load_state_dict(step_ring.state_dict())

    context = 64
    args_roll = list(example_inputs(roll, context))
    k_hist, v_hist = args_roll[7], args_roll[8]
    # Ring sees the same history minus the entry roll is about to discard.
    args_ring = list(args_roll)
    args_ring[5] = torch.zeros(1, 1, 1, ring.cache_len(SLIDING, context) + 1)
    args_ring[7] = k_hist[:, :, 1:]
    args_ring[8] = v_hist[:, :, 1:]

    with torch.no_grad():
        out_ring = step_ring(*args_ring)
        out_roll = step_roll(*args_roll)
    assert torch.allclose(out_ring[0], out_roll[0], atol=2e-4, rtol=2e-4)


@pytest.mark.core_model
@pytest.mark.cpu
def test_roll_layout_drops_the_oldest_entry():
    """Roll's rolled window is the old cache shifted by one with the new key
    appended -- what makes a masked pad slot safe to prepend."""
    cfg = _cfg(cache_layout="roll", layer_slice=slice(0, 1))
    step = SparkDecodeStep(cfg).eval()
    args = list(example_inputs(cfg, context=1024))
    k_cache = args[7]
    with torch.no_grad():
        out = step(*args)
    k_new = out[1]
    assert torch.allclose(k_new[:, :, :-1], k_cache[:, :, 1:])


@pytest.mark.core_model
@pytest.mark.cpu
def test_ring_sliding_cache_holds_one_less_than_the_window():
    ring, roll = _cfg(cache_layout="ring"), _cfg(cache_layout="roll")
    assert ring.cache_len(SLIDING, 1024) == HF_CONFIG["sliding_window"] - 1
    assert roll.cache_len(SLIDING, 1024) == HF_CONFIG["sliding_window"]
    # Full-attention layers are unaffected by the layout.
    assert ring.cache_len(FULL, 1024) == roll.cache_len(FULL, 1024) == 1024


@pytest.mark.core_model
@pytest.mark.cpu
def test_layer_slice_selects_one_layer_type():
    types = HF_CONFIG["layer_types"]
    sliding = _cfg(layer_slice=slice(0, 1))
    assert sliding.layer_types == (SLIDING,)
    full = _cfg(layer_slice=slice(types.index(FULL), types.index(FULL) + 1))
    assert full.layer_types == (FULL,)
    assert full.first_layer == types.index(FULL)


@pytest.mark.core_model
@pytest.mark.cpu
def test_unknown_layer_type_is_rejected():
    bad = dict(HF_CONFIG, layer_types=[SLIDING, "linear_attention", SLIDING, FULL])
    with pytest.raises(ValueError, match="unknown layer types"):
        config_from_spark(bad)


@pytest.mark.core_model
@pytest.mark.cpu
def test_gelu_modes_stay_close_to_the_exact_form():
    """The approximations are offered because erf GELU is ~12% of a layer on
    device, but they are only adoptable if they track the trained activation."""
    x = torch.linspace(-6, 6, 4096)
    exact = torch.nn.functional.gelu(x)
    from vllm_omni.edge.spark_export import _gelu

    tanh_err = (_gelu(x, "tanh") - exact).abs().max().item()
    sigmoid_err = (_gelu(x, "sigmoid") - exact).abs().max().item()
    assert torch.equal(_gelu(x, "exact"), exact)
    assert tanh_err < 2e-3
    # The sigmoid form is markedly looser; it is kept for measurement only.
    assert 1e-2 < sigmoid_err < 3e-2
    assert tanh_err < sigmoid_err


@pytest.mark.core_model
@pytest.mark.cpu
def test_unknown_gelu_mode_is_rejected():
    with pytest.raises(ValueError, match="unknown gelu_mode"):
        _cfg(gelu_mode="relu")


@pytest.mark.core_model
@pytest.mark.cpu
def test_unknown_cache_layout_is_rejected():
    with pytest.raises(ValueError, match="unknown cache_layout"):
        _cfg(cache_layout="linked_list")


@pytest.mark.core_model
@pytest.mark.cpu
def test_auto_layout_picks_ring_for_sliding_and_roll_for_full():
    """Ring pays off only on sliding layers, and only there does it compile at
    w4a16, so the default mixes them."""
    cfg = _cfg()  # default cache_layout="auto"
    assert cfg.layout_for(SLIDING) == "ring"
    assert cfg.layout_for(FULL) == "roll"
    assert cfg.cache_len(SLIDING, 1024) == HF_CONFIG["sliding_window"] - 1
    forced = _cfg(cache_layout="roll")
    assert forced.layout_for(SLIDING) == forced.layout_for(FULL) == "roll"


@pytest.mark.core_model
@pytest.mark.cpu
def test_auto_layout_returns_one_entry_for_sliding_and_a_window_for_full():
    cfg = _cfg()
    step = SparkDecodeStep(cfg).eval()
    with torch.no_grad():
        out = step(*example_inputs(cfg, 1024))
    caches = out[1:]
    for i, layer_type in enumerate(cfg.layer_types):
        # Both layouts hand back exactly one new entry here: ring by design,
        # roll because a full layer only ever appends.
        assert caches[2 * i].shape[2] == 1, layer_type