forked from Karylab-cklius/vllm
ecf4aa5ce2
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Config-only resolution of DFlash draft attention causality.
|
|
|
|
``dflash_has_any_non_causal`` decides pre-build whether the draft needs a
|
|
non-causal-capable backend, so its branch table (explicit override, SWA-derived
|
|
per-layer causality, and the no-``layer_types`` fallback) is worth pinning.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm.model_executor.models.qwen3_dflash import (
|
|
_dflash_layer_causal,
|
|
dflash_has_any_non_causal,
|
|
)
|
|
|
|
|
|
def _config(num_hidden_layers, layer_types=None, causal_override=None):
|
|
dflash_config = None if causal_override is None else {"causal": causal_override}
|
|
return SimpleNamespace(
|
|
num_hidden_layers=num_hidden_layers,
|
|
layer_types=layer_types,
|
|
dflash_config=dflash_config,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"config,expected",
|
|
[
|
|
# Override forces causality on every layer, ignoring layer_types.
|
|
(_config(2, layer_types=["full_attention"] * 2, causal_override=True), False),
|
|
# Override forces non-causal on every layer.
|
|
(
|
|
_config(2, layer_types=["sliding_attention"] * 2, causal_override=False),
|
|
True,
|
|
),
|
|
# SWA-derived: full-attention layers are non-causal.
|
|
(_config(2, layer_types=["sliding_attention", "full_attention"]), True),
|
|
# SWA-derived: all-sliding is fully causal.
|
|
(_config(2, layer_types=["sliding_attention", "sliding_attention"]), False),
|
|
# No layer_types -> non-causal fallback.
|
|
(_config(2, layer_types=None), True),
|
|
(_config(2, layer_types=[]), True),
|
|
],
|
|
)
|
|
def test_dflash_has_any_non_causal(config, expected):
|
|
assert dflash_has_any_non_causal(config) is expected
|
|
|
|
|
|
def test_dflash_layer_causal_is_per_layer():
|
|
config = _config(2, layer_types=["sliding_attention", "full_attention"])
|
|
assert _dflash_layer_causal(config, 0) is True
|
|
assert _dflash_layer_causal(config, 1) is False
|